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

use std::{cell::Cell, io::IoSlice, mem, mem::MaybeUninit, ops::Range};

use crate::{
    buffer::{Buffer, CellBuffer, Drain, Resize},
    loom::sync::atomic::{AtomicUsize, Ordering},
    write_vectored::{VectoredSlice, EMPTY_SLICE},
};

/// A buffer of [`IoSlice`]
pub struct WriteVectoredVecBuffer<T> {
    owned: Box<[Cell<MaybeUninit<T>>]>,
    slices: Box<[Cell<IoSlice<'static>>]>,
    total_size: AtomicUsize,
}

impl<T> Default for WriteVectoredVecBuffer<T> {
    fn default() -> Self {
        Self {
            owned: Default::default(),
            slices: Default::default(),
            total_size: Default::default(),
        }
    }
}

// SAFETY: `WriteVectoredVecBuffer::clear` does clear the inserted range from the buffer
unsafe impl<T> Buffer for WriteVectoredVecBuffer<T>
where
    T: AsRef<[u8]>,
{
    type Slice<'a> = VectoredSlice<'a>
    where
        T: 'a;

    #[inline]
    fn capacity(&self) -> usize {
        self.owned.len()
    }

    #[inline]
    unsafe fn slice(&mut self, range: Range<usize>) -> Self::Slice<'_> {
        // SAFETY: [Cell<IoSlice>] has the same layout as [IoSlice]
        // and function contract guarantees that the range is initialized
        let slices = unsafe {
            &mut *(&mut self.slices[range.start..range.end + 2] as *mut _
                as *mut [IoSlice<'static>])
        };
        // SAFETY: slices are never read and live along their owner in the buffer, as they are
        // inserted and removed together
        unsafe { VectoredSlice::new(slices, self.total_size.load(Ordering::Acquire)) }
    }

    #[inline]
    unsafe fn clear(&mut self, range: Range<usize>) {
        *self.total_size.get_mut() = 0;
        for index in range {
            // SAFETY: function contract guarantees that the range is initialized
            unsafe { self.remove(index) };
        }
    }
}

// SAFETY: `insert` does initialize the index in the buffer
unsafe impl<T> CellBuffer<T> for WriteVectoredVecBuffer<T>
where
    T: AsRef<[u8]>,
{
    unsafe fn insert(&self, index: usize, value: T) {
        // SAFETY: slice is never read with static lifetime, it will only be used as a reference
        // with the same lifetime than the slice owner
        let slice = unsafe { mem::transmute(IoSlice::new(value.as_ref())) };
        self.slices[index + 1].set(slice);
        self.owned[index].set(MaybeUninit::new(value));
        self.total_size.fetch_add(slice.len(), Ordering::AcqRel);
    }
}

impl<T> Resize for WriteVectoredVecBuffer<T>
where
    T: AsRef<[u8]>,
{
    fn resize(&mut self, capacity: usize) {
        self.owned = (0..capacity)
            .map(|_| Cell::new(MaybeUninit::uninit()))
            .collect();
        self.slices = (0..capacity + 2)
            .map(|_| Cell::new(IoSlice::new(EMPTY_SLICE)))
            .collect();
    }
}

// SAFETY: `WriteVectoredVecBuffer::remove` does remove the index from the buffer
unsafe impl<T> Drain for WriteVectoredVecBuffer<T>
where
    T: AsRef<[u8]>,
{
    type Value = T;

    #[inline]
    unsafe fn remove(&mut self, index: usize) -> Self::Value {
        // SAFETY: function contract guarantees that the index has been inserted and is then initialized
        let value = unsafe {
            self.owned[index]
                .replace(MaybeUninit::uninit())
                .assume_init()
        };
        self.total_size
            .fetch_sub(value.as_ref().len(), Ordering::Release);
        value
    }
}