feat(ser): add probe length

feat(ser): now to dtb will return total size

Signed-off-by: Woshiluo Luo <woshiluo.luo@outlook.com>

Author: woshiluo

Cargo.toml

@@ -18,7 +18,7 @@ serde = { version = "1.0", default-features = false, features = ["derive"] }
 dyn_serde = { version = "=1.1.2", default-features = false, optional = true }
 
 [features]
-default = ["std", "ser"]
+default = ["std", "ser", "alloc"]
 
 ser = ["dep:dyn_serde"]
 std = ["serde/std"]

src/de.rs

@@ -518,8 +518,6 @@ mod tests {
     use alloc::format;
     #[cfg(any(feature = "std", feature = "alloc"))]
     use serde::Deserialize;
-    #[cfg(feature = "std")]
-    use std::format;
 
     #[cfg(any(feature = "std", feature = "alloc"))]
     #[test]

src/ser/mod.rs

@@ -1,6 +1,7 @@
 pub mod patch;
 pub mod pointer;
 pub mod serializer;
+
 pub mod string_block;
 
 use crate::common::*;
@@ -10,13 +11,17 @@ use crate::ser::patch::Patch;
 const RSVMAP_LEN: usize = 16;
 
 /// Make dtb header with structure block and string block length.
-fn make_header<'se>(writer: &'se mut [u8], structure_length: u32, string_block_length: u32) {
+fn make_header<'se>(
+    writer: &'se mut [u8],
+    structure_length: u32,
+    string_block_length: u32,
+) -> usize {
     let (header, _) = writer.split_at_mut(HEADER_LEN as usize);
     let header = unsafe { &mut *(header.as_mut_ptr() as *mut Header) };
+    let total_size =
+        HEADER_PADDING_LEN + RSVMAP_LEN as u32 + structure_length + string_block_length;
     header.magic = u32::from_be(DEVICE_TREE_MAGIC);
-    header.total_size = u32::from_be(
-        HEADER_PADDING_LEN + RSVMAP_LEN as u32 + structure_length + string_block_length,
-    );
+    header.total_size = u32::from_be(total_size);
     assert_eq!(header.total_size % 8, 0);
     header.off_dt_struct = u32::from_be(HEADER_PADDING_LEN + RSVMAP_LEN as u32);
     header.off_dt_strings = u32::from_be(HEADER_PADDING_LEN + RSVMAP_LEN as u32 + structure_length);
@@ -26,61 +31,64 @@ fn make_header<'se>(writer: &'se mut [u8], structure_length: u32, string_block_l
     header.boot_cpuid_phys = 0; // TODO
     header.size_dt_strings = u32::from_be(string_block_length as u32);
     header.size_dt_struct = u32::from_be(structure_length as u32);
+
+    total_size as usize
 }
 
 /// Serialize the data to dtb, with a list fof Patch, write to the `writer`.
 ///
 /// We do run-twice on convert, first time to generate string block, second time todo real
 /// structure.
-pub fn to_dtb<'se, T>(data: &T, list: &'se [Patch<'se>], writer: &'se mut [u8]) -> Result<(), Error>
+pub fn to_dtb<'se, T>(
+    data: &T,
+    list: &'se [Patch<'se>],
+    writer: &'se mut [u8],
+) -> Result<usize, Error>
 where
     T: serde::ser::Serialize,
 {
     writer.iter_mut().for_each(|x| *x = 0);
 
     let string_block_length = {
         let mut offset: usize = 0;
+        let mut block = crate::ser::string_block::StringBlock::new(Some(writer), &mut offset);
+        let mut dst = crate::ser::pointer::Pointer::new(None);
+        let mut patch_list = crate::ser::patch::PatchList::new(list);
+        let mut ser =
+            crate::ser::serializer::SerializerInner::new(&mut dst, &mut block, &mut patch_list);
+        let ser = crate::ser::serializer::Serializer::new(&mut ser);
+        data.serialize(ser)?;
         {
-            let mut dst = crate::ser::pointer::Pointer::new(None);
-            let mut patch_list = crate::ser::patch::PatchList::new(list);
-            let mut block = crate::ser::string_block::StringBlock::new(writer, &mut offset);
-            let mut ser =
-                crate::ser::serializer::SerializerInner::new(&mut dst, &mut block, &mut patch_list);
-            let ser = crate::ser::serializer::Serializer::new(&mut ser);
-            data.serialize(ser)?;
-            offset
-        };
-        {
-            let mut block = crate::ser::string_block::StringBlock::new(writer, &mut offset);
+            let mut block = crate::ser::string_block::StringBlock::new(Some(writer), &mut offset);
             block.align();
         };
         offset
     };
 
     list.iter().for_each(|patch| patch.init());
     let bottom_string_block_start = writer.len() - string_block_length;
-    // Write to bottom, avoid overlap.
-    unsafe {
-        core::ptr::copy(
-            core::ptr::addr_of!(writer[0]),
-            core::ptr::addr_of_mut!(writer[bottom_string_block_start]),
-            string_block_length,
-        );
-    }
+
+    // Clear string block
     writer[0..string_block_length].fill(0);
 
     let structure_length = {
         let (data_block, string_block) = writer.split_at_mut(writer.len() - string_block_length);
         let (_, data_block) = data_block.split_at_mut(HEADER_PADDING_LEN as usize + RSVMAP_LEN);
 
         let mut patch_list = crate::ser::patch::PatchList::new(list);
-        let mut temp_length = string_block_length;
-        let mut block = crate::ser::string_block::StringBlock::new(string_block, &mut temp_length);
+        let mut temp_length = 0;
+        let mut block =
+            crate::ser::string_block::StringBlock::new(Some(string_block), &mut temp_length);
         let mut dst = crate::ser::pointer::Pointer::new(Some(data_block));
         let mut ser =
             crate::ser::serializer::SerializerInner::new(&mut dst, &mut block, &mut patch_list);
         let ser = crate::ser::serializer::Serializer::new(&mut ser);
         let struct_len = data.serialize(ser)?.1;
+        {
+            let mut block =
+                crate::ser::string_block::StringBlock::new(Some(writer), &mut temp_length);
+            block.align();
+        };
         assert_eq!(struct_len % 4, 0); // As spec, structure block align with 4 bytes.
         assert_eq!(temp_length, string_block_length); // StringBlock should be same with first run.
         struct_len
@@ -97,9 +105,32 @@ where
     }
     writer[bottom_string_block_start..].fill(0);
 
-    make_header(writer, structure_length as u32, string_block_length as u32);
+    let result = make_header(writer, structure_length as u32, string_block_length as u32);
+
+    Ok(result)
+}
+
+#[cfg(feature = "alloc")]
+pub fn probe_dtb_length<'se, T>(data: &T, list: &'se [Patch<'se>]) -> Result<usize, Error>
+where
+    T: serde::ser::Serialize,
+{
+    let mut offset: usize = 0;
+    let structure_length = {
+        let mut dst = crate::ser::pointer::Pointer::new(None);
+        let mut patch_list = crate::ser::patch::PatchList::new(list);
+        let mut block = crate::ser::string_block::StringBlock::new(None, &mut offset);
+        let mut ser =
+            crate::ser::serializer::SerializerInner::new(&mut dst, &mut block, &mut patch_list);
+        let ser = crate::ser::serializer::Serializer::new(&mut ser);
+        data.serialize(ser)?.1
+    };
+    {
+        let mut block = crate::ser::string_block::StringBlock::new(None, &mut offset);
+        block.align();
+    };
 
-    Ok(())
+    Ok(structure_length + offset + HEADER_PADDING_LEN as usize + RSVMAP_LEN)
 }
 
 #[derive(Debug)]

src/ser/serializer.rs

@@ -763,4 +763,39 @@ mod tests {
         // println!("{:x?}", buf1);
         // assert!(false);
     }
+    #[test]
+    #[cfg(feature = "alloc")]
+    fn probe_length() {
+        #[derive(Serialize)]
+        struct Base {
+            pub hello: u32,
+            pub base1: Base1,
+            pub hello2: u32,
+            pub base2: Base1,
+        }
+        #[derive(Serialize)]
+        struct Base1 {
+            pub hello: &'static str,
+        }
+        let mut buf1 = [0u8; MAX_SIZE];
+
+        let new_base = Base1 { hello: "added" };
+        let patch = crate::ser::patch::Patch::new(
+            "/base3",
+            &new_base as _,
+            crate::ser::serializer::ValueType::Node,
+        );
+        let list = [patch];
+        let base = Base {
+            hello: 0xdeedbeef,
+            base1: Base1 {
+                hello: "Hello, World!",
+            },
+            hello2: 0x11223344,
+            base2: Base1 { hello: "Roger" },
+        };
+        let probe_legnth = crate::ser::probe_dtb_length(&base, &list).unwrap();
+        let result = crate::ser::to_dtb(&base, &list, &mut buf1).unwrap();
+        assert_eq!(probe_legnth, result);
+    }
 }

src/ser/string_block.rs

@@ -1,17 +1,27 @@
+#[cfg(feature = "alloc")]
+use alloc::collections::BTreeMap;
+
 /// StringBlock
 /// As spec said, dtb have a block called string block for saving prop names.
 pub struct StringBlock<'se> {
     end: &'se mut usize,
-    data: &'se mut [u8],
+    data: Option<&'se mut [u8]>,
+    #[cfg(feature = "alloc")]
+    tree: BTreeMap<&'se str, usize>,
 }
 
 impl<'se> StringBlock<'se> {
     /// Make a new string block.
     ///
     /// For get how long is string block, we make `end` as a mut ref.
     #[inline(always)]
-    pub fn new(dst: &'se mut [u8], end: &'se mut usize) -> StringBlock<'se> {
-        StringBlock { data: dst, end }
+    pub fn new(dst: Option<&'se mut [u8]>, end: &'se mut usize) -> StringBlock<'se> {
+        StringBlock {
+            data: dst,
+            #[cfg(feature = "alloc")]
+            tree: BTreeMap::new(),
+            end,
+        }
     }
 
     // TODO: show as error
@@ -24,26 +34,39 @@ impl<'se> StringBlock<'se> {
         if offset > *self.end {
             panic!("invalid read");
         }
-        let current_slice = &self.data[offset..];
-        let pos = current_slice
-            .iter()
-            .position(|&x| x == b'\0')
-            .unwrap_or(self.data.len());
-        let (a, _) = current_slice.split_at(pos + 1);
-        let result = unsafe { core::str::from_utf8_unchecked(&a[..a.len() - 1]) };
-        (result, pos + offset + 1)
+        if let Some(data) = &self.data {
+            let current_slice = &data[offset..];
+            let pos = current_slice
+                .iter()
+                .position(|&x| x == b'\0')
+                .unwrap_or(data.len());
+            let (a, _) = current_slice.split_at(pos + 1);
+            let result = unsafe { core::str::from_utf8_unchecked(&a[..a.len() - 1]) };
+            (result, pos + offset + 1)
+        } else {
+            panic!("must have writer when no alloc");
+        }
+    }
+
+    #[inline(always)]
+    fn write_u8(&mut self, index: usize, data: u8) {
+        if let Some(buffer) = &mut self.data {
+            buffer[index] = data;
+        }
     }
 
     #[inline(always)]
     fn insert_u8(&mut self, data: u8) {
-        self.data[*self.end] = data;
+        self.write_u8(*self.end, data);
         *self.end += 1;
     }
 
     /// Return the start offset of inserted string.
     #[inline(always)]
-    pub fn insert_str(&mut self, name: &str) -> usize {
+    pub fn insert_str(&mut self, name: &'se str) -> usize {
         let result = *self.end;
+        #[cfg(feature = "alloc")]
+        self.tree.insert(name, result);
         name.bytes().for_each(|x| {
             self.insert_u8(x);
         });
@@ -55,14 +78,14 @@ impl<'se> StringBlock<'se> {
     #[inline(always)]
     pub fn align(&mut self) {
         while (*self.end & 0b111) != 0 {
-            self.data[*self.end] = 0;
-            *self.end += 1;
+            self.insert_u8(0);
         }
     }
 
     /// Find a string. If not found, insert it.
     #[inline(always)]
-    pub fn find_or_insert(&mut self, name: &str) -> usize {
+    #[cfg(not(feature = "alloc"))]
+    pub fn find_or_insert(&mut self, name: &'se str) -> usize {
         let mut current_pos = 0;
         while current_pos < *self.end {
             let (result, new_pos) = self.get_str_by_offset(current_pos);
@@ -74,4 +97,15 @@ impl<'se> StringBlock<'se> {
 
         self.insert_str(name)
     }
+
+    #[inline(always)]
+    #[cfg(feature = "alloc")]
+    pub fn find_or_insert(&mut self, name: &'se str) -> usize {
+        let result = self.tree.get(name);
+        if result.is_some() {
+            *result.unwrap()
+        } else {
+            self.insert_str(name)
+        }
+    }
 }