Merge #252: Word jets with word values

f4e259a feat: Optimize scribe (Christian Lewe)
3fd665d refactor: Simplify encode_value (Christian Lewe)
33fcd9f refactor: Use words in word jets (Christian Lewe)
67e1246 feat: Add Word (Christian Lewe)

Pull request description:

  Word jets can only host word values. Currently, our code has some assertions here and there that will fail at runtime if a non-word is passed. The assertions are not fully rigorous, especially regarding the maximum word type of `TWO^(2^31)`.

  This PR introduces the wrapper type `Word` which is guaranteed to be a valid word. `Word` is used everywhere where a word should be used.

  This PR also optimized scribe to use as many word jets as possible. Fixes #155

ACKs for top commit:
  apoelstra:
    ACK f4e259a successfully ran local tests

Tree-SHA512: 5af47f28cf1a9aa56ae07fbfdfe2b4e2caf540bfe041db5538a6995bcb1ebbd94475383c785e078d7571f7b2ad09b83d78bd333f9946afeab3ef62575ab52307

src/analysis.rs

@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: CC0-1.0
 
 use crate::jet::Jet;
-use crate::Value;
 use std::{cmp, fmt};
 
+use crate::value::Word;
 #[cfg(feature = "elements")]
 use elements::encode::Encodable;
 #[cfg(feature = "elements")]
@@ -360,11 +360,11 @@ impl NodeBounds {
     }
 
     /// Node bounds for an arbitrary constant word node
-    pub fn const_word(value: &Value) -> NodeBounds {
+    pub fn const_word(word: &Word) -> NodeBounds {
         NodeBounds {
             extra_cells: 0,
             extra_frames: 0,
-            cost: Cost::OVERHEAD + Cost::of_type(value.padded_len()),
+            cost: Cost::OVERHEAD + Cost::of_type(word.len()),
         }
     }
 

src/bit_encoding/decode.rs

@@ -13,7 +13,8 @@ use crate::node::{
     WitnessConstructible,
 };
 use crate::types;
-use crate::{BitIter, FailEntropy, Value};
+use crate::value::Word;
+use crate::{BitIter, FailEntropy};
 use std::collections::HashSet;
 use std::sync::Arc;
 use std::{error, fmt};
@@ -129,7 +130,7 @@ enum DecodeNode<J: Jet> {
     Fail(FailEntropy),
     Hidden(Cmr),
     Jet(J),
-    Word(Value),
+    Word(Word),
 }
 
 impl<'d, J: Jet> DagLike for (usize, &'d [DecodeNode<J>]) {
@@ -265,8 +266,8 @@ fn decode_node<I: Iterator<Item = u8>, J: Jet>(
         if bits.read_bit()? {
             J::decode(bits).map(|jet| DecodeNode::Jet(jet))
         } else {
-            let depth = bits.read_natural(Some(32))?;
-            let word = decode_power_of_2(bits, 1 << (depth - 1))?;
+            let n = bits.read_natural(Some(32))? as u32; // cast safety: decoded number is at most the number 32
+            let word = Word::from_bits(bits, n - 1)?;
             Ok(DecodeNode::Word(word))
         }
     } else {
@@ -320,43 +321,6 @@ fn decode_node<I: Iterator<Item = u8>, J: Jet>(
     }
 }
 
-/// Decode a value from bits, of the form 2^exp
-///
-/// # Panics
-///
-/// Panics if `exp` itself is not a power of 2
-pub fn decode_power_of_2<I: Iterator<Item = bool>>(
-    iter: &mut I,
-    exp: usize,
-) -> Result<Value, Error> {
-    struct StackElem {
-        value: Value,
-        width: usize,
-    }
-
-    assert_eq!(exp.count_ones(), 1, "exp must be a power of 2");
-
-    let mut stack = Vec::with_capacity(32);
-    for _ in 0..exp {
-        // Read next bit
-        let bit = Value::u1(u8::from(iter.next().ok_or(Error::EndOfStream)?));
-        stack.push(StackElem {
-            value: bit,
-            width: 1,
-        });
-
-        while stack.len() >= 2 && stack[stack.len() - 1].width == stack[stack.len() - 2].width {
-            let right = stack.pop().unwrap();
-            let left = stack.pop().unwrap();
-            stack.push(StackElem {
-                value: Value::product(left.value, right.value),
-                width: left.width * 2,
-            });
-        }
-    }
-    Ok(stack.pop().unwrap().value)
-}
-
 /// Decode a natural number from bits.
 /// If a bound is specified, then the decoding terminates before trying to decode a larger number.
 pub fn decode_natural<I: Iterator<Item = bool>>(

src/bit_encoding/encode.rs

@@ -256,13 +256,11 @@ fn encode_node<W: io::Write, N: node::Marker>(
                 w.write_bit(true)?; // jet
                 jet.encode(w)?;
             }
-            node::Inner::Word(val) => {
+            node::Inner::Word(word) => {
                 w.write_bit(true)?; // jet or word
                 w.write_bit(false)?; // word
-                assert_eq!(val.compact_len().count_ones(), 1);
-                let depth = val.compact_len().trailing_zeros();
-                encode_natural(1 + depth as usize, w)?;
-                encode_value(val, w)?;
+                encode_natural(1 + word.n() as usize, w)?;
+                encode_value(word.as_value(), w)?;
             }
             _ => unreachable!(),
         }
@@ -288,18 +286,9 @@ where
 /// Encode a value to bits.
 pub fn encode_value<W: io::Write>(value: &Value, w: &mut BitWriter<W>) -> io::Result<usize> {
     let n_start = w.n_total_written();
-
-    if let Some(left) = value.as_left() {
-        w.write_bit(false)?;
-        encode_value(left, w)?;
-    } else if let Some(right) = value.as_right() {
-        w.write_bit(true)?;
-        encode_value(right, w)?;
-    } else if let Some((left, right)) = value.as_product() {
-        encode_value(left, w)?;
-        encode_value(right, w)?;
+    for bit in value.iter_compact() {
+        w.write_bit(bit)?;
     }
-
     Ok(w.n_total_written() - n_start)
 }
 

src/bit_machine/mod.rs

@@ -334,7 +334,7 @@ impl BitMachine {
                 }
                 node::Inner::Witness(value) => self.write_value(value),
                 node::Inner::Jet(jet) => self.exec_jet(*jet, env)?,
-                node::Inner::Word(value) => self.write_value(value),
+                node::Inner::Word(value) => self.write_value(value.as_value()),
                 node::Inner::Fail(entropy) => {
                     return Err(ExecutionError::ReachedFailNode(*entropy))
                 }

src/human_encoding/mod.rs

@@ -9,7 +9,6 @@
 mod error;
 mod named_node;
 mod parse;
-mod serialize;
 
 use crate::dag::{DagLike, MaxSharing};
 use crate::jet::Jet;
@@ -150,8 +149,8 @@ impl<J: Jet> Forest<J> {
                     node::Inner::AssertR(cmr, _) => format!("{} := assertr #{}", name, cmr),
                     node::Inner::Fail(entropy) => format!("{} := fail {}", name, entropy),
                     node::Inner::Jet(ref j) => format!("{} := jet_{}", name, j),
-                    node::Inner::Word(ref v) => {
-                        format!("{} := const {}", name, serialize::DisplayWord(v))
+                    node::Inner::Word(ref word) => {
+                        format!("{} := const {}", name, word)
                     }
                     inner => format!("{} := {}", name, inner),
                 };

src/human_encoding/parse/ast.rs

@@ -7,7 +7,8 @@ use std::sync::Arc;
 
 use crate::human_encoding::{Error, ErrorSet, Position, WitnessOrHole};
 use crate::jet::Jet;
-use crate::{node, types, Value};
+use crate::value::Word;
+use crate::{node, types};
 use crate::{BitIter, Cmr, FailEntropy};
 use santiago::grammar::{Associativity, Grammar};
 use santiago::lexer::{Lexeme, LexerRules};
@@ -644,12 +645,11 @@ fn grammar<J: Jet + 'static>() -> Grammar<Ast<J>> {
                     Error::BadWordLength { bit_length },
                 ));
             }
-            let ty = types::Final::two_two_n(bit_length.trailing_zeros() as usize);
             // unwrap ok here since literally every sequence of bits is a valid
             // value for the given type
-            let value = Value::from_compact_bits(&mut iter, &ty).unwrap();
+            let word = Word::from_bits(&mut iter, bit_length.trailing_zeros()).unwrap();
             Ast::Expression(Expression {
-                inner: ExprInner::Inline(node::Inner::Word(value)),
+                inner: ExprInner::Inline(node::Inner::Word(word)),
                 position,
             })
         };

src/human_encoding/parse/mod.rs

@@ -575,6 +575,7 @@ mod tests {
     use crate::human_encoding::Forest;
     use crate::jet::{Core, Jet};
     use crate::node::Inner;
+    use crate::value::Word;
     use crate::{BitMachine, Value};
 
     fn assert_cmr_witness<J: Jet>(
@@ -619,15 +620,15 @@ mod tests {
         }
     }
 
-    fn assert_const<J: Jet>(s: &str, value: Value) {
+    fn assert_const<J: Jet>(s: &str, word: Word) {
         match parse::<J>(s) {
             Ok(forest) => {
                 assert_eq!(forest.len(), 1);
                 let main = &forest["main"];
 
                 for data in main.clone().post_order_iter::<MaxSharing<_>>() {
-                    if let Inner::Word(parsed_value) = data.node.inner() {
-                        assert_eq!(&value, parsed_value);
+                    if let Inner::Word(parsed_word) = data.node.inner() {
+                        assert_eq!(&word, parsed_word);
                     }
                 }
             }
@@ -743,32 +744,34 @@ mod tests {
 
     #[test]
     fn const_word() {
-        let human_values = [
-            ("0b0", Value::u1(0b0)),
-            ("0b1", Value::u1(0b1)),
-            ("0b00", Value::u2(0b00)),
-            ("0b11", Value::u2(0b11)),
-            ("0b0000", Value::u4(0b0000)),
-            ("0b1111", Value::u4(0b1111)),
-            ("0b00000000", Value::u8(0b00000000)),
-            ("0b11111111", Value::u8(0b11111111)),
+        let human_words = [
+            ("0b0", Word::u1(0b0)),
+            ("0b1", Word::u1(0b1)),
+            ("0b00", Word::u2(0b00)),
+            ("0b11", Word::u2(0b11)),
+            ("0b0000", Word::u4(0b0000)),
+            ("0b1111", Word::u4(0b1111)),
+            ("0b00000000", Word::u8(0b00000000)),
+            ("0b11111111", Word::u8(0b11111111)),
             (
                 "0b00000001001000110100010101100111",
-                Value::from_byte_array([0b00000001, 0b00100011, 0b01000101, 0b01100111]),
+                Word::u32(u32::from_be_bytes([
+                    0b00000001, 0b00100011, 0b01000101, 0b01100111,
+                ])),
             ),
-            ("0x0", Value::u4(0x0)),
-            ("0xf", Value::u4(0xf)),
-            ("0x00", Value::u8(0x00)),
-            ("0xff", Value::u8(0xff)),
+            ("0x0", Word::u4(0x0)),
+            ("0xf", Word::u4(0xf)),
+            ("0x00", Word::u8(0x00)),
+            ("0xff", Word::u8(0xff)),
             (
                 "0xdeadbeef",
-                Value::from_byte_array([0xde, 0xad, 0xbe, 0xef]),
+                Word::u32(u32::from_be_bytes([0xde, 0xad, 0xbe, 0xef])),
             ),
         ];
 
-        for (human, value) in human_values {
+        for (human, word) in human_words {
             let s = format!("main := comp const {human} unit");
-            assert_const::<Core>(s.as_str(), value);
+            assert_const::<Core>(s.as_str(), word);
         }
     }
 

src/jet/mod.rs

@@ -94,6 +94,7 @@ mod tests {
     use crate::jet::Core;
     use crate::node::{ConstructNode, CoreConstructible, JetConstructible};
     use crate::types;
+    use crate::value::Word;
     use crate::{BitMachine, Value};
     use std::sync::Arc;
 
@@ -102,8 +103,8 @@ mod tests {
         let ctx = types::Context::new();
         let two_words = Arc::<ConstructNode<_>>::comp(
             &Arc::<ConstructNode<_>>::pair(
-                &Arc::<ConstructNode<_>>::const_word(&ctx, Value::u32(2)),
-                &Arc::<ConstructNode<_>>::const_word(&ctx, Value::u32(16)),
+                &Arc::<ConstructNode<_>>::const_word(&ctx, Word::u32(2)),
+                &Arc::<ConstructNode<_>>::const_word(&ctx, Word::u32(16)),
             )
             .unwrap(),
             &Arc::<ConstructNode<_>>::jet(&ctx, Core::Add32),
@@ -122,8 +123,8 @@ mod tests {
     fn test_simple() {
         let ctx = types::Context::new();
         let two_words = Arc::<ConstructNode<Core>>::pair(
-            &Arc::<ConstructNode<_>>::const_word(&ctx, Value::u32(2)),
-            &Arc::<ConstructNode<_>>::const_word(&ctx, Value::u16(16)),
+            &Arc::<ConstructNode<_>>::const_word(&ctx, Word::u32(2)),
+            &Arc::<ConstructNode<_>>::const_word(&ctx, Word::u16(16)),
         )
         .unwrap();
         assert_eq!(

src/merkle/amr.rs

@@ -4,6 +4,7 @@ use crate::impl_midstate_wrapper;
 use crate::jet::Jet;
 use crate::merkle::compact_value;
 use crate::types::arrow::FinalArrow;
+use crate::value::Word;
 use crate::{Cmr, Tmr, Value};
 use hashes::sha256::Midstate;
 
@@ -167,8 +168,8 @@ impl Amr {
     ///
     /// This is equal to the IMR of the equivalent scribe, converted to a CMR in
     /// the usual way for jets.
-    pub fn const_word(v: &Value) -> Self {
-        Cmr::const_word(v).into()
+    pub fn const_word(word: &Word) -> Self {
+        Cmr::const_word(word).into()
     }
 
     #[rustfmt::skip]