Add cloning support for Snapshot in Memdb (#12)

* reverse iterator init * some fixes * fix reverse iterator seaklowerbound * fix tests * fix track channels * longest prefix on txn * some fixes * revisit * todo revisit * major code refactor * some minor fixes * add lru * fix bugs * added walk func * add more tests * some prog * some progress track mutate * some progress * fix tests * some optimizations * some memory optimizations * fix tests * minor fixes * some minor fixes * some fixes * fix delete prefix * some fixes * init stack * fix seek lower bound * clone * fix watch * fix clone bug * some fixes for go mem db * merge absl * Fix race * fix go mod * fix race in add child * fix last memdb test * clone * fix memdb last tests
absolutelightning · May 21, 2024 · fcb7feb · fcb7feb
1 parent 643e521
commit fcb7feb
Show file tree

Hide file tree

Showing 9 changed files with 51 additions and 19 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)