Merge pull request #532 from Chia-Network/node_from_stream_backrefs_o…

…ptimisation

Node from stream backrefs optimisation

benches/deserialize.rs

@@ -1,7 +1,8 @@
 use clvmr::allocator::Allocator;
 use clvmr::serde::{
-    node_from_bytes, node_from_bytes_backrefs, node_to_bytes_backrefs,
-    serialized_length_from_bytes, serialized_length_from_bytes_trusted, tree_hash_from_stream,
+    node_from_bytes, node_from_bytes_backrefs, node_from_bytes_backrefs_old,
+    node_to_bytes_backrefs, serialized_length_from_bytes, serialized_length_from_bytes_trusted,
+    tree_hash_from_stream,
 };
 use criterion::{criterion_group, criterion_main, Criterion};
 use std::include_bytes;
@@ -64,6 +65,15 @@ fn deserialize_benchmark(c: &mut Criterion) {
                 start.elapsed()
             })
         });
+
+        group.bench_function(format!("node_from_bytes_backrefs_old{name_suffix}"), |b| {
+            b.iter(|| {
+                a.restore_checkpoint(&iter_checkpoint);
+                let start = Instant::now();
+                node_from_bytes_backrefs_old(&mut a, bl).expect("node_from_bytes_backrefs_old");
+                start.elapsed()
+            })
+        });
     }
 
     let mut a = Allocator::new();

fuzz/Cargo.toml

@@ -10,7 +10,7 @@ cargo-fuzz = true
 
 [dependencies]
 libfuzzer-sys = { workspace = true }
-clvmr = { workspace = true }
+clvmr = { workspace = true, features = ["counters"] }
 chia-sha2 = { workspace = true }
 hex = { workspace = true }
 arbitrary = { workspace = true }

fuzz/fuzz_targets/deserialize_br.rs

@@ -1,23 +1,29 @@
 #![no_main]
+
+mod node_eq;
+
 use clvmr::allocator::Allocator;
 use clvmr::serde::node_from_bytes_backrefs;
-use clvmr::serde::node_to_bytes_backrefs;
+use clvmr::serde::node_from_bytes_backrefs_old;
 use libfuzzer_sys::fuzz_target;
 
 fuzz_target!(|data: &[u8]| {
     let mut allocator = Allocator::new();
-    let program = match node_from_bytes_backrefs(&mut allocator, data) {
-        Err(_) => {
+    let res1 = node_from_bytes_backrefs(&mut allocator, data);
+    let node_count = allocator.pair_count();
+    let res2 = node_from_bytes_backrefs_old(&mut allocator, data);
+    // check that the new implementation creates the same number of pair nodes as the old one
+    assert_eq!(node_count * 2, allocator.pair_count());
+    match (res1, res2) {
+        (Err(_e1), Err(_e2)) => {
+            // both failed, that's fine
             return;
         }
-        Ok(r) => r,
-    };
-
-    let b1 = node_to_bytes_backrefs(&allocator, program).unwrap();
-
-    let mut allocator = Allocator::new();
-    let program = node_from_bytes_backrefs(&mut allocator, &b1).unwrap();
-
-    let b2 = node_to_bytes_backrefs(&allocator, program).unwrap();
-    assert_eq!(b1, b2);
+        (Ok(n1), Ok(n2)) => {
+            assert!(node_eq::node_eq(&allocator, n1, n2));
+        }
+        _ => {
+            panic!("mismatching results");
+        }
+    }
 });

src/allocator.rs

@@ -107,6 +107,7 @@ pub struct Checkpoint {
     pairs: usize,
     atoms: usize,
     small_atoms: usize,
+    ghost_pairs: usize,
 }
 
 pub enum NodeVisitor<'a> {
@@ -177,6 +178,10 @@ pub struct Allocator {
     // the number of small atoms we've allocated. We keep track of these to ensure the limit on the
     // number of atoms is identical to what it was before the small-atom optimization
     small_atoms: usize,
+
+    // this tracks the pairs that are being skipped from optimisations
+    // we track this to simulate a compatible maximum with older versions
+    num_ghost_pairs: usize,
 }
 
 impl Default for Allocator {
@@ -244,6 +249,7 @@ impl Allocator {
             // initialize this to 2 to behave as if we had allocated atoms for
             // nil() and one(), like we used to
             small_atoms: 2,
+            num_ghost_pairs: 0,
         };
         r.u8_vec.reserve(1024 * 1024);
         r.atom_vec.reserve(256);
@@ -260,6 +266,7 @@ impl Allocator {
             pairs: self.pair_vec.len(),
             atoms: self.atom_vec.len(),
             small_atoms: self.small_atoms,
+            ghost_pairs: self.num_ghost_pairs,
         }
     }
 
@@ -275,6 +282,7 @@ impl Allocator {
         self.pair_vec.truncate(cp.pairs);
         self.atom_vec.truncate(cp.atoms);
         self.small_atoms = cp.small_atoms;
+        self.num_ghost_pairs = cp.ghost_pairs;
     }
 
     pub fn new_atom(&mut self, v: &[u8]) -> Result<NodePtr, EvalErr> {
@@ -332,13 +340,32 @@ impl Allocator {
 
     pub fn new_pair(&mut self, first: NodePtr, rest: NodePtr) -> Result<NodePtr, EvalErr> {
         let idx = self.pair_vec.len();
-        if idx == MAX_NUM_PAIRS {
+        if idx >= MAX_NUM_PAIRS - self.num_ghost_pairs {
             return err(self.nil(), "too many pairs");
         }
         self.pair_vec.push(IntPair { first, rest });
         Ok(NodePtr::new(ObjectType::Pair, idx))
     }
 
+    // this code is used when we are simulating pairs with a vec locally
+    // in the deserialize_br code
+    // we must maintain parity with the old deserialize_br code so need to track the skipped pairs
+    pub fn add_ghost_pair(&mut self, amount: usize) -> Result<(), EvalErr> {
+        if MAX_NUM_PAIRS - self.num_ghost_pairs - self.pair_vec.len() < amount {
+            return err(self.nil(), "too many pairs");
+        }
+        self.num_ghost_pairs += amount;
+        Ok(())
+    }
+
+    // this code is used when we actually create the pairs that were previously skipped ghost pairs
+    pub fn remove_ghost_pair(&mut self, amount: usize) -> Result<(), EvalErr> {
+        // currently let this panic with overflow if we go below 0 to debug if/where it happens
+        debug_assert!(self.num_ghost_pairs >= amount);
+        self.num_ghost_pairs -= amount;
+        Ok(())
+    }
+
     pub fn new_substr(&mut self, node: NodePtr, start: u32, end: u32) -> Result<NodePtr, EvalErr> {
         self.check_atom_limit()?;
 
@@ -669,6 +696,11 @@ impl Allocator {
 
     #[cfg(feature = "counters")]
     pub fn pair_count(&self) -> usize {
+        self.pair_vec.len() + self.num_ghost_pairs
+    }
+
+    #[cfg(feature = "counters")]
+    pub fn pair_count_no_ghosts(&self) -> usize {
         self.pair_vec.len()
     }
 
@@ -1008,6 +1040,19 @@ mod tests {
         }
 
         assert_eq!(a.new_pair(atom, atom).unwrap_err().1, "too many pairs");
+        assert_eq!(a.add_ghost_pair(1).unwrap_err().1, "too many pairs");
+    }
+
+    #[test]
+    fn test_ghost_pair_limit() {
+        let mut a = Allocator::new();
+        let atom = a.new_atom(b"foo").unwrap();
+        // one pair is OK
+        let _pair1 = a.new_pair(atom, atom).unwrap();
+        a.add_ghost_pair(MAX_NUM_PAIRS - 1).unwrap();
+
+        assert_eq!(a.new_pair(atom, atom).unwrap_err().1, "too many pairs");
+        assert_eq!(a.add_ghost_pair(1).unwrap_err().1, "too many pairs");
     }
 
     #[test]