1
use std::cell::Cell;
2
use std::marker::PhantomData;
3
use std::ptr::NonNull;
4

            
5
/// Intrusive node requirements for use in [`FreeList`].
6
///
7
/// # Safety
8
///
9
/// Implementors must guarantee that `get_next` and `set_next` read and write
10
/// the same link field, and that the field is valid for all nodes managed by
11
/// a corresponding freelist.
12
pub unsafe trait FreeListEntry: Sized {
13
    /// Returns the next pointer for `ptr` (or null).
14
    ///
15
    /// # Safety
16
    ///
17
    /// `ptr` must be a valid pointer to a node managed by the corresponding
18
    /// freelist.
19
    unsafe fn get_next(ptr: *mut Self) -> *mut Self;
20

            
21
    /// Sets the next pointer for `ptr`.
22
    ///
23
    /// # Safety
24
    ///
25
    /// `ptr` must be a valid pointer to a node managed by the corresponding
26
    /// freelist, and `next` must be either null or a valid pointer to a node
27
    /// managed by the same freelist.
28
    unsafe fn set_next(ptr: *mut Self, next: *mut Self);
29
}
30

            
31
/// Intrusive freelist based on a stack.
32
pub struct FreeList<T: FreeListEntry> {
33
    /// The head of the freelist. Null means empty.
34
    head: Cell<*mut T>,
35

            
36
    /// We implement Send and Sync manually.
37
    _marker: PhantomData<*mut T>,
38
}
39

            
40
// Safety: Transferring a FreeList to another thread is safe if T is Send.
41
unsafe impl<T: FreeListEntry + Send> Send for FreeList<T> {}
42

            
43
impl<T: FreeListEntry> Default for FreeList<T> {
44
    fn default() -> Self {
45
        Self::new()
46
    }
47
}
48

            
49
impl<T: FreeListEntry> FreeList<T> {
50
4928
    pub fn new() -> Self {
51
4928
        Self {
52
4928
            head: Cell::new(std::ptr::null_mut()),
53
4928
            _marker: PhantomData,
54
4928
        }
55
4928
    }
56

            
57
    /// Pops one entry from the freelist.
58
14853164
    pub fn try_pop(&self) -> Option<NonNull<T>> {
59
14853164
        let node = NonNull::new(self.head.get())?;
60

            
61
        // Safety: `head` is non-null.
62
13695071
        let next = unsafe { T::get_next(node.as_ptr()) };
63
13695071
        self.head.set(next);
64

            
65
13695071
        Some(node)
66
14853164
    }
67

            
68
    /// Pushes an entry onto the freelist.
69
13830029
    pub fn push(&self, entry: NonNull<T>) {
70
13830029
        let ptr = entry.as_ptr();
71
13830029
        let head = self.head.get();
72

            
73
        // Safety: caller transfers ownership of a valid node; writing its next link is valid.
74
13830029
        unsafe { T::set_next(ptr, head) };
75
13830029
        self.head.set(ptr);
76
13830029
    }
77

            
78
    /// Returns an iterator over freelist entries.
79
    ///
80
    /// # Safety
81
    ///
82
    /// The freelist uses a non-atomic [`Cell`] head.
83
1991656
    pub unsafe fn iter(&self) -> FreeListIterator<T> {
84
1991656
        FreeListIterator {
85
1991656
            current: NonNull::new(self.head.get()),
86
1991656
        }
87
1991656
    }
88

            
89
    /// Returns a mutable iterator over freelist entries.
90
    ///
91
    /// # Safety
92
    ///
93
    /// The freelist uses a non-atomic [`Cell`] head.
94
    pub unsafe fn iter_mut(&mut self) -> FreeListIteratorMut<'_, T> {
95
        FreeListIteratorMut {
96
            current: NonNull::new(self.head.get()),
97
            marker: PhantomData,
98
        }
99
    }
100

            
101
    /// Clears the freelist head.
102
1991656
    pub fn clear(&mut self) {
103
1991656
        self.head.set(std::ptr::null_mut());
104
1991656
    }
105

            
106
    /// Returns whether the freelist is currently empty.
107
680812
    pub fn is_empty(&self) -> bool {
108
680812
        self.head.get().is_null()
109
680812
    }
110

            
111
    /// Replaces the freelist head with `head`.
112
    ///
113
    /// # Safety
114
    ///
115
    /// `head` must be either null or a valid linked list of nodes
116
    /// managed by this freelist.
117
680812
    pub unsafe fn set_head(&self, head: *mut T) {
118
680812
        self.head.set(head);
119
680812
    }
120
}
121

            
122
/// Iterator over entries in a [`FreeList`].
123
pub struct FreeListIterator<T: FreeListEntry> {
124
    current: Option<NonNull<T>>,
125
}
126

            
127
impl<T: FreeListEntry> Iterator for FreeListIterator<T> {
128
    type Item = NonNull<T>;
129

            
130
320393380
    fn next(&mut self) -> Option<Self::Item> {
131
320393380
        if let Some(current) = self.current {
132
            // Safety: `current` is a freelist node; its link field is valid to read.
133
318401724
            unsafe {
134
318401724
                self.current = NonNull::new(T::get_next(current.as_ptr()));
135
318401724
            }
136
318401724
            Some(current)
137
        } else {
138
1991656
            None
139
        }
140
320393380
    }
141
}
142

            
143
/// Mutable iterator over entries in a [`FreeList`].
144
pub struct FreeListIteratorMut<'a, T: FreeListEntry> {
145
    current: Option<NonNull<T>>,
146
    marker: PhantomData<&'a mut T>,
147
}
148

            
149
impl<'a, T: FreeListEntry> Iterator for FreeListIteratorMut<'a, T> {
150
    type Item = &'a mut T;
151

            
152
    fn next(&mut self) -> Option<Self::Item> {
153
        if let Some(current) = self.current {
154
            // Safety: `current` is a freelist node; its link field is valid to read.
155
            unsafe {
156
                let current_ptr = current.as_ptr();
157
                self.current = NonNull::new(T::get_next(current_ptr));
158
                Some(&mut *current_ptr)
159
            }
160
        } else {
161
            None
162
        }
163
    }
164
}