feat: add path generation for nmt subroots (#621)

* feat: add path generation for nmt subroots

* delete unused struct

* linter

* use specific no lint comments

Co-authored-by: Rootul Patel <[email protected]>

* docs typo

Co-authored-by: Rootul Patel <[email protected]>

* docs typo

Co-authored-by: Rootul Patel <[email protected]>

* docs typo

Co-authored-by: Rootul Patel <[email protected]>

* add ascii tree

* cherry pick changes from non-interactive defaults branch

* add extra condition for climbing right

Co-authored-by: Rootul Patel <[email protected]>

pkg/inclusion/paths.go

@@ -0,0 +1,125 @@
+package inclusion
+
+// genSubTreeRootPath calculates the path to a given subtree root of a node, given the
+// depth and position of the node. note: the root of the tree is depth 0.
+// The following nolint can be removed after this function is used.
+//nolint:unused,deadcode
+func genSubTreeRootPath(depth int, pos uint) []bool {
+	path := make([]bool, depth)
+	counter := 0
+	for i := depth - 1; i >= 0; i-- {
+		if (pos & (1 << i)) == 0 {
+			path[counter] = false
+		} else {
+			path[counter] = true
+		}
+		counter++
+	}
+	return path
+}
+
+// coord identifies a tree node using the depth and position
+// Depth       Position
+// 0              0
+//               / \
+//              /   \
+// 1           0     1
+//            /\     /\
+// 2         0  1   2  3
+//          /\  /\ /\  /\
+// 3       0 1 2 3 4 5 6 7
+type coord struct {
+	// depth is the typical depth of a tree, 0 being the root
+	depth uint64
+	// position is the index of a node of a given depth, 0 being the left most
+	// node
+	position uint64
+}
+
+// climb is a state transition function to simulate climbing a balanced binary
+// tree, using the current node as input and the next highest node as output.
+func (c coord) climb() coord {
+	return coord{
+		depth:    c.depth - 1,
+		position: c.position / 2,
+	}
+}
+
+// canClimbRight uses the current position to calculate the direction of the next
+// climb. Returns true if the next climb is right (if the position (index) is
+// even). please see depth and position example map in docs for coord.
+func (c coord) canClimbRight() bool {
+	return c.position%2 == 0 && c.depth > 0
+}
+
+// calculateSubTreeRootCoordinates generates the sub tree root coordinates of a
+// set of shares for a balanced binary tree of a given depth. It assumes that
+// end does not exceed the range of a tree of the provided depth, and that end
+// >= start. This function works by starting at the first index of the msg and
+// working our way right.
+func calculateSubTreeRootCoordinates(maxDepth, start, end uint64) []coord {
+	cds := []coord{}
+	// leafCursor keeps track of the current leaf that we are starting with when
+	// finding the subtree root for some set. When leafCursor == end, we are
+	// finished calculating sub tree roots
+	leafCursor := start
+	// nodeCursor keeps track of the current tree node when finding sub
+	// tree roots
+	nodeCursor := coord{
+		depth:    maxDepth,
+		position: start,
+	}
+	// lastNodeCursor keeps track of the last node cursor so that when we climb
+	// too high, we can use this node as a sub tree root
+	lastNodeCursor := nodeCursor
+	lastLeafCursor := leafCursor
+	// nodeRangeCursor keeps track of the number of leaves that are under the
+	// current tree node. We could calculate this each time, but this acts as a
+	// cache
+	nodeRangeCursor := uint64(1)
+	// reset is used to reset the above state after finding a subtree root. We
+	// reset by setting the node cursors to the values equal to the next leaf
+	// node.
+	reset := func() {
+		lastNodeCursor = nodeCursor
+		lastLeafCursor = leafCursor
+		nodeCursor = coord{
+			depth:    maxDepth,
+			position: leafCursor,
+		}
+		nodeRangeCursor = uint64(1)
+	}
+	// recursively climb the tree starting at the left most leaf node (the
+	// starting leaf), and save each subtree root as we find it. After finding a
+	// subtree root, if there's still leaves left in the message, then restart
+	// the process from that leaf.
+	for {
+		switch {
+		// check if we're finished, if so add the last coord and return
+		case leafCursor == end:
+			cds = append(cds, nodeCursor)
+			return cds
+		// check if we've climbed too high in the tree. if so, add the last
+		// highest node and proceed.
+		case leafCursor > end:
+			cds = append(cds, lastNodeCursor)
+			leafCursor = lastLeafCursor + 1
+			reset()
+		// check if can climb right again (only even positions will climb
+		// right). If not, we want to record this coord as it is a subtree
+		// root, then adjust the cursor and proceed.
+		case !nodeCursor.canClimbRight():
+			cds = append(cds, nodeCursor)
+			leafCursor++
+			reset()
+		// proceed to climb higher by incrementing the relevant state and
+		// progressing through the loop.
+		default:
+			lastLeafCursor = leafCursor
+			lastNodeCursor = nodeCursor
+			leafCursor = leafCursor + nodeRangeCursor
+			nodeRangeCursor = nodeRangeCursor * 2
+			nodeCursor = nodeCursor.climb()
+		}
+	}
+}

pkg/inclusion/paths_test.go

@@ -0,0 +1,249 @@
+package inclusion
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_calculateSubTreeRootCoordinates(t *testing.T) {
+	type test struct {
+		name                 string
+		start, end, maxDepth uint64
+		expected             []coord
+	}
+	tests := []test{
+		{
+			name:     "first four shares of an 8 leaf tree",
+			start:    0,
+			end:      3,
+			maxDepth: 3,
+			expected: []coord{
+				{
+					depth:    1,
+					position: 0,
+				},
+			},
+		},
+		{
+			name:     "second set of four shares of an 8 leaf tree",
+			start:    4,
+			end:      7,
+			maxDepth: 3,
+			expected: []coord{
+				{
+					depth:    1,
+					position: 1,
+				},
+			},
+		},
+		{
+			name:     "middle 2 shares of an 8 leaf tree",
+			start:    3,
+			end:      4,
+			maxDepth: 3,
+			expected: []coord{
+				{
+					depth:    3,
+					position: 3,
+				},
+				{
+					depth:    3,
+					position: 4,
+				},
+			},
+		},
+		{
+			name:     "third lone share of an 8 leaf tree",
+			start:    3,
+			end:      3,
+			maxDepth: 3,
+			expected: []coord{
+				{
+					depth:    3,
+					position: 3,
+				},
+			},
+		},
+		{
+			name:     "middle 3 shares of an 8 leaf tree",
+			start:    3,
+			end:      5,
+			maxDepth: 3,
+			expected: []coord{
+				{
+					depth:    3,
+					position: 3,
+				},
+				{
+					depth:    2,
+					position: 2,
+				},
+			},
+		},
+		{
+			name:     "middle 6 shares of an 8 leaf tree",
+			start:    1,
+			end:      6,
+			maxDepth: 3,
+			expected: []coord{
+				{
+					depth:    3,
+					position: 1,
+				},
+				{
+					depth:    2,
+					position: 1,
+				},
+				{
+					depth:    2,
+					position: 2,
+				},
+				{
+					depth:    3,
+					position: 6,
+				},
+			},
+		},
+		{
+			name:     "first 5 shares of an 8 leaf tree",
+			start:    0,
+			end:      4,
+			maxDepth: 3,
+			expected: []coord{
+				{
+					depth:    1,
+					position: 0,
+				},
+				{
+					depth:    3,
+					position: 4,
+				},
+			},
+		},
+		{
+			name:     "first 7 shares of an 8 leaf tree",
+			start:    0,
+			end:      6,
+			maxDepth: 3,
+			expected: []coord{
+				{
+					depth:    1,
+					position: 0,
+				},
+				{
+					depth:    2,
+					position: 2,
+				},
+				{
+					depth:    3,
+					position: 6,
+				},
+			},
+		},
+		{
+			name:     "all shares of an 8 leaf tree",
+			start:    0,
+			end:      7,
+			maxDepth: 3,
+			expected: []coord{
+				{
+					depth:    0,
+					position: 0,
+				},
+			},
+		},
+		{
+			name:     "first 32 shares of a 128 leaf tree",
+			start:    0,
+			end:      31,
+			maxDepth: 7,
+			expected: []coord{
+				{
+					depth:    2,
+					position: 0,
+				},
+			},
+		},
+		{
+			name:     "first 33 shares of a 128 leaf tree",
+			start:    0,
+			end:      32,
+			maxDepth: 7,
+			expected: []coord{
+				{
+					depth:    2,
+					position: 0,
+				},
+				{
+					depth:    7,
+					position: 32,
+				},
+			},
+		},
+		{
+			name:     "first 31 shares of a 128 leaf tree",
+			start:    0,
+			end:      30,
+			maxDepth: 7,
+			expected: []coord{
+				{
+					depth:    3,
+					position: 0,
+				},
+				{
+					depth:    4,
+					position: 2,
+				},
+				{
+					depth:    5,
+					position: 6,
+				},
+				{
+					depth:    6,
+					position: 14,
+				},
+				{
+					depth:    7,
+					position: 30,
+				},
+			},
+		},
+		{
+			name:     "first 64 shares of a 128 leaf tree",
+			start:    0,
+			end:      63,
+			maxDepth: 7,
+			expected: []coord{
+				{
+					depth:    1,
+					position: 0,
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		res := calculateSubTreeRootCoordinates(tt.maxDepth, tt.start, tt.end)
+		assert.Equal(t, tt.expected, res, tt.name)
+	}
+}
+
+func Test_genSubTreeRootPath(t *testing.T) {
+	type test struct {
+		depth    int
+		pos      uint
+		expected []bool
+	}
+	tests := []test{
+		{2, 0, []bool{false, false}},
+		{0, 0, []bool{}},
+		{3, 0, []bool{false, false, false}},
+		{3, 1, []bool{false, false, true}},
+		{3, 2, []bool{false, true, false}},
+		{5, 16, []bool{true, false, false, false, false}},
+	}
+	for _, tt := range tests {
+		path := genSubTreeRootPath(tt.depth, tt.pos)
+		assert.Equal(t, tt.expected, path)
+	}
+}