path: root/src/session/segment_buffer.rs

use std::collections::BTreeMap;
use std::num::Wrapping;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SegmentBufferErr {
    TooManySegments,
    TooLargeData,
}

#[derive(Debug, Clone)]
struct Segment {
    offset: Wrapping<u32>,
    payload: Vec<u8>,
}

impl Segment {
    fn offset_part(&self, offset: Wrapping<u32>) -> &[u8] {
        let this_right = self.offset + Wrapping(self.payload.len() as u32);
        if this_right <= offset {
            return &[];
        } else {
            let overlap_size = (offset - self.offset).0 as usize;
            return &self.payload[overlap_size..];
        }
    }
}

#[derive(Debug)]
pub struct SegmentBuffer {
    next_offset: Wrapping<u32>,
    segments: BTreeMap<u32, Segment>,

    max_n_segments: usize,
    max_buf_len: usize,
    used_buf_len: usize,
}

impl SegmentBuffer {
    pub fn new(max_buf_len: usize, max_n_segments: usize) -> Self {
        SegmentBuffer {
            next_offset: Wrapping(0),
            segments: BTreeMap::new(),
            max_n_segments,
            max_buf_len,
            used_buf_len: 0,
        }
    }
    pub fn update(&mut self, offset: u32, payload: &[u8]) -> Result<(), SegmentBufferErr> {
        if self.segments.len() >= self.max_n_segments {
            return Err(SegmentBufferErr::TooManySegments);
        }

        if self.used_buf_len + payload.len()> self.max_buf_len {
            return Err(SegmentBufferErr::TooLargeData);
        }

        if payload.is_empty() {
            return Ok(());
        }
        let segment = Segment {
            offset: Wrapping(offset),
            payload: payload.to_vec(),
        };
        self.insert_sorted(segment);

        Ok(())
    }
    pub fn pullup(&mut self) -> Option<Vec<u8>> {
        if self.segments.is_empty() {
            println!("No segment to send");
            return None;
        }
        let mut ret = Vec::new();

        while !self.segments.is_empty() {
            if self.segments.first_entry().unwrap().get().offset > self.next_offset { // there is a gap
                println!("Gap detected, the offset is {}, expecting: {}", self.segments.first_entry().unwrap().get().offset.0, self.next_offset.0);
                break;
            }

            let segment = self.segments.pop_first().unwrap().1;
            let extension = segment.offset_part(self.next_offset);
            ret.extend_from_slice(extension);
            self.next_offset += Wrapping(extension.len() as u32);

            self.used_buf_len -= segment.payload.len();
        }

        if ret.is_empty() {
            return None;
        } else {
            return Some(ret);
        }
    }
    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 {
                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);
    
        ret
    }

    fn insert_sorted(&mut self, s: Segment) {
        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)]
mod tests {
    use super::*;

    #[test]
    fn several_ordered_segments() {
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(0, &[1, 2, 3]), Ok(()));
        assert_eq!(s.update(3, &[4, 5]), Ok(()));
        assert_eq!(s.update(5, &[6]), Ok(()));
        
        assert_eq!(s.pullup().unwrap(), &[1, 2, 3, 4, 5, 6]);
        assert_eq!(s.clear(), Vec::<Vec<u8>>::new());
    }

    #[test]
    fn several_unordered_segments() {
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(5, &[6]), Ok(()));
        assert_eq!(s.update(0, &[1, 2, 3]), Ok(()));
        assert_eq!(s.update(3, &[4, 5]), Ok(()));

        let expected_clear = vec![vec![1, 2, 3, 4, 5, 6]];
        assert_eq!(s.clear(), expected_clear);
    }

    #[test]
    fn with_hole() {
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(0, &[1, 2]), Ok(())); // miss 3
        assert_eq!(s.update(3, &[4, 5]), Ok(()));
        assert_eq!(s.update(5, &[6]), Ok(()));  // miss 7
        assert_eq!(s.update(7, &[8, 9]), Ok(()));

        let expected_clear = vec![vec![1, 2], vec![4, 5, 6], vec![8, 9]];
        assert_eq!(s.clear(), expected_clear);
    }

    #[test]
    fn duplicate_packet() {
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(0, &[1, 2, 3]), Ok(()));
        assert_eq!(s.update(0, &[1, 2, 3]), Ok(()));
        assert_eq!(s.update(3, &[4, 5]), Ok(()));

        let expected_clear = vec![vec![1, 2, 3, 4, 5]];
        assert_eq!(s.clear(), expected_clear);
    }

    #[test]
    fn pop_empty() {
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.pullup(), None);
    }

    #[test]
    fn pop_blocked_by_hole() {
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(0, &[1, 2]), Ok(())); // miss 3
        assert_eq!(s.update(3, &[4, 5]), Ok(()));

        assert_eq!(s.pullup().unwrap(), &[1, 2]);
        assert_eq!(s.pullup(), None);
        assert_eq!(s.pullup(), None);

        assert_eq!(s.update(2, &[3]), Ok(()));
        assert_eq!(s.pullup().unwrap(), &[3, 4, 5]);
    }

    #[test]
    fn overlap_when_popped() {
        // [1,2,3] -> popped
        // [3,4,5,6] -> [4,5,6]
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(0, &[1, 2, 3]), Ok(()));
        assert_eq!(s.pullup().unwrap(), &[1, 2, 3]);
        assert_eq!(s.update(2, &[3, 4, 5, 6]), Ok(()));
        assert_eq!(s.pullup().unwrap(), &[4, 5, 6]);
    }

    #[test]
    fn overlap_as_old_packet() {
        // [1,2,3,4,5,6] -> popped
        // [2,3,4] -> drop(description: old packet)
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(0, &[1, 2, 3, 4, 5, 6]), Ok(()));
        assert_eq!(s.pullup().unwrap(), &[1, 2, 3, 4, 5, 6]);
        assert_eq!(s.update(1, &[2, 3, 4]), Ok(()));
        assert_eq!(s.pullup(), None);
    }

    #[test]
    fn overlap_change_next_one() {
        // [2,3,4] -> wait
        // [3,4,5] -> wait(join the first as [2,3,4,5])
        // [1] -> send [1 2,3,4,5]
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(1, &[2, 3, 4]), Ok(()));
        assert_eq!(s.update(2, &[3, 4, 5]), Ok(()));
        assert!(s.pullup().is_none());

        assert_eq!(s.update(0, &[1]), Ok(()));
        assert_eq!(s.pullup().unwrap(), &[1, 2, 3, 4, 5]);
    }

    #[test]
    fn overlap_del_next_one() {
        // [2,3,4,5] -> wait
        // [3,4,5] -> del(overlapped)
        // [1,2,3] -> send [1,2,3,4,5]
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(1, &[2, 3, 4, 5]), Ok(()));
        assert_eq!(s.update(2, &[3, 4, 5]), Ok(()));
        assert!(s.pullup().is_none());

        assert_eq!(s.update(0, &[1, 2, 3]), Ok(()));
        assert_eq!(s.pullup().unwrap(), &[1, 2, 3, 4, 5]);
    }

    #[test]
    fn overlap_surpass_all() {
        // [2] -> wait
        // [3] -> wait
        // [1,2,3,4] -> send [1,2,3,4]
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(1, &[2]), Ok(()));
        assert_eq!(s.update(2, &[3]), Ok(()));
        assert!(s.pullup().is_none());

        assert_eq!(s.update(0, &[1, 2, 3, 4]), Ok(()));
        assert_eq!(s.pullup().unwrap(), &[1, 2, 3, 4]);   
    }

    #[test]
    fn clear_and_restart() {
        // [2,3,4] -> wait
        // [3,4,5] -> clear ( return [2,3,4,5] )
        // [1] -> send [1]
        let mut s = SegmentBuffer::new(10000, 100);
        assert_eq!(s.update(1, &[2, 3, 4]), Ok(()));
        assert_eq!(s.update(2, &[3, 4, 5]), Ok(()));
        assert_eq!(s.clear(), vec![vec![2, 3, 4, 5]]);

        assert_eq!(s.update(0, &[1]), Ok(()));
        assert_eq!(s.pullup().unwrap(), &[1]);
    }

    #[test]
    fn full_window() {
        let mut s = SegmentBuffer::new(10, 100);

        assert_eq!(s.update(0, &[1, 2, 3, 4]), Ok(()));
        assert_eq!(s.update(4, &[5, 6, 7, 8]), Ok(()));
        assert_eq!(s.update(8, &[2, 3, 4, 5]), Err(SegmentBufferErr::TooLargeData));

        assert_eq!(s.pullup().unwrap(), &[1, 2, 3, 4, 5, 6, 7, 8]); // pop will clear window size

        let v:Vec<_> = (1..=10).collect();
        assert_eq!(s.update(8, v.as_slice()), Ok(()));
        assert_eq!(s.update(18, &[1]), Err(SegmentBufferErr::TooLargeData));
    }

    #[test]
    fn too_many_packet() {
        let mut s = SegmentBuffer::new(10000, 2);

        assert_eq!(s.update(0, &[1, 2, 3, 4]), Ok(()));
        assert_eq!(s.update(5, &[6, 7, 8]), Ok(())); // hole
        assert_eq!(s.update(8, &[2, 3, 4, 5]), Err(SegmentBufferErr::TooManySegments));

        assert_eq!(s.pullup().unwrap(), &[1, 2, 3, 4]); // pop 1

        assert_eq!(s.update(8, &[11]), Ok(()));
        assert_eq!(s.update(9, &[12]), Err(SegmentBufferErr::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]), 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]), 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]), 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]), 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]), Ok(()));
    //     assert_eq!(s.update(3, &[4, 5]), Ok(()));
    //     s.pullup();

    //     s.prepend();
    //     assert_eq!(s.update(2, &[42]), Ok(()));
    //     assert_eq!(s.pullup().unwrap(), &[4, 5, 42]);
    // }
}