From 815471d6d1ca91f2005ff8b112404d4e0f341a40 Mon Sep 17 00:00:00 2001
From: haskal <haskal@bepis.xyz>
Date: Sat, 14 Mar 2020 14:42:03 -0400
Subject: [PATCH] Implement more efficient scheme for ack packet coding

---
 hptp/src/msg.rs | 90 +++++++++++++++++++++++++++++++++++++++++++------
 hptp/src/seg.rs |  3 +-
 2 files changed, 80 insertions(+), 13 deletions(-)

diff --git a/hptp/src/msg.rs b/hptp/src/msg.rs
index 0ac02fb..d2bf8a2 100644
--- a/hptp/src/msg.rs
+++ b/hptp/src/msg.rs
@@ -1,5 +1,5 @@
 use super::seg::{SegData, SegIdx};
-pub use super::seg::{DOWN_HEADER_SIZE, MAX_TOTAL_PACKET_SIZE, UP_HEADER_SIZE};
+pub use super::seg::{DOWN_HEADER_SIZE, MAX_TOTAL_PACKET_SIZE, UP_HEADER_SIZE, MAX_SEG_SIZE};
 use byteorder::ByteOrder;
 
 #[derive(Clone)]
@@ -78,29 +78,97 @@ impl SerDes for UpMsg {
 
 impl SerDes for DownMsg {
     fn des(buf: &[u8]) -> Result<Self, DesError> {
-        let mut idxs = vec![];
-        for (i, b) in buf.iter().cloned().enumerate() {
-            for j in 0..8 {
-                if b & (1 << j) != 0 {
-                    idxs.push((i * 8 + j) as SegIdx);
+        if buf.len() < DOWN_HEADER_SIZE {
+            Err(DesError)
+        } else {
+            let first_unacked = BO::read_u32(&buf[0..4]);
+            let mut idxs = vec![];
+            for i in 0..first_unacked {
+                idxs.push(i);
+            }
+            for (i, b) in buf.iter().skip(4).cloned().enumerate() {
+                for j in 0..8 {
+                    if b & (1 << j) != 0 {
+                        idxs.push((i * 8 + j + (first_unacked as usize) + 1) as SegIdx);
+                    }
                 }
             }
+            Ok(DownMsg::Ack { idxs })
         }
-        Ok(DownMsg::Ack { idxs })
     }
 
     fn ser_to(&self, buf: &mut [u8]) -> usize {
         match self {
             DownMsg::Ack { idxs } => {
+                let mut first_unacked: u32 = 0;
+                for idx in idxs.iter().cloned() {
+                    if idx == first_unacked {
+                        first_unacked += 1;
+                    } else {
+                        break;
+                    }
+                }
+                BO::write_u32(&mut buf[0..4], first_unacked);
                 let mut max_i = 0;
                 for seg_idx in idxs.iter().cloned() {
-                    let i = seg_idx as usize / 8;
-                    let j = (seg_idx % 8) as u8;
-                    buf[i] |= 1 << j;
+                    if seg_idx < first_unacked {
+                        continue;
+                    }
+                    let seg_idx_delta: usize = (seg_idx - first_unacked - 1) as usize;
+                    let i = seg_idx_delta / 8;
+                    let j = (seg_idx_delta % 8) as u8;
+                    buf[i + 4] |= 1 << j;
                     max_i = std::cmp::max(i, max_i);
                 }
-                max_i + 1
+                max_i + 4 + 1
             }
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn serdes_downmsg(init_idxs: &Vec<SegIdx>, expected_hdr: u32) {
+        let mut buf = [0u8; MAX_SEG_SIZE];
+        DownMsg::Ack{idxs: init_idxs.clone()}.ser_to(&mut buf);
+        assert_eq!(BO::read_u32(&buf[0..4]), expected_hdr);
+        let des = DownMsg::des(&buf).unwrap();
+        match des {
+            DownMsg::Ack { idxs } => {
+                assert_eq!(*init_idxs, idxs);
+            }
+        }
+    }
+
+    #[test]
+    fn test_serdes_downmsg() {
+        let init_idxs: Vec<SegIdx> = vec![0, 1, 2, 3];
+        serdes_downmsg(&init_idxs, 4);
+    }
+
+    #[test]
+    fn test_serdes_downmsg_empty() {
+        let init_idxs: Vec<SegIdx> = vec![];
+        serdes_downmsg(&init_idxs, 0);
+    }
+
+    #[test]
+    fn test_serdes_downmsg_missing0() {
+        let init_idxs: Vec<SegIdx> = vec![1, 2, 4];
+        serdes_downmsg(&init_idxs, 0);
+    }
+
+    #[test]
+    fn test_serdes_downmsg_missing1() {
+        let init_idxs: Vec<SegIdx> = vec![0, 2, 7];
+        serdes_downmsg(&init_idxs, 1);
+    }
+
+    #[test]
+    fn test_serdes_downmsg_missing2() {
+        let init_idxs: Vec<SegIdx> = vec![0, 1, 2, 3, 4, 7, 8, 10];
+        serdes_downmsg(&init_idxs, 5);
+    }
+}
diff --git a/hptp/src/seg.rs b/hptp/src/seg.rs
index 6839959..5066093 100644
--- a/hptp/src/seg.rs
+++ b/hptp/src/seg.rs
@@ -1,9 +1,8 @@
 /// Per the assignment spec, `1472` is the maximum size packet we're allowed to send.
 pub const MAX_TOTAL_PACKET_SIZE: usize = 1472;
 
-// TODO: change these based off the decoders
 pub const UP_HEADER_SIZE: usize = 4;
-pub const DOWN_HEADER_SIZE: usize = 1;
+pub const DOWN_HEADER_SIZE: usize = 4;
 
 /// This is the maximum amount of segment data we can fit into a packet.
 pub const MAX_SEG_SIZE: usize = MAX_TOTAL_PACKET_SIZE - UP_HEADER_SIZE;