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
132
133
134
135
136
137
138
139
140
141
142
143

//! Buffer management for same-process client<->server communication.

use std::io::{self, Write};
use std::mem;
use std::ops::{Deref, DerefMut};
use std::slice;

#[repr(C)]
pub struct Buffer<T: Copy> {
    data: *mut T,
    len: usize,
    capacity: usize,
    reserve: extern "C" fn(Buffer<T>, usize) -> Buffer<T>,
    drop: extern "C" fn(Buffer<T>),
}

unsafe impl<T: Copy + Sync> Sync for Buffer<T> {}
unsafe impl<T: Copy + Send> Send for Buffer<T> {}

impl<T: Copy> Default for Buffer<T> {
    fn default() -> Self {
        Self::from(vec![])
    }
}

impl<T: Copy> Deref for Buffer<T> {
    type Target = [T];
    fn deref(&self) -> &[T] {
        unsafe { slice::from_raw_parts(self.data as *const T, self.len) }
    }
}

impl<T: Copy> DerefMut for Buffer<T> {
    fn deref_mut(&mut self) -> &mut [T] {
        unsafe { slice::from_raw_parts_mut(self.data, self.len) }
    }
}

impl<T: Copy> Buffer<T> {
    pub(super) fn new() -> Self {
        Self::default()
    }

    pub(super) fn clear(&mut self) {
        self.len = 0;
    }

    pub(super) fn take(&mut self) -> Self {
        mem::take(self)
    }

    // We have the array method separate from extending from a slice. This is
    // because in the case of small arrays, codegen can be more efficient
    // (avoiding a memmove call). With extend_from_slice, LLVM at least
    // currently is not able to make that optimization.
    pub(super) fn extend_from_array<const N: usize>(&mut self, xs: &[T; N]) {
        if xs.len() > (self.capacity - self.len) {
            let b = self.take();
            *self = (b.reserve)(b, xs.len());
        }
        unsafe {
            xs.as_ptr().copy_to_nonoverlapping(self.data.add(self.len), xs.len());
            self.len += xs.len();
        }
    }

    pub(super) fn extend_from_slice(&mut self, xs: &[T]) {
        if xs.len() > (self.capacity - self.len) {
            let b = self.take();
            *self = (b.reserve)(b, xs.len());
        }
        unsafe {
            xs.as_ptr().copy_to_nonoverlapping(self.data.add(self.len), xs.len());
            self.len += xs.len();
        }
    }

    pub(super) fn push(&mut self, v: T) {
        // The code here is taken from Vec::push, and we know that reserve()
        // will panic if we're exceeding isize::MAX bytes and so there's no need
        // to check for overflow.
        if self.len == self.capacity {
            let b = self.take();
            *self = (b.reserve)(b, 1);
        }
        unsafe {
            *self.data.add(self.len) = v;
            self.len += 1;
        }
    }
}

impl Write for Buffer<u8> {
    fn write(&mut self, xs: &[u8]) -> io::Result<usize> {
        self.extend_from_slice(xs);
        Ok(xs.len())
    }

    fn write_all(&mut self, xs: &[u8]) -> io::Result<()> {
        self.extend_from_slice(xs);
        Ok(())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl<T: Copy> Drop for Buffer<T> {
    fn drop(&mut self) {
        let b = self.take();
        (b.drop)(b);
    }
}

impl<T: Copy> From<Vec<T>> for Buffer<T> {
    fn from(mut v: Vec<T>) -> Self {
        let (data, len, capacity) = (v.as_mut_ptr(), v.len(), v.capacity());
        mem::forget(v);

        // This utility function is nested in here because it can *only*
        // be safely called on `Buffer`s created by *this* `proc_macro`.
        fn to_vec<T: Copy>(b: Buffer<T>) -> Vec<T> {
            unsafe {
                let Buffer { data, len, capacity, .. } = b;
                mem::forget(b);
                Vec::from_raw_parts(data, len, capacity)
            }
        }

        extern "C" fn reserve<T: Copy>(b: Buffer<T>, additional: usize) -> Buffer<T> {
            let mut v = to_vec(b);
            v.reserve(additional);
            Buffer::from(v)
        }

        extern "C" fn drop<T: Copy>(b: Buffer<T>) {
            mem::drop(to_vec(b));
        }

        Buffer { data, len, capacity, reserve, drop }
    }
}