abstract raw byte writers from strict writers and io::Write

rust/src/embedded.rs

@@ -36,8 +36,8 @@ use crate::stl::AsciiSym;
 use crate::{
     DecodeError, DefineUnion, ReadRaw, ReadTuple, ReadUnion, RestrictedCharacter, RestrictedString,
     Sizing, StrictDecode, StrictDumb, StrictEncode, StrictProduct, StrictStruct, StrictSum,
-    StrictTuple, StrictType, StrictUnion, TypeName, TypedRead, TypedWrite, WriteTuple, WriteUnion,
-    LIB_EMBEDDED,
+    StrictTuple, StrictType, StrictUnion, TypeName, TypedRead, TypedWrite, WriteRaw, WriteTuple,
+    WriteUnion, LIB_EMBEDDED,
 };
 
 pub trait DecodeRawLe: Sized {
@@ -54,8 +54,12 @@ pub trait DecodeRawLe: Sized {
 pub struct Byte(u8);
 
 impl StrictEncode for Byte {
-    fn strict_encode<W: TypedWrite>(&self, writer: W) -> io::Result<W> {
-        unsafe { writer.register_primitive(BYTE)._write_raw::<1>([self.0]) }
+    fn strict_encode<W: TypedWrite>(&self, mut writer: W) -> io::Result<W> {
+        unsafe {
+            writer = writer.register_primitive(BYTE);
+            writer.raw_writer().write_raw::<1>([self.0])?;
+        }
+        Ok(writer)
     }
 }
 
@@ -65,12 +69,12 @@ macro_rules! encode_num {
             const STRICT_LIB_NAME: &'static str = $crate::LIB_EMBEDDED;
         }
         impl $crate::StrictEncode for $ty {
-            fn strict_encode<W: TypedWrite>(&self, writer: W) -> io::Result<W> {
+            fn strict_encode<W: TypedWrite>(&self, mut writer: W) -> io::Result<W> {
                 unsafe {
-                    writer
-                        .register_primitive($id)
-                        ._write_raw_array(self.to_le_bytes())
+                    writer = writer.register_primitive($id);
+                    writer.raw_writer().write_raw_array(self.to_le_bytes())?;
                 }
+                Ok(writer)
             }
         }
         impl $crate::DecodeRawLe for $ty {
@@ -93,12 +97,14 @@ macro_rules! encode_nonzero {
             const STRICT_LIB_NAME: &'static str = $crate::LIB_EMBEDDED;
         }
         impl $crate::StrictEncode for $ty {
-            fn strict_encode<W: TypedWrite>(&self, writer: W) -> io::Result<W> {
+            fn strict_encode<W: TypedWrite>(&self, mut writer: W) -> io::Result<W> {
                 unsafe {
+                    writer = writer.register_primitive($id);
                     writer
-                        .register_primitive($id)
-                        ._write_raw_array(self.get().to_le_bytes())
+                        .raw_writer()
+                        .write_raw_array(self.get().to_le_bytes())?;
                 }
+                Ok(writer)
             }
         }
         impl $crate::StrictDecode for $ty {
@@ -123,10 +129,14 @@ macro_rules! encode_float {
         }
         #[cfg(feature = "float")]
         impl $crate::StrictEncode for $ty {
-            fn strict_encode<W: TypedWrite>(&self, writer: W) -> io::Result<W> {
+            fn strict_encode<W: TypedWrite>(&self, mut writer: W) -> io::Result<W> {
                 let mut be = [0u8; $len];
                 be.copy_from_slice(&self.to_bits().to_le_bytes()[..$len]);
-                unsafe { writer.register_primitive($id)._write_raw_array(be) }
+                unsafe {
+                    writer = writer.register_primitive($id);
+                    writer.raw_writer().write_raw_array(be)?;
+                }
+                Ok(writer)
             }
         }
         #[cfg(feature = "float")]
@@ -549,7 +559,7 @@ impl<
 {
     fn strict_encode<W: TypedWrite>(&self, mut writer: W) -> io::Result<W> {
         unsafe {
-            writer = writer._write_raw_len::<MAX_LEN>(self.len())?;
+            writer.raw_writer().write_raw_len::<MAX_LEN>(self.len())?;
         }
         for (k, v) in self {
             writer = k.strict_encode(writer)?;

rust/src/lib.rs

@@ -69,7 +69,9 @@ pub use stl::{Bool, RestrictedCharacter, RestrictedString, U4, U5};
 pub use traits::*;
 pub use types::*;
 pub use util::{Sizing, Variant};
-pub use writer::{SplitParent, StrictParent, StrictWriter, StructWriter, UnionWriter};
+pub use writer::{
+    SplitParent, StreamWriter, StrictParent, StrictWriter, StructWriter, UnionWriter,
+};
 
 #[deprecated(since = "2.2.0", note = "use LIB_EMBEDDED")]
 pub const NO_LIB: &str = LIB_EMBEDDED;

rust/src/reader.rs

@@ -123,19 +123,19 @@ pub struct StrictReader<R: ReadRaw>(R);
 
 impl<T: AsRef<[u8]>> StrictReader<StreamReader<io::Cursor<T>>> {
     pub fn in_memory<const MAX: usize>(data: T) -> Self {
-        StrictReader(StreamReader::new::<MAX>(io::Cursor::new(data)))
+        Self(StreamReader::new::<MAX>(io::Cursor::new(data)))
     }
     pub fn into_cursor(self) -> io::Cursor<T> { self.0.unconfine() }
 }
 
 impl StrictReader<StreamReader<ReadCounter>> {
     pub fn counter<const MAX: usize>() -> Self {
-        StrictReader(StreamReader::new::<MAX>(ReadCounter::default()))
+        Self(StreamReader::new::<MAX>(ReadCounter::default()))
     }
 }
 
 impl<R: ReadRaw> StrictReader<R> {
-    pub fn with(reader: R) -> Self { StrictReader(reader) }
+    pub fn with(reader: R) -> Self { Self(reader) }
 
     pub fn unbox(self) -> R { self.0 }
 }

rust/src/test.rs

@@ -29,8 +29,8 @@ use crate::{StrictDecode, StrictEncode, StrictReader, StrictWriter};
 pub fn encode<T: StrictEncode + Debug + Eq>(val: &T) -> Vec<u8> {
     const MAX: usize = u16::MAX as usize;
 
-    let ast_data = StrictWriter::in_memory(MAX);
-    let data = val.strict_encode(ast_data).unwrap().unbox();
+    let ast_data = StrictWriter::in_memory::<MAX>();
+    let data = val.strict_encode(ast_data).unwrap().unbox().unconfine();
     Confined::<Vec<u8>, 0, MAX>::try_from(data)
         .unwrap()
         .into_inner()

rust/src/traits.rs

@@ -29,18 +29,48 @@ use amplify::Wrapper;
 
 use super::{DecodeError, DecodeRawLe, VariantName};
 use crate::reader::StreamReader;
+use crate::writer::StreamWriter;
 use crate::{
     DeserializeError, FieldName, Primitive, SerializeError, Sizing, StrictDumb, StrictEnum,
     StrictReader, StrictStruct, StrictSum, StrictTuple, StrictType, StrictUnion, StrictWriter,
 };
 
 pub trait TypedParent: Sized {}
 
+pub trait WriteRaw {
+    fn write_raw<const MAX_LEN: usize>(&mut self, bytes: impl AsRef<[u8]>) -> io::Result<()>;
+    fn write_raw_array<const LEN: usize>(&mut self, raw: [u8; LEN]) -> io::Result<()> {
+        self.write_raw::<LEN>(raw)
+    }
+    fn write_raw_len<const MAX_LEN: usize>(&mut self, len: usize) -> io::Result<()> {
+        match MAX_LEN {
+            tiny if tiny <= u8::MAX as usize => self.write_raw_array((len as u8).to_le_bytes()),
+            small if small <= u16::MAX as usize => self.write_raw_array((len as u16).to_le_bytes()),
+            medium if medium <= u24::MAX.into_usize() => {
+                self.write_raw_array((u24::with(len as u32)).to_le_bytes())
+            }
+            large if large <= u32::MAX as usize => self.write_raw_array((len as u32).to_le_bytes()),
+            huge if huge <= u64::MAX as usize => self.write_raw_array((len as u64).to_le_bytes()),
+            _ => unreachable!("confined collections larger than u64::MAX must not exist"),
+        }
+    }
+}
+
+impl<T: WriteRaw> WriteRaw for &mut T {
+    fn write_raw<const MAX_LEN: usize>(&mut self, bytes: impl AsRef<[u8]>) -> io::Result<()> {
+        (*self).write_raw::<MAX_LEN>(bytes)
+    }
+}
+
 #[allow(unused_variables)]
 pub trait TypedWrite: Sized {
     type TupleWriter: WriteTuple<Parent = Self>;
     type StructWriter: WriteStruct<Parent = Self>;
     type UnionDefiner: DefineUnion<Parent = Self>;
+    type RawWriter: WriteRaw;
+
+    #[doc(hidden)]
+    unsafe fn raw_writer(&mut self) -> &mut Self::RawWriter;
 
     fn write_union<T: StrictUnion>(
         self,
@@ -85,38 +115,16 @@ pub trait TypedWrite: Sized {
         self
     }
 
-    #[doc(hidden)]
-    unsafe fn _write_raw<const MAX_LEN: usize>(self, bytes: impl AsRef<[u8]>) -> io::Result<Self>;
-    #[doc(hidden)]
-    unsafe fn _write_raw_array<const LEN: usize>(self, raw: [u8; LEN]) -> io::Result<Self> {
-        self._write_raw::<LEN>(raw)
-    }
-    #[doc(hidden)]
-    unsafe fn _write_raw_len<const MAX_LEN: usize>(self, len: usize) -> io::Result<Self> {
-        match MAX_LEN {
-            tiny if tiny <= u8::MAX as usize => self._write_raw_array((len as u8).to_le_bytes()),
-            small if small <= u16::MAX as usize => {
-                self._write_raw_array((len as u16).to_le_bytes())
-            }
-            medium if medium <= u24::MAX.into_usize() => {
-                self._write_raw_array((u24::with(len as u32)).to_le_bytes())
-            }
-            large if large <= u32::MAX as usize => {
-                self._write_raw_array((len as u32).to_le_bytes())
-            }
-            huge if huge <= u64::MAX as usize => self._write_raw_array((len as u64).to_le_bytes()),
-            _ => unreachable!("confined collections larger than u64::MAX must not exist"),
-        }
-    }
-
     /// Used by unicode strings, ASCII strings and restricted char set strings.
     #[doc(hidden)]
     unsafe fn write_string<const MAX_LEN: usize>(
-        self,
+        mut self,
         bytes: impl AsRef<[u8]>,
     ) -> io::Result<Self> {
-        self._write_raw_len::<MAX_LEN>(bytes.as_ref().len())?
-            ._write_raw::<MAX_LEN>(bytes)
+        self.raw_writer()
+            .write_raw_len::<MAX_LEN>(bytes.as_ref().len())?;
+        self.raw_writer().write_raw::<MAX_LEN>(bytes)?;
+        Ok(self)
     }
 
     /// Vec and sets - excluding strings, written by [`Self::write_string`].
@@ -129,7 +137,7 @@ pub trait TypedWrite: Sized {
         for<'a> &'a C: IntoIterator,
         for<'a> <&'a C as IntoIterator>::Item: StrictEncode,
     {
-        self = self._write_raw_len::<MAX_LEN>(col.len())?;
+        self.raw_writer().write_raw_len::<MAX_LEN>(col.len())?;
         for item in col {
             self = item.strict_encode(self)?;
         }
@@ -336,9 +344,10 @@ pub trait ReadUnion: Sized {
 
 pub trait StrictEncode: StrictType {
     fn strict_encode<W: TypedWrite>(&self, writer: W) -> io::Result<W>;
-    fn strict_write(&self, lim: usize, writer: impl io::Write) -> io::Result<usize> {
-        let counter = StrictWriter::with(lim, writer);
-        Ok(self.strict_encode(counter)?.count())
+    fn strict_write(&self, writer: impl WriteRaw) -> io::Result<()> {
+        let counter = StrictWriter::with(writer);
+        self.strict_encode(counter)?;
+        Ok(())
     }
 }
 
@@ -367,22 +376,24 @@ impl<T> StrictDecode for PhantomData<T> {
 pub trait StrictSerialize: StrictEncode {
     fn strict_serialized_len(&self) -> io::Result<usize> {
         let counter = StrictWriter::counter();
-        Ok(self.strict_encode(counter)?.unbox().count)
+        Ok(self.strict_encode(counter)?.unbox().unconfine().count)
     }
 
     fn to_strict_serialized<const MAX: usize>(
         &self,
     ) -> Result<Confined<Vec<u8>, 0, MAX>, SerializeError> {
-        let ast_data = StrictWriter::in_memory(MAX);
-        let data = self.strict_encode(ast_data)?.unbox();
+        let ast_data = StrictWriter::in_memory::<MAX>();
+        let data = self.strict_encode(ast_data)?.unbox().unconfine();
         Confined::<Vec<u8>, 0, MAX>::try_from(data).map_err(SerializeError::from)
     }
 
     fn strict_serialize_to_file<const MAX: usize>(
         &self,
         path: impl AsRef<std::path::Path>,
     ) -> Result<(), SerializeError> {
-        let file = StrictWriter::with(MAX, fs::File::create(path)?);
+        let file = fs::File::create(path)?;
+        // TODO: Do FileReader
+        let file = StrictWriter::with(StreamWriter::new::<MAX>(file));
         self.strict_encode(file)?;
         Ok(())
     }