Merge branch 'main' into absl

absolutelightning · Jun 14, 2024 · 22574a5 · 22574a5
2 parents 81fbd5f + dac3d3e
commit 22574a5
Showing 1 changed file with 37 additions and 16 deletions.
diff --git a/iterator.go b/iterator.go
@@ -60,8 +60,8 @@ func (i *Iterator[T]) Next() ([]byte, T, bool) {
 			return nil, zero, false
 		}
 
-		switch node.getArtNodeType() {
-		case leafType:
+		switch node.(type) {
+		case *NodeLeaf[T]:
 			leafCh := node.(*NodeLeaf[T])
 			if i.lowerBound {
 				if bytes.Compare(getKey(leafCh.key), getKey(i.path)) >= 0 {
@@ -73,8 +73,9 @@ func (i *Iterator[T]) Next() ([]byte, T, bool) {
 				continue
 			}
 			return getKey(leafCh.key), leafCh.value, true
-		case node4:
+		case *Node4[T]:
 			n4 := node.(*Node4[T])
+			hasLeaf := n4.leaf != nil
 			for itr := int(n4.numChildren) - 1; itr >= 0; itr-- {
 				nodeCh := n4.children[itr]
 				if nodeCh == nil {
@@ -91,17 +92,18 @@ func (i *Iterator[T]) Next() ([]byte, T, bool) {
 					i.stack = append(i.stack, NodeWrapper[T]{nodeCh, nodeW.d + int(n4.partialLen) + 1})
 				}
 			}
-			if i.lowerBound && n4.leaf != nil {
+			if hasLeaf && i.lowerBound {
 				if bytes.Compare(n4.leaf.key, i.path) >= 0 {
 					return getKey(n4.leaf.key), n4.leaf.value, true
 				}
 				continue
 			}
-			if n4.leaf != nil && i.seeKPrefixWatch && hasPrefix(getKey(n4.leaf.key), i.path) {
+			if hasLeaf && i.seeKPrefixWatch {
 				return getKey(n4.leaf.key), n4.leaf.value, true
 			}
-		case node16:
+		case *Node16[T]:
 			n16 := node.(*Node16[T])
+			hasLeaf := n16.leaf != nil
 			for itr := int(n16.numChildren) - 1; itr >= 0; itr-- {
 				nodeCh := n16.children[itr]
 				if nodeCh == nil {
@@ -118,17 +120,18 @@ func (i *Iterator[T]) Next() ([]byte, T, bool) {
 					i.stack = append(i.stack, NodeWrapper[T]{nodeCh, nodeW.d + int(n16.partialLen) + 1})
 				}
 			}
-			if i.lowerBound && n16.leaf != nil {
+			if hasLeaf && i.lowerBound {
 				if bytes.Compare(n16.leaf.key, i.path) >= 0 {
 					return getKey(n16.leaf.key), n16.leaf.value, true
 				}
 				continue
 			}
-			if n16.leaf != nil && i.seeKPrefixWatch && hasPrefix(getKey(n16.leaf.key), i.path) {
+			if hasLeaf && i.seeKPrefixWatch {
 				return getKey(n16.leaf.key), n16.leaf.value, true
 			}
-		case node48:
+		case *Node48[T]:
 			n48 := node.(*Node48[T])
+			hasLeaf := n48.leaf != nil
 			for itr := 255; itr >= 0; itr-- {
 				idx := n48.keys[itr]
 				if idx == 0 {
@@ -149,31 +152,32 @@ func (i *Iterator[T]) Next() ([]byte, T, bool) {
 					i.stack = append(i.stack, NodeWrapper[T]{nodeCh, nodeW.d + int(n48.partialLen) + 1})
 				}
 			}
-			if i.lowerBound && n48.leaf != nil {
+			if hasLeaf && i.lowerBound {
 				if bytes.Compare(n48.leaf.key, i.path) >= 0 {
 					return getKey(n48.leaf.key), n48.leaf.value, true
 				}
 				continue
 			}
-			if n48.leaf != nil && i.seeKPrefixWatch && hasPrefix(getKey(n48.leaf.key), i.path) {
+			if hasLeaf && i.seeKPrefixWatch {
 				return getKey(n48.leaf.key), n48.leaf.value, true
 			}
-		case node256:
+		case *Node256[T]:
 			n256 := node.(*Node256[T])
+			hasLeaf := n256.leaf != nil
 			for itr := 255; itr >= 0; itr-- {
 				nodeCh := n256.children[itr]
 				if nodeCh == nil {
 					continue
 				}
 				i.stack = append(i.stack, NodeWrapper[T]{nodeCh, nodeW.d + int(n256.partialLen) + 1})
 			}
-			if i.lowerBound && n256.leaf != nil {
+			if hasLeaf && i.lowerBound {
 				if bytes.Compare(n256.leaf.key, i.path) >= 0 {
 					return getKey(n256.leaf.key), n256.leaf.value, true
 				}
 				continue
 			}
-			if n256.leaf != nil && i.seeKPrefixWatch && hasPrefix(getKey(n256.leaf.key), i.path) {
+			if hasLeaf && i.seeKPrefixWatch {
 				return getKey(n256.leaf.key), n256.leaf.value, true
 			}
 		}
@@ -215,10 +219,23 @@ func (i *Iterator[T]) SeekPrefixWatch(prefix []byte) (watch <-chan struct{}) {
 		if node.getPartialLen() > 0 {
 			// If the node has a prefix, compare it with the prefix
 			mismatchIdx := prefixMismatch[T](node, prefix, len(prefix), depth)
+			if mismatchIdx < int(node.getPartialLen()) || mismatchIdx >= len(prefix) {
+				i.stack = nil
+				i.node = node
+				i.depth = depth
+				if node.isLeaf() && !hasPrefix((getKey(node.getNodeLeaf().getKey())), prefix) {
+					i.node = nil
+				}
+				return node.getMutateCh()
+			}
 			if mismatchIdx < int(node.getPartialLen()) {
 				// If there's a mismatch, set the node to nil to break the loop
 				i.stack = nil
-				i.node = node
+				if node.getNodeLeaf() != nil && hasPrefix(getKey(node.getNodeLeaf().getKey()), prefix) {
+					i.node = node
+				} else {
+					i.node = nil
+				}
 				i.depth = depth
 				return node.getMutateCh()
 			}
@@ -238,7 +255,11 @@ func (i *Iterator[T]) SeekPrefixWatch(prefix []byte) (watch <-chan struct{}) {
 		if child == nil {
 			// If the child node doesn't exist, break the loop
 			i.stack = nil
-			i.node = node
+			if node.getNodeLeaf() != nil && hasPrefix(getKey(node.getNodeLeaf().getKey()), prefix) {
+				i.node = node
+			} else {
+				i.node = nil
+			}
 			return node.getMutateCh()
 		}