Kasavinal DA

chain/block.go

@@ -4,8 +4,11 @@
 package chain
 
 import (
+	"bytes"
 	"context"
+	"crypto/sha256"
 	"encoding/binary"
+	"encoding/json"
 	"fmt"
 	"time"
 
@@ -15,6 +18,8 @@ import (
 	"github.com/ava-labs/avalanchego/snow/engine/snowman/block"
 	"github.com/ava-labs/avalanchego/utils/set"
 	"github.com/ava-labs/avalanchego/x/merkledb"
+	"github.com/celestiaorg/merkletree"
+	"github.com/manojkgorle/rsmt2d"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
@@ -50,6 +55,11 @@ type StatefulBlock struct {
 	// starting the verification of another block, etc.
 	StateRoot ids.ID `json:"stateRoot"`
 
+	// DataRoot, RowRoots, ColumnRoots are produced asynchronously
+	DataRoot    []byte   `json:"dataRoot"`
+	RowRoots    [][]byte `json:"rowRoots"`
+	ColumnRoots [][]byte `json:"columnRoots"`
+
 	size int
 
 	// authCounts can be used by batch signature verification
@@ -70,6 +80,18 @@ func (b *StatefulBlock) ID() (ids.ID, error) {
 }
 
 func NewGenesisBlock(root ids.ID) *StatefulBlock {
+	data := make([]byte, ShareSize)
+	eds, _ := rsmt2d.ComputeExtendedDataSquare([][]byte{data}, rsmt2d.NewLeoRSCodec(), rsmt2d.NewDefaultTree)
+	rroots, _ := eds.RowRoots()
+	croots, _ := eds.ColRoots()
+	tree := merkletree.NewFromTreehasher(merkletree.NewDefaultHasher(sha256.New()))
+	for _, root := range rroots {
+		tree.Push(root)
+	}
+	for _, root := range croots {
+		tree.Push(root)
+	}
+	mroot := tree.Root()
 	return &StatefulBlock{
 		// We set the genesis block timestamp to be after the ProposerVM fork activation.
 		//
@@ -82,7 +104,10 @@ func NewGenesisBlock(root ids.ID) *StatefulBlock {
 		Tmstmp: time.Date(2023, time.January, 1, 0, 0, 0, 0, time.UTC).UnixMilli(),
 
 		// StateRoot should include all allocates made when loading the genesis file
-		StateRoot: root,
+		StateRoot:   root,
+		DataRoot:    mroot,
+		RowRoots:    rroots,
+		ColumnRoots: croots,
 	}
 }
 
@@ -101,6 +126,8 @@ type StatelessBlock struct {
 	results    []*Result
 	feeManager *fees.Manager
 
+	eds *rsmt2d.ExtendedDataSquare
+
 	vm   VM
 	view merkledb.View
 
@@ -472,7 +499,44 @@ func (b *StatelessBlock) innerVerify(ctx context.Context, vctx VerifyContext) er
 			b.StateRoot,
 		)
 	}
-
+	eds, prr, pcc, pdr := b.vm.GetBlockDataEdsAsync(b.Parent())
+	for i, root := range prr {
+		if !bytes.Equal(root, b.RowRoots[i]) {
+			return fmt.Errorf("row root mismatch %d", i)
+		}
+	}
+	for i, root := range pcc {
+		if !bytes.Equal(root, b.ColumnRoots[i]) {
+			return fmt.Errorf("column root mismatch %d", i)
+		}
+	}
+	if !bytes.Equal(pdr, b.DataRoot) {
+		return fmt.Errorf("data root mismatch")
+	}
+	b.eds = eds
+	txData := make([]byte, 0)
+	for _, tx := range b.Txs {
+		if tx.Action.GetTypeID() == 1 {
+			txData = append(txData, tx.Action.Data()...)
+		}
+	}
+	if len(txData) == 0 {
+		txData = append(txData, make([]byte, ShareSize)...)
+	}
+	if len(txData)%ShareSize != 0 {
+		diff := ShareSize - len(txData)%ShareSize
+		txData = append(txData, make([]byte, diff)...)
+	}
+	var holder [][]byte
+	for i := 0; i < len(txData)/ShareSize; i++ {
+		holder = append(holder, txData[i*ShareSize:(i+1)*ShareSize])
+	}
+	if !isSquareNumber(len(holder)) {
+		diff := nextSquareOffset(len(holder))
+		for i := 0; i < diff; i++ {
+			holder = append(holder, make([]byte, ShareSize))
+		}
+	}
 	// Ensure signatures are verified
 	_, sspan := b.vm.Tracer().Start(ctx, "StatelessBlock.Verify.WaitSignatures")
 	start = time.Now()
@@ -492,6 +556,32 @@ func (b *StatelessBlock) innerVerify(ctx context.Context, vctx VerifyContext) er
 	}
 	b.view = view
 
+	go func() {
+		rcodec := rsmt2d.NewLeoRSCodec()
+		eds, err := rsmt2d.ComputeExtendedDataSquare(holder, rcodec, rsmt2d.NewDefaultTree)
+		if err != nil {
+			log.Error("failed to compute extended data square", zap.Error(err))
+			return
+		}
+		rroots, err := eds.RowRoots()
+		if err != nil {
+			log.Error("failed to get row roots", zap.Error(err))
+			return
+		}
+		croots, err := eds.ColRoots()
+		if err != nil {
+			log.Error("failed to get column roots", zap.Error(err))
+			return
+		}
+		tree := merkletree.NewFromTreehasher(merkletree.NewDefaultHasher(sha256.New()))
+		for _, root := range rroots {
+			tree.Push(root)
+		}
+		for _, root := range croots {
+			tree.Push(root)
+		}
+		b.vm.WriteBlockDataEdsAsync(b.ID(), eds, rroots, croots, tree.Root())
+	}()
 	// Kickoff root generation
 	go func() {
 		start := time.Now()
@@ -610,6 +700,10 @@ func (b *StatelessBlock) Processed() bool {
 	return b.view != nil
 }
 
+func (b *StatelessBlock) EDS() *rsmt2d.ExtendedDataSquare {
+	return b.eds
+}
+
 // View returns the [merkledb.TrieView] of the block (representing the state
 // post-execution) or returns the accepted state if the block is accepted or
 // is height 0 (genesis).
@@ -820,6 +914,17 @@ func (b *StatefulBlock) Marshal() ([]byte, error) {
 	}
 
 	p.PackID(b.StateRoot)
+	p.PackBytes(b.DataRoot)
+	bt, err := json.Marshal(&b.RowRoots)
+	if err != nil {
+		return nil, err
+	}
+	p.PackBytes(bt)
+	ct, err := json.Marshal(&b.ColumnRoots)
+	if err != nil {
+		return nil, err
+	}
+	p.PackBytes(ct)
 	bytes := p.Bytes()
 	if err := p.Err(); err != nil {
 		return nil, err
@@ -854,7 +959,20 @@ func UnmarshalBlock(raw []byte, parser Parser) (*StatefulBlock, error) {
 	}
 
 	p.UnpackID(false, &b.StateRoot)
-
+	p.UnpackBytes(-1, false, &b.DataRoot)
+	var bt []byte
+	p.UnpackBytes(-1, false, &bt)
+	err := json.Unmarshal(bt, &b.RowRoots)
+	if err != nil {
+		return nil, fmt.Errorf("row root unmarshall err, %w", err)
+	}
+	var ct []byte
+	p.UnpackBytes(-1, false, &ct)
+	err = json.Unmarshal(ct, &b.ColumnRoots)
+	if err != nil {
+		// fmt.Println("row roots")
+		return nil, fmt.Errorf("column root unmarshall err, %w", err)
+	}
 	// Ensure no leftover bytes
 	if !p.Empty() {
 		return nil, fmt.Errorf("%w: remaining=%d", ErrInvalidObject, len(raw)-p.Offset())

chain/builder.go

@@ -5,16 +5,20 @@ package chain
 
 import (
 	"context"
+	"crypto/sha256"
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"math"
 	"sync"
 	"time"
 
 	"github.com/ava-labs/avalanchego/database"
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/utils/logging"
 	"github.com/ava-labs/avalanchego/utils/set"
+	"github.com/celestiaorg/merkletree"
+	"github.com/manojkgorle/rsmt2d"
 	"go.opentelemetry.io/otel/attribute"
 	"go.uber.org/zap"
 
@@ -133,6 +137,7 @@ func BuildBlock(
 		// asynchronously.
 		prepareStreamLock sync.Mutex
 
+		daTx bool
 		// stop is used to trigger that we should stop building, assuming we are no longer executing
 		stop bool
 	)
@@ -141,6 +146,9 @@ func BuildBlock(
 	mempool.StartStreaming(ctx)
 	b.Txs = []*Transaction{}
 	for time.Since(start) < vm.GetTargetBuildDuration() && !stop {
+		if daTx && time.Since(start) > vm.GetTargetBuildDuration()-600*time.Millisecond {
+			break
+		}
 		prepareStreamLock.Lock()
 		txs := mempool.Stream(ctx, streamBatch)
 		prepareStreamLock.Unlock()
@@ -259,7 +267,9 @@ func BuildBlock(
 						cacheLock.Unlock()
 					}()
 				}
-
+				if tx.Action.GetTypeID() == 1 {
+					daTx = true
+				}
 				// Execute block
 				tsv := ts.NewView(stateKeys, storage)
 				if err := tx.PreExecute(ctx, feeManager, sm, r, tsv, nextTime); err != nil {
@@ -425,18 +435,71 @@ func BuildBlock(
 	}
 	b.StateRoot = root
 
+	txData := make([]byte, 0)
+	for _, tx := range b.Txs {
+		if tx.Action.GetTypeID() == 1 {
+			txData = append(txData, tx.Action.Data()...)
+		}
+	}
+	if len(txData) == 0 {
+		txData = append(txData, make([]byte, ShareSize)...)
+	}
+	if len(txData)%ShareSize != 0 {
+		diff := ShareSize - len(txData)%ShareSize
+		txData = append(txData, make([]byte, diff)...)
+	}
+	var holder [][]byte
+	for i := 0; i < len(txData)/ShareSize; i++ {
+		holder = append(holder, txData[i*ShareSize:(i+1)*ShareSize])
+	}
+	if !isSquareNumber(len(holder)) {
+		diff := nextSquareOffset(len(holder))
+		for i := 0; i < diff; i++ {
+			holder = append(holder, make([]byte, ShareSize))
+		}
+	}
 	// Get view from [tstate] after writing all changed keys
 	view, err := ts.ExportMerkleDBView(ctx, vm.Tracer(), parentView)
 	if err != nil {
 		return nil, err
 	}
-
+	eds, rr, cr, dr := vm.GetBlockDataEdsAsync(parent.id)
+	b.DataRoot = dr
+	b.RowRoots = rr
+	b.ColumnRoots = cr
+	b.eds = eds
 	// Compute block hash and marshaled representation
 	if err := b.initializeBuilt(ctx, view, results, feeManager); err != nil {
 		log.Warn("block failed", zap.Int("txs", len(b.Txs)), zap.Any("consumed", feeManager.UnitsConsumed()))
 		return nil, err
 	}
 
+	go func() {
+		rcodec := rsmt2d.NewLeoRSCodec()
+		eds, err := rsmt2d.ComputeExtendedDataSquare(holder, rcodec, rsmt2d.NewDefaultTree)
+		if err != nil {
+			log.Error("extended data square computation failed", zap.Error(err))
+			return
+		}
+		rroots, err := eds.RowRoots()
+		if err != nil {
+			log.Error("row roots computation failed", zap.Error(err))
+			return
+		}
+		croots, err := eds.ColRoots()
+		if err != nil {
+			log.Error("column roots computation failed", zap.Error(err))
+			return
+		}
+		tree := merkletree.NewFromTreehasher(merkletree.NewDefaultHasher(sha256.New()))
+		for _, root := range rroots {
+			tree.Push(root)
+		}
+		for _, root := range croots {
+			tree.Push(root)
+		}
+		vm.WriteBlockDataEdsAsync(b.id, eds, rroots, croots, tree.Root())
+	}()
 	// Kickoff root generation
 	go func() {
 		start := time.Now()
@@ -465,3 +528,22 @@ func BuildBlock(
 	)
 	return b, nil
 }
+
+func isSquareNumber(n int) bool {
+	if n < 0 {
+		return false // Negative numbers are not square numbers
+	}
+	sqrt := int(math.Sqrt(float64(n)))
+	return sqrt*sqrt == n
+}
+
+// nextSquareOffset calculates how far n is from the next square number.
+func nextSquareOffset(n int) int {
+	if n < 0 {
+		return -1 // Negative numbers cannot be squares
+	}
+	sqrt := math.Sqrt(float64(n))
+	nextSquare := int(math.Ceil(sqrt)) // Ceiling to get the next integer
+	nextSquareNumber := nextSquare * nextSquare
+	return nextSquareNumber - n
+}

chain/consts.go

@@ -26,6 +26,8 @@ const (
 	MaxKeyDependencies = 100_000_000
 )
 
+const ShareSize = 512
+
 func HeightKey(prefix []byte) []byte {
 	return keys.EncodeChunks(prefix, HeightKeyChunks)
 }

chain/dependencies.go

@@ -14,6 +14,7 @@ import (
 	"github.com/ava-labs/avalanchego/utils/logging"
 	"github.com/ava-labs/avalanchego/utils/set"
 	"github.com/ava-labs/avalanchego/x/merkledb"
+	"github.com/manojkgorle/rsmt2d"
 
 	"github.com/ava-labs/hypersdk/codec"
 	"github.com/ava-labs/hypersdk/executor"
@@ -80,6 +81,9 @@ type VM interface {
 	GetTransactionExecutionCores() int
 	GetStateFetchConcurrency() int
 
+	WriteBlockDataEdsAsync(ids.ID, *rsmt2d.ExtendedDataSquare, [][]byte, [][]byte, []byte)
+	GetBlockDataEdsAsync(ids.ID) (*rsmt2d.ExtendedDataSquare, [][]byte, [][]byte, []byte)
+
 	Verified(context.Context, *StatelessBlock)
 	Rejected(context.Context, *StatelessBlock)
 	Accepted(context.Context, *StatelessBlock)
@@ -256,6 +260,8 @@ type Action interface {
 		actor codec.Address,
 		txID ids.ID,
 	) (success bool, computeUnits uint64, output []byte, err error)
+
+	Data() []byte
 }
 
 type Auth interface {