From 4ccd209df7f13511f443229d2ff2057cd812c56b Mon Sep 17 00:00:00 2001
From: Hailey Somerville <hailey@hailey.lol>
Date: Wed, 27 Dec 2023 15:57:58 +1100
Subject: [PATCH] modify decode trait to allow for signalling missing packets

---
 bark-core/src/decode/mod.rs  |  9 +++++----
 bark-core/src/decode/opus.rs | 15 ++++++++++++---
 bark-core/src/decode/pcm.rs  | 13 ++++++++++---
 bark/src/receive.rs          |  2 +-
 4 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/bark-core/src/decode/mod.rs b/bark-core/src/decode/mod.rs
index 043c19b..2e4b7a2 100644
--- a/bark-core/src/decode/mod.rs
+++ b/bark-core/src/decode/mod.rs
@@ -47,13 +47,14 @@ impl Decoder {
         &self.decode as &dyn Display
     }
 
-    pub fn decode(&mut self, packet: &Audio, out: &mut SampleBuffer) -> Result<(), DecodeError> {
-        self.decode.decode_packet(packet.buffer_bytes(), out)
+    pub fn decode(&mut self, packet: Option<&Audio>, out: &mut SampleBuffer) -> Result<(), DecodeError> {
+        let bytes = packet.map(|packet| packet.buffer_bytes());
+        self.decode.decode_packet(bytes, out)
     }
 }
 
 trait Decode: Display {
-    fn decode_packet(&mut self, bytes: &[u8], out: &mut SampleBuffer) -> Result<(), DecodeError>;
+    fn decode_packet(&mut self, bytes: Option<&[u8]>, out: &mut SampleBuffer) -> Result<(), DecodeError>;
 }
 
 enum DecodeFormat {
@@ -63,7 +64,7 @@ enum DecodeFormat {
 }
 
 impl Decode for DecodeFormat {
-    fn decode_packet(&mut self, bytes: &[u8], out: &mut SampleBuffer) -> Result<(), DecodeError> {
+    fn decode_packet(&mut self, bytes: Option<&[u8]>, out: &mut SampleBuffer) -> Result<(), DecodeError> {
         match self {
             DecodeFormat::S16LE(dec) => dec.decode_packet(bytes, out),
             DecodeFormat::F32LE(dec) => dec.decode_packet(bytes, out),
diff --git a/bark-core/src/decode/opus.rs b/bark-core/src/decode/opus.rs
index be86825..79653c1 100644
--- a/bark-core/src/decode/opus.rs
+++ b/bark-core/src/decode/opus.rs
@@ -26,11 +26,20 @@ impl Display for OpusDecoder {
 }
 
 impl Decode for OpusDecoder {
-    fn decode_packet(&mut self, bytes: &[u8], out: &mut SampleBuffer) -> Result<(), DecodeError> {
+    fn decode_packet(&mut self, bytes: Option<&[u8]>, out: &mut SampleBuffer) -> Result<(), DecodeError> {
         let expected = out.len() / 2;
-        let length = self.opus.decode_float(bytes, out, false)?;
+        let length;
 
-        log::debug!("opus decode: bytes={} -> frames={}", bytes.len(), length);
+        match bytes {
+            Some(bytes) => {
+                length = self.opus.decode_float(bytes, out, false)?;
+                log::debug!("opus decode: bytes={} -> frames={}", bytes.len(), length);
+            }
+            None => {
+                length = self.opus.decode_float(&[], out, false)?;
+                log::debug!("opus decode: correcting packet loss! frames={}", length);
+            }
+        };
 
         if expected != length {
             return Err(DecodeError::WrongLength { length, expected });
diff --git a/bark-core/src/decode/pcm.rs b/bark-core/src/decode/pcm.rs
index 348408c..52f76a4 100644
--- a/bark-core/src/decode/pcm.rs
+++ b/bark-core/src/decode/pcm.rs
@@ -11,7 +11,7 @@ impl Display for S16LEDecoder {
 }
 
 impl Decode for S16LEDecoder {
-    fn decode_packet(&mut self, bytes: &[u8], out: &mut SampleBuffer) -> Result<(), DecodeError> {
+    fn decode_packet(&mut self, bytes: Option<&[u8]>, out: &mut SampleBuffer) -> Result<(), DecodeError> {
         decode_packed(bytes, out, |bytes| {
             let input = i16::from_le_bytes(bytes);
             let scale = i16::MAX as f32;
@@ -29,16 +29,23 @@ impl Display for F32LEDecoder {
 }
 
 impl Decode for F32LEDecoder {
-    fn decode_packet(&mut self, bytes: &[u8], out: &mut SampleBuffer) -> Result<(), DecodeError> {
+    fn decode_packet(&mut self, bytes: Option<&[u8]>, out: &mut SampleBuffer) -> Result<(), DecodeError> {
         decode_packed(bytes, out, f32::from_le_bytes)
     }
 }
 
 fn decode_packed<const N: usize>(
-    bytes: &[u8],
+    bytes: Option<&[u8]>,
     out: &mut SampleBuffer,
     func: impl Fn([u8; N]) -> f32,
 ) -> Result<(), DecodeError> {
+    let Some(bytes) = bytes else {
+        // PCM codecs have no packet loss correction
+        // just zero fill and return
+        out.fill(0.0);
+        return Ok(());
+    };
+
     check_length(bytes, out.len() * N)?;
 
     for (input, output) in bytes.chunks_exact(N).zip(out) {
diff --git a/bark/src/receive.rs b/bark/src/receive.rs
index 55a43d2..e2222a4 100644
--- a/bark/src/receive.rs
+++ b/bark/src/receive.rs
@@ -203,7 +203,7 @@ impl Receiver {
         // decode packet
         let mut decode_buffer: SampleBuffer = array::from_fn(|_| 0.0);
         if let Some(decoder) = stream.decoder.as_mut() {
-            match decoder.decode(&packet, &mut decode_buffer) {
+            match decoder.decode(Some(&packet), &mut decode_buffer) {
                 Ok(()) => {}
                 Err(e) => {
                     log::warn!("error in decoder, skipping packet: {e}");