Subrootpaths (#49)

* initial unoptimized subrooth paths traversal function * cleanup and optimization * fix mistakenly omitted preprocessed roots * make comments more helpful, and special case check clearer * undo error for special case check * add more comprehensive tests, fix some issues, make GetSubrootPaths return a 3d list * add test case for 100% coverage * change to table tests, assign named errors * replace recursive subdivide with bitwise extraction * generalize prune function and fix edge case, add test * clean up comments, add clarifying comment on new case * more idiomatic errors Co-authored-by: Ismail Khoffi <[email protected]> * better test descriptions, lowercase errors, add entire-square test * add infinite loop bug in case of submitting a span for last two rows * remove trailing debug statement * significantly simplify implementation, use cleaner code * remove ineffectual assign to pass go lint * final cleanup * add fixes suggested by john Co-authored-by: Ismail Khoffi <[email protected]>
celestiaorg · Oct 21, 2021 · 56714f3 · 56714f3
1 parent 1d72cff
commit 56714f3
Show file tree

Hide file tree

Showing 2 changed files with 314 additions and 0 deletions.
diff --git a/subrootpaths.go b/subrootpaths.go
@@ -0,0 +1,172 @@
+package nmt
+
+import (
+	"errors"
+	"math"
+	"math/bits"
+)
+
+var (
+	srpNotPowerOf2       = errors.New("GetSubrootPaths: Supplied square size is not a power of 2")
+	srpInvalidShareCount = errors.New("GetSubrootPaths: Can't compute path for 0 share count slice")
+	srpPastSquareSize    = errors.New("GetSubrootPaths: Share slice can't be past the square size")
+)
+
+// merkle path to a node is equivalent to the index's binary representation
+// this is just a quick function to return that representation as a list of ints
+func subdivide(idxStart uint, width uint) []int {
+	var path []int
+	pathlen := int(math.Log2(float64(width)))
+	for i := pathlen - 1; i >= 0; i-- {
+		if (idxStart & (1 << i)) == 0 {
+			path = append(path, 0)
+		} else {
+			path = append(path, 1)
+		}
+	}
+	return path
+}
+
+// this function takes a path, and returns a copy of that path with path[index] set to branch,
+// and cuts off the list at path[:index+offset] - used to create inclusion branches during traversal
+func extractBranch(path []int, index int, offset int, branch int) []int {
+	rightCapture := make([]int, len(path))
+	copy(rightCapture, path)
+	rightCapture[index] = branch
+	return rightCapture[:index+offset]
+}
+
+func prune(idxStart uint, idxEnd uint, maxWidth uint) [][]int {
+
+	var prunedPaths [][]int
+	var preprocessedPaths [][]int
+
+	pathStart := subdivide(idxStart, maxWidth)
+	pathEnd := subdivide(idxEnd, maxWidth)
+
+	// special case of two-share path, just return one or two paths
+	if idxStart+1 >= idxEnd {
+		if idxStart%2 == 1 {
+			return [][]int{pathStart, pathEnd}
+		} else {
+			return [][]int{pathStart[:len(pathStart)-1]}
+		}
+	}
+
+	// if starting share is on an odd index, add that single path and shift it right 1
+	if idxStart%2 == 1 {
+		idxStart++
+		preprocessedPaths = append(preprocessedPaths, pathStart)
+		pathStart = subdivide(idxStart, maxWidth)
+	}
+
+	// if ending share is on an even index, add that single index and shift it left 1
+	if idxEnd%2 == 0 {
+		idxEnd--
+		preprocessedPaths = append(preprocessedPaths, pathEnd)
+	}
+
+	treeDepth := len(pathStart)
+	capturedSpan := uint(0)
+	rightTraversed := false
+
+	for i := treeDepth - 1; i >= 0 && capturedSpan < idxEnd; i-- {
+		nodeSpan := uint(math.Pow(float64(2), float64(treeDepth-i)))
+		if pathStart[i] == 0 {
+			// if nodespan is less than end index, continue traversing upwards
+			lastNode := nodeSpan + idxStart - 1
+			if lastNode <= idxEnd {
+				capturedSpan = lastNode
+				// if a right path has been encountered, we want to return the right
+				// branch one level down
+				if rightTraversed {
+					prunedPaths = append(prunedPaths, extractBranch(pathStart, i, 1, 1))
+				} else {
+					// else add *just* the current root node
+					prunedPaths = [][]int{pathStart[:i]}
+				}
+			} else {
+				// else if it's greater than the end index, break out of the left-capture loop
+				break
+			}
+		} else {
+			// on a right upwards traverse, we skip processing
+			// besides adjusting the idxStart for span calculation
+			// and modifying the previous path calculations to not include
+			// containing roots as they would span beyond the start index
+			idxStart = idxStart - nodeSpan/2
+			rightTraversed = true
+		}
+	}
+
+	combined := append(preprocessedPaths, prunedPaths...)
+	// if the process captured the span to the end, return the results
+	if capturedSpan == idxEnd {
+		return combined
+	}
+	// else recurse into the leftover span
+	return append(combined, prune(capturedSpan+1, idxEnd, maxWidth)...)
+}
+
+// GetSubrootPaths is a pure function that takes arguments: square size, share index start,
+// and share Count, and returns a minimal set of paths to the subtree roots that
+// encompasses that entire range of shares, with each top level entry in the list
+// starting from the nearest row root.
+//
+// An empty entry in the top level list means the shares span that entire row and so
+// the root for that segment of shares is equivalent to the row root.
+func GetSubrootPaths(squareSize uint, idxStart uint, shareCount uint) ([][][]int, error) {
+
+	var paths [][]int
+	var top [][][]int
+
+	shares := squareSize * squareSize
+
+	// check squareSize is at least 2 and that it's
+	// a power of 2 by checking that only 1 bit is on
+	if squareSize < 2 || bits.OnesCount(squareSize) != 1 {
+		return nil, srpNotPowerOf2
+	}
+
+	// no path exists for 0 count slice
+	if shareCount == 0 {
+		return nil, srpInvalidShareCount
+	}
+
+	// sanity checking
+	if idxStart >= shares || idxStart+shareCount > shares {
+		return nil, srpPastSquareSize
+	}
+
+	// adjust for 0 index
+	shareCount = shareCount - 1
+
+	startRow := int(math.Floor(float64(idxStart) / float64(squareSize)))
+	closingRow := int(math.Ceil(float64(idxStart+shareCount) / float64(squareSize)))
+
+	shareStart := idxStart % squareSize
+	shareEnd := (idxStart + shareCount) % squareSize
+
+	// if the count is one, just return the subdivided start path
+	if shareCount == 0 {
+		return append(top, append(paths, subdivide(shareStart, squareSize))), nil
+	}
+
+	// if the shares are all in one row, do the normal case
+	if startRow == closingRow-1 {
+		top = append(top, prune(shareStart, shareEnd, squareSize))
+	} else {
+		// if the shares span multiple rows, treat it as 2 different path generations,
+		// one from left-most root to end of a row, and one from start of a row to right-most root,
+		// and returning nil lists for the fully covered rows in between
+		left, _ := GetSubrootPaths(squareSize, shareStart, squareSize-shareStart)
+		right, _ := GetSubrootPaths(squareSize, 0, shareEnd+1)
+		top = append(top, left[0])
+		for i := 1; i < (closingRow-startRow)-1; i++ {
+			top = append(top, [][]int{{}})
+		}
+		top = append(top, right[0])
+	}
+
+	return top, nil
+}
diff --git a/subrootpaths_test.go b/subrootpaths_test.go
@@ -0,0 +1,142 @@
+package nmt
+
+import (
+	"reflect"
+	"testing"
+)
+
+type pathSpan struct {
+	squareSize uint
+	startNode  uint
+	length     uint
+}
+
+type pathResult [][][]int
+
+func TestArgValidation(t *testing.T) {
+
+	type test struct {
+		input pathSpan
+		want  error
+	}
+
+	tests := []test{
+		{input: pathSpan{squareSize: 0, startNode: 0, length: 0}, want: srpNotPowerOf2},
+		{input: pathSpan{squareSize: 1, startNode: 0, length: 1}, want: srpNotPowerOf2},
+		{input: pathSpan{squareSize: 20, startNode: 0, length: 1}, want: srpNotPowerOf2},
+		{input: pathSpan{squareSize: 4, startNode: 0, length: 17}, want: srpPastSquareSize},
+		{input: pathSpan{squareSize: 4, startNode: 0, length: 0}, want: srpInvalidShareCount},
+	}
+
+	for _, tc := range tests {
+		paths, err := GetSubrootPaths(tc.input.squareSize, tc.input.startNode, tc.input.length)
+		if err != tc.want {
+			t.Fatalf(`GetSubrootPaths(%v) = %v, %v, want %v`, tc.input, paths, err, tc.want)
+		}
+	}
+}
+
+func TestPathGeneration(t *testing.T) {
+
+	type test struct {
+		input pathSpan
+		want  pathResult
+		desc  string
+	}
+
+	tests := []test{
+		{
+			input: pathSpan{squareSize: 2, startNode: 0, length: 2},
+			want:  pathResult{{{}}},
+			desc:  "Single row span, should return empty to signify one row root",
+		},
+		{
+			input: pathSpan{squareSize: 2, startNode: 0, length: 1},
+			want:  pathResult{{{0}}},
+			desc:  "Single left-most node span, should return left-most branch",
+		},
+		{
+			input: pathSpan{squareSize: 2, startNode: 1, length: 1},
+			want:  pathResult{{{1}}},
+			desc:  "Single right-most node span on first row, should return single-row right-most branch",
+		},
+		{
+			input: pathSpan{squareSize: 4, startNode: 1, length: 2},
+			want:  pathResult{{{0, 1}, {1, 0}}},
+			desc:  "2-node span on unaligned start, should return two branch paths leading to two nodes in the middle of first row's tree",
+		},
+		{
+			input: pathSpan{squareSize: 8, startNode: 1, length: 6},
+			want:  pathResult{{{0, 0, 1}, {1, 1, 0}, {0, 1}, {1, 0}}},
+			desc:  "Single row span, taking whole row minus start and end nodes, unaligned start and end. Should return two offset paths, two internal paths, in one row",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 16, length: 16},
+			want:  pathResult{{{1}}},
+			desc:  "Single row span, taking the right half of the first row, should return right (1) branch of one row",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 0, length: 32},
+			want:  pathResult{{{}}},
+			desc:  "Whole row span of a larger square, should return empty to signify one row root",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 0, length: 64},
+			want:  pathResult{{{}}, {{}}},
+			desc:  "Whole row span of 2 rows, should return two empty lists to signify two row roots",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 0, length: 96},
+			want:  pathResult{{{}}, {{}}, {{}}},
+			desc:  "Whole row span of 3 rows, should return three empty lists to signify three row roots",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 18, length: 11},
+			want:  pathResult{{{1, 1, 1, 0, 0}, {1, 0, 0, 1}, {1, 0, 1}, {1, 1, 0}}},
+			desc:  "Span starting on right side of first row's tree, on an even-index start but not on a power-of-two alignment, ending on an even-index. Should return 4 paths: branch spanning 18-19, branch spanning 20-23, branch spanning 24-28, and single-node path to 29",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 14, length: 18},
+			want:  pathResult{{{0, 1, 1, 1}, {1}}},
+			desc:  "Span starting on left side of first row's tree, spanning until end of tree. Should return two paths in one row: right-most branch on left side of tree, and whole right side of tree",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 14, length: 17},
+			want:  pathResult{{{1, 1, 1, 1, 0}, {0, 1, 1, 1}, {1, 0}, {1, 1, 0}, {1, 1, 1, 0}}},
+			desc:  "Span starting on the last branch of the left side of the first row's tree, starting on an even index, ending at the second-to-last branch of the first row's tree, on an even index. Should return 5 paths: branch spanning 14-15, branch spanning 16-23, branch spanning 24-27, branch spanning 28-29, single-node path to 30",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 48, length: 16},
+			want:  pathResult{{{1}}},
+			desc:  "Span for right side of second row in square. Should return a single branch in a single list, pointing to the first right path of the row within that starting index",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 0, length: 1024},
+			want:  pathResult{{{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}, {{}}},
+			desc:  "Span for the entire square. Should return 32 empty lists to signify span covers every row in the square",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 988, length: 32},
+			want:  pathResult{{{1, 1, 1}}, {{0}, {1, 0}, {1, 1, 0}}},
+			desc:  "Span for last two rows in square, should return last branch of second to last row, left half of last row, and two branches on right half of last row",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 992, length: 32},
+			want:  pathResult{{{}}},
+			desc:  "Span for last row in the square, should return empty list.",
+		},
+		{
+			input: pathSpan{squareSize: 32, startNode: 1023, length: 1},
+			want:  pathResult{{{1, 1, 1, 1, 1}}},
+			desc:  "Span for last node in the last row in the square, should return a path of 1s",
+		},
+	}
+
+	for _, tc := range tests {
+		paths, err := GetSubrootPaths(tc.input.squareSize, tc.input.startNode, tc.input.length)
+		if !reflect.DeepEqual(pathResult(paths), tc.want) {
+			t.Fatalf(`GetSubrootPaths(%v) = %v, %v, want %v - rationale: %v`, tc.input, paths, err, tc.want, tc.desc)
+		}
+	}
+
+}