linked-buffer 0.0.2

Yet another linked buffer implemention.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use std::{cell::UnsafeCell, mem::MaybeUninit, ptr::NonNull};

pub(crate) const BLOCK_CAP: usize = 4 * 1024;
thread_local! {
    pub(crate) static BLOCK_MANAGER: UnsafeCell<BlockManager>  = UnsafeCell::new(BlockManager::new());
}

#[derive(Debug)]
pub(crate) struct Block {
    /// Reference count of this block.
    /// Buffer hold one ref_count if it use its data for length > 0.
    pub(crate) ref_count: usize,

    /// Block data.
    pub(crate) buf: [MaybeUninit<u8>; BLOCK_CAP],
}

impl Block {
    /// Create a new Block.
    pub(crate) const fn new() -> Self {
        Block {
            ref_count: 0,
            buf: unsafe { MaybeUninit::uninit().assume_init() },
        }
    }

    /// Create a new Block and convert it to a ptr.
    /// You should manually release memory.
    pub(crate) fn new_ptr() -> NonNull<Self> {
        let block = Self::new();
        unsafe { NonNull::new_unchecked(Box::leak(Box::new(block))) }
    }

    pub(crate) fn get_buf(&self) -> &[u8] {
        unsafe { std::mem::transmute(self.buf.as_ref()) }
    }

    pub(crate) fn get_mut_buf(&mut self) -> &mut [u8] {
        unsafe { std::mem::transmute(self.buf.as_mut()) }
    }
}

#[derive(Debug)]
pub(crate) struct BlockNode {
    pub(crate) block: NonNull<Block>,
    pub(crate) next: Option<NonNull<BlockNode>>,
}

impl BlockNode {
    pub(crate) fn new_ptr() -> NonNull<Self> {
        let block = Block::new_ptr();
        let node = Self { block, next: None };
        unsafe { NonNull::new_unchecked(Box::leak(Box::new(node))) }
    }

    pub(crate) fn alloc() -> NonNull<Self> {
        let mut node = BLOCK_MANAGER.with(|m| unsafe { (&mut *m.get()).alloc_node() });
        unsafe { node.as_mut() }.ref_inc();
        node
    }

    pub(crate) fn ref_inc(&mut self) {
        unsafe { self.block.as_mut() }.ref_count += 1;
    }
}

pub(crate) struct BlockManager {
    head: Option<NonNull<BlockNode>>,
}

impl Drop for BlockManager {
    fn drop(&mut self) {
        let mut maybe_node = self.head.take();
        while let Some(mut node) = maybe_node {
            let mut block = unsafe { node.as_ref() }.block;
            maybe_node = unsafe { node.as_ref() }.next;

            // free block
            unsafe { std::ptr::drop_in_place(block.as_mut()) };
            // free node
            unsafe { std::ptr::drop_in_place(node.as_mut()) };
        }
    }
}

impl BlockManager {
    /// Create a new block manager.
    pub(crate) const fn new() -> Self {
        Self { head: None }
    }

    /// Alloc a BlockNode.
    pub(crate) fn alloc_node(&mut self) -> NonNull<BlockNode> {
        if let Some(mut node) = self.head {
            self.head = unsafe { node.as_mut() }.next.take();
            node
        } else {
            BlockNode::new_ptr()
        }
    }

    /// Free a BlockNode to free list.
    /// The block ref_count must be zero.
    pub(crate) unsafe fn free_node(&mut self, mut node: NonNull<BlockNode>) {
        debug_assert!(node.as_ref().block.as_ref().ref_count == 0);
        node.as_mut().next = self.head;
        self.head = Some(node);
    }
}

/// Free block node to tls pool.
pub(crate) unsafe fn free_block_node(node: NonNull<BlockNode>) {
    BLOCK_MANAGER.with(|m| unsafe { (&mut *m.get()).free_node(node) });
}

/// Free node only.
pub(crate) unsafe fn free_node(mut node: NonNull<BlockNode>) {
    std::ptr::drop_in_place(node.as_mut());
}

/// Drop node.
/// Ref dec and free.
pub(crate) unsafe fn drop_node(mut node: NonNull<BlockNode>) {
    let block_rc = &mut node.as_mut().block.as_mut().ref_count;
    *block_rc -= 1;
    if *block_rc == 0 {
        free_block_node(node);
    } else {
        free_node(node);
    }
}