chore: update to rsmt2d v0.7.0 (#926)

Updates dependency of rsmt2d to v0.7.0
Changes relevant constructors/methods to update to new definition of
rsmt2d.TreeConstructorFn

- [x] Closes #712
- [x] Closes #754

Co-authored-by: evan-forbes <[email protected]>
Co-authored-by: Evan Forbes <[email protected]>
Co-authored-by: CHAMI Rachid <[email protected]>

go.mod

@@ -27,7 +27,7 @@ require (
 require (
 	cosmossdk.io/errors v1.0.0-beta.7
 	cosmossdk.io/math v1.0.0-beta.2
-	github.com/celestiaorg/rsmt2d v0.6.0
+	github.com/celestiaorg/rsmt2d v0.7.0
 	github.com/cosmos/cosmos-sdk v0.0.0-00010101000000-000000000000
 	github.com/regen-network/cosmos-proto v0.3.1
 	github.com/tendermint/tendermint v0.34.20
@@ -52,7 +52,6 @@ require (
 	github.com/bgentry/speakeasy v0.1.0 // indirect
 	github.com/btcsuite/btcd v0.22.1 // indirect
 	github.com/btcsuite/btcd/btcec/v2 v2.1.2 // indirect
-	github.com/celestiaorg/go-leopard v0.1.0 // indirect
 	github.com/celestiaorg/merkletree v0.0.0-20210714075610-a84dc3ddbbe4 // indirect
 	github.com/cenkalti/backoff/v4 v4.1.3 // indirect
 	github.com/cespare/xxhash v1.1.0 // indirect
@@ -115,6 +114,8 @@ require (
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/jmhodges/levigo v1.0.0 // indirect
 	github.com/klauspost/compress v1.15.6 // indirect
+	github.com/klauspost/cpuid/v2 v2.1.1 // indirect
+	github.com/klauspost/reedsolomon v1.11.1 // indirect
 	github.com/lib/pq v1.10.6 // indirect
 	github.com/libp2p/go-buffer-pool v0.1.0 // indirect
 	github.com/magiconair/properties v1.8.6 // indirect
@@ -124,6 +125,7 @@ require (
 	github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369 // indirect
 	github.com/mimoo/StrobeGo v0.0.0-20210601165009-122bf33a46e0 // indirect
 	github.com/minio/highwayhash v1.0.2 // indirect
+	github.com/minio/sha256-simd v1.0.0 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
 	github.com/mitchellh/go-testing-interface v1.0.0 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
@@ -163,6 +165,7 @@ require (
 	go.etcd.io/bbolt v1.3.6 // indirect
 	go.opencensus.io v0.23.0 // indirect
 	golang.org/x/oauth2 v0.0.0-20220411215720-9780585627b5 // indirect
+	golang.org/x/sync v0.0.0-20220929204114-8fcdb60fdcc0 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	golang.org/x/xerrors v0.0.0-20220517211312-f3a8303e98df // indirect
 	google.golang.org/api v0.81.0 // indirect

go.sum

@@ -182,16 +182,14 @@ github.com/celestiaorg/celestia-core v1.5.0-tm-v0.34.20 h1:BqlcOQqL2UqdDTcdCtrOL
 github.com/celestiaorg/celestia-core v1.5.0-tm-v0.34.20/go.mod h1:f4R8qNJrP1CDH0SNwj4jA3NymBLQM4lNdx6Ijmfllbw=
 github.com/celestiaorg/cosmos-sdk v1.4.0-sdk-v0.46.0 h1:65gnQ92mfz+9XNVTHeVwMp+SZuBqmToEnz8+WdDRmQ8=
 github.com/celestiaorg/cosmos-sdk v1.4.0-sdk-v0.46.0/go.mod h1:ByQ2rOrZs7s2OnPfeaiTMC8IOlcrT195xIRPgevydCI=
-github.com/celestiaorg/go-leopard v0.1.0 h1:28z2EkvKJIez5J9CEaiiUEC+OxalRLtTGJJ1oScfE1g=
-github.com/celestiaorg/go-leopard v0.1.0/go.mod h1:NtO/rjlB8dw2aq7jr06vZFKGvryQcTDXaNHelmPNOAM=
 github.com/celestiaorg/merkletree v0.0.0-20210714075610-a84dc3ddbbe4 h1:CJdIpo8n5MFP2MwK0gSRcOVlDlFdQJO1p+FqdxYzmvc=
 github.com/celestiaorg/merkletree v0.0.0-20210714075610-a84dc3ddbbe4/go.mod h1:fzuHnhzj1pUygGz+1ZkB3uQbEUL4htqCGJ4Qs2LwMZA=
 github.com/celestiaorg/nmt v0.11.0 h1:iqTaNwnVzM3njBmPklpHzb3A4Xy/JKahoclRPbAzxNc=
 github.com/celestiaorg/nmt v0.11.0/go.mod h1:NN3W8EEoospv8EHCw50DDNWwPLpJkFHoEFiqCEcNCH4=
 github.com/celestiaorg/quantum-gravity-bridge v1.2.0 h1:l/LEEUP+x8MhhXB8rrWkyUVFZgQj1Ur/TAwUpnyLK38=
 github.com/celestiaorg/quantum-gravity-bridge v1.2.0/go.mod h1:6WOajINTDEUXpSj5UZzod16UZ96ZVB/rFNKyM+Mt1gI=
-github.com/celestiaorg/rsmt2d v0.6.0 h1:32Eq5t7lPNbhftPFFjxwCUeEjWg/yGgeMbshxnGw03c=
-github.com/celestiaorg/rsmt2d v0.6.0/go.mod h1:EZ+O2KdCq8xI7WFwjATLdhtMdrdClmAs2w7zENDr010=
+github.com/celestiaorg/rsmt2d v0.7.0 h1:r8fybOWhE2/VJ2XEJ6BncnYTSlLYx2c7dQDUD+5hBqg=
+github.com/celestiaorg/rsmt2d v0.7.0/go.mod h1:hhlsTi6G3+X5jOP/8Lb/d7i5y2XNFmnyMddYbFSmrgo=
 github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
 github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM=
 github.com/cenkalti/backoff/v4 v4.1.1/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw=
@@ -669,8 +667,13 @@ github.com/klauspost/compress v1.12.3/go.mod h1:8dP1Hq4DHOhN9w426knH3Rhby4rFm6D8
 github.com/klauspost/compress v1.15.6 h1:6D9PcO8QWu0JyaQ2zUMmu16T1T+zjjEpP91guRsvDfY=
 github.com/klauspost/compress v1.15.6/go.mod h1:PhcZ0MbTNciWF3rruxRgKxI5NkcHHrHUDtV4Yw2GlzU=
 github.com/klauspost/cpuid v0.0.0-20170728055534-ae7887de9fa5/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
+github.com/klauspost/cpuid/v2 v2.0.4/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
+github.com/klauspost/cpuid/v2 v2.1.1 h1:t0wUqjowdm8ezddV5k0tLWVklVuvLJpoHeb4WBdydm0=
+github.com/klauspost/cpuid/v2 v2.1.1/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
 github.com/klauspost/crc32 v0.0.0-20161016154125-cb6bfca970f6/go.mod h1:+ZoRqAPRLkC4NPOvfYeR5KNOrY6TD+/sAC3HXPZgDYg=
 github.com/klauspost/pgzip v1.0.2-0.20170402124221-0bf5dcad4ada/go.mod h1:Ch1tH69qFZu15pkjo5kYi6mth2Zzwzt50oCQKQE9RUs=
+github.com/klauspost/reedsolomon v1.11.1 h1:0gCWQXOB8pVe1Y5SGozDA5t2qoVxX3prsV+qHgI/Fik=
+github.com/klauspost/reedsolomon v1.11.1/go.mod h1:FXLZzlJIdfqEnQLdUKWNRuMZg747hZ4oYp2Ml60Lb/k=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/kr/fs v0.1.0/go.mod h1:FFnZGqtBN9Gxj7eW1uZ42v5BccTP0vu6NEaFoC2HwRg=
@@ -734,6 +737,8 @@ github.com/mimoo/StrobeGo v0.0.0-20210601165009-122bf33a46e0 h1:QRUSJEgZn2Snx0Em
 github.com/mimoo/StrobeGo v0.0.0-20210601165009-122bf33a46e0/go.mod h1:43+3pMjjKimDBf5Kr4ZFNGbLql1zKkbImw+fZbw3geM=
 github.com/minio/highwayhash v1.0.2 h1:Aak5U0nElisjDCfPSG79Tgzkn2gl66NxOMspRrKnA/g=
 github.com/minio/highwayhash v1.0.2/go.mod h1:BQskDq+xkJ12lmlUUi7U0M5Swg3EWR+dLTk+kldvVxY=
+github.com/minio/sha256-simd v1.0.0 h1:v1ta+49hkWZyvaKwrQB8elexRqm6Y0aMLjCNsrYxo6g=
+github.com/minio/sha256-simd v1.0.0/go.mod h1:OuYzVNI5vcoYIAmbIvHPl3N3jUzVedXbKy5RFepssQM=
 github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=
 github.com/mitchellh/go-homedir v1.0.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
@@ -1233,6 +1238,8 @@ golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220513210516-0976fa681c29/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20220929204114-8fcdb60fdcc0 h1:cu5kTvlzcw1Q5S9f5ip1/cpiB4nXvw1XYzFPGgzLUOY=
+golang.org/x/sync v0.0.0-20220929204114-8fcdb60fdcc0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180823144017-11551d06cbcc/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
@@ -1334,6 +1341,7 @@ golang.org/x/sys v0.0.0-20220412211240-33da011f77ad/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20220502124256-b6088ccd6cba/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220517195934-5e4e11fc645e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f h1:v4INt8xihDGvnrfjMDVXGxw9wrfxYyCjk0KbXjhR55s=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=

pkg/da/data_availability_header.go

@@ -67,8 +67,8 @@ func ExtendShares(squareSize uint64, shares [][]byte) (*rsmt2d.ExtendedDataSquar
 			squareSize*squareSize,
 		)
 	}
-	tree := wrapper.NewErasuredNamespacedMerkleTree(squareSize)
-	return rsmt2d.ComputeExtendedDataSquare(shares, appconsts.DefaultCodec(), tree.Constructor)
+	// here we construct a tree
+	return rsmt2d.ComputeExtendedDataSquare(shares, appconsts.DefaultCodec(), wrapper.NewConstructor(squareSize))
 }
 
 // String returns hex representation of merkle hash of the DAHeader.

pkg/inclusion/nmt_caching.go

@@ -2,6 +2,7 @@ package inclusion
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/celestiaorg/celestia-app/pkg/da"
 	"github.com/celestiaorg/celestia-app/pkg/wrapper"
@@ -78,40 +79,36 @@ func (strc subTreeRootCacher) walk(root []byte, path []WalkInstruction) ([]byte,
 // threadsafe, but with a future refactor, we could simply read from rsmt2d and
 // not use the tree constructor which would fix both of these issues.
 type EDSSubTreeRootCacher struct {
-	caches     []*subTreeRootCacher
+	mut        *sync.RWMutex
+	caches     map[uint]*subTreeRootCacher
 	squareSize uint64
-	// counter is used to ignore columns NOTE: this is a leaky abstraction that
-	// we make because rsmt2d is used to generate the roots for us, so we have
-	// to assume that it will generate a row root every other tree contructed.
-	// This is also one of the reasons this implementation is not thread safe.
-	// Please see note above on a better refactor.
-	counter int
 }
 
 func NewSubtreeCacher(squareSize uint64) *EDSSubTreeRootCacher {
 	return &EDSSubTreeRootCacher{
-		caches:     []*subTreeRootCacher{},
+		mut:        &sync.RWMutex{},
+		caches:     make(map[uint]*subTreeRootCacher),
 		squareSize: squareSize,
 	}
 }
 
 // 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() rsmt2d.Tree {
+func (stc *EDSSubTreeRootCacher) Constructor(axis rsmt2d.Axis, index 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
 	var newTree wrapper.ErasuredNamespacedMerkleTree
-	switch stc.counter % 2 {
-	case 0:
+	switch axis {
+	case rsmt2d.Row:
 		strc := newSubTreeRootCacher()
-		stc.caches = append(stc.caches, strc)
-		newTree = wrapper.NewErasuredNamespacedMerkleTree(stc.squareSize, nmt.NodeVisitor(strc.Visit))
+		stc.mut.Lock()
+		stc.caches[index] = strc
+		stc.mut.Unlock()
+		newTree = wrapper.NewErasuredNamespacedMerkleTree(stc.squareSize, index, nmt.NodeVisitor(strc.Visit))
 	default:
-		newTree = wrapper.NewErasuredNamespacedMerkleTree(stc.squareSize)
+		newTree = wrapper.NewErasuredNamespacedMerkleTree(stc.squareSize, index)
 	}
-
-	stc.counter++
 	return &newTree
 }
 
@@ -124,5 +121,8 @@ func (stc *EDSSubTreeRootCacher) getSubTreeRoot(dah da.DataAvailabilityHeader, r
 	if row >= len(stc.caches) {
 		return nil, fmt.Errorf("row exceeds range of cache: max %d got %d", len(stc.caches), row)
 	}
-	return stc.caches[row].walk(dah.RowsRoots[row], path)
+	stc.mut.RLock()
+	sbt, err := stc.caches[uint(row)].walk(dah.RowsRoots[row], path)
+	stc.mut.RUnlock()
+	return sbt, err
 }

pkg/inclusion/nmt_caching_test.go

@@ -19,33 +19,24 @@ func TestWalkCachedSubTreeRoot(t *testing.T) {
 	// create the first main tree
 	strc := newSubTreeRootCacher()
 	oss := uint64(8)
-	tr := wrapper.NewErasuredNamespacedMerkleTree(oss, nmt.NodeVisitor(strc.Visit))
+	tr := wrapper.NewErasuredNamespacedMerkleTree(oss, 0, nmt.NodeVisitor(strc.Visit))
 	d := []byte{0, 0, 0, 0, 0, 0, 0, 1, 1, 2, 3, 4, 5, 6, 7, 8}
 	for i := 0; i < 8; i++ {
-		tr.Push(d, rsmt2d.SquareIndex{
-			Axis: uint(rsmt2d.Row),
-			Cell: uint(i),
-		})
+		tr.Push(d)
 	}
 	highestRoot := tr.Root()
 
 	// create a short sub tree
-	shortSubTree := wrapper.NewErasuredNamespacedMerkleTree(oss)
+	shortSubTree := wrapper.NewErasuredNamespacedMerkleTree(oss, 0)
 	for i := 0; i < 2; i++ {
-		shortSubTree.Push(d, rsmt2d.SquareIndex{
-			Axis: uint(rsmt2d.Row),
-			Cell: uint(i),
-		})
+		shortSubTree.Push(d)
 	}
 	shortSTR := shortSubTree.Root()
 
 	// create a tall sub tree root
-	tallSubTree := wrapper.NewErasuredNamespacedMerkleTree(oss)
+	tallSubTree := wrapper.NewErasuredNamespacedMerkleTree(oss, 0)
 	for i := 0; i < 4; i++ {
-		tallSubTree.Push(d, rsmt2d.SquareIndex{
-			Axis: uint(rsmt2d.Row),
-			Cell: uint(i),
-		})
+		tallSubTree.Push(d)
 	}
 	tallSTR := tallSubTree.Root()
 
@@ -154,13 +145,10 @@ func calculateSubTreeRoots(row [][]byte, depth int) [][]byte {
 	subTreeRoots := make([][]byte, count)
 	chunks := chunkSlice(row, subLeafRange)
 	for i, rowChunk := range chunks {
-		tr := wrapper.NewErasuredNamespacedMerkleTree(uint64(len(row)))
-		for j, r := range rowChunk {
-			c := (i * subLeafRange) + j
-			tr.Push(r, rsmt2d.SquareIndex{
-				Axis: uint(rsmt2d.Row),
-				Cell: uint(c),
-			})
+		tr := wrapper.NewErasuredNamespacedMerkleTree(uint64(len(row)), 0)
+		for _, r := range rowChunk {
+			// c := (i * subLeafRange) + j
+			tr.Push(r)
 		}
 		subTreeRoots[i] = tr.Root()
 	}

pkg/prove/proof.go

@@ -40,14 +40,10 @@ func TxInclusion(codec rsmt2d.Codec, data types.Data, txIndex uint64) (types.TxP
 	var rowRoots []tmbytes.HexBytes //nolint:prealloc // rarely will this contain more than a single root
 	for i, row := range rowShares {
 		// create an nmt to use to generate a proof
-		tree := wrapper.NewErasuredNamespacedMerkleTree(data.OriginalSquareSize)
-		for j, share := range row {
+		tree := wrapper.NewErasuredNamespacedMerkleTree(data.OriginalSquareSize, uint(i))
+		for _, share := range row {
 			tree.Push(
 				share,
-				rsmt2d.SquareIndex{
-					Axis: uint(i),
-					Cell: uint(j),
-				},
 			)
 		}
 

pkg/wrapper/nmt_wrapper.go

@@ -9,8 +9,10 @@ import (
 )
 
 // Fulfills the rsmt2d.Tree interface and rsmt2d.TreeConstructorFn function
-var _ rsmt2d.TreeConstructorFn = ErasuredNamespacedMerkleTree{}.Constructor
-var _ rsmt2d.Tree = &ErasuredNamespacedMerkleTree{}
+var (
+	_ rsmt2d.TreeConstructorFn = NewConstructor(0)
+	_ rsmt2d.Tree              = &ErasuredNamespacedMerkleTree{}
+)
 
 // ErasuredNamespacedMerkleTree wraps NamespaceMerkleTree to conform to the
 // rsmt2d.Tree interface while also providing the correct namespaces to the
@@ -23,21 +25,38 @@ 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
 }
 
 // NewErasuredNamespacedMerkleTree issues a new ErasuredNamespacedMerkleTree. squareSize must be greater than zero
-func NewErasuredNamespacedMerkleTree(squareSize uint64, setters ...nmt.Option) ErasuredNamespacedMerkleTree {
+func NewErasuredNamespacedMerkleTree(squareSize uint64, axisIndex uint, setters ...nmt.Option) ErasuredNamespacedMerkleTree {
 	if squareSize == 0 {
 		panic("cannot create a ErasuredNamespacedMerkleTree of squareSize == 0")
 	}
 	tree := nmt.New(appconsts.NewBaseHashFunc(), setters...)
-	return ErasuredNamespacedMerkleTree{squareSize: squareSize, options: setters, tree: tree}
+	return ErasuredNamespacedMerkleTree{squareSize: squareSize, options: setters, tree: tree, idx: &rsmt2d.SquareIndex{Axis: axisIndex, Cell: 0}}
+}
+
+type constructor struct {
+	squareSize uint64
+	opts       []nmt.Option
+}
+
+// NewConstructor creates a tree constructor function as required by rsmt2d to
+// calculate the data root. It creates that tree using the
+// wrapper.ErasuredNamespacedMerkleTree.
+func NewConstructor(squareSize uint64, opts ...nmt.Option) rsmt2d.TreeConstructorFn {
+	return constructor{
+		squareSize: squareSize,
+		opts:       opts,
+	}.NewTree
 }
 
-// Constructor acts as the rsmt2d.TreeConstructorFn for
-// ErasuredNamespacedMerkleTree
-func (w ErasuredNamespacedMerkleTree) Constructor() rsmt2d.Tree {
-	newTree := NewErasuredNamespacedMerkleTree(w.squareSize, w.options...)
+// 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...)
 	return &newTree
 }
 
@@ -46,14 +65,14 @@ func (w ErasuredNamespacedMerkleTree) Constructor() rsmt2d.Tree {
 // namespace unless the data pushed to the second half of the tree. Fulfills the
 // 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, idx rsmt2d.SquareIndex) {
-	if idx.Axis+1 > 2*uint(w.squareSize) || idx.Cell+1 > 2*uint(w.squareSize) {
-		panic(fmt.Sprintf("pushed past predetermined square size: boundary at %d index at %+v", 2*w.squareSize, idx))
+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))
 	}
 	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 idx.Axis+1 > uint(w.squareSize) || idx.Cell+1 > uint(w.squareSize) {
+	if w.idx.Axis+1 > uint(w.squareSize) || w.idx.Cell+1 > uint(w.squareSize) {
 		copy(nidAndData[:appconsts.NamespaceSize], appconsts.ParitySharesNamespaceID)
 	} else {
 		copy(nidAndData[:appconsts.NamespaceSize], data[:appconsts.NamespaceSize])
@@ -64,6 +83,7 @@ func (w *ErasuredNamespacedMerkleTree) Push(data []byte, idx rsmt2d.SquareIndex)
 	if err != nil {
 		panic(err)
 	}
+	w.idx.Cell++
 }
 
 // Root fulfills the rsmt.Tree interface by generating and returning the