refactor: cliarify a few names and add some docs to the wrapper

pkg/inclusion/nmt_caching.go

@@ -94,7 +94,7 @@ func NewSubtreeCacher(squareSize uint64) *EDSSubTreeRootCacher {
 
 // Constructor fullfills the rsmt2d.TreeCreatorFn by keeping a pointer to the
 // cache and embedding it as a nmt.NodeVisitor into a new wrapped nmt.
-func (stc *EDSSubTreeRootCacher) Constructor(axis rsmt2d.Axis, index uint) rsmt2d.Tree {
+func (stc *EDSSubTreeRootCacher) Constructor(axis rsmt2d.Axis, axisIndex uint) rsmt2d.Tree {
 	// see docs of counter field for more
 	// info. if the counter is even or == 0, then we make the assumption that we
 	// are creating a tree for a row
@@ -103,11 +103,11 @@ func (stc *EDSSubTreeRootCacher) Constructor(axis rsmt2d.Axis, index uint) rsmt2
 	case rsmt2d.Row:
 		strc := newSubTreeRootCacher()
 		stc.mut.Lock()
-		stc.caches[index] = strc
+		stc.caches[axisIndex] = strc
 		stc.mut.Unlock()
-		newTree = wrapper.NewErasuredNamespacedMerkleTree(stc.squareSize, index, nmt.NodeVisitor(strc.Visit))
+		newTree = wrapper.NewErasuredNamespacedMerkleTree(stc.squareSize, axisIndex, nmt.NodeVisitor(strc.Visit))
 	default:
-		newTree = wrapper.NewErasuredNamespacedMerkleTree(stc.squareSize, index)
+		newTree = wrapper.NewErasuredNamespacedMerkleTree(stc.squareSize, axisIndex)
 	}
 	return &newTree
 }

pkg/wrapper/nmt_wrapper.go

@@ -25,19 +25,28 @@ type ErasuredNamespacedMerkleTree struct {
 	squareSize uint64 // note: this refers to the width of the original square before erasure-coded
 	options    []nmt.Option
 	tree       *nmt.NamespacedMerkleTree
-	idx        *rsmt2d.SquareIndex
+	// axisIndex is the index of the axis (row or column) that this tree is on. This is passed
+	// by rsmt2d and used to help determine which quadrant each leaf belongs to.
+	axisIndex uint64
+	// shareIndex is the index of the share in a row or column that is being
+	// pushed to the tree. It is expected to be in the range: 0 <= shareIndex <
+	// 2*squareSize. shareIndex is used to help determine which quadrant each
+	// leaf belongs to, along with keeping track of how many leaves have been
+	// addeded to the tree so far.
+	shareIndex uint64
 }
 
 // NewErasuredNamespacedMerkleTree creates a new ErasuredNamespacedMerkleTree
-// with an underlying NMT of namespace size `appconsts.NamespaceSize`.
-// squareSize must be greater than zero
+// with an underlying NMT of namespace size `appconsts.NamespaceSize`. axisIndex
+// is the index of the row or column that this tree is committing to. squareSize
+// must be greater than zero.
 func NewErasuredNamespacedMerkleTree(squareSize uint64, axisIndex uint, setters ...nmt.Option) ErasuredNamespacedMerkleTree {
 	if squareSize == 0 {
 		panic("cannot create a ErasuredNamespacedMerkleTree of squareSize == 0")
 	}
 	setters = append(setters, nmt.NamespaceIDSize(appconsts.NamespaceSize))
 	tree := nmt.New(appconsts.NewBaseHashFunc(), setters...)
-	return ErasuredNamespacedMerkleTree{squareSize: squareSize, options: setters, tree: tree, idx: &rsmt2d.SquareIndex{Axis: axisIndex, Cell: 0}}
+	return ErasuredNamespacedMerkleTree{squareSize: squareSize, options: setters, tree: tree, axisIndex: uint64(axisIndex), shareIndex: 0}
 }
 
 type constructor struct {
@@ -58,8 +67,8 @@ func NewConstructor(squareSize uint64, opts ...nmt.Option) rsmt2d.TreeConstructo
 // NewTree creates a new rsmt2d.Tree using the
 // wrapper.ErasuredNamespacedMerkleTree with predefined square size and
 // nmt.Options
-func (c constructor) NewTree(_ rsmt2d.Axis, index uint) rsmt2d.Tree {
-	newTree := NewErasuredNamespacedMerkleTree(c.squareSize, index, c.opts...)
+func (c constructor) NewTree(_ rsmt2d.Axis, axisIndex uint) rsmt2d.Tree {
+	newTree := NewErasuredNamespacedMerkleTree(c.squareSize, axisIndex, c.opts...)
 	return &newTree
 }
 
@@ -69,24 +78,24 @@ func (c constructor) NewTree(_ rsmt2d.Axis, index uint) rsmt2d.Tree {
 // rsmt.Tree interface. NOTE: panics if an error is encountered while pushing or
 // if the tree size is exceeded.
 func (w *ErasuredNamespacedMerkleTree) Push(data []byte) {
-	if w.idx.Axis+1 > 2*uint(w.squareSize) || w.idx.Cell+1 > 2*uint(w.squareSize) {
-		panic(fmt.Sprintf("pushed past predetermined square size: boundary at %d index at %+v", 2*w.squareSize, w.idx))
+	if w.axisIndex+1 > 2*w.squareSize || w.shareIndex+1 > 2*w.squareSize {
+		panic(fmt.Sprintf("pushed past predetermined square size: boundary at %d index at %d %d", 2*w.squareSize, w.axisIndex, w.shareIndex))
 	}
 	nidAndData := make([]byte, appconsts.NamespaceSize+len(data))
 	copy(nidAndData[appconsts.NamespaceSize:], data)
 	// use the parity namespace if the cell is not in Q0 of the extended data square
-	if w.idx.Axis+1 > uint(w.squareSize) || w.idx.Cell+1 > uint(w.squareSize) {
-		copy(nidAndData[:appconsts.NamespaceSize], appconsts.ParitySharesNamespaceID)
-	} else {
+	if w.isQuadrantZero() {
 		copy(nidAndData[:appconsts.NamespaceSize], data[:appconsts.NamespaceSize])
+	} else {
+		copy(nidAndData[:appconsts.NamespaceSize], appconsts.ParitySharesNamespaceID)
 	}
 	// push to the underlying tree
 	err := w.tree.Push(nidAndData)
 	// panic on error
 	if err != nil {
 		panic(err)
 	}
-	w.idx.Cell++
+	w.incrementShareIndex()
 }
 
 // Root fulfills the rsmt.Tree interface by generating and returning the
@@ -103,3 +112,14 @@ func (w *ErasuredNamespacedMerkleTree) Prove(ind int) (nmt.Proof, error) {
 func (w *ErasuredNamespacedMerkleTree) Tree() *nmt.NamespacedMerkleTree {
 	return w.tree
 }
+
+// incrementShareIndex increments the share index by one.
+func (w *ErasuredNamespacedMerkleTree) incrementShareIndex() {
+	w.shareIndex++
+}
+
+// isQuadrantZero returns true if the current share index and axis index are both
+// in the original data square.
+func (w *ErasuredNamespacedMerkleTree) isQuadrantZero() bool {
+	return w.shareIndex < w.squareSize && w.axisIndex < w.squareSize
+}