Encode QoL

src/bit_encoding/bitwriter.rs

@@ -109,3 +109,38 @@ impl<W: io::Write> BitWriter<W> {
         Ok(len)
     }
 }
+
+/// Write the result of a bit operation into a byte vector and return the vector.
+///
+/// I/O to a vector never fails.
+pub fn write_to_vec<F>(f: F) -> Vec<u8>
+where
+    F: FnOnce(&mut BitWriter<&mut Vec<u8>>) -> io::Result<usize>,
+{
+    let mut bytes = Vec::new();
+    let mut bits = BitWriter::new(&mut bytes);
+    f(&mut bits).expect("I/O to vector never fails");
+    bits.flush_all().expect("I/O to vector never fails");
+    bytes
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::jet::Core;
+    use crate::node::CoreConstructible;
+    use crate::ConstructNode;
+    use std::sync::Arc;
+
+    #[test]
+    fn vec() {
+        let program = Arc::<ConstructNode<Core>>::unit();
+        let _ = write_to_vec(|w| program.encode(w));
+    }
+
+    #[test]
+    fn empty_vec() {
+        let vec = write_to_vec(|_| Ok(0));
+        assert!(vec.is_empty());
+    }
+}

src/bit_encoding/encode.rs

@@ -173,14 +173,14 @@ pub fn encode_program<W: io::Write, N: node::Marker>(
     let iter = EncodeNode::Node(program).post_order_iter::<EncodeSharing<N>>();
 
     let len = iter.clone().count();
-    let start_n = w.n_total_written();
+    let n_start = w.n_total_written();
     encode_natural(len, w)?;
 
     for node in iter {
         encode_node(node, w)?;
     }
 
-    Ok(w.n_total_written() - start_n)
+    Ok(w.n_total_written() - n_start)
 }
 
 /// Encode a node to bits.
@@ -291,7 +291,7 @@ where
     I: Iterator<Item = &'a Value> + Clone,
 {
     let mut bit_len = 0;
-    let start_n = w.n_total_written();
+    let n_start = w.n_total_written();
 
     for value in witness.clone() {
         bit_len += value.len();
@@ -308,11 +308,13 @@ where
         }
     }
 
-    Ok(w.n_total_written() - start_n)
+    Ok(w.n_total_written() - n_start)
 }
 
 /// Encode a value to bits.
-pub fn encode_value<W: io::Write>(value: &Value, w: &mut BitWriter<W>) -> io::Result<()> {
+pub fn encode_value<W: io::Write>(value: &Value, w: &mut BitWriter<W>) -> io::Result<usize> {
+    let n_start = w.n_total_written();
+
     match value {
         Value::Unit => {}
         Value::SumL(left) => {
@@ -329,32 +331,49 @@ pub fn encode_value<W: io::Write>(value: &Value, w: &mut BitWriter<W>) -> io::Re
         }
     }
 
-    Ok(())
+    Ok(w.n_total_written() - n_start)
 }
 
 /// Encode a hash to bits.
-pub fn encode_hash<W: io::Write>(h: &[u8], w: &mut BitWriter<W>) -> io::Result<()> {
+pub fn encode_hash<W: io::Write>(h: &[u8], w: &mut BitWriter<W>) -> io::Result<usize> {
     for byte in h {
         w.write_bits_be(u64::from(*byte), 8)?;
     }
 
-    Ok(())
+    Ok(h.len() * 8)
 }
 
-/// Encode a natural number to bits.
-pub fn encode_natural<W: io::Write>(n: usize, w: &mut BitWriter<W>) -> io::Result<()> {
-    assert_ne!(n, 0); // Cannot encode zero
-    let len = 8 * mem::size_of::<usize>() - n.leading_zeros() as usize - 1;
+/// Encode a positive integer to bits.
+pub fn encode_natural<W: io::Write>(mut n: usize, w: &mut BitWriter<W>) -> io::Result<usize> {
+    assert!(n > 0, "Zero cannot be encoded");
+    let n_start = w.n_total_written();
 
-    if len == 0 {
-        w.write_bit(false)?;
-    } else {
-        w.write_bit(true)?;
-        encode_natural(len, w)?;
-        w.write_bits_be(n as u64, len)?;
+    /// Minimum number of bits to represent `n` minus the most-significant bit
+    fn truncated_bit_len(n: usize) -> usize {
+        8 * mem::size_of::<usize>() - n.leading_zeros() as usize - 1
     }
 
-    Ok(())
+    let mut suffix = Vec::new();
+
+    loop {
+        debug_assert!(n > 0);
+        let len = truncated_bit_len(n);
+        if len == 0 {
+            w.write_bit(false)?;
+            break;
+        } else {
+            w.write_bit(true)?;
+            suffix.push((n, len));
+            n = len;
+        }
+    }
+
+    while let Some((bits, len)) = suffix.pop() {
+        let bits = bits as u64; // Case safety: assuming 64-bit machine or lower
+        w.write_bits_be(bits, len)?;
+    }
+
+    Ok(w.n_total_written() - n_start)
 }
 
 #[cfg(test)]

src/bit_encoding/mod.rs

@@ -28,4 +28,4 @@ pub mod decode;
 pub mod encode;
 
 pub use bititer::{BitIter, EarlyEndOfStreamError};
-pub use bitwriter::BitWriter;
+pub use bitwriter::{write_to_vec, BitWriter};

src/lib.rs

@@ -51,7 +51,7 @@ mod value;
 
 pub use bit_encoding::decode;
 pub use bit_encoding::encode;
-pub use bit_encoding::BitWriter;
+pub use bit_encoding::{write_to_vec, BitWriter};
 pub use bit_encoding::{BitIter, EarlyEndOfStreamError};
 
 #[cfg(feature = "elements")]

src/node/redeem.rs

@@ -16,7 +16,7 @@ use crate::analysis::NodeBounds;
 use crate::dag::{DagLike, InternalSharing, MaxSharing, PostOrderIterItem};
 use crate::jet::Jet;
 use crate::types::{self, arrow::FinalArrow};
-use crate::{encode, WitnessNode};
+use crate::{encode, write_to_vec, WitnessNode};
 use crate::{Amr, BitIter, BitWriter, Cmr, Error, FirstPassImr, Imr, Value};
 
 use super::{
@@ -364,7 +364,9 @@ impl<J: Jet> RedeemNode<J> {
         Ok(program)
     }
 
-    /// Encode a Simplicity program to bits, including the witness data.
+    /// Encode the program to bits.
+    ///
+    /// Includes witness data. Returns the number of written bits.
     pub fn encode<W: io::Write>(&self, w: &mut BitWriter<W>) -> io::Result<usize> {
         let sharing_iter = self.post_order_iter::<MaxSharing<Redeem<J>>>();
         let program_bits = encode::encode_program(self, w)?;
@@ -374,15 +376,11 @@ impl<J: Jet> RedeemNode<J> {
         Ok(program_bits + witness_bits)
     }
 
-    /// Encode a Simplicity program to a vector of bytes, including the witness data.
+    /// Encode the program to a byte vector.
+    ///
+    /// Includes witness data.
     pub fn encode_to_vec(&self) -> Vec<u8> {
-        let mut program_and_witness_bytes = Vec::<u8>::new();
-        let mut writer = BitWriter::new(&mut program_and_witness_bytes);
-        self.encode(&mut writer)
-            .expect("write to vector never fails");
-        debug_assert!(!program_and_witness_bytes.is_empty());
-
-        program_and_witness_bytes
+        write_to_vec(|w| self.encode(w))
     }
 }