[perf] Use Vec instead of BTreeMap in tree_from_memory (#1338)

* Build memory tree presumably faster

* Construct all-zero nodes in advance

* Revert "Construct all-zero nodes in advance"

This reverts commit 9f7dc7f.

* Cosmetic change

* Calculate the capacity first, then reserve exactly this much

* Revert "Calculate the capacity first, then reserve exactly this much"

This reverts commit 3ee49b3.

* Make pagedvec iter return &T, not T

* Calculate the capacity first, then reserve exactly this much

* Revert "Make pagedvec iter return &T, not T"

This reverts commit f1e8403.

* Revert "Calculate the capacity first, then reserve exactly this much"

This reverts commit c719be1.

* Cache only all-zero leaf

* Avoid building `Vec<(_, [F; CHUNK])>`

* Revert "Cache only all-zero leaf"

This reverts commit 98ae517.

* Revert "Revert "Cache only all-zero leaf""

This reverts commit e247f5b.

---------

Co-authored-by: Jonathan Wang <[email protected]>

Author: Golovanov399

crates/vm/src/system/memory/tree/mod.rs

@@ -1,6 +1,6 @@
 pub mod public_values;
 
-use std::{collections::BTreeMap, sync::Arc};
+use std::{ops::Range, sync::Arc};
 
 use openvm_stark_backend::{p3_field::PrimeField32, p3_maybe_rayon::prelude::*};
 use MemoryNode::*;
@@ -66,23 +66,71 @@ impl<const CHUNK: usize, F: PrimeField32> MemoryNode<CHUNK, F> {
     }
 
     fn from_memory(
-        memory: &BTreeMap<u64, [F; CHUNK]>,
-        height: usize,
+        memory: &[(u64, F)],
+        lookup_range: Range<usize>,
+        length: u64,
         from: u64,
         hasher: &(impl Hasher<CHUNK, F> + Sync),
+        zero_leaf: &MemoryNode<CHUNK, F>,
     ) -> MemoryNode<CHUNK, F> {
-        let mut range = memory.range(from..from + (1 << height));
-        if height == 0 {
-            let values = *memory.get(&from).unwrap_or(&[F::ZERO; CHUNK]);
-            MemoryNode::new_leaf(hasher.hash(&values))
-        } else if range.next().is_none() {
+        if length == CHUNK as u64 {
+            if lookup_range.is_empty() {
+                zero_leaf.clone()
+            } else {
+                debug_assert_eq!(memory[lookup_range.start].0, from);
+                let mut values = [F::ZERO; CHUNK];
+                for (index, value) in memory[lookup_range].iter() {
+                    values[(index % CHUNK as u64) as usize] = *value;
+                }
+                MemoryNode::new_leaf(hasher.hash(&values))
+            }
+        } else if lookup_range.is_empty() {
             let leaf_value = hasher.hash(&[F::ZERO; CHUNK]);
-            MemoryNode::construct_uniform(height, leaf_value, hasher)
+            MemoryNode::construct_uniform(
+                (length / CHUNK as u64).trailing_zeros() as usize,
+                leaf_value,
+                hasher,
+            )
         } else {
-            let midpoint = from + (1 << (height - 1));
+            let midpoint = from + length / 2;
+            let mid = {
+                let mut left = lookup_range.start;
+                let mut right = lookup_range.end;
+                if memory[left].0 >= midpoint {
+                    left
+                } else {
+                    while left + 1 < right {
+                        let mid = left + (right - left) / 2;
+                        if memory[mid].0 < midpoint {
+                            left = mid;
+                        } else {
+                            right = mid;
+                        }
+                    }
+                    right
+                }
+            };
             let (left, right) = join(
-                || Self::from_memory(memory, height - 1, from, hasher),
-                || Self::from_memory(memory, height - 1, midpoint, hasher),
+                || {
+                    Self::from_memory(
+                        memory,
+                        lookup_range.start..mid,
+                        length >> 1,
+                        from,
+                        hasher,
+                        zero_leaf,
+                    )
+                },
+                || {
+                    Self::from_memory(
+                        memory,
+                        mid..lookup_range.end,
+                        length >> 1,
+                        midpoint,
+                        hasher,
+                        zero_leaf,
+                    )
+                },
             );
             NonLeaf {
                 hash: hasher.compress(&left.hash(), &right.hash()),
@@ -97,22 +145,33 @@ impl<const CHUNK: usize, F: PrimeField32> MemoryNode<CHUNK, F> {
         memory: &MemoryImage<F>,
         hasher: &(impl Hasher<CHUNK, F> + Sync),
     ) -> MemoryNode<CHUNK, F> {
-        // Construct a BTreeMap that includes the address space in the label calculation,
+        // Construct a Vec that includes the address space in the label calculation,
         // representing the entire memory tree.
-        let mut memory_partition = BTreeMap::new();
-        for ((address_space, pointer), value) in memory.items() {
-            let label = (address_space, pointer / CHUNK as u32);
-            let index = memory_dimensions.label_to_index(label);
-            let chunk = memory_partition
-                .entry(index)
-                .or_insert_with(|| [F::ZERO; CHUNK]);
-            chunk[(pointer % CHUNK as u32) as usize] = value;
-        }
+        let memory_items = memory
+            .items()
+            .filter(|((_, ptr), _)| *ptr as usize / CHUNK < (1 << memory_dimensions.address_height))
+            .map(|((address_space, pointer), value)| {
+                (
+                    memory_dimensions.label_to_index((address_space, pointer / CHUNK as u32))
+                        * CHUNK as u64
+                        + (pointer % CHUNK as u32) as u64,
+                    value,
+                )
+            })
+            .collect::<Vec<_>>();
+        debug_assert!(memory_items.is_sorted_by_key(|(addr, _)| addr));
+        debug_assert!(
+            memory_items.last().map_or(0, |(addr, _)| *addr)
+                < ((CHUNK as u64) << memory_dimensions.overall_height())
+        );
+        let zero_leaf = MemoryNode::new_leaf(hasher.hash(&[F::ZERO; CHUNK]));
         Self::from_memory(
-            &memory_partition,
-            memory_dimensions.overall_height(),
+            &memory_items,
+            0..memory_items.len(),
+            (CHUNK as u64) << memory_dimensions.overall_height(),
             0,
             hasher,
+            &zero_leaf,
         )
     }
 }