Handle edge case of extension with modified prefix

Author: BrianBland

src/storage/overlay_root.rs

@@ -180,17 +180,39 @@ impl StorageEngine {
                         stats,
                         updated_storage_branch_nodes,
                     )?;
-                    let hash = node_hash.unwrap();
-                    println!("Adding unaffected branch: {:?}", path_prefix);
-                    hash_builder.add_branch(path_prefix.clone(), hash, true);
-                    stats.account.branch_nodes += 1;
-                    self.process_nonoverlapping_overlay_state(
-                        context,
-                        &post_overlay,
-                        hash_builder,
-                        stats,
-                        updated_storage_branch_nodes,
-                    )?;
+                    if !node.prefix().is_empty() {
+                        // The node is actually an extension + branch node. We cannot add it by hash
+                        // as the extension prefix may have changed, which
+                        // would in turn change the hash. We need to
+                        // traverse the node via its full path.
+                        println!(
+                            "Processing extension + branch: {:?} / {:?}",
+                            path_prefix, full_path
+                        );
+                        self.handle_affected_node_with_overlay(
+                            context,
+                            &with_prefix,
+                            hash_builder,
+                            slotted_page,
+                            &node,
+                            &full_path,
+                            stats,
+                            updated_storage_branch_nodes,
+                        )?;
+                    } else {
+                        // This is a true branch node, so we can add it by hash.
+                        let hash = node_hash.unwrap();
+                        println!("Adding unaffected branch: {:?}", path_prefix);
+                        hash_builder.add_branch(path_prefix.clone(), hash, true);
+                        stats.account.branch_nodes += 1;
+                        self.process_nonoverlapping_overlay_state(
+                            context,
+                            &post_overlay,
+                            hash_builder,
+                            stats,
+                            updated_storage_branch_nodes,
+                        )?;
+                    }
                     return Ok(());
                 } else {
                     let (pre_overlay, with_prefix, post_overlay) =
@@ -309,6 +331,7 @@ impl StorageEngine {
                 }
                 None => {
                     // Tombstone - skip
+                    println!("Skipping tombstone: {:?}", path);
                     last_processed_path = Some(path);
                 }
             }
@@ -355,11 +378,18 @@ impl StorageEngine {
                 )?;
             }
             NodeKind::StorageLeaf { .. } => {
-                if let Some((_, Some(OverlayValue::Storage(value)))) = overlay.get(0) {
-                    hash_builder.add_leaf(full_path.clone(), &self.encode_storage(value)?);
-                    stats.storage.overlay_leaves += 1;
-                } else {
-                    panic!("Storage leaf does not have a valid overlay: {full_path:?}");
+                match overlay.get(0) {
+                    Some((_, Some(OverlayValue::Storage(value)))) => {
+                        hash_builder.add_leaf(full_path.clone(), &self.encode_storage(value)?);
+                        stats.storage.overlay_leaves += 1;
+                    }
+                    Some((_, None)) => {
+                        println!("Skipping deleted storage leaf: {:?}", full_path);
+                        stats.storage.overlay_leaves += 1;
+                    }
+                    _ => {
+                        panic!("Storage leaf does not have a valid overlay: {full_path:?}");
+                    }
                 }
             }
         }
@@ -396,22 +426,28 @@ impl StorageEngine {
             children[i as usize] = (child_path, child_overlay, node.child(i as u8).unwrap());
         }
 
-        let num_children = children
+        // This conservatively counts the minimum number of children that will exist. It may
+        // undercount.
+        let min_num_children = children
             .iter()
             .filter(|(_, child_overlay, child_pointer)| {
                 if let Some((_, value)) = child_overlay.get(0) {
+                    // If the overlayed value is none, this branch path might be removed.
                     return value.is_some();
                 } else if child_pointer.is_some() {
+                    // If there is no overlayed value and an existing child, then this path will
+                    // definitely exist.
                     return true;
                 }
 
                 false
             })
             .count();
 
-        if num_children > 1 {
-            // We have at least 2 children, so we can add non-overlayed children by hash instead of
-            // traversing them. This can save many lookups, hashes, and page reads.
+        if min_num_children > 1 {
+            // We are guaranteed to have at least 2 children, keeping this branch node alive.
+            // We can add non-overlayed children by hash instead of traversing them. This can save
+            // many lookups, hashes, and page reads.
             for (child_path, child_overlay, child_pointer) in children.iter() {
                 if let Some((path, _)) = child_overlay.get(0) {
                     if &path == child_path {
@@ -476,8 +512,8 @@ impl StorageEngine {
             return Ok(());
         }
 
-        // Otherwise, this branch may no longer exist, so we need to traverse and add the children
-        // directly
+        // Otherwise, it is possible that this branch might be removed.
+        // We need to traverse and add the children directly in order to be safe.
         for (child_path, child_overlay, child_pointer) in children.iter() {
             if let Some((path, _)) = child_overlay.get(0) {
                 if &path == child_path {
@@ -573,7 +609,7 @@ impl StorageEngine {
                 }
             } else {
                 // If overlay_value is None, it's a deletion - skip
-                // println!("Skipping deletion: {full_path:?}");
+                println!("Skipping deleted account: {full_path:?}");
                 None
             }
         } else {
@@ -2002,8 +2038,6 @@ mod tests {
             db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap();
         assert_ne!(overlay_root, initial_root);
 
-        // println!("Overlay branches: {:?}", overlay_branches);
-
         for (path, branch) in overlay_branches.iter() {
             if let Some(root_hash) = branch.root_hash {
                 overlay_mut_with_branches.insert(path.clone(), Some(OverlayValue::Hash(root_hash)));