src/libproc_macro/bridge/buffer.rs

   1 //! Buffer management for same-process client<->server communication.
   2
   3 use std::io::{self, Write};
   4 use std::mem;
   5 use std::ops::{Deref, DerefMut};
   6 use std::slice;
   7
   8 #[repr(C)]
   9 struct Slice<'a, T> {
  10     data: &'a [T; 0],
  11     len: usize,
  12 }
  13
  14 unsafe impl<'a, T: Sync> Sync for Slice<'a, T> {}
  15 unsafe impl<'a, T: Sync> Send for Slice<'a, T> {}
  16
  17 impl<T> Copy for Slice<'a, T> {}
  18 impl<T> Clone for Slice<'a, T> {
  19     fn clone(&self) -> Self {
  20         *self
  21     }
  22 }
  23
  24 impl<T> From<&'a [T]> for Slice<'a, T> {
  25     fn from(xs: &'a [T]) -> Self {
  26         Slice { data: unsafe { &*(xs.as_ptr() as *const [T; 0]) }, len: xs.len() }
  27     }
  28 }
  29
  30 impl<T> Deref for Slice<'a, T> {
  31     type Target = [T];
  32     fn deref(&self) -> &[T] {
  33         unsafe { slice::from_raw_parts(self.data.as_ptr(), self.len) }
  34     }
  35 }
  36
  37 #[repr(C)]
  38 pub struct Buffer<T: Copy> {
  39     data: *mut T,
  40     len: usize,
  41     capacity: usize,
  42     extend_from_slice: extern "C" fn(Buffer<T>, Slice<'_, T>) -> Buffer<T>,
  43     drop: extern "C" fn(Buffer<T>),
  44 }
  45
  46 unsafe impl<T: Copy + Sync> Sync for Buffer<T> {}
  47 unsafe impl<T: Copy + Send> Send for Buffer<T> {}
  48
  49 impl<T: Copy> Default for Buffer<T> {
  50     fn default() -> Self {
  51         Self::from(vec![])
  52     }
  53 }
  54
  55 impl<T: Copy> Deref for Buffer<T> {
  56     type Target = [T];
  57     fn deref(&self) -> &[T] {
  58         unsafe { slice::from_raw_parts(self.data as *const T, self.len) }
  59     }
  60 }
  61
  62 impl<T: Copy> DerefMut for Buffer<T> {
  63     fn deref_mut(&mut self) -> &mut [T] {
  64         unsafe { slice::from_raw_parts_mut(self.data, self.len) }
  65     }
  66 }
  67
  68 impl<T: Copy> Buffer<T> {
  69     pub(super) fn new() -> Self {
  70         Self::default()
  71     }
  72
  73     pub(super) fn clear(&mut self) {
  74         self.len = 0;
  75     }
  76
  77     pub(super) fn take(&mut self) -> Self {
  78         mem::take(self)
  79     }
  80
  81     pub(super) fn extend_from_slice(&mut self, xs: &[T]) {
  82         // Fast path to avoid going through an FFI call.
  83         if let Some(final_len) = self.len.checked_add(xs.len()) {
  84             if final_len <= self.capacity {
  85                 let dst = unsafe { slice::from_raw_parts_mut(self.data, self.capacity) };
  86                 dst[self.len..][..xs.len()].copy_from_slice(xs);
  87                 self.len = final_len;
  88                 return;
  89             }
  90         }
  91         let b = self.take();
  92         *self = (b.extend_from_slice)(b, Slice::from(xs));
  93     }
  94 }
  95
  96 impl Write for Buffer<u8> {
  97     fn write(&mut self, xs: &[u8]) -> io::Result<usize> {
  98         self.extend_from_slice(xs);
  99         Ok(xs.len())
 100     }
 101
 102     fn write_all(&mut self, xs: &[u8]) -> io::Result<()> {
 103         self.extend_from_slice(xs);
 104         Ok(())
 105     }
 106
 107     fn flush(&mut self) -> io::Result<()> {
 108         Ok(())
 109     }
 110 }
 111
 112 impl<T: Copy> Drop for Buffer<T> {
 113     fn drop(&mut self) {
 114         let b = self.take();
 115         (b.drop)(b);
 116     }
 117 }
 118
 119 impl<T: Copy> From<Vec<T>> for Buffer<T> {
 120     fn from(mut v: Vec<T>) -> Self {
 121         let (data, len, capacity) = (v.as_mut_ptr(), v.len(), v.capacity());
 122         mem::forget(v);
 123
 124         // This utility function is nested in here because it can *only*
 125         // be safely called on `Buffer`s created by *this* `proc_macro`.
 126         fn to_vec<T: Copy>(b: Buffer<T>) -> Vec<T> {
 127             unsafe {
 128                 let Buffer { data, len, capacity, .. } = b;
 129                 mem::forget(b);
 130                 Vec::from_raw_parts(data, len, capacity)
 131             }
 132         }
 133
 134         extern "C" fn extend_from_slice<T: Copy>(b: Buffer<T>, xs: Slice<'_, T>) -> Buffer<T> {
 135             let mut v = to_vec(b);
 136             v.extend_from_slice(&xs);
 137             Buffer::from(v)
 138         }
 139
 140         extern "C" fn drop<T: Copy>(b: Buffer<T>) {
 141             mem::drop(to_vec(b));
 142         }
 143
 144         Buffer { data, len, capacity, extend_from_slice, drop }
 145     }
 146 }