Browse Source

aya-log-common: generalize TagLenValue

This allows logging values backed by generators.
Tamir Duberstein 1 year ago
parent
commit
9a1a720
1 changed files with 49 additions and 44 deletions
  1. 49 44
      aya-log-common/src/lib.rs

+ 49 - 44
aya-log-common/src/lib.rs

@@ -1,6 +1,6 @@
 #![no_std]
 
-use core::{mem, num, ptr, slice};
+use core::{mem, num, ptr};
 
 use num_enum::IntoPrimitive;
 
@@ -98,58 +98,65 @@ pub enum DisplayHint {
     UpperMac,
 }
 
-struct TagLenValue<'a, T> {
-    tag: T,
-    value: &'a [u8],
+struct TagLenValue<T, V> {
+    pub tag: T,
+    pub value: V,
 }
 
-impl<'a, T> TagLenValue<'a, T>
+impl<T, V> TagLenValue<T, V>
 where
-    T: Copy,
+    V: IntoIterator<Item = u8>,
+    <V as IntoIterator>::IntoIter: ExactSizeIterator,
 {
-    #[inline(always)]
-    pub(crate) fn new(tag: T, value: &'a [u8]) -> TagLenValue<'a, T> {
-        TagLenValue { tag, value }
-    }
-
-    pub(crate) fn write(&self, mut buf: &mut [u8]) -> Result<usize, ()> {
+    pub(crate) fn write(self, mut buf: &mut [u8]) -> Result<usize, ()> {
         // Break the abstraction to please the verifier.
         if buf.len() > LOG_BUF_CAPACITY {
             buf = &mut buf[..LOG_BUF_CAPACITY];
         }
         let Self { tag, value } = self;
+        let value = value.into_iter();
         let len = value.len();
         let wire_len: LogValueLength = value
             .len()
             .try_into()
             .map_err(|num::TryFromIntError { .. }| ())?;
-        let size = mem::size_of_val(tag) + mem::size_of_val(&wire_len) + len;
+        let size = mem::size_of_val(&tag) + mem::size_of_val(&wire_len) + len;
         if size > buf.len() {
             return Err(());
         }
 
-        unsafe { ptr::write_unaligned(buf.as_mut_ptr() as *mut _, *tag) };
-        buf = &mut buf[mem::size_of_val(tag)..];
+        let tag_size = mem::size_of_val(&tag);
+        unsafe { ptr::write_unaligned(buf.as_mut_ptr() as *mut _, tag) };
+        buf = &mut buf[tag_size..];
 
         unsafe { ptr::write_unaligned(buf.as_mut_ptr() as *mut _, wire_len) };
         buf = &mut buf[mem::size_of_val(&wire_len)..];
 
-        unsafe { ptr::copy_nonoverlapping(value.as_ptr(), buf.as_mut_ptr(), len) };
+        buf.iter_mut().zip(value).for_each(|(dst, src)| {
+            *dst = src;
+        });
 
         Ok(size)
     }
 }
 
+impl<T, V> TagLenValue<T, V> {
+    #[inline(always)]
+    pub(crate) fn new(tag: T, value: V) -> TagLenValue<T, V> {
+        TagLenValue { tag, value }
+    }
+}
+
 pub trait WriteToBuf {
     #[allow(clippy::result_unit_err)]
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()>;
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()>;
 }
 
 macro_rules! impl_write_to_buf {
     ($type:ident, $arg_type:expr) => {
         impl WriteToBuf for $type {
-            fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-                TagLenValue::<Argument>::new($arg_type, &self.to_ne_bytes()).write(buf)
+            fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+                TagLenValue::new($arg_type, self.to_ne_bytes()).write(buf)
             }
         }
     };
@@ -171,35 +178,34 @@ impl_write_to_buf!(f32, Argument::F32);
 impl_write_to_buf!(f64, Argument::F64);
 
 impl WriteToBuf for [u8; 16] {
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-        TagLenValue::<Argument>::new(Argument::ArrU8Len16, self).write(buf)
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        TagLenValue::new(Argument::ArrU8Len16, self).write(buf)
     }
 }
 
 impl WriteToBuf for [u16; 8] {
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-        let ptr = self.as_ptr().cast::<u8>();
-        let bytes = unsafe { slice::from_raw_parts(ptr, 16) };
-        TagLenValue::<Argument>::new(Argument::ArrU16Len8, bytes).write(buf)
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        let bytes = unsafe { core::mem::transmute::<_, [u8; 16]>(self) };
+        TagLenValue::new(Argument::ArrU16Len8, bytes).write(buf)
     }
 }
 
 impl WriteToBuf for [u8; 6] {
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-        TagLenValue::<Argument>::new(Argument::ArrU8Len6, self).write(buf)
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        TagLenValue::new(Argument::ArrU8Len6, self).write(buf)
     }
 }
 
-impl WriteToBuf for str {
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-        TagLenValue::<Argument>::new(Argument::Str, self.as_bytes()).write(buf)
+impl WriteToBuf for &str {
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        TagLenValue::new(Argument::Str, self.as_bytes().iter().copied()).write(buf)
     }
 }
 
 impl WriteToBuf for DisplayHint {
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-        let v: u8 = (*self).into();
-        TagLenValue::<Argument>::new(Argument::DisplayHint, &v.to_ne_bytes()).write(buf)
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        let v: u8 = self.into();
+        TagLenValue::new(Argument::DisplayHint, v.to_ne_bytes()).write(buf)
     }
 }
 
@@ -217,17 +223,16 @@ pub fn write_record_header(
 ) -> Result<usize, ()> {
     let level: u8 = level.into();
     let mut size = 0;
-    for attr in [
-        TagLenValue::<RecordField>::new(RecordField::Target, target.as_bytes()),
-        TagLenValue::<RecordField>::new(RecordField::Level, &level.to_ne_bytes()),
-        TagLenValue::<RecordField>::new(RecordField::Module, module.as_bytes()),
-        TagLenValue::<RecordField>::new(RecordField::File, file.as_bytes()),
-        TagLenValue::<RecordField>::new(RecordField::Line, &line.to_ne_bytes()),
-        TagLenValue::<RecordField>::new(RecordField::NumArgs, &num_args.to_ne_bytes()),
-    ] {
-        size += attr.write(&mut buf[size..])?;
-    }
-
+    size += TagLenValue::new(RecordField::Target, target.as_bytes().iter().copied())
+        .write(&mut buf[size..])?;
+    size += TagLenValue::new(RecordField::Level, level.to_ne_bytes()).write(&mut buf[size..])?;
+    size += TagLenValue::new(RecordField::Module, module.as_bytes().iter().copied())
+        .write(&mut buf[size..])?;
+    size += TagLenValue::new(RecordField::File, file.as_bytes().iter().copied())
+        .write(&mut buf[size..])?;
+    size += TagLenValue::new(RecordField::Line, line.to_ne_bytes()).write(&mut buf[size..])?;
+    size +=
+        TagLenValue::new(RecordField::NumArgs, num_args.to_ne_bytes()).write(&mut buf[size..])?;
     Ok(size)
 }