linked-buffer 0.0.2

use std::ptr::NonNull;

use crate::{
    block::{drop_node, BlockNode, BLOCK_CAP},
    Buffer,
};

/// BufferMut
// Implemention notice:
// It is allowed that pointer is at prev block and offset is BLOCK_CAP,
// it is also allowed that the pointer at next block and offset is 0.
// These two form is the same.
pub struct BufferMut {
    // the first block pointer
    pub(crate) head: Option<NonNull<BlockNode>>,
    // the write block pointer(not the last)
    pub(crate) tail: Option<NonNull<BlockNode>>,

    // read offset
    pub(crate) read_offset: usize,
    // write offset
    pub(crate) write_offset: usize,
    // whole capacity
    pub(crate) cap: usize,
    // whole length
    pub(crate) len: usize,
}

impl Drop for BufferMut {
    /// Loop to get every BlockNode and free it.
    #[inline]
    fn drop(&mut self) {
        let mut maybe_node = self.head;
        while let Some(mut node) = maybe_node {
            maybe_node = unsafe { node.as_mut() }.next.take();
            unsafe { drop_node(node) };
        }
    }
}

impl BufferMut {
    /// Creates a new `BufferMut` with at least specified capacity.
    #[inline]
    pub fn with_capacity(capacity: usize) -> Self {
        let mut allocated = 0;
        let mut head = None;
        while allocated < capacity {
            let mut node = BlockNode::alloc();
            allocated += BLOCK_CAP;

            let node_mut = unsafe { node.as_mut() };
            node_mut.next = head;
            head = Some(node);
        }
        Self {
            head,
            tail: head,
            read_offset: 0,
            write_offset: 0,
            cap: allocated,
            len: 0,
        }
    }

    /// Creates a new `BufferMut` with default capacity.
    #[inline]
    pub const fn new() -> Self {
        Self {
            head: None,
            tail: None,
            read_offset: 0,
            write_offset: 0,
            cap: 0,
            len: 0,
        }
    }

    /// Returns the number of bytes contained in this `BufferMut`.
    #[inline]
    pub fn len(&self) -> usize {
        self.len
    }

    /// Returns true if the `BufferMut` has a length of 0.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Returns the number of bytes the `BufferMut` can hold.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.cap
    }

    /// Clears the buffer, removing all data. Existing capacity is preserved.
    #[inline]
    pub fn clear(&mut self) {
        self.len = 0;
        self.tail = self.head;
        self.write_offset = self.read_offset;
    }

    /// Sets the length of the buffer.
    ///
    /// # Safety
    /// The caller must ensure that the data has been initialized.
    #[inline]
    pub unsafe fn set_len(&mut self, len: usize) {
        debug_assert!(len <= self.cap, "set_len out of bounds");
        self.len = len;

        let mut target_offset = self.read_offset + len;
        let mut new_tail = self.head;

        while target_offset > BLOCK_CAP {
            target_offset -= BLOCK_CAP;
            // new_tail must be Some since the length <= capacity
            new_tail = new_tail.unwrap().as_mut().next;
        }

        // if current offset is at the last of the block, and next block is not nullptr
        // we will move to next block.
        if target_offset == BLOCK_CAP && new_tail.is_some() {
            // new_tail is Some, we checked it already
            if let Some(next) = new_tail.unwrap_unchecked().as_mut().next {
                target_offset = 0;
                new_tail = Some(next);
            }
        }

        self.tail = new_tail;
        self.write_offset = target_offset;
    }

    /// Reserves capacity for at least `additional` more bytes to be inserted
    /// into the given `BufferMut`.
    #[inline]
    pub fn reserve(&mut self, additional: usize) {
        let want = self.len() + additional;
        while self.cap < want {
            self.reserve_block();
        }
    }

    /// Advance read offset.
    #[inline]
    pub fn advance(&mut self, cnt: usize) {
        assert!(
            cnt <= self.len(),
            "cannot advance past `remaining`: {:?} <= {:?}",
            cnt,
            self.len(),
        );
        self.len -= cnt;
        self.cap -= cnt;

        let mut offset = self.read_offset + cnt;
        while offset > BLOCK_CAP {
            offset -= BLOCK_CAP;

            // we have to move to next block, it must exists since we checked the cnt.
            let mut to_drop = self.head.unwrap();
            self.head = unsafe { to_drop.as_mut() }.next;

            unsafe { drop_node(to_drop) };
        }
        if offset == BLOCK_CAP && self.head.is_some() {
            // move to next if the next exists
            // head is Some, we checked it already
            if let Some(next) = unsafe { self.head.unwrap_unchecked().as_mut() }.next {
                self.head = Some(next);
                offset = 0;
            }
        }
        self.read_offset = offset;
    }

    /// Advance write offset.
    ///
    /// # Safety
    /// User must ensure the data is initialized.
    #[inline]
    pub unsafe fn advance_mut(&mut self, cnt: usize) {
        let new_len = self.len() + cnt;
        assert!(
            new_len <= self.cap,
            "new_len = {}; capacity = {}",
            new_len,
            self.cap
        );
        self.len = new_len;

        let mut offset = self.write_offset + cnt;
        while offset > BLOCK_CAP {
            offset -= BLOCK_CAP;

            // we have to move to next block, it must exists since we checked the cnt.
            self.tail = self.tail.unwrap().as_ref().next;
        }
        self.write_offset = offset;
    }

    /// Freeze and convert to Buffer
    #[inline]
    pub fn freeze(mut self) -> Buffer {
        if self.len == 0 {
            // release all head chain
            let mut head = self.head;
            while let Some(node) = head {
                head = unsafe { node.as_ref() }.next;
                unsafe { drop_node(node) };
            }

            return Buffer {
                head: None,
                read_offset: 0,
                len: 0,
            };
        }

        // length is not 0, head must exists
        if self.read_offset == BLOCK_CAP {
            self.read_offset = 0;
            let mut h = self.head.unwrap();
            self.head = unsafe { h.as_mut() }.next;
            unsafe { drop_node(h) };
        }

        let mut node = self.head.unwrap();
        let mut cnt = 0;
        let target = self.len + self.read_offset;
        loop {
            cnt += BLOCK_CAP;
            if cnt >= target {
                break;
            }
            node = unsafe { node.as_mut() }.next.unwrap();
        }

        let mut to_drop = unsafe { node.as_mut() }.next.take();
        while let Some(node) = to_drop {
            to_drop = unsafe { node.as_ref() }.next;
            unsafe { drop_node(node) };
        }
        Buffer {
            head: self.head,
            read_offset: self.read_offset,
            len: self.len,
        }
    }

    pub(crate) fn reserve_block(&mut self) {
        let mut node = BlockNode::alloc();
        if let Some(mut tail) = self.tail {
            unsafe { node.as_mut() }.next = unsafe { tail.as_mut() }.next.take();
            unsafe { tail.as_mut() }.next = Some(node);
        } else {
            if self.head.is_none() {
                self.head = Some(node);
            }
            self.tail = Some(node);
            self.read_offset = 0;
            self.write_offset = 0;
        }

        self.cap += BLOCK_CAP;

        // there must be a new block, so we can move it when the current one is full.
        if self.write_offset == BLOCK_CAP {
            self.write_offset = 0;
            // we have to move to next block, it must exists since we allocated it.
            self.tail = unsafe { self.tail.unwrap().as_ref() }.next;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new() {
        let mut b = BufferMut::new();
        assert_eq!(b.len(), 0);
        assert!(b.is_empty());
        assert_eq!(b.capacity(), 0);

        b.reserve(8);
        unsafe { b.set_len(4) };
        assert_eq!(b.len(), 4, "length should be 4");
        assert_eq!(b.capacity(), BLOCK_CAP, "one block should be allocated");
        assert_eq!(b.write_offset, 4);
    }

    #[test]
    fn set_len() {
        let b = BufferMut::with_capacity(BLOCK_CAP * 3);
        assert_eq!(b.len(), 0);
        assert_eq!(b.capacity(), BLOCK_CAP * 3);

        let mut b = BufferMut::with_capacity(BLOCK_CAP * 3 + 1);
        assert_eq!(b.len(), 0);
        assert_eq!(b.capacity(), BLOCK_CAP * 4);

        unsafe { b.set_len(BLOCK_CAP) };
        assert_eq!(b.len(), BLOCK_CAP, "length should equals to BLOCK_CAP");
        assert_eq!(b.capacity(), BLOCK_CAP * 4);
        assert_eq!(b.write_offset, 0);

        unsafe { b.set_len(1) };
        assert_eq!(b.len(), 1, "length should equals to 1");
        assert_eq!(b.capacity(), BLOCK_CAP * 4);
        assert_eq!(b.write_offset, 1);
        assert_eq!(b.head, b.tail);
    }

    #[test]
    fn clear() {
        let mut b = BufferMut::with_capacity(BLOCK_CAP + 1);
        unsafe { b.set_len(BLOCK_CAP) };
        b.clear();
        assert!(b.is_empty());
        assert_eq!(b.capacity(), BLOCK_CAP * 2);
    }

    #[test]
    #[should_panic]
    fn advance_panic() {
        let mut b = BufferMut::new();
        b.advance(4);
    }

    #[test]
    fn advance() {
        let mut b = BufferMut::with_capacity(BLOCK_CAP);
        unsafe { b.set_len(8) };
        b.advance(4);

        assert_eq!(b.capacity(), BLOCK_CAP - 4);
        assert_eq!(b.len(), 4);

        unsafe { b.set_len(2) };
        assert_eq!(b.capacity(), BLOCK_CAP - 4);
        assert_eq!(b.len(), 2);
    }
}