Refactor I/O traits

rust/src/embedded.rs

@@ -34,11 +34,16 @@
 use crate::constants::*;
 use crate::stl::AsciiSym;
 use crate::{
-    DecodeError, DefineUnion, ReadTuple, ReadUnion, RestrictedCharacter, RestrictedString, Sizing,
-    StrictDecode, StrictDumb, StrictEncode, StrictProduct, StrictStruct, StrictSum, StrictTuple,
-    StrictType, StrictUnion, TypeName, TypedRead, TypedWrite, WriteTuple, WriteUnion, LIB_EMBEDDED,
+    DecodeError, DefineUnion, ReadRaw, ReadTuple, ReadUnion, RestrictedCharacter, RestrictedString,
+    Sizing, StrictDecode, StrictDumb, StrictEncode, StrictProduct, StrictStruct, StrictSum,
+    StrictTuple, StrictType, StrictUnion, TypeName, TypedRead, TypedWrite, WriteRaw, WriteTuple,
+    WriteUnion, LIB_EMBEDDED,
 };
 
+pub trait DecodeRawLe: Sized {
+    fn decode_raw_le(reader: &mut (impl ReadRaw + ?Sized)) -> Result<Self, DecodeError>;
+}
+
 #[derive(
     Wrapper, WrapperMut, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Default, From
 )]
@@ -49,8 +54,12 @@
 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)
     }
 }
 
@@ -60,18 +69,23 @@
             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 {
+            fn decode_raw_le(reader: &mut (impl ReadRaw + ?Sized)) -> Result<Self, DecodeError> {
+                let buf = reader.read_raw_array::<{ Self::BITS as usize / 8 }>()?;
+                Ok(Self::from_le_bytes(buf))
             }
         }
         impl $crate::StrictDecode for $ty {
             fn strict_decode(reader: &mut impl TypedRead) -> Result<Self, DecodeError> {
-                let buf = unsafe { reader._read_raw_array::<{ Self::BITS as usize / 8 }>()? };
-                Ok(Self::from_le_bytes(buf))
+                Self::decode_raw_le(unsafe { reader.raw_reader() })
             }
         }
     };
@@ -83,17 +97,23 @@
             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 {
             fn strict_decode(reader: &mut impl TypedRead) -> Result<Self, DecodeError> {
-                let buf = unsafe { reader._read_raw_array::<{ Self::BITS as usize / 8 }>()? };
+                let buf = unsafe {
+                    reader
+                        .raw_reader()
+                        .read_raw_array::<{ Self::BITS as usize / 8 }>()?
+                };
                 let v = <$p>::from_le_bytes(buf);
                 Self::new(v).ok_or(DecodeError::ZeroNatural)
             }
@@ -109,18 +129,22 @@
         }
         #[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")]
         impl $crate::StrictDecode for $ty {
             fn strict_decode(reader: &mut impl TypedRead) -> Result<Self, DecodeError> {
                 const BYTES: usize = <$ty>::BITS / 8;
                 let mut inner = [0u8; 32];
-                let buf = unsafe { reader._read_raw_array::<BYTES>()? };
+                let buf = unsafe { reader.raw_reader().read_raw_array::<BYTES>()? };
                 inner[..BYTES].copy_from_slice(&buf[..]);
                 let bits = u256::from_le_bytes(inner);
                 Ok(Self::from_bits(bits))
@@ -474,7 +498,7 @@
     for Confined<Vec<T>, MIN_LEN, MAX_LEN>
 {
     fn strict_decode(reader: &mut impl TypedRead) -> Result<Self, DecodeError> {
-        let len = unsafe { reader._read_raw_len::<MAX_LEN>()? };
+        let len = unsafe { reader.raw_reader().read_raw_len::<MAX_LEN>()? };
         let mut col = Vec::<T>::with_capacity(len);
         for _ in 0..len {
             col.push(StrictDecode::strict_decode(reader)?);
@@ -505,7 +529,7 @@
     for Confined<BTreeSet<T>, MIN_LEN, MAX_LEN>
 {
     fn strict_decode(reader: &mut impl TypedRead) -> Result<Self, DecodeError> {
-        let len = unsafe { reader._read_raw_len::<MAX_LEN>()? };
+        let len = unsafe { reader.raw_reader().read_raw_len::<MAX_LEN>()? };
         let mut col = BTreeSet::<T>::new();
         for _ in 0..len {
             let item = StrictDecode::strict_decode(reader)?;
@@ -535,7 +559,7 @@
 {
     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)?;
@@ -558,7 +582,7 @@
 > StrictDecode for Confined<BTreeMap<K, V>, MIN_LEN, MAX_LEN>
 {
     fn strict_decode(reader: &mut impl TypedRead) -> Result<Self, DecodeError> {
-        let len = unsafe { reader._read_raw_len::<MAX_LEN>()? };
+        let len = unsafe { reader.raw_reader().read_raw_len::<MAX_LEN>()? };
         let mut col = BTreeMap::new();
         for _ in 0..len {
             let key = StrictDecode::strict_decode(reader)?;

rust/src/lib.rs

@@ -60,16 +60,18 @@ pub mod stl;
 #[cfg(test)]
 pub(crate) mod test;
 
-pub use embedded::Byte;
+pub use embedded::{Byte, DecodeRawLe};
 pub use error::{DecodeError, DeserializeError, SerializeError};
 pub use ident::{FieldName, Ident, InvalidIdent, LibName, TypeName, VariantName};
 pub use primitives::{constants, NumCls, NumInfo, NumSize, Primitive};
-pub use reader::StrictReader;
+pub use reader::{ConfinedReader, StreamReader, StrictReader};
 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

@@ -22,7 +22,7 @@
 use std::io;
 
 use crate::{
-    DecodeError, FieldName, ReadStruct, ReadTuple, ReadUnion, StrictDecode, StrictEnum,
+    DecodeError, FieldName, ReadRaw, ReadStruct, ReadTuple, ReadUnion, StrictDecode, StrictEnum,
     StrictStruct, StrictSum, StrictTuple, StrictUnion, TypedRead, VariantName,
 };
 
@@ -46,13 +46,13 @@
 
 // TODO: Move to amplify crate
 #[derive(Clone, Debug)]
-pub struct CountingReader<R: io::Read> {
+pub struct ConfinedReader<R: io::Read> {
     count: usize,
     limit: usize,
     reader: R,
 }
 
-impl<R: io::Read> From<R> for CountingReader<R> {
+impl<R: io::Read> From<R> for ConfinedReader<R> {
     fn from(reader: R) -> Self {
         Self {
             count: 0,
@@ -62,7 +62,7 @@
     }
 }
 
-impl<R: io::Read> CountingReader<R> {
+impl<R: io::Read> ConfinedReader<R> {
     pub fn with(limit: usize, reader: R) -> Self {
         Self {
             count: 0,
@@ -73,10 +73,10 @@
 
     pub fn count(&self) -> usize { self.count }
 
-    pub fn unbox(self) -> R { self.reader }
+    pub fn unconfine(self) -> R { self.reader }
 }
 
-impl<R: io::Read> io::Read for CountingReader<R> {
+impl<R: io::Read> io::Read for ConfinedReader<R> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         let len = self.reader.read(buf)?;
         match self.count.checked_add(len) {
@@ -88,31 +88,72 @@
     }
 }
 
-#[derive(Clone, Debug, From)]
-pub struct StrictReader<R: io::Read>(CountingReader<R>);
+#[derive(Clone, Debug)]
+pub struct StreamReader<R: io::Read>(ConfinedReader<R>);
+
+impl<R: io::Read> StreamReader<R> {
+    pub fn new<const MAX: usize>(inner: R) -> Self { Self(ConfinedReader::with(MAX, inner)) }
+    pub fn unconfine(self) -> R { self.0.unconfine() }
+}
+
+impl<T: AsRef<[u8]>> StreamReader<io::Cursor<T>> {
+    pub fn cursor<const MAX: usize>(inner: T) -> Self {
+        Self(ConfinedReader::with(MAX, io::Cursor::new(inner)))
+    }
+}
+
+impl<R: io::Read> ReadRaw for StreamReader<R> {
+    fn read_raw<const MAX_LEN: usize>(&mut self, len: usize) -> io::Result<Vec<u8>> {
+        use io::Read;
+        let mut buf = vec![0u8; len];
+        self.0.read_exact(&mut buf)?;
+        Ok(buf)
+    }
 
-impl StrictReader<io::Cursor<Vec<u8>>> {
-    pub fn in_memory(data: Vec<u8>, limit: usize) -> Self {
-        StrictReader(CountingReader::with(limit, io::Cursor::new(data)))
+    fn read_raw_array<const LEN: usize>(&mut self) -> io::Result<[u8; LEN]> {
+        use io::Read;
+        let mut buf = [0u8; LEN];
+        self.0.read_exact(&mut buf)?;
+        Ok(buf)
     }
 }
 
-impl StrictReader<ReadCounter> {
-    pub fn counter() -> Self { StrictReader(CountingReader::from(ReadCounter::default())) }
+impl<T: AsRef<[u8]>> StreamReader<io::Cursor<T>> {
+    pub fn in_memory<const MAX: usize>(data: T) -> Self { Self::new::<MAX>(io::Cursor::new(data)) }
+    pub fn into_cursor(self) -> io::Cursor<T> { self.0.unconfine() }
 }
 
-impl<R: io::Read> StrictReader<R> {
-    pub fn with(limit: usize, reader: R) -> Self {
-        StrictReader(CountingReader::with(limit, reader))
+impl StreamReader<ReadCounter> {
+    pub fn counter<const MAX: usize>() -> Self { Self::new::<MAX>(ReadCounter::default()) }
+}
+
+#[derive(Clone, Debug, From)]
+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 {
+        Self(StreamReader::in_memory::<MAX>(data))
     }
+    pub fn into_cursor(self) -> io::Cursor<T> { self.0.into_cursor() }
+}
+
+impl StrictReader<StreamReader<ReadCounter>> {
+    pub fn counter<const MAX: usize>() -> Self { Self(StreamReader::counter::<MAX>()) }
+}
 
-    pub fn unbox(self) -> R { self.0.unbox() }
+impl<R: ReadRaw> StrictReader<R> {
+    pub fn with(reader: R) -> Self { Self(reader) }
+
+    pub fn unbox(self) -> R { self.0 }
 }
 
-impl<R: io::Read> TypedRead for StrictReader<R> {
+impl<R: ReadRaw> TypedRead for StrictReader<R> {
     type TupleReader<'parent> = TupleReader<'parent, R> where Self: 'parent;
     type StructReader<'parent> = StructReader<'parent, R> where Self: 'parent;
     type UnionReader = Self;
+    type RawReader = R;
+
+    unsafe fn raw_reader(&mut self) -> &mut Self::RawReader { &mut self.0 }
 
     fn read_union<T: StrictUnion>(
         &mut self,
@@ -183,49 +224,35 @@
         assert!(reader.named_fields.is_empty(), "excessive fields are read for {}", name);
         Ok(res)
     }
-
-    unsafe fn _read_raw<const MAX_LEN: usize>(&mut self, len: usize) -> io::Result<Vec<u8>> {
-        use io::Read;
-        let mut buf = vec![0u8; len];
-        self.0.read_exact(&mut buf)?;
-        Ok(buf)
-    }
-
-    unsafe fn _read_raw_array<const LEN: usize>(&mut self) -> io::Result<[u8; LEN]> {
-        use io::Read;
-        let mut buf = [0u8; LEN];
-        self.0.read_exact(&mut buf)?;
-        Ok(buf)
-    }
 }
 
 #[derive(Debug)]
-pub struct TupleReader<'parent, R: io::Read> {
+pub struct TupleReader<'parent, R: ReadRaw> {
     read_fields: u8,
     parent: &'parent mut StrictReader<R>,
 }
 
-impl<'parent, R: io::Read> ReadTuple for TupleReader<'parent, R> {
+impl<'parent, R: ReadRaw> ReadTuple for TupleReader<'parent, R> {
     fn read_field<T: StrictDecode>(&mut self) -> Result<T, DecodeError> {
         self.read_fields += 1;
         T::strict_decode(self.parent)
     }
 }
 
 #[derive(Debug)]
-pub struct StructReader<'parent, R: io::Read> {
+pub struct StructReader<'parent, R: ReadRaw> {
     named_fields: Vec<FieldName>,
     parent: &'parent mut StrictReader<R>,
 }
 
-impl<'parent, R: io::Read> ReadStruct for StructReader<'parent, R> {
+impl<'parent, R: ReadRaw> ReadStruct for StructReader<'parent, R> {
     fn read_field<T: StrictDecode>(&mut self, field: FieldName) -> Result<T, DecodeError> {
         self.named_fields.push(field);
         T::strict_decode(self.parent)
     }
 }
 
-impl<R: io::Read> ReadUnion for StrictReader<R> {
+impl<R: ReadRaw> ReadUnion for StrictReader<R> {
     type TupleReader<'parent> = TupleReader<'parent, R> where Self: 'parent;
     type StructReader<'parent> = StructReader<'parent, R> where Self: 'parent;
 

rust/src/test.rs

@@ -20,7 +20,6 @@
 // limitations under the License.
 
 use std::fmt::Debug;
-use std::io;
 use std::io::BufRead;
 
 use amplify::confinement::Confined;
@@ -30,8 +29,8 @@
 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()
@@ -40,16 +39,15 @@
 pub fn decode<T: StrictDecode + Debug + Eq>(data: impl AsRef<[u8]>) -> T {
     const MAX: usize = u16::MAX as usize;
 
-    let cursor = io::Cursor::new(data);
-    let mut reader = StrictReader::with(MAX, cursor);
+    let mut reader = StrictReader::in_memory::<MAX>(data);
     let val2 = T::strict_decode(&mut reader).unwrap();
-    let mut cursor = reader.unbox();
+    let mut cursor = reader.into_cursor();
     assert!(!cursor.fill_buf().unwrap().is_empty(), "data not entirely consumed");
 
     val2
 }
 
 #[allow(dead_code)]
 pub fn encoding_roundtrip<T: StrictEncode + StrictDecode + Debug + Eq>(val: &T) {
     let data = encode(val);
     let val2: T = decode(data);