Deduce nodes created from descendants diff

internal/trie/node/branch.go

@@ -68,6 +68,7 @@ func (b *Branch) StringNode() (stringNode *gotree.Node) {
 	stringNode.Appendf("Dirty: %t", b.Dirty)
 	stringNode.Appendf("Key: " + bytesToString(b.Key))
 	stringNode.Appendf("Value: " + bytesToString(b.Value))
+	stringNode.Appendf("Descendants: %d", b.Descendants)
 	stringNode.Appendf("Calculated encoding: " + bytesToString(b.Encoding))
 	stringNode.Appendf("Calculated digest: " + bytesToString(b.HashDigest))
 

internal/trie/node/branch_test.go

@@ -82,14 +82,16 @@ func Test_Branch_String(t *testing.T) {
 ├── Dirty: false
 ├── Key: nil
 ├── Value: nil
+├── Descendants: 0
 ├── Calculated encoding: nil
 └── Calculated digest: nil`,
 		},
 		"branch with value smaller than 1024": {
 			branch: &Branch{
-				Key:   []byte{1, 2},
-				Value: []byte{3, 4},
-				Dirty: true,
+				Key:         []byte{1, 2},
+				Value:       []byte{3, 4},
+				Dirty:       true,
+				Descendants: 3,
 				Children: [16]Node{
 					nil, nil, nil,
 					&Leaf{},
@@ -105,6 +107,7 @@ func Test_Branch_String(t *testing.T) {
 ├── Dirty: true
 ├── Key: 0x0102
 ├── Value: 0x0304
+├── Descendants: 3
 ├── Calculated encoding: nil
 ├── Calculated digest: nil
 ├── Child 3
@@ -121,6 +124,7 @@ func Test_Branch_String(t *testing.T) {
 |       ├── Dirty: false
 |       ├── Key: nil
 |       ├── Value: nil
+|       ├── Descendants: 0
 |       ├── Calculated encoding: nil
 |       └── Calculated digest: nil
 └── Child 11
@@ -134,9 +138,10 @@ func Test_Branch_String(t *testing.T) {
 		},
 		"branch with value higher than 1024": {
 			branch: &Branch{
-				Key:   []byte{1, 2},
-				Value: make([]byte, 1025),
-				Dirty: true,
+				Key:         []byte{1, 2},
+				Value:       make([]byte, 1025),
+				Dirty:       true,
+				Descendants: 3,
 				Children: [16]Node{
 					nil, nil, nil,
 					&Leaf{},
@@ -152,6 +157,7 @@ func Test_Branch_String(t *testing.T) {
 ├── Dirty: true
 ├── Key: 0x0102
 ├── Value: 0x0000000000000000...0000000000000000
+├── Descendants: 3
 ├── Calculated encoding: nil
 ├── Calculated digest: nil
 ├── Child 3
@@ -168,6 +174,7 @@ func Test_Branch_String(t *testing.T) {
 |       ├── Dirty: false
 |       ├── Key: nil
 |       ├── Value: nil
+|       ├── Descendants: 0
 |       ├── Calculated encoding: nil
 |       └── Calculated digest: nil
 └── Child 11

internal/trie/node/children.go

@@ -3,6 +3,63 @@
 
 package node
 
+func (b *Branch) AssignChild(index int, newChild Node) {
+	oldChild := b.Children[index]
+	// fmt.Println(oldChild)
+	b.Children[index] = newChild
+	// fmt.Println(newChild)
+
+	b.SetDirty(true)
+
+	descendantsDelta := calculateDescendantsDelta(oldChild, newChild)
+	b.HandleDelta(descendantsDelta)
+}
+
+func (b *Branch) HandleDelta(descendantsDelta int) {
+	if descendantsDelta < 0 {
+		descendantsDelta = -descendantsDelta
+		b.Descendants -= uint32(descendantsDelta)
+	} else {
+		b.Descendants += uint32(descendantsDelta)
+	}
+}
+
+func calculateDescendantsDelta(oldChild, newChild Node) int {
+	switch {
+	case oldChild == nil:
+		switch {
+		case newChild == nil:
+			return 0
+		case newChild.Type() == LeafType:
+			return 1
+		default: // branch
+			newBranchChild := newChild.(*Branch)
+			return 1 + int(newBranchChild.Descendants)
+		}
+	case oldChild.Type() == LeafType:
+		switch {
+		case newChild == nil:
+			return -1
+		case newChild.Type() == LeafType:
+			return 0
+		default: // branch
+			newBranchChild := newChild.(*Branch)
+			return int(newBranchChild.Descendants)
+		}
+	default: // branch old child
+		oldBranchChild := oldChild.(*Branch)
+		switch {
+		case newChild == nil:
+			return -1 - int(oldBranchChild.Descendants)
+		case newChild.Type() == LeafType:
+			return -int(oldBranchChild.Descendants)
+		default: // branch
+			newBranchChild := newChild.(*Branch)
+			return int(newBranchChild.Descendants) - int(oldBranchChild.Descendants)
+		}
+	}
+}
+
 // ChildrenBitmap returns the 16 bit bitmap
 // of the children in the branch.
 func (b *Branch) ChildrenBitmap() (bitmap uint16) {

lib/trie/print_test.go

@@ -39,8 +39,9 @@ func Test_Trie_String(t *testing.T) {
 		"branch root": {
 			trie: Trie{
 				root: &node.Branch{
-					Key:   nil,
-					Value: []byte{1, 2},
+					Key:         nil,
+					Value:       []byte{1, 2},
+					Descendants: 2,
 					Children: [16]node.Node{
 						&node.Leaf{
 							Key:        []byte{1, 2, 3},
@@ -61,6 +62,7 @@ func Test_Trie_String(t *testing.T) {
 ├── Dirty: false
 ├── Key: nil
 ├── Value: 0x0102
+├── Descendants: 2
 ├── Calculated encoding: nil
 ├── Calculated digest: nil
 ├── Child 0

lib/trie/trie.go

@@ -326,20 +326,19 @@ func (t *Trie) Put(keyLE, value []byte) {
 }
 
 func (t *Trie) put(key, value []byte) {
-	t.root, _ = t.insert(t.root, key, value)
+	t.root = t.insert(t.root, key, value)
 }
 
 // insert inserts a value in the trie at the key specified.
 // It may create one or more new nodes or update an existing node.
-func (t *Trie) insert(parent Node, key, value []byte) (newParent Node, nodesCreated uint32) {
+func (t *Trie) insert(parent Node, key, value []byte) (newParent Node) {
 	if parent == nil {
-		const nodesCreated = 1
 		return &node.Leaf{
 			Key:        key,
 			Value:      value,
 			Generation: t.generation,
 			Dirty:      true,
-		}, nodesCreated
+		}
 	}
 
 	// TODO ensure all values have dirty set to true
@@ -354,17 +353,15 @@ func (t *Trie) insert(parent Node, key, value []byte) (newParent Node, nodesCrea
 	}
 }
 
-func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, value []byte) (
-	newParent Node, nodesCreated uint32) {
+func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, value []byte) (newParent Node) {
 	if bytes.Equal(parentLeaf.Key, key) {
-		nodesCreated = 0
 		if bytes.Equal(value, parentLeaf.Value) {
-			return parentLeaf, nodesCreated
+			return parentLeaf
 		}
 
 		parentLeaf = t.prepLeafForMutation(parentLeaf)
 		parentLeaf.Value = value
-		return parentLeaf, nodesCreated
+		return parentLeaf
 	}
 
 	commonPrefixLength := lenCommonPrefix(key, parentLeaf.Key)
@@ -384,14 +381,12 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, value []byte) (
 		if len(key) < len(parentLeafKey) {
 			// Move the current leaf parent as a child to the new branch.
 			parentLeaf = t.prepLeafForMutation(parentLeaf)
-			childIndex := parentLeafKey[commonPrefixLength]
+			childIndex := int(parentLeafKey[commonPrefixLength])
 			parentLeaf.Key = parentLeaf.Key[commonPrefixLength+1:]
-			newBranchParent.Children[childIndex] = parentLeaf
-			newBranchParent.AddDescendants(1)
-			nodesCreated++
+			newBranchParent.AssignChild(childIndex, parentLeaf)
 		}
 
-		return newBranchParent, nodesCreated
+		return newBranchParent
 	}
 
 	if len(parentLeaf.Key) == commonPrefixLength {
@@ -400,38 +395,35 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, value []byte) (
 	} else {
 		// make the leaf a child of the new branch
 		parentLeaf = t.prepLeafForMutation(parentLeaf)
-		childIndex := parentLeafKey[commonPrefixLength]
+		childIndex := int(parentLeafKey[commonPrefixLength])
 		parentLeaf.Key = parentLeaf.Key[commonPrefixLength+1:]
-		newBranchParent.Children[childIndex] = parentLeaf
-		newBranchParent.AddDescendants(1)
-		nodesCreated++
+		newBranchParent.AssignChild(childIndex, parentLeaf)
 	}
-	childIndex := key[commonPrefixLength]
-	newBranchParent.Children[childIndex] = &node.Leaf{
+	childIndex := int(key[commonPrefixLength])
+	newChildLeaf := &node.Leaf{
 		Key:        key[commonPrefixLength+1:],
 		Value:      value,
 		Generation: t.generation,
 		Dirty:      true,
 	}
-	newBranchParent.AddDescendants(1)
-	nodesCreated++
+	newBranchParent.AssignChild(childIndex, newChildLeaf)
 
-	return newBranchParent, nodesCreated
+	return newBranchParent
 }
 
-func (t *Trie) insertInBranch(parentBranch *node.Branch, key, value []byte) (
-	newParent Node, nodesCreated uint32) {
+func (t *Trie) insertInBranch(parentBranch *node.Branch,
+	key, value []byte) (newParent *node.Branch) {
 	parentBranch = t.prepBranchForMutation(parentBranch)
 
 	if bytes.Equal(key, parentBranch.Key) {
 		parentBranch.Value = value
-		return parentBranch, 0
+		return parentBranch
 	}
 
 	if bytes.HasPrefix(key, parentBranch.Key) {
 		// key is included in parent branch key
 		commonPrefixLength := lenCommonPrefix(key, parentBranch.Key)
-		childIndex := key[commonPrefixLength]
+		childIndex := int(key[commonPrefixLength])
 		remainingKey := key[commonPrefixLength+1:]
 		child := parentBranch.Children[childIndex]
 
@@ -442,45 +434,63 @@ func (t *Trie) insertInBranch(parentBranch *node.Branch, key, value []byte) (
 				Generation: t.generation,
 				Dirty:      true,
 			}
-			nodesCreated = 1
 		} else {
-			child, nodesCreated = t.insert(child, remainingKey, value)
+			var childWasBranch bool
+			var oldChildDescendants uint32
+			if isNodeBranch(child) {
+				// we need this since the child branch may be modified in place by the call
+				// to t.insert(child, remainingKey, value) and its statistics updated,
+				// making the comparisons in AssignChild invalid.
+				childWasBranch = true
+				oldChildDescendants = child.(*node.Branch).Descendants
+			}
+
+			child = t.insert(child, remainingKey, value)
+			child.SetDirty(true)
+
+			if childWasBranch && isNodeBranch(child) {
+				// new child is a branch
+				// then check the descendants count against the descendants count
+				// of the previously existing child branch.
+				newChildDescendants := child.(*node.Branch).Descendants
+				delta := int(newChildDescendants) - int(oldChildDescendants)
+				parentBranch.HandleDelta(delta)
+			}
 		}
 
-		parentBranch.Children[childIndex] = child
-		parentBranch.AddDescendants(nodesCreated)
-		return parentBranch, nodesCreated
+		parentBranch.AssignChild(childIndex, child)
+		parentBranch.SetDirty(true)
+		parentBranch.Generation = t.generation
+
+		return parentBranch
 	}
 
 	// we need to branch out at the point where the keys diverge
 	// update partial keys, new branch has key up to matching length
-	nodesCreated = 1
 	commonPrefixLength := lenCommonPrefix(key, parentBranch.Key)
 	newParentBranch := &node.Branch{
 		Key:        key[:commonPrefixLength],
 		Generation: t.generation,
 		Dirty:      true,
 	}
 
-	oldParentIndex := parentBranch.Key[commonPrefixLength]
+	oldParentIndex := int(parentBranch.Key[commonPrefixLength])
 	remainingOldParentKey := parentBranch.Key[commonPrefixLength+1:]
 
 	parentBranch.Key = remainingOldParentKey
-	newParentBranch.Children[oldParentIndex] = parentBranch
-	newParentBranch.AddDescendants(1 + parentBranch.GetDescendants())
+	parentBranch.Generation = t.generation
+	newParentBranch.AssignChild(oldParentIndex, parentBranch)
 
 	if len(key) <= commonPrefixLength {
 		newParentBranch.Value = value
 	} else {
-		childIndex := key[commonPrefixLength]
+		childIndex := int(key[commonPrefixLength])
 		remainingKey := key[commonPrefixLength+1:]
-		var additionalNodesCreated uint32
-		newParentBranch.Children[childIndex], additionalNodesCreated = t.insert(nil, remainingKey, value)
-		nodesCreated += additionalNodesCreated
-		newParentBranch.AddDescendants(additionalNodesCreated)
+		newChild := t.insert(nil, remainingKey, value)
+		newParentBranch.AssignChild(childIndex, newChild)
 	}
 
-	return newParentBranch, nodesCreated
+	return newParentBranch
 }
 
 // LoadFromMap loads the given data mapping of key to value into the trie.
@@ -1115,3 +1125,7 @@ func concatenateSlices(sliceOne, sliceTwo []byte, otherSlices ...[]byte) (concat
 func intToByteSlice(n int) (slice []byte) {
 	return []byte{byte(n)}
 }
+
+func isNodeBranch(n Node) (isBranch bool) {
+	return n != nil && n.Type() != node.LeafType
+}