feat: [MR-144] Reuse unchanged canister digests across HashTrees

rs/canonical_state/src/lazy_tree_conversion.rs

@@ -9,9 +9,10 @@ use crate::{
 };
 use LazyTree::Blob;
 use ic_canonical_state_tree_hash::{
-    hash_tree::HashTree,
+    hash_tree::{HashTree, HashTreeError, hash_lazy_tree},
     lazy_tree::{
-        Lazy, LazyFork, LazyTree, blob, fork, materialize::materialize_partial, num, string,
+        Lazy, LazyFork, LazyTree, SubtreeExpander, SubtreeSource, blob, fork,
+        materialize::materialize_partial, num, string,
     },
 };
 use ic_crypto_tree_hash::{Label, Witness, sparse_labeled_tree_from_paths};
@@ -786,7 +787,7 @@ const CANISTER_NO_MODULE_LABELS: [&[u8]; 1] = [CONTROLLERS_LABEL];
 
 #[derive(Clone)]
 struct CanisterFork<'a> {
-    canister: &'a CanisterState,
+    canister: &'a Arc<CanisterState>,
     version: CertificationVersion,
 }
 
@@ -849,6 +850,54 @@ impl<'a> LazyFork<'a> for CanisterFork<'a> {
     fn len(&self) -> usize {
         self.valid_labels().len()
     }
+
+    /// A canister's certified subtree is stored as a reusable stub identified by
+    /// the backing `Arc<CanisterState>` and the version-specific expander. An
+    /// unchanged canister keeps the same `Arc` (copy-on-write) and the same
+    /// expander, so its precomputed digest is reused from the baseline; any
+    /// mutation or version change yields a mismatched [`SubtreeSource`] and a
+    /// rebuild.
+    fn subtree_source(&self) -> Option<SubtreeSource> {
+        Some(SubtreeSource::new(
+            self.canister,
+            select_canister_expander(self.version),
+        ))
+    }
+}
+
+/// Rebuilds a canister's stubbed [subtree](`NodeKind::Stub`) for witness
+/// generation, by recovering the `Arc<CanisterState>` from the stub's
+/// [`SubtreeSource`] and traversing its [`CanisterFork`].
+///
+/// The certification version (which the canonical encoding depends on) is baked
+/// in as the const parameter `V`, so the stored function pointer alone fully
+/// determines the expansion — see [`select_canister_expander`].
+fn expand_canister<const V: u32>(source: &SubtreeSource) -> Result<HashTree, HashTreeError> {
+    let canister = source.downcast::<CanisterState>();
+    let version = CertificationVersion::try_from(V)
+        .expect("const version parameter is a valid certification version");
+    // `canister` (and thus the borrow below) outlives `hash_lazy_tree`, which
+    // returns an owned `HashTree`; no borrow escapes.
+    hash_lazy_tree(&fork(CanisterFork {
+        canister: &canister,
+        version,
+    }))
+}
+
+/// Selects the [`expand_canister`] monomorphization for `version`, so the
+/// resulting [`SubtreeExpander`] function pointer carries the version with it
+/// (rather than replicating it in every stub).
+fn select_canister_expander(version: CertificationVersion) -> SubtreeExpander {
+    match version {
+        CertificationVersion::V19 => expand_canister::<{ CertificationVersion::V19 as u32 }>,
+        CertificationVersion::V20 => expand_canister::<{ CertificationVersion::V20 as u32 }>,
+        CertificationVersion::V21 => expand_canister::<{ CertificationVersion::V21 as u32 }>,
+        CertificationVersion::V22 => expand_canister::<{ CertificationVersion::V22 as u32 }>,
+        CertificationVersion::V23 => expand_canister::<{ CertificationVersion::V23 as u32 }>,
+        CertificationVersion::V24 => expand_canister::<{ CertificationVersion::V24 as u32 }>,
+        CertificationVersion::V25 => expand_canister::<{ CertificationVersion::V25 as u32 }>,
+        CertificationVersion::V26 => expand_canister::<{ CertificationVersion::V26 as u32 }>,
+    }
 }
 
 fn api_boundary_nodes_as_tree(

rs/canonical_state/tree_hash/BUILD.bazel

@@ -10,6 +10,7 @@ rust_library(
     deps = [
         # Keep sorted.
         "//rs/crypto/tree_hash",
+        "//rs/utils",
         "@crate_index//:itertools",
         "@crate_index//:leb128",
         "@crate_index//:scoped_threadpool",

rs/canonical_state/tree_hash/Cargo.toml

@@ -8,6 +8,7 @@ documentation.workspace = true
 
 [dependencies]
 ic-crypto-tree-hash = { path = "../../crypto/tree_hash" }
+ic-utils = { path = "../../utils" }
 itertools = { workspace = true }
 leb128 = "0.2.1"
 scoped_threadpool = "0.1.*"