feat: write data length to first compact share (#781)

Closes #725

pkg/prove/proof.go

@@ -2,6 +2,7 @@ package prove
 
 import (
 	"errors"
+	"fmt"
 
 	"github.com/celestiaorg/celestia-app/pkg/appconsts"
 	"github.com/celestiaorg/celestia-app/pkg/shares"
@@ -97,20 +98,39 @@ func txSharePosition(txs types.Txs, txIndex uint64) (startSharePos, endSharePos
 		return startSharePos, endSharePos, errors.New("transaction index is greater than the number of txs")
 	}
 
-	totalLen := 0
+	prevTxTotalLen := 0
 	for i := uint64(0); i < txIndex; i++ {
 		txLen := len(txs[i])
-		totalLen += (shares.DelimLen(uint64(txLen)) + txLen)
+		prevTxTotalLen += (shares.DelimLen(uint64(txLen)) + txLen)
 	}
 
-	txLen := len(txs[txIndex])
-
-	startSharePos = uint64((totalLen) / appconsts.ContinuationCompactShareContentSize)
-	endSharePos = uint64((totalLen + txLen + shares.DelimLen(uint64(txLen))) / appconsts.ContinuationCompactShareContentSize)
+	currentTxLen := len(txs[txIndex])
+	currentTxTotalLen := shares.DelimLen(uint64(currentTxLen)) + currentTxLen
+	endOfCurrentTxLen := prevTxTotalLen + currentTxTotalLen
 
+	startSharePos = txShareIndex(prevTxTotalLen)
+	endSharePos = txShareIndex(endOfCurrentTxLen)
+	fmt.Printf("prevTxTotalLen: %d, endOfCurrentTxLen: %d, startSharePos: %d, endSharePos %d, currentTxTotalLen %d\n", prevTxTotalLen, endOfCurrentTxLen, startSharePos, endSharePos, currentTxTotalLen)
 	return startSharePos, endSharePos, nil
 }
 
+// txShareIndex returns the index of the compact share that would contain
+// transactions with totalTxLen
+func txShareIndex(totalTxLen int) (index uint64) {
+	if totalTxLen <= appconsts.FirstCompactShareContentSize {
+		return 0
+	}
+
+	index++
+	totalTxLen -= appconsts.FirstCompactShareContentSize
+
+	for totalTxLen > appconsts.ContinuationCompactShareContentSize {
+		index++
+		totalTxLen -= appconsts.ContinuationCompactShareContentSize
+	}
+	return index
+}
+
 // genRowShares progessively generates data square rows from block data
 func genRowShares(codec rsmt2d.Codec, data types.Data, startRow, endRow uint64) ([][][]byte, error) {
 	if endRow > data.OriginalSquareSize {

pkg/shares/compact_shares_test.go

@@ -62,7 +62,7 @@ func TestFuzz_processCompactShares(t *testing.T) {
 
 func Test_processCompactShares(t *testing.T) {
 	// exactTxShareSize is the length of tx that will fit exactly into a single
-	// share, accounting for namespace id and the length delimiter prepended to
+	// share, accounting for the tx length delimiter prepended to
 	// each tx. Note that the length delimiter can be 1 to 10 bytes (varint) but
 	// this test assumes it is 1 byte.
 	const exactTxShareSize = appconsts.ContinuationCompactShareContentSize - 1
@@ -81,7 +81,7 @@ func Test_processCompactShares(t *testing.T) {
 		{"single big tx", appconsts.ContinuationCompactShareContentSize * 4, 1},
 		{"many big txs", appconsts.ContinuationCompactShareContentSize * 4, 10},
 		{"single exact size tx", exactTxShareSize, 1},
-		{"many exact size txs", exactTxShareSize, 10},
+		{"many exact size txs", exactTxShareSize, 100},
 	}
 
 	for _, tc := range tests {

pkg/shares/parse_compact_shares.go

@@ -9,7 +9,8 @@ import (
 
 // parseCompactShares takes raw shares and extracts out transactions,
 // intermediate state roots, or evidence. The returned [][]byte do not have
-// namespaces, info bytes, or length delimiters and are ready to be unmarshalled
+// namespaces, info bytes, data length delimiter, or unit length
+// delimiters and are ready to be unmarshalled
 func parseCompactShares(shares [][]byte) (data [][]byte, err error) {
 	if len(shares) == 0 {
 		return nil, nil
@@ -44,7 +45,7 @@ func (ss *shareStack) resolve() ([][]byte, error) {
 	if !infoByte.IsMessageStart() {
 		return nil, errors.New("first share is not a message start")
 	}
-	err = ss.peel(ss.shares[0][appconsts.NamespaceSize+appconsts.ShareInfoBytes+appconsts.CompactShareReservedBytes:], true)
+	err = ss.peel(ss.shares[0][appconsts.NamespaceSize+appconsts.ShareInfoBytes+appconsts.FirstCompactShareDataLengthBytes+appconsts.CompactShareReservedBytes:], true)
 	return ss.data, err
 }
 

pkg/shares/share_splitting_test.go

@@ -26,11 +26,12 @@ func TestSplitTxs(t *testing.T) {
 			want: [][]uint8{
 				append([]uint8{
 					0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, // namespace id
-					0x1, // info byte
-					0x0, // reserved byte
+					0x1,                // info byte
+					0x2, 0x0, 0x0, 0x0, // 1 byte (unit)  + 1 byte (unit length) = 2 bytes message length
+					0x0, // BUG: reserved byte should be non-zero
 					0x1, // unit length of first transaction
 					0xa, // data of first transaction
-				}, bytes.Repeat([]byte{0}, 244)...), // padding
+				}, bytes.Repeat([]byte{0}, 240)...), // padding
 			},
 		},
 		{
@@ -39,25 +40,27 @@ func TestSplitTxs(t *testing.T) {
 			want: [][]uint8{
 				append([]uint8{
 					0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, // namespace id
-					0x1, // info byte
-					0x0, // reserved byte
+					0x1,                // info byte
+					0x4, 0x0, 0x0, 0x0, // 2 bytes (first transaction) + 2 bytes (second transaction) = 4 bytes message length
+					0x0, // BUG: reserved byte should be non-zero
 					0x1, // unit length of first transaction
 					0xa, // data of first transaction
 					0x1, // unit length of second transaction
 					0xb, // data of second transaction
-				}, bytes.Repeat([]byte{0}, 242)...), // padding
+				}, bytes.Repeat([]byte{0}, 238)...), // padding
 			},
 		},
 		{
 			name: "one large tx that spans two shares",
-			txs:  coretypes.Txs{bytes.Repeat([]byte{0xC}, 245)},
+			txs:  coretypes.Txs{bytes.Repeat([]byte{0xC}, 241)},
 			want: [][]uint8{
 				append([]uint8{
 					0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, // namespace id
-					0x1,    // info byte
-					0x0,    // BUG reserved byte should be non-zero see https://github.com/celestiaorg/celestia-app/issues/802
-					245, 1, // unit length of first transaction is 245
-				}, bytes.Repeat([]byte{0xc}, 244)...), // data of first transaction
+					0x1,          // info byte
+					243, 1, 0, 0, // 241 (unit) + 2 (unit length) = 243 message length
+					0x0,    // BUG: reserved byte should be non-zero
+					241, 1, // unit length of first transaction is 241
+				}, bytes.Repeat([]byte{0xc}, 240)...), // data of first transaction
 				append([]uint8{
 					0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, // namespace id
 					0x0, // info byte
@@ -68,22 +71,23 @@ func TestSplitTxs(t *testing.T) {
 		},
 		{
 			name: "one small tx then one large tx that spans two shares",
-			txs:  coretypes.Txs{coretypes.Tx{0xd}, bytes.Repeat([]byte{0xe}, 243)},
+			txs:  coretypes.Txs{coretypes.Tx{0xd}, bytes.Repeat([]byte{0xe}, 241)},
 			want: [][]uint8{
 				append([]uint8{
 					0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, // namespace id
-					0x1,    // info byte
-					0x0,    // BUG reserved byte should be non-zero see https://github.com/celestiaorg/celestia-app/issues/802
+					0x1,          // info byte
+					245, 1, 0, 0, // 2 bytes (first transaction) + 243 bytes (second transaction) = 245 bytes message length
+					0x0,    // BUG: reserved byte should be non-zero
 					1,      // unit length of first transaction
 					0xd,    // data of first transaction
-					243, 1, // unit length of second transaction is 243
-				}, bytes.Repeat([]byte{0xe}, 242)...), // data of first transaction
+					241, 1, // unit length of second transaction is 241
+				}, bytes.Repeat([]byte{0xe}, 238)...), // data of first transaction
 				append([]uint8{
 					0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, // namespace id
-					0x0, // info byte
-					0x0, // reserved byte
-					0xe, // continuation data of second transaction
-				}, bytes.Repeat([]byte{0x0}, 245)...), // padding
+					0x0,           // info byte
+					0x0,           // reserved byte
+					0xe, 0xe, 0xe, // continuation data of second transaction
+				}, bytes.Repeat([]byte{0x0}, 243)...), // padding
 			},
 		},
 	}

pkg/shares/split_compact_shares.go

@@ -29,8 +29,10 @@ func NewCompactShareSplitter(ns namespace.ID, version uint8) *CompactShareSplitt
 	if err != nil {
 		panic(err)
 	}
+	placeholderDataLength := make([]byte, appconsts.FirstCompactShareDataLengthBytes)
 	pendingShare.Share = append(pendingShare.Share, ns...)
 	pendingShare.Share = append(pendingShare.Share, byte(infoByte))
+	pendingShare.Share = append(pendingShare.Share, placeholderDataLength...)
 	return &CompactShareSplitter{pendingShare: pendingShare, namespace: ns}
 }
 
@@ -59,7 +61,7 @@ func (css *CompactShareSplitter) WriteEvidence(evd coretypes.Evidence) error {
 func (css *CompactShareSplitter) WriteBytes(rawData []byte) {
 	// if this is the first time writing to a pending share, we must add the
 	// reserved bytes
-	if len(css.pendingShare.Share) == appconsts.NamespaceSize+appconsts.ShareInfoBytes {
+	if css.isEmptyPendingShare() {
 		reservedBytes := make([]byte, appconsts.CompactShareReservedBytes)
 		css.pendingShare.Share = append(css.pendingShare.Share, reservedBytes...)
 	}
@@ -130,14 +132,28 @@ func (css *CompactShareSplitter) stackPending() {
 
 // Export finalizes and returns the underlying compact shares.
 func (css *CompactShareSplitter) Export() NamespacedShares {
+	if css.isEmpty() {
+		return []NamespacedShare{}
+	}
+
+	var bytesOfPadding int
 	// add the pending share to the current shares before returning
-	if len(css.pendingShare.Share) > appconsts.NamespaceSize+appconsts.ShareInfoBytes {
-		css.pendingShare.Share, _ = zeroPadIfNecessary(css.pendingShare.Share, appconsts.ShareSize)
+	if !css.isEmptyPendingShare() {
+		css.pendingShare.Share, bytesOfPadding = zeroPadIfNecessary(css.pendingShare.Share, appconsts.ShareSize)
 		css.shares = append(css.shares, css.pendingShare)
 	}
-	// force the last share to have a reserve byte of zero
+
+	dataLengthVarint := css.dataLengthVarint(bytesOfPadding)
+	css.writeDataLengthVarintToFirstShare(dataLengthVarint)
+	css.forceLastShareReserveByteToZero()
+	return css.shares
+}
+
+// forceLastShareReserveByteToZero overwrites the reserve byte of the last share
+// with zero. See https://github.com/celestiaorg/celestia-app/issues/779
+func (css *CompactShareSplitter) forceLastShareReserveByteToZero() {
 	if len(css.shares) == 0 {
-		return css.shares
+		return
 	}
 	lastShare := css.shares[len(css.shares)-1]
 	rawLastShare := lastShare.Data()
@@ -147,21 +163,99 @@ func (css *CompactShareSplitter) Export() NamespacedShares {
 		// confusion for light clients parsing these shares, as the rest of the
 		// data after transaction is padding. See
 		// https://github.com/celestiaorg/celestia-specs/blob/master/src/specs/data_structures.md#share
-		rawLastShare[appconsts.NamespaceSize+appconsts.ShareInfoBytes+i] = byte(0)
+		if len(css.shares) == 1 {
+			// the reserved byte is after the namespace, info byte, and data length varint
+			rawLastShare[appconsts.NamespaceSize+appconsts.ShareInfoBytes+appconsts.FirstCompactShareDataLengthBytes+i] = byte(0)
+		} else {
+			// the reserved byte is after the namespace, info byte
+			rawLastShare[appconsts.NamespaceSize+appconsts.ShareInfoBytes+i] = byte(0)
+		}
 	}
 
 	newLastShare := NamespacedShare{
 		Share: rawLastShare,
 		ID:    lastShare.NamespaceID(),
 	}
 	css.shares[len(css.shares)-1] = newLastShare
-	return css.shares
 }
 
-// Count returns the current number of shares that will be made if exporting.
+// dataLengthVarint returns a varint of the data length written to this compact
+// share splitter.
+func (css *CompactShareSplitter) dataLengthVarint(bytesOfPadding int) []byte {
+	if css.isEmpty() {
+		return []byte{}
+	}
+
+	// declare and initialize the data length
+	dataLengthVarint := make([]byte, appconsts.FirstCompactShareDataLengthBytes)
+	binary.PutUvarint(dataLengthVarint, css.dataLength(bytesOfPadding))
+	zeroPadIfNecessary(dataLengthVarint, appconsts.FirstCompactShareDataLengthBytes)
+
+	return dataLengthVarint
+}
+
+func (css *CompactShareSplitter) writeDataLengthVarintToFirstShare(dataLengthVarint []byte) {
+	if css.isEmpty() {
+		return
+	}
+
+	// write the data length varint to the first share
+	firstShare := css.shares[0]
+	rawFirstShare := firstShare.Data()
+	for i := 0; i < appconsts.FirstCompactShareDataLengthBytes; i++ {
+		rawFirstShare[appconsts.NamespaceSize+appconsts.ShareInfoBytes+i] = dataLengthVarint[i]
+	}
+
+	// replace existing first share with new first share
+	newFirstShare := NamespacedShare{
+		Share: rawFirstShare,
+		ID:    firstShare.NamespaceID(),
+	}
+	css.shares[0] = newFirstShare
+}
+
+// dataLength returns the total length in bytes of all units (transactions,
+// intermediate state roots, or evidence) written to this splitter.
+// dataLength does not include the # of bytes occupied by the namespace ID or
+// the share info byte in each share. dataLength does include the reserved
+// byte in each share and the unit length delimiter prefixed to each unit.
+func (css *CompactShareSplitter) dataLength(bytesOfPadding int) uint64 {
+	if len(css.shares) == 0 {
+		return 0
+	}
+	if len(css.shares) == 1 {
+		return uint64(appconsts.FirstCompactShareContentSize) - uint64(bytesOfPadding)
+	}
+
+	continuationSharesCount := len(css.shares) - 1
+	continuationSharesDataLength := uint64(continuationSharesCount) * appconsts.ContinuationCompactShareContentSize
+	return uint64(appconsts.FirstCompactShareContentSize) + continuationSharesDataLength - uint64(bytesOfPadding)
+}
+
+// isEmptyPendingShare returns true if the pending share is empty, false otherwise.
+func (css *CompactShareSplitter) isEmptyPendingShare() bool {
+	if css.isPendingShareTheFirstShare() {
+		return len(css.pendingShare.Share) == appconsts.NamespaceSize+appconsts.ShareInfoBytes+appconsts.FirstCompactShareDataLengthBytes
+	}
+	return len(css.pendingShare.Share) == appconsts.NamespaceSize+appconsts.ShareInfoBytes
+}
+
+// isPendingShareTheFirstShare returns true if the pending share is the first
+// share of this compact share splitter and false otherwise.
+func (css *CompactShareSplitter) isPendingShareTheFirstShare() bool {
+	return len(css.shares) == 0
+}
+
+// isEmpty returns whether this compact share splitter is empty.
+func (css *CompactShareSplitter) isEmpty() bool {
+	return len(css.shares) == 0 && css.isEmptyPendingShare()
+}
+
+// Count returns the number of shares that would be made if `Export` was invoked
+// on this compact share splitter.
 func (css *CompactShareSplitter) Count() (shareCount int) {
-	if len(css.pendingShare.Share) > appconsts.NamespaceSize+appconsts.ShareInfoBytes {
-		// pending share is non-empty, so we must add one to the count
+	if !css.isEmptyPendingShare() {
+		// pending share is non-empty, so it will be zero padded and added to shares during export
 		return len(css.shares) + 1
 	}
 	return len(css.shares)