Ref counting (#95)

* Reference counting * reference counting with memdb tests passing * removed atomic add * fixes in watch
absolutelightning · Jun 21, 2024 · d03bc52 · d03bc52
1 parent 2a927cd
commit d03bc52
Show file tree

Hide file tree

Showing 10 changed files with 338 additions and 97 deletions.
diff --git a/node.go b/node.go
@@ -16,8 +16,11 @@ type Node[T any] interface {
 	isLeaf() bool
 	matchPrefix([]byte) bool
 	getChild(int) Node[T]
+	incrementLazyRefCount(delta int64)
+	getRefCount() int64
+	processRefCount()
 	setChild(int, Node[T])
-	clone(bool) Node[T]
+	clone(bool, bool) Node[T]
 	setMutateCh(chan struct{})
 	getKey() []byte
 	getValue() T

diff --git a/node_16.go b/node_16.go
@@ -10,14 +10,16 @@ import (
 )
 
 type Node16[T any] struct {
-	id          uint64
-	partialLen  uint32
-	numChildren uint8
-	partial     []byte
-	keys        [16]byte
-	children    [16]Node[T]
-	mutateCh    atomic.Pointer[chan struct{}]
-	leaf        *NodeLeaf[T]
+	id           uint64
+	partialLen   uint32
+	numChildren  uint8
+	partial      []byte
+	keys         [16]byte
+	children     [16]Node[T]
+	mutateCh     atomic.Pointer[chan struct{}]
+	leaf         *NodeLeaf[T]
+	refCount     int64
+	lazyRefCount int64
 }
 
 func (n *Node16[T]) getId() uint64 {
@@ -88,27 +90,45 @@ 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] {
+	n.processRefCount()
 	newNode := &Node16[T]{
 		partialLen:  n.getPartialLen(),
 		numChildren: n.getNumChildren(),
+		refCount:    n.getRefCount(),
 	}
 	if keepWatch {
 		newNode.setMutateCh(n.getMutateCh())
 	}
 	newPartial := make([]byte, maxPrefixLen)
-	newNode.setNodeLeaf(n.getNodeLeaf())
+	if deep {
+		if n.getNodeLeaf() != nil {
+			newNode.setNodeLeaf(n.getNodeLeaf().clone(true, true).(*NodeLeaf[T]))
+		}
+	} else {
+		newNode.setNodeLeaf(n.getNodeLeaf())
+	}
 	copy(newPartial, n.partial)
 	newNode.setPartial(newPartial)
 	newNode.setId(n.getId())
 	copy(newNode.keys[:], n.keys[:])
-	cpy := make([]Node[T], len(n.children))
-	copy(cpy, n.children[:])
-	for i := 0; i < 16; i++ {
-		newNode.setChild(i, cpy[i])
+	if deep {
+		cpy := make([]Node[T], len(n.children))
+		copy(cpy, n.children[:])
+		for i := 0; i < 16; i++ {
+			if cpy[i] == nil {
+				continue
+			}
+			newNode.setChild(i, cpy[i].clone(keepWatch, true))
+		}
+	} else {
+		cpy := make([]Node[T], len(n.children))
+		copy(cpy, n.children[:])
+		for i := 0; i < 16; i++ {
+			newNode.setChild(i, cpy[i])
+		}
 	}
-	nodeT := Node[T](newNode)
-	return nodeT
+	return newNode
 }
 
 func (n *Node16[T]) getKeyLen() uint32 {
@@ -210,3 +230,28 @@ func (n *Node16[T]) LowerBoundIterator() *LowerBoundIterator[T] {
 		node: n,
 	}
 }
+
+func (n *Node16[T]) incrementLazyRefCount(inc int64) {
+	n.lazyRefCount += inc
+}
+
+func (n *Node16[T]) processRefCount() {
+	if n.lazyRefCount == 0 {
+		return
+	}
+	n.refCount += n.lazyRefCount
+	if n.getNodeLeaf() != nil {
+		n.getNodeLeaf().incrementLazyRefCount(n.lazyRefCount)
+	}
+	for _, child := range n.children {
+		if child != nil {
+			child.incrementLazyRefCount(n.lazyRefCount)
+		}
+	}
+	n.lazyRefCount = 0
+}
+
+func (n *Node16[T]) getRefCount() int64 {
+	n.processRefCount()
+	return n.refCount
+}
diff --git a/node_256.go b/node_256.go
@@ -8,13 +8,15 @@ import (
 )
 
 type Node256[T any] struct {
-	id          uint64
-	partialLen  uint32
-	numChildren uint8
-	partial     []byte
-	children    [256]Node[T]
-	mutateCh    atomic.Pointer[chan struct{}]
-	leaf        *NodeLeaf[T]
+	id           uint64
+	partialLen   uint32
+	numChildren  uint8
+	partial      []byte
+	children     [256]Node[T]
+	mutateCh     atomic.Pointer[chan struct{}]
+	leaf         *NodeLeaf[T]
+	lazyRefCount int64
+	refCount     int64
 }
 
 func (n *Node256[T]) getId() uint64 {
@@ -96,23 +98,42 @@ 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, deep bool) Node[T] {
+	n.processRefCount()
 	newNode := &Node256[T]{
 		partialLen:  n.getPartialLen(),
 		numChildren: n.getNumChildren(),
+		refCount:    n.getRefCount(),
 	}
 	if keepWatch {
 		newNode.setMutateCh(n.getMutateCh())
 	}
-	newNode.setNodeLeaf(n.getNodeLeaf())
+	if deep {
+		if n.getNodeLeaf() != nil {
+			newNode.setNodeLeaf(n.getNodeLeaf().clone(true, true).(*NodeLeaf[T]))
+		}
+	} else {
+		newNode.setNodeLeaf(n.getNodeLeaf())
+	}
 	newPartial := make([]byte, maxPrefixLen)
 	newNode.setId(n.getId())
 	copy(newPartial, n.partial)
 	newNode.setPartial(newPartial)
-	cpy := make([]Node[T], len(n.children))
-	copy(cpy, n.children[:])
-	for i := 0; i < 256; i++ {
-		newNode.setChild(i, cpy[i])
+	if deep {
+		cpy := make([]Node[T], len(n.children))
+		copy(cpy, n.children[:])
+		for i := 0; i < 256; i++ {
+			if cpy[i] == nil {
+				continue
+			}
+			newNode.setChild(i, cpy[i].clone(keepWatch, true))
+		}
+	} else {
+		cpy := make([]Node[T], len(n.children))
+		copy(cpy, n.children[:])
+		for i := 0; i < 256; i++ {
+			newNode.setChild(i, cpy[i])
+		}
 	}
 	return newNode
 }
@@ -206,3 +227,28 @@ func (n *Node256[T]) LowerBoundIterator() *LowerBoundIterator[T] {
 		node: nodeT,
 	}
 }
+
+func (n *Node256[T]) incrementLazyRefCount(inc int64) {
+	n.lazyRefCount += inc
+}
+
+func (n *Node256[T]) processRefCount() {
+	if n.lazyRefCount == 0 {
+		return
+	}
+	n.refCount += n.lazyRefCount
+	if n.getNodeLeaf() != nil {
+		n.getNodeLeaf().incrementLazyRefCount(n.lazyRefCount)
+	}
+	for _, child := range n.children {
+		if child != nil {
+			child.incrementLazyRefCount(n.lazyRefCount)
+		}
+	}
+	n.lazyRefCount = 0
+}
+
+func (n *Node256[T]) getRefCount() int64 {
+	n.processRefCount()
+	return n.refCount
+}
diff --git a/node_4.go b/node_4.go
@@ -10,14 +10,16 @@ import (
 )
 
 type Node4[T any] struct {
-	id          uint64
-	partialLen  uint32
-	numChildren uint8
-	partial     []byte
-	keys        [4]byte
-	children    [4]Node[T]
-	mutateCh    atomic.Pointer[chan struct{}]
-	leaf        *NodeLeaf[T]
+	id           uint64
+	partialLen   uint32
+	numChildren  uint8
+	partial      []byte
+	keys         [4]byte
+	children     [4]Node[T]
+	mutateCh     atomic.Pointer[chan struct{}]
+	leaf         *NodeLeaf[T]
+	refCount     int64
+	lazyRefCount int64
 }
 
 func (n *Node4[T]) getId() uint64 {
@@ -87,24 +89,43 @@ 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] {
+	n.processRefCount()
 	newNode := &Node4[T]{
 		partialLen:  n.getPartialLen(),
 		numChildren: n.getNumChildren(),
+		refCount:    n.getRefCount(),
 	}
 	newNode.setId(n.getId())
 	if keepWatch {
 		newNode.setMutateCh(n.getMutateCh())
 	}
-	newNode.setNodeLeaf(n.getNodeLeaf())
+	if deep {
+		if n.getNodeLeaf() != nil {
+			newNode.setNodeLeaf(n.getNodeLeaf().clone(true, true).(*NodeLeaf[T]))
+		}
+	} else {
+		newNode.setNodeLeaf(n.getNodeLeaf())
+	}
 	newPartial := make([]byte, maxPrefixLen)
 	copy(newPartial, n.partial)
 	newNode.setPartial(newPartial)
 	copy(newNode.keys[:], n.keys[:])
-	cpy := make([]Node[T], len(n.children))
-	copy(cpy, n.children[:])
-	for i := 0; i < 4; i++ {
-		newNode.setChild(i, cpy[i])
+	if deep {
+		cpy := make([]Node[T], len(n.children))
+		copy(cpy, n.children[:])
+		for i := 0; i < 4; i++ {
+			if cpy[i] == nil {
+				continue
+			}
+			newNode.setChild(i, cpy[i].clone(keepWatch, true))
+		}
+	} else {
+		cpy := make([]Node[T], len(n.children))
+		copy(cpy, n.children[:])
+		for i := 0; i < 4; i++ {
+			newNode.setChild(i, cpy[i])
+		}
 	}
 	return newNode
 }
@@ -209,3 +230,28 @@ func (n *Node4[T]) LowerBoundIterator() *LowerBoundIterator[T] {
 		node: n,
 	}
 }
+
+func (n *Node4[T]) incrementLazyRefCount(inc int64) {
+	n.lazyRefCount += inc
+}
+
+func (n *Node4[T]) processRefCount() {
+	if n.lazyRefCount == 0 {
+		return
+	}
+	n.refCount += n.lazyRefCount
+	if n.getNodeLeaf() != nil {
+		n.getNodeLeaf().incrementLazyRefCount(n.lazyRefCount)
+	}
+	for _, child := range n.children {
+		if child != nil {
+			child.incrementLazyRefCount(n.lazyRefCount)
+		}
+	}
+	n.lazyRefCount = 0
+}
+
+func (n *Node4[T]) getRefCount() int64 {
+	n.processRefCount()
+	return n.refCount
+}