sealable-trie: add del operation

Cargo.toml

@@ -26,6 +26,7 @@ anchor-lang = {version = "0.28.0", features = ["init-if-needed"]}
 base64 = { version = "0.21", default-features = false, features = ["alloc"] }
 borsh = { version = "0.10.3", default-features = false }
 derive_more = "0.99.17"
+hex-literal = "0.4.1"
 ibc = { version = "0.45.0", default-features = false, features = ["serde", "borsh"] }
 ibc-proto = { version = "0.35.0", default-features = false, features = ["serde"] }
 pretty_assertions = "1.4.0"

common/sealable-trie/Cargo.toml

@@ -16,6 +16,7 @@ memory.workspace = true
 stdx.workspace = true
 
 [dev-dependencies]
+hex-literal.workspace = true
 pretty_assertions.workspace = true
 rand.workspace = true
 

common/sealable-trie/src/trie.rs

@@ -6,6 +6,7 @@ use memory::Ptr;
 use crate::nodes::{Node, NodeRef, RawNode, Reference};
 use crate::{bits, proof};
 
+mod del;
 mod seal;
 mod set;
 #[cfg(test)]
@@ -244,9 +245,8 @@ impl<A: memory::Allocator<Value = Value>> Trie<A> {
     pub fn set(&mut self, key: &[u8], value_hash: &CryptoHash) -> Result<()> {
         let (ptr, hash) = (self.root_ptr, self.root_hash.clone());
         let key = bits::Slice::from_bytes(key).ok_or(Error::KeyTooLong)?;
-        let (ptr, hash) =
-            set::SetContext::new(&mut self.alloc, key, value_hash)
-                .set(ptr, &hash)?;
+        let (ptr, hash) = set::Context::new(&mut self.alloc, key, value_hash)
+            .set(ptr, &hash)?;
         self.root_ptr = Some(ptr);
         self.root_hash = hash;
         Ok(())
@@ -264,19 +264,35 @@ impl<A: memory::Allocator<Value = Value>> Trie<A> {
     /// an error.
     // TODO(mina86): Add seal_with_proof.
     pub fn seal(&mut self, key: &[u8]) -> Result<()> {
-        let key = bits::Slice::from_bytes(key).ok_or(Error::KeyTooLong)?;
         if self.root_hash == EMPTY_TRIE_ROOT {
             return Err(Error::NotFound);
         }
-
-        let seal = seal::SealContext::new(&mut self.alloc, key)
+        let key = bits::Slice::from_bytes(key).ok_or(Error::KeyTooLong)?;
+        let removed = seal::Context::new(&mut self.alloc, key)
             .seal(NodeRef::new(self.root_ptr, &self.root_hash))?;
-        if seal {
+        if removed {
             self.root_ptr = None;
         }
         Ok(())
     }
 
+    /// Deletes value at given key.  Returns `false` if key was not found.
+    pub fn del(&mut self, key: &[u8]) -> Result<bool> {
+        let key = bits::Slice::from_bytes(key).ok_or(Error::KeyTooLong)?;
+        let res = del::Context::new(&mut self.alloc, key)
+            .del(self.root_ptr, &self.root_hash);
+        match res {
+            Ok(res) => {
+                let (ptr, hash) = res.unwrap_or((None, EMPTY_TRIE_ROOT));
+                self.root_ptr = ptr;
+                self.root_hash = hash;
+                Ok(true)
+            }
+            Err(Error::NotFound) => Ok(false),
+            Err(err) => Err(err),
+        }
+    }
+
     /// Prints the trie.  Used for testing and debugging only.
     #[cfg(test)]
     pub(crate) fn print(&self) {

common/sealable-trie/src/trie/del.rs

@@ -0,0 +1,296 @@
+use lib::hash::CryptoHash;
+use memory::Ptr;
+
+use super::{Error, Result};
+use crate::bits;
+use crate::nodes::{Node, NodeRef, RawNode, Reference, ValueRef};
+
+/// Context for [`Trie::del`] operation.
+pub(super) struct Context<'a, A: memory::Allocator<Value = super::Value>> {
+    /// Part of the key yet to be traversed.
+    ///
+    /// It starts as the key user provided and as trie is traversed bits are
+    /// removed from its front.
+    key: bits::Slice<'a>,
+
+    /// Allocator used to allocate new nodes.
+    wlog: memory::WriteLog<'a, A>,
+}
+
+impl<'a, A: memory::Allocator<Value = super::Value>> Context<'a, A> {
+    pub(super) fn new(alloc: &'a mut A, key: bits::Slice<'a>) -> Self {
+        let wlog = memory::WriteLog::new(alloc);
+        Self { key, wlog }
+    }
+
+    /// Inserts value hash into the trie.
+    pub(super) fn del(
+        mut self,
+        root_ptr: Option<Ptr>,
+        root_hash: &CryptoHash,
+    ) -> Result<Option<(Option<Ptr>, CryptoHash)>> {
+        if *root_hash == super::EMPTY_TRIE_ROOT {
+            return Err(Error::NotFound);
+        };
+        let action =
+            self.handle(NodeRef { ptr: root_ptr, hash: root_hash }, false)?;
+        let res = self.ref_from_action(action)?.map(|child| match child {
+            OwnedRef::Node(ptr, hash) => (ptr, hash),
+            _ => unreachable!(),
+        });
+        self.wlog.commit();
+        Ok(res)
+    }
+
+    /// Processes a reference which may be either node or value reference.
+    fn handle_reference(
+        &mut self,
+        child: Reference,
+        from_ext: bool,
+    ) -> Result<Action> {
+        match child {
+            Reference::Value(vref) => {
+                if vref.is_sealed {
+                    Err(Error::Sealed)
+                } else if self.key.is_empty() {
+                    Ok(Action::Drop)
+                } else {
+                    Err(Error::NotFound)
+                }
+            }
+            Reference::Node(nref) => self.handle(nref, from_ext),
+        }
+    }
+
+    /// Processes a node.
+    fn handle(&mut self, nref: NodeRef, from_ext: bool) -> Result<Action> {
+        let ptr = nref.ptr.ok_or(Error::Sealed)?;
+        let node = RawNode(*self.wlog.allocator().get(ptr));
+        let node = node.decode();
+        debug_assert_eq!(*nref.hash, node.hash());
+
+        match node {
+            Node::Branch { children } => self.handle_branch(ptr, children),
+            Node::Extension { key, child } => {
+                self.handle_extension(ptr, key, child)
+            }
+            Node::Value { value, child } => {
+                self.handle_value(ptr, value, child, from_ext)
+            }
+        }
+    }
+
+    /// Processes a Branch node.
+    fn handle_branch(
+        &mut self,
+        ptr: Ptr,
+        children: [Reference; 2],
+    ) -> Result<Action> {
+        let key_offset = self.key.offset;
+
+        let side = usize::from(self.key.pop_front().ok_or(Error::NotFound)?);
+        let action = self.handle_reference(children[side], false)?;
+
+        // If the branch changed but wasn’t deleted, we just need to replace the
+        // reference.  Otherwise, we’ll need to convert the Branch into an
+        // Extension.
+        if let Some(child) = self.ref_from_action(action)? {
+            let child = child.to_ref();
+            let (left, right) = if side == 0 {
+                (child, children[1])
+            } else {
+                (children[0], child)
+            };
+            let node = RawNode::branch(left, right);
+            return Ok(Action::Ref(self.set_node(ptr, node)));
+        }
+
+        // The child has been deleted.  We need to convert this Branch into an
+        // Extension with a single-bit key and the other child.  However, if the
+        // other child already is also an Extension, we need to merge them.
+
+        self.del_node(ptr);
+        let child = children[1 - side];
+        Ok(self
+            .maybe_pop_extension(child, &|key| {
+                bits::Owned::unshift(side == 0, key).unwrap()
+            })
+            .unwrap_or_else(|| {
+                Action::Ext(
+                    bits::Owned::bit(side == 0, key_offset),
+                    OwnedRef::from(child),
+                )
+            }))
+    }
+
+    /// Processes an Extension node.
+    fn handle_extension(
+        &mut self,
+        ptr: Ptr,
+        key: bits::Slice,
+        child: Reference,
+    ) -> Result<Action> {
+        if !self.key.strip_prefix(key) {
+            return Err(Error::NotFound);
+        }
+        self.del_node(ptr);
+        Ok(match self.handle_reference(child, true)? {
+            Action::Drop => Action::Drop,
+            Action::Ref(child) => Action::Ext(key.into(), child),
+            Action::Ext(suffix, child) => {
+                let key = bits::Owned::concat(key, suffix).unwrap();
+                Action::Ext(key, child)
+            }
+        })
+    }
+
+    /// Processes a Branch node.
+    fn handle_value(
+        &mut self,
+        ptr: Ptr,
+        value: ValueRef<'_, ()>,
+        child: NodeRef,
+        from_ext: bool,
+    ) -> Result<Action> {
+        // We’ve reached the value we want to delete.  Drop the Value node and
+        // replace parent’s reference with child we’re pointing at.  The one
+        // complication is that if our parent is an Extension, we need to fetch
+        // the child to check if it’s an Extension as well.
+        if self.key.is_empty() {
+            self.del_node(ptr);
+            if from_ext {
+                let action = self
+                    .maybe_pop_extension(Reference::Node(child), &|key| {
+                        key.into()
+                    });
+                if let Some(action) = action {
+                    return Ok(action);
+                }
+            }
+            return Ok(Action::Ref(child.into()));
+        }
+
+        // Traverse into the child and handle that.
+        let action = self.handle(child, false)?;
+        Ok(match self.ref_from_action(action)? {
+            None => {
+                // We’re deleting the child which means we need to delete the
+                // Value node and replace parent’s reference to ValueRef.
+                self.del_node(ptr);
+                let value = ValueRef::new(false, value.hash);
+                Action::Ref(value.into())
+            }
+            Some(OwnedRef::Node(child_ptr, hash)) => {
+                let child = NodeRef::new(child_ptr, &hash);
+                let node = RawNode::value(value, child);
+                Action::Ref(self.set_node(ptr, node))
+            }
+            Some(OwnedRef::Value(..)) => unreachable!(),
+        })
+    }
+
+    /// If `child` is a node reference pointing at an Extension node, pops that
+    /// node and returns corresponding `Action::Ext` action.
+    fn maybe_pop_extension(
+        &mut self,
+        child: Reference,
+        make_key: &dyn Fn(bits::Slice) -> bits::Owned,
+    ) -> Option<Action> {
+        if let Reference::Node(NodeRef { ptr: Some(ptr), hash }) = child {
+            let node = RawNode(*self.wlog.allocator().get(ptr));
+            let node = node.decode();
+            debug_assert_eq!(*hash, node.hash());
+
+            if let Node::Extension { key, child } = node {
+                // Drop the child Extension and merge keys.
+                self.del_node(ptr);
+                return Some(Action::Ext(make_key(key), OwnedRef::from(child)));
+            }
+        }
+        None
+    }
+
+    /// Sets value of a node cell at given address and returns an [`OwnedRef`]
+    /// pointing at the node.
+    fn set_node(&mut self, ptr: Ptr, node: RawNode) -> OwnedRef {
+        let hash = node.decode().hash();
+        self.wlog.set(ptr, *node);
+        OwnedRef::Node(Some(ptr), hash)
+    }
+
+    /// Frees a node.
+    fn del_node(&mut self, ptr: Ptr) { self.wlog.free(ptr); }
+
+    /// Converts an [`Action`] into an [`OwnedRef`] if it’s not a `Drop` action.
+    ///
+    /// If action is [`Action::Ext`] allocates a new node or sequence of nodes
+    /// and adds them eventually converting the action to an `OwnedRef`.  If
+    /// action is [`Action::Ref`] already, simply returns the reference.
+    fn ref_from_action(&mut self, action: Action) -> Result<Option<OwnedRef>> {
+        let (key, mut child) = match action {
+            Action::Ext(key, child) => (key, child),
+            Action::Ref(owned) => return Ok(Some(owned)),
+            Action::Drop => return Ok(None),
+        };
+
+        for chunk in key.as_slice().chunks().rev() {
+            let node = RawNode::extension(chunk, child.to_ref()).unwrap();
+            let ptr = self.wlog.alloc(node.0)?;
+            child = OwnedRef::Node(Some(ptr), node.decode().hash());
+        }
+
+        Ok(Some(child))
+    }
+}
+
+/// An internal representation of results of handling of a node.
+enum Action {
+    /// The node has been deleted.
+    ///
+    /// Deletion should propagate upstream.
+    Drop,
+
+    /// The node needs to be replaced with given Extension node.
+    ///
+    /// This may propagate upstream through Extension nodes that may need to be
+    /// merged or split.
+    Ext(bits::Owned, OwnedRef),
+
+    /// The reference has been replaced by the given owned reference.
+    Ref(OwnedRef),
+}
+
+enum OwnedRef {
+    Node(Option<Ptr>, CryptoHash),
+    Value(bool, CryptoHash),
+}
+
+impl OwnedRef {
+    fn to_ref(&self) -> Reference {
+        match self {
+            Self::Node(ptr, hash) => Reference::node(*ptr, hash),
+            Self::Value(is_sealed, hash) => Reference::value(*is_sealed, hash),
+        }
+    }
+}
+
+impl From<NodeRef<'_>> for OwnedRef {
+    fn from(nref: NodeRef) -> OwnedRef {
+        Self::Node(nref.ptr, nref.hash.clone())
+    }
+}
+
+impl From<ValueRef<'_>> for OwnedRef {
+    fn from(vref: ValueRef) -> OwnedRef {
+        Self::Value(vref.is_sealed, vref.hash.clone())
+    }
+}
+
+impl From<Reference<'_>> for OwnedRef {
+    fn from(rf: Reference<'_>) -> Self {
+        match rf {
+            Reference::Node(nref) => nref.into(),
+            Reference::Value(vref) => vref.into(),
+        }
+    }
+}