diff --git a/iterator.go b/iterator.go index 9ef92b5..337aa2c 100644 --- a/iterator.go +++ b/iterator.go @@ -141,7 +141,7 @@ func (i *Iterator[T]) SeekPrefixWatch(prefixKey []byte) (watch <-chan struct{}) if prefixKey == nil { i.node = node i.stack = []Node[T]{node} - return + return watch } for { @@ -151,7 +151,7 @@ func (i *Iterator[T]) SeekPrefixWatch(prefixKey []byte) (watch <-chan struct{}) watch = node.getMutateCh() if node.isLeaf() { - return + return watch } // Determine the child index to proceed based on the next byte of the prefix @@ -184,7 +184,7 @@ func (i *Iterator[T]) SeekPrefixWatch(prefixKey []byte) (watch <-chan struct{}) node = child depth++ } - return + return watch } func (i *Iterator[T]) SeekPrefix(prefixKey []byte) { diff --git a/node.go b/node.go index 7aeae8b..79945e6 100644 --- a/node.go +++ b/node.go @@ -15,7 +15,7 @@ type Node[T any] interface { matchPrefix([]byte) bool getChild(int) Node[T] setChild(int, Node[T]) - clone(bool) Node[T] + clone(bool, bool) Node[T] getKey() []byte getValue() T setValue(T) diff --git a/node_16.go b/node_16.go index 4e6aa28..bd8fceb 100644 --- a/node_16.go +++ b/node_16.go @@ -78,7 +78,7 @@ func (n *Node16[T]) getChild(index int) Node[T] { return n.children[index] } -func (n *Node16[T]) clone(keepWatch bool) Node[T] { +func (n *Node16[T]) clone(keepWatch, deep bool) Node[T] { newNode := &Node16[T]{ partialLen: n.getPartialLen(), numChildren: n.getNumChildren(), @@ -90,7 +90,15 @@ func (n *Node16[T]) clone(keepWatch bool) Node[T] { newNode.mutateCh = make(chan struct{}) } copy(newNode.keys[:], n.keys[:]) - copy(newNode.children[:], n.children[:]) + if deep { + for i := 0; i < 16; i++ { + if n.children[i] != nil { + newNode.children[i] = n.children[i].clone(keepWatch, deep) + } + } + } else { + copy(newNode.children[:], n.children[:]) + } nodeT := Node[T](newNode) return nodeT } diff --git a/node_256.go b/node_256.go index 5c3a9b1..97cc322 100644 --- a/node_256.go +++ b/node_256.go @@ -83,7 +83,7 @@ func (n *Node256[T]) getChild(index int) Node[T] { return n.children[index] } -func (n *Node256[T]) clone(keepWatch bool) Node[T] { +func (n *Node256[T]) clone(keepWatch bool, deep bool) Node[T] { newNode := &Node256[T]{ partialLen: n.getPartialLen(), numChildren: n.getNumChildren(), @@ -94,7 +94,15 @@ func (n *Node256[T]) clone(keepWatch bool) Node[T] { } else { newNode.mutateCh = make(chan struct{}) } - copy(newNode.children[:], n.children[:]) + if deep { + for i := 0; i < 256; i++ { + if n.children[i] != nil { + newNode.children[i] = n.children[i].clone(keepWatch, deep) + } + } + } else { + copy(newNode.children[:], n.children[:]) + } return newNode } diff --git a/node_4.go b/node_4.go index 53c3d86..8899f86 100644 --- a/node_4.go +++ b/node_4.go @@ -77,7 +77,7 @@ func (n *Node4[T]) getChild(index int) Node[T] { return n.children[index] } -func (n *Node4[T]) clone(keepWatch bool) Node[T] { +func (n *Node4[T]) clone(keepWatch, deep bool) Node[T] { newNode := &Node4[T]{ partialLen: n.getPartialLen(), numChildren: n.getNumChildren(), @@ -89,7 +89,15 @@ func (n *Node4[T]) clone(keepWatch bool) Node[T] { newNode.mutateCh = make(chan struct{}) } copy(newNode.keys[:], n.keys[:]) - copy(newNode.children[:], n.children[:]) + if deep { + for i := 0; i < 4; i++ { + if n.children[i] != nil { + newNode.children[i] = n.children[i].clone(keepWatch, deep) + } + } + } else { + copy(newNode.children[:], n.children[:]) + } return newNode } diff --git a/node_48.go b/node_48.go index bc6c959..563fa53 100644 --- a/node_48.go +++ b/node_48.go @@ -87,7 +87,7 @@ func (n *Node48[T]) getChild(index int) Node[T] { return n.children[index] } -func (n *Node48[T]) clone(keepWatch bool) Node[T] { +func (n *Node48[T]) clone(keepWatch, deep bool) Node[T] { newNode := &Node48[T]{ partialLen: n.getPartialLen(), numChildren: n.getNumChildren(), @@ -99,7 +99,15 @@ func (n *Node48[T]) clone(keepWatch bool) Node[T] { newNode.mutateCh = make(chan struct{}) } copy(newNode.keys[:], n.keys[:]) - copy(newNode.children[:], n.children[:]) + if deep { + for i := 0; i < 48; i++ { + if n.children[i] != nil { + newNode.children[i] = n.children[i].clone(keepWatch, deep) + } + } + } else { + copy(newNode.children[:], n.children[:]) + } return newNode } diff --git a/node_leaf.go b/node_leaf.go index 7658c0a..ce8da0f 100644 --- a/node_leaf.go +++ b/node_leaf.go @@ -114,7 +114,7 @@ func (n *NodeLeaf[T]) getChild(index int) Node[T] { return nil } -func (n *NodeLeaf[T]) clone(keepWatch bool) Node[T] { +func (n *NodeLeaf[T]) clone(keepWatch, deep bool) Node[T] { newNode := &NodeLeaf[T]{ keyLen: n.getKeyLen(), key: make([]byte, len(n.getKey())), diff --git a/tree.go b/tree.go index 0db76cb..8112f95 100644 --- a/tree.go +++ b/tree.go @@ -41,8 +41,8 @@ func (t *RadixTree[T]) Len() int { } // Clone is used to return the clone of tree -func (t *RadixTree[T]) Clone() *RadixTree[T] { - return &RadixTree[T]{root: t.root.clone(true), size: t.size} +func (t *RadixTree[T]) Clone(deep bool) *RadixTree[T] { + return &RadixTree[T]{root: t.root.clone(true, deep), size: t.size} } func (t *RadixTree[T]) GetPathIterator(path []byte) *PathIterator[T] { diff --git a/txn.go b/txn.go index e25c4ff..14136dc 100644 --- a/txn.go +++ b/txn.go @@ -41,7 +41,7 @@ type Txn[T any] struct { // Txn starts a new transaction that can be used to mutate the tree func (t *RadixTree[T]) Txn() *Txn[T] { - treeClone := t.Clone() + treeClone := t.Clone(false) txn := &Txn[T]{ size: t.size, snap: treeClone.root, @@ -56,8 +56,8 @@ func (t *Txn[T]) Clone() *Txn[T] { // reset the writable node cache to avoid leaking future writes into the clone txn := &Txn[T]{ - tree: t.tree.Clone(), - snap: t.snap.clone(false), + tree: t.tree.Clone(false), + snap: t.snap.clone(false, false), size: t.size, } return txn @@ -527,7 +527,7 @@ func (t *Txn[T]) writeNode(n Node[T]) Node[T] { // safe to replace this leaf with another after you get your node for // writing. You MUST replace it, because the channel associated with // this leaf will be closed when this transaction is committed. - nc := n.clone(false) + nc := n.clone(false, false) // Mark this node as writable. t.writable.Add(nc, nil)