proof: Avoid all reallocations (#23)

This change introduces pre-allocation into the Proof generation types, thereby
reducing the total number of slice allocations to one. A new helper for computing
the compact range size compact.RangeSize is introduced, and the signature of
the previous compact.RangeNodes now requires providing a slice to append the
result to.

Author: pav-kv

compact/nodes.go

@@ -56,10 +56,11 @@ func (id NodeID) Coverage() (uint64, uint64) {
 	return id.Index << id.Level, (id.Index + 1) << id.Level
 }
 
-// RangeNodes returns node IDs that comprise the [begin, end) compact range.
-func RangeNodes(begin, end uint64) []NodeID {
+// RangeNodes appends the IDs of the nodes that comprise the [begin, end)
+// compact range to the given slice, and returns the new slice. The caller may
+// pre-allocate space with the help of the RangeSize function.
+func RangeNodes(begin, end uint64, ids []NodeID) []NodeID {
 	left, right := Decompose(begin, end)
-	ids := make([]NodeID, 0, bits.OnesCount64(left)+bits.OnesCount64(right))
 
 	pos := begin
 	// Iterate over perfect subtrees along the left border of the range, ordered
@@ -80,3 +81,9 @@ func RangeNodes(begin, end uint64) []NodeID {
 
 	return ids
 }
+
+// RangeSize returns the number of nodes in the [begin, end) compact range.
+func RangeSize(begin, end uint64) int {
+	left, right := Decompose(begin, end)
+	return bits.OnesCount64(left) + bits.OnesCount64(right)
+}

compact/nodes_test.go

@@ -21,7 +21,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 )
 
-func TestRangeNodes(t *testing.T) {
+func TestRangeNodesAndSize(t *testing.T) {
 	n := func(level uint, index uint64) NodeID {
 		return NewNodeID(level, index)
 	}
@@ -31,9 +31,9 @@ func TestRangeNodes(t *testing.T) {
 		want  []NodeID
 	}{
 		// Empty ranges.
-		{end: 0, want: []NodeID{}},
-		{begin: 10, end: 10, want: []NodeID{}},
-		{begin: 1024, end: 1024, want: []NodeID{}},
+		{end: 0, want: nil},
+		{begin: 10, end: 10, want: nil},
+		{begin: 1024, end: 1024, want: nil},
 		// One entry.
 		{begin: 10, end: 11, want: []NodeID{n(0, 10)}},
 		{begin: 1024, end: 1025, want: []NodeID{n(0, 1024)}},
@@ -67,19 +67,35 @@ func TestRangeNodes(t *testing.T) {
 		{begin: 1, end: 17, want: []NodeID{n(0, 1), n(1, 1), n(2, 1), n(3, 1), n(0, 16)}},
 	} {
 		t.Run(fmt.Sprintf("range:%d:%d", tc.begin, tc.end), func(t *testing.T) {
-			got := RangeNodes(tc.begin, tc.end)
+			got := RangeNodes(tc.begin, tc.end, nil)
 			if diff := cmp.Diff(got, tc.want); diff != "" {
 				t.Fatalf("RangeNodes: diff(-want +got):\n%s", diff)
 			}
+			if got, want := RangeSize(tc.begin, tc.end), len(tc.want); got != want {
+				t.Errorf("RangeSize: got %d, want %d", got, want)
+			}
 		})
 	}
 }
 
+func TestRangeNodesAppend(t *testing.T) {
+	prefix := []NodeID{NewNodeID(0, 0), NewNodeID(10, 0), NewNodeID(11, 5)}
+	nodes := RangeNodes(123, 456, prefix)
+
+	if got, min := len(nodes), len(prefix); got < min {
+		t.Fatalf("RangeNodes returned %d IDs, want >= %d", got, min)
+	}
+	got := nodes[:len(prefix)]
+	if diff := cmp.Diff(got, prefix); diff != "" {
+		t.Fatalf("RangeNodes: diff(-prefix +got):\n%s", diff)
+	}
+}
+
 func TestGenRangeNodes(t *testing.T) {
 	const size = uint64(512)
 	for begin := uint64(0); begin <= size; begin++ {
 		for end := begin; end <= size; end++ {
-			got := RangeNodes(begin, end)
+			got := RangeNodes(begin, end, nil)
 			want := refRangeNodes(NewNodeID(63, 0), begin, end)
 			if diff := cmp.Diff(got, want); diff != "" {
 				t.Fatalf("RangeNodes(%d, %d): diff(-want +got):\n%s", begin, end, diff)
@@ -89,11 +105,11 @@ func TestGenRangeNodes(t *testing.T) {
 }
 
 // refRangeNodes returns node IDs that comprise the [begin, end) compact range.
-// This is a reference implementation for cross-cthehecking.
+// This is a reference implementation for cross-checking.
 func refRangeNodes(root NodeID, begin, end uint64) []NodeID {
 	b, e := root.Coverage()
 	if end <= b || begin >= e {
-		return []NodeID{}
+		return nil
 	}
 	if b >= begin && e <= end {
 		return []NodeID{root}

compact/range.go

@@ -41,10 +41,8 @@ func (f *RangeFactory) NewRange(begin, end uint64, hashes [][]byte) (*Range, err
 	if end < begin {
 		return nil, fmt.Errorf("invalid range: end=%d, want >= %d", end, begin)
 	}
-	left, right := Decompose(begin, end)
-	ones := bits.OnesCount64(left) + bits.OnesCount64(right)
-	if ln := len(hashes); ln != ones {
-		return nil, fmt.Errorf("invalid hashes: got %d values, want %d", ln, ones)
+	if got, want := len(hashes), RangeSize(begin, end); got != want {
+		return nil, fmt.Errorf("invalid hashes: got %d values, want %d", got, want)
 	}
 	return &Range{f: f, begin: begin, end: end, hashes: hashes}, nil
 }

compact/range_test.go

@@ -368,7 +368,7 @@ func TestNewRangeWithStorage(t *testing.T) {
 		}); err != nil {
 			t.Fatalf("%d: Append: %v", i, err)
 		}
-		hashes := getHashes(RangeNodes(0, i+1))
+		hashes := getHashes(RangeNodes(0, i+1, nil))
 		var err error
 		if cr, err = factory.NewRange(0, i+1, hashes); err != nil {
 			t.Fatalf("%d: NewRange: %v", i+1, err)

proof/proof.go

@@ -49,7 +49,7 @@ func Inclusion(index, size uint64) (Nodes, error) {
 	if index >= size {
 		return Nodes{}, fmt.Errorf("index %d out of bounds for tree size %d", index, size)
 	}
-	return nodes(index, 0, size), nil
+	return nodes(index, 0, size).skipFirst(), nil
 }
 
 // Consistency returns the information on how to fetch and construct a
@@ -72,57 +72,59 @@ func Consistency(size1, size2 uint64) (Nodes, error) {
 	// into this node in the tree of size2.
 	p := nodes(index, level, size2)
 
-	// Handle the case when size1 is not a power of 2.
-	if index != 0 {
-		// Prepend the earlier computed node to the proof.
-		// TODO(pavelkalinnikov): This code path is invoked almost always. Avoid
-		// the extra allocation that append does.
-		p.IDs = append(p.IDs, compact.NodeID{})
-		copy(p.IDs[1:], p.IDs)
-		p.IDs[0] = compact.NewNodeID(level, index)
-
-		// Fixup the indices into the IDs slice.
-		if p.begin < p.end {
-			p.begin++
-			p.end++
-		}
+	// Handle the case when size1 is a power of 2.
+	if index == 0 {
+		return p.skipFirst(), nil
 	}
-
 	return p, nil
 }
 
 // nodes returns the node IDs necessary to prove that the (level, index) node
 // is included in the Merkle tree of the given size.
 func nodes(index uint64, level uint, size uint64) Nodes {
-	node := compact.NewNodeID(level, index)
-	begin, _ := node.Coverage()
+	// Compute the `fork` node, where the path from root to (level, index) node
+	// diverges from the path to (0, size).
+	//
+	// The sibling of this node is the ephemeral node which represents a subtree
+	// that is not complete in the tree of the given size. To compute the hash
+	// of the ephemeral node, we need all the non-ephemeral nodes that cover the
+	// same range of leaves.
+	//
+	// The `inner` variable is how many layers up from (level, index) the `fork`
+	// and the ephemeral nodes are.
+	inner := bits.Len64(index^(size>>level)) - 1
+	fork := compact.NewNodeID(level+uint(inner), index>>inner)
 
-	// Compute the level at which the path to leaf `begin` diverges from the path
-	// to `size`. This is where the ephemeral node is located. The ephemeral node
-	// represents a subtree that is not complete in the tree of the given size,
-	// so we instead provide the minimal list of non-ephemeral nodes which cover
-	// the same range of leaves.
-	ephemLevel := uint(bits.Len64(begin^size) - 1)
+	begin, end := fork.Coverage()
+	left := compact.RangeSize(0, begin)
+	right := compact.RangeSize(end, size)
+
+	node := compact.NewNodeID(level, index)
+	// Pre-allocate the exact number of nodes for the proof, in order:
+	// - The seed node for which we are building the proof.
+	// - The `inner` nodes at each level up to the fork node.
+	// - The `right` nodes, comprising the ephemeral node.
+	// - The `left` nodes, completing the coverage of the whole [0, size) range.
+	nodes := append(make([]compact.NodeID, 0, 1+inner+right+left), node)
 
 	// The first portion of the proof consists of the siblings for nodes of the
 	// path going up to the level at which the ephemeral node appears.
-	// TODO(pavelkalinnikov): Pre-allocate the full capacity.
-	nodes := make([]compact.NodeID, 0, ephemLevel-level)
-	for ; node.Level < ephemLevel; node = node.Parent() {
+	for ; node.Level < fork.Level; node = node.Parent() {
 		nodes = append(nodes, node.Sibling())
 	}
-	// This portion of the proof covers the range under the reached node. The
+	// This portion of the proof covers the range [begin, end) under it. The
 	// ranges to the left and to the right from it remain to be covered.
-	begin, end := node.Coverage()
 
 	// Add all the nodes (potentially none) that cover the right range, and
 	// represent the ephemeral node. Reverse them so that the Rehash method can
 	// process hashes in the convenient order, from lower to upper levels.
 	len1 := len(nodes)
-	nodes = append(nodes, reverse(compact.RangeNodes(end, size))...)
+	nodes = compact.RangeNodes(end, size, nodes)
+	reverse(nodes[len(nodes)-right:])
 	len2 := len(nodes)
 	// Add the nodes that cover the left range, ordered increasingly by level.
-	nodes = append(nodes, reverse(compact.RangeNodes(0, begin))...)
+	nodes = compact.RangeNodes(0, begin, nodes)
+	reverse(nodes[len(nodes)-left:])
 
 	// nodes[len1:len2] contains the nodes representing the ephemeral node. If
 	// it's empty or only has one node, make it zero.
@@ -133,7 +135,7 @@ func nodes(index uint64, level uint, size uint64) Nodes {
 		len1, len2 = 0, 0
 	}
 
-	return Nodes{IDs: nodes, begin: len1, end: len2, ephem: node.Sibling()}
+	return Nodes{IDs: nodes, begin: len1, end: len2, ephem: fork.Sibling()}
 }
 
 // Ephem returns the ephemeral node, and indices begin and end, such that
@@ -173,9 +175,18 @@ func (n Nodes) Rehash(h [][]byte, hc func(left, right []byte) []byte) ([][]byte,
 	return h[:cursor], nil
 }
 
-func reverse(ids []compact.NodeID) []compact.NodeID {
+func (n Nodes) skipFirst() Nodes {
+	n.IDs = n.IDs[1:]
+	// Fixup the indices into the IDs slice.
+	if n.begin < n.end {
+		n.begin--
+		n.end--
+	}
+	return n
+}
+
+func reverse(ids []compact.NodeID) {
 	for i, j := 0, len(ids)-1; i < j; i, j = i+1, j-1 {
 		ids[i], ids[j] = ids[j], ids[i]
 	}
-	return ids
 }