pullup_with_len

1 files changed, 91 insertions, 23 deletions

diff --git a/src/session/segment_buffer.rs b/src/session/segment_buffer.rs
index 1590241..f9f00b6 100644
--- a/src/session/segment_buffer.rs
+++ b/src/session/segment_buffer.rs
@@ -34,8 +34,6 @@ pub struct SegmentBuffer {
     max_n_segments: usize,
     max_buf_len: usize,
     used_buf_len: usize,
-
-    // todo: 支持有截止的输出（输出一部分）
 }
 
 impl SegmentBuffer {
@@ -77,7 +75,7 @@ impl SegmentBuffer {
 
         while !self.segments.is_empty() {
             if self.segments.first_entry().unwrap().get().offset > self.next_offset { // there is a gap
-                println!("Gap detected");
+                println!("Gap detected, the offset is {}, expecting: {}", self.segments.first_entry().unwrap().get().offset.0, self.next_offset.0);
                 break;
             }
 
@@ -95,29 +93,33 @@ impl SegmentBuffer {
             return Some(ret);
         }
     }
-
-    pub fn clear(&mut self) -> Vec<Vec<u8>> {
-        let mut ret = Vec::new();
-        
-        while !self.segments.is_empty() {
-            let segment = self.segments.pop_first().unwrap().1;
-
-            if segment.offset + Wrapping(segment.payload.len() as u32) <= self.next_offset {
-                continue;
-            }
-
-            let rel = segment.payload.len() as u32;
-            if segment.offset > self.next_offset || ret.is_empty() {
-                ret.push(segment.payload);
-                self.next_offset = segment.offset + Wrapping(rel);
+    pub fn pullup_with_len(&mut self, len: u32) -> Option<Vec<u8>> {
+        let mut ret = None;
+        if let Some(s) = self.pullup() {
+            let s_len = s.len() as u32;
+            if s_len > len {
+                let (s1, s2) = s.split_at(len as usize);
+                ret = Some(s1.to_vec());
+
+                self.next_offset -= Wrapping(s2.len() as u32); // s2 is not sent
+                self.insert_sorted(Segment {
+                    offset: self.next_offset,
+                    payload: s2.to_vec(),
+                });
             } else {
-                if let Some(last) = ret.last_mut() {
-                    let extension = segment.offset_part(self.next_offset);
-                    last.extend_from_slice(extension);
-                    self.next_offset += Wrapping(extension.len() as u32);
-                }
+                ret = Some(s);
             }
         }
+        
+        ret
+    }
+    pub fn clear(&mut self) -> Vec<Vec<u8>> {
+        let mut ret = Vec::new();
+        self.fill_hole();
+        while let Some(s) = self.pullup() {
+            ret.push(s);
+            self.fill_hole();
+        }
 
         self.used_buf_len = 0;
         self.next_offset = Wrapping(0);
@@ -129,6 +131,38 @@ impl SegmentBuffer {
         self.used_buf_len += s.payload.len();
         self.segments.insert(s.offset.0, s);
     }
+    fn fill_hole(&mut self) -> u32 {
+        if self.segments.is_empty() {
+            return 0;
+        }
+        let start_at = self.segments.first_entry().unwrap().get().offset;
+        if start_at <= self.next_offset {
+            return 0;
+        }
+        let hole_len = (start_at - self.next_offset).0;
+        self.next_offset = start_at;
+
+        hole_len
+    }
+    // remove the hole. Let the data to the beginning of the buffer.
+    // pub fn prepend(&mut self) {
+    //     let hole_len = self.fill_hole();
+    //     println!("Hole length: {}", hole_len);
+    //     let start_at = self.segments.first_entry().unwrap().get().offset;
+    //     println!("Prepend {} bytes", start_at.0);
+
+    //     for (_, s) in self.segments.iter_mut() {
+    //         s.offset -= start_at;
+    //         todo: 光改这儿不行，还要改btree的key，这样就有点复杂了，要不然改成加一个intial offset，把所有以后输入的都加一个值？总之暂时不实现了
+    //     }
+    //     self.next_offset -= start_at;
+
+    //     for (_, s) in self.segments.iter_mut() {
+    //         println!("Offset: {}, payload: {:?}", s.offset.0, s.payload);
+    //     }
+
+    //     println!("Next offset: {}", self.next_offset.0)
+    // }
 }
 
 #[cfg(test)]
@@ -307,4 +341,38 @@ mod tests {
         assert_eq!(s.update(9, &[12]), Description::TooManySegments);
         assert_eq!(s.clear(), vec![vec![6, 7, 8, 11]]);
     }
+
+    #[test]
+    fn pullup_given_different_buffer() {
+        let mut s = SegmentBuffer::new(10000, 100);
+        assert_eq!(s.update(0, &[1, 2, 3, 4, 5, 6]), Description::Ok);
+        assert_eq!(s.pullup_with_len(5).unwrap(), &[1, 2, 3, 4, 5]);
+        assert_eq!(s.pullup_with_len(5).unwrap(), &[6]);
+
+        assert_eq!(s.update(6, &[1, 2, 3, 4, 5, 6]), Description::Ok);
+        assert_eq!(s.pullup_with_len(10).unwrap(), &[1, 2, 3, 4, 5, 6]); // big enough
+        assert_eq!(s.pullup_with_len(10), None);
+    }
+
+    #[test]
+    fn pulluplen_and_continue_to_add() {
+        let mut s = SegmentBuffer::new(10000, 100);
+        assert_eq!(s.update(0, &[1, 2, 3, 4, 5, 6]), Description::Ok);
+        assert_eq!(s.pullup_with_len(5).unwrap(), &[1, 2, 3, 4, 5]);
+
+        assert_eq!(s.update(6, &[7, 8, 9, 10, 11, 12]), Description::Ok);
+        assert_eq!(s.pullup_with_len(10).unwrap(), &[6, 7, 8, 9, 10, 11, 12]);
+    }
+
+    // #[test]
+    // fn prepend_when_has_hole() {
+    //     let mut s = SegmentBuffer::new(10000, 100);
+    //     assert_eq!(s.update(0, &[1, 2]), Description::Ok);
+    //     assert_eq!(s.update(3, &[4, 5]), Description::Ok);
+    //     s.pullup();
+
+    //     s.prepend();
+    //     assert_eq!(s.update(2, &[42]), Description::Ok);
+    //     assert_eq!(s.pullup().unwrap(), &[4, 5, 42]);
+    // }
 }