chunk.go

package flow

import (
	"fmt"
	"io"
	"log"

	"github.com/ipfs/go-cid"
	"github.com/onflow/go-ethereum/rlp"
	"github.com/vmihailenco/msgpack/v4"
)

var EmptyEventCollectionID Identifier

func init() {
	// Convert hexadecimal string to a byte slice.
	var err error
	emptyEventCollectionHex := "0e5751c026e543b2e8ab2eb06099daa1d1e5df47778f7787faab45cdf12fe3a8"
	EmptyEventCollectionID, err = HexStringToIdentifier(emptyEventCollectionHex)
	if err != nil {
		log.Fatalf("Failed to decode hex: %v", err)
	}
}

type ChunkBody struct {
	CollectionIndex uint

	// execution info
	StartState      StateCommitment // start state when starting executing this chunk
	EventCollection Identifier      // Events generated by executing results
	// ServiceEventCount defines how many service events were emitted in this chunk.
	// By reading these fields from the prior chunks in the same ExecutionResult, we can
	// compute exactly what service events were emitted in this chunk.
	//
	// Let C be this chunk, K be the set of chunks in the ExecutionResult containing C.
	// Then the service event indices for C are given by:
	//    StartIndex = ∑Ci.ServiceEventCount : Ci ∈ K, Ci.Index < C.Index
	//    EndIndex = StartIndex + C.ServiceEventCount
	// The service events for C are given by:
	//    ExecutionResult.ServiceEvents[StartIndex:EndIndex]
	//
	ServiceEventCount uint16
	BlockID           Identifier // Block id of the execution result this chunk belongs to

	// Computation consumption info
	TotalComputationUsed uint64 // total amount of computation used by running all txs in this chunk
	NumberOfTransactions uint64 // number of transactions inside the collection
}

//structwrite:immutable - mutations allowed only within the constructor
type Chunk struct {
	ChunkBody

	Index uint64 // chunk index inside the ER (starts from zero)
	// EndState inferred from next chunk or from the ER
	EndState StateCommitment
}

// UntrustedChunk is an untrusted input-only representation of an Chunk,
// used for construction.
//
// This type exists to ensure that constructor functions are invoked explicitly
// with named fields, which improves clarity and reduces the risk of incorrect field
// ordering during construction.
//
// An instance of UntrustedChunk should be validated and converted into
// a trusted Chunk using NewChunk constructor.
type UntrustedChunk Chunk

// NewChunk returns a Chunk compliant with Protocol Version 2 and later.
// Construction Chunk allowed only within the constructor.
//
// All errors indicate a valid Chunk cannot be constructed from the input.
func NewChunk(untrusted UntrustedChunk) (*Chunk, error) {
	if untrusted.BlockID == ZeroID {
		return nil, fmt.Errorf("BlockID must not be empty")
	}

	if untrusted.StartState == (StateCommitment{}) {
		return nil, fmt.Errorf("StartState must not be zero-value")
	}

	if untrusted.EventCollection == ZeroID {
		return nil, fmt.Errorf("EventCollection must not be empty")
	}

	if untrusted.EndState == (StateCommitment{}) {
		return nil, fmt.Errorf("EndState must not be zero-value")
	}

	return &Chunk{
		ChunkBody: ChunkBody{
			BlockID:              untrusted.BlockID,
			CollectionIndex:      untrusted.CollectionIndex,
			StartState:           untrusted.StartState,
			NumberOfTransactions: untrusted.NumberOfTransactions,
			EventCollection:      untrusted.EventCollection,
			ServiceEventCount:    untrusted.ServiceEventCount,
			TotalComputationUsed: untrusted.TotalComputationUsed,
		},
		Index:    untrusted.Index,
		EndState: untrusted.EndState,
	}, nil
}

// NewRootChunk creates a chunk whose final state is the given commit, with all other fields set to zero.
// This is a special kind of chunk used only as the sole chunk of a root execution result, which forms
// a part of the root protocol state snapshot used as the trusted root for a spork.
func NewRootChunk(
	commit StateCommitment,
) *Chunk {
	return &Chunk{
		ChunkBody: ChunkBody{
			BlockID:              Identifier{},
			CollectionIndex:      0,
			StartState:           StateCommitment{},
			EventCollection:      Identifier{},
			ServiceEventCount:    0,
			TotalComputationUsed: 0,
			NumberOfTransactions: 0,
		},
		Index:    0,
		EndState: commit,
	}
}

// ID returns the unique identifier of the Chunk
func (ch *Chunk) ID() Identifier {
	return MakeID(ch)
}

// ChunkDataPackHeader is a reduced representation of ChunkDataPack. In a nutshell, we substitute
// the larger [ChunkDataPack.Proof] and [ChunkDataPack.Collection] with their collision-resistant hashes.
// Note, ChunkDataPackHeader.ID() is the same as ChunkDataPack.ID().
//
//structwrite:immutable - mutations allowed only within the constructor
type ChunkDataPackHeader struct {
	ChunkID    Identifier      // ID of the chunk this data pack is for
	StartState StateCommitment // commitment for starting state
	Proof      Identifier      // Hash of the proof for all registers touched (read or written) during the chunk execution
	Collection Identifier      // ID of collection executed in this chunk; [flow.ZeroID] for system chunk

	// ExecutionDataRoot is the root data structure of an execution_data.BlockExecutionData.
	// It contains the necessary information for a verification node to validate that the
	// BlockExecutionData produced is valid.
	ExecutionDataRoot BlockExecutionDataRoot
}

// NewChunkDataPackHeader instantiates an "immutable"  ChunkDataPackHeader.
// The `CollectionID` field is set to [flow.ZeroID] for system chunks.
func NewChunkDataPackHeader(ChunkID Identifier, StartState StateCommitment, ProofID Identifier, CollectionID Identifier, ExecutionDataRoot BlockExecutionDataRoot) *ChunkDataPackHeader {
	return &ChunkDataPackHeader{
		ChunkID:           ChunkID,
		StartState:        StartState,
		Proof:             ProofID,
		Collection:        CollectionID,
		ExecutionDataRoot: ExecutionDataRoot,
	}
}

func (c *ChunkDataPackHeader) ID() Identifier {
	return MakeID(c)
}

// ChunkDataPack holds all register touches (any read, or write).
//
// Note that we have to include merkle paths as storage proof for all registers touched (read or written) for
// the _starting_ state of the chunk (i.e. before the chunk computation updates the registers).
// For instance, if an execution state contains three registers: { A: 1, B: 2, C: 3}, and a certain
// chunk has a tx that assigns A = A + B, then its chunk data pack should include the merkle
// paths for { A: 1, B: 2 } as storage proof.
// C is not included because it's neither read or written by the chunk.
// B is included because it's read by the chunk.
// A is included because it's updated by the chunk, and its value 1 is included because it's
// the value before the chunk computation.
// This is necessary for Verification Nodes to (i) check that the read register values are
// consistent with the starting state's root hash and (ii) verify the correctness of the resulting
// state after the chunk computation. `Proof` includes merkle proofs for all touched registers
// during the execution of the chunk.
// Register proofs order must not be correlated to the order of register reads during
// the chunk execution in order to enforce the SPoCK secret high entropy.
//
//structwrite:immutable - mutations allowed only within the constructor
type ChunkDataPack struct {
	ChunkID    Identifier      // ID of the chunk this data pack is for
	StartState StateCommitment // commitment for starting state
	Proof      StorageProof    // proof for all registers touched (read or written) during the chunk execution
	Collection *Collection     // collection executed in this chunk; nil for system chunk

	// ExecutionDataRoot is the root data structure of an execution_data.BlockExecutionData.
	// It contains the necessary information for a verification node to validate that the
	// BlockExecutionData produced is valid.
	ExecutionDataRoot BlockExecutionDataRoot
}

// UntrustedChunkDataPack is an untrusted input-only representation of an ChunkDataPack,
// used for construction.
//
// This type exists to ensure that constructor functions are invoked explicitly
// with named fields, which improves clarity and reduces the risk of incorrect field
// ordering during construction.
//
// An instance of UntrustedChunkDataPack should be validated and converted into
// a trusted ChunkDataPack using NewChunkDataPack constructor.
type UntrustedChunkDataPack ChunkDataPack

// NewChunkDataPack converts a chunk data pack from an untrusted source
// into its canonical representation. Here, basic structural validation is performed.
// Construction of ChunkDataPacks is ONLY allowed via THIS CONSTRUCTOR.
//
// All errors indicate a valid ChunkDataPack cannot be constructed from the input.
func NewChunkDataPack(untrusted UntrustedChunkDataPack) (*ChunkDataPack, error) {
	if untrusted.ChunkID == ZeroID {
		return nil, fmt.Errorf("ChunkID must not be empty")
	}

	if untrusted.StartState == (StateCommitment{}) {
		return nil, fmt.Errorf("StartState must not be zero-value")
	}

	if len(untrusted.Proof) == 0 {
		return nil, fmt.Errorf("Proof must not be empty")
	}

	if untrusted.ExecutionDataRoot.BlockID == ZeroID {
		return nil, fmt.Errorf("ExecutionDataRoot.BlockID must not be empty")
	}

	if len(untrusted.ExecutionDataRoot.ChunkExecutionDataIDs) == 0 {
		return nil, fmt.Errorf("ExecutionDataRoot.ChunkExecutionDataIDs must not be empty")
	}

	return &ChunkDataPack{
		ChunkID:           untrusted.ChunkID,
		StartState:        untrusted.StartState,
		Proof:             untrusted.Proof,
		Collection:        untrusted.Collection,
		ExecutionDataRoot: untrusted.ExecutionDataRoot,
	}, nil
}

// ID returns a collision-resistant hash of the ChunkDataPack struct.
func (c *ChunkDataPack) ID() Identifier {
	var collectionID Identifier
	if c.Collection != nil {
		collectionID = c.Collection.ID()
	} else {
		collectionID = ZeroID
	}

	return NewChunkDataPackHeader(c.ChunkID, c.StartState, MakeID(c.Proof), collectionID, c.ExecutionDataRoot).ID()
}

// TODO: This is the basic version of the list, we need to substitute it with something like Merkle tree at some point
type ChunkList []*Chunk

func (cl ChunkList) Fingerprint() Identifier {
	return MerkleRoot(GetIDs(cl)...)
}

func (cl *ChunkList) Insert(ch *Chunk) {
	*cl = append(*cl, ch)
}

func (cl ChunkList) Items() []*Chunk {
	return cl
}

// Empty returns true if the chunk list is empty. Otherwise it returns false.
func (cl ChunkList) Empty() bool {
	return len(cl) == 0
}

func (cl ChunkList) Indices() []uint64 {
	indices := make([]uint64, len(cl))
	for i, chunk := range cl {
		indices[i] = chunk.Index
	}

	return indices
}

// ByIndex returns an entity from the list by index
// if requested chunk is within range of list, it returns chunk and true
// if requested chunk is out of the range, it returns nil and false
// boolean return value indicates whether requested chunk is within range
func (cl ChunkList) ByIndex(i uint64) (*Chunk, bool) {
	if i >= uint64(len(cl)) {
		// index out of range
		return nil, false
	}
	return cl[i], true
}

// Len returns the number of Chunks in the list. It is also part of the sort
// interface that makes ChunkList sortable
func (cl ChunkList) Len() int {
	return len(cl)
}

// BlockExecutionDataRoot represents the root of a serialized execution_data.BlockExecutionData.
// The hash of the serialized BlockExecutionDataRoot is the ExecutionDataID used within an
// flow.ExecutionResult.
// Context:
//   - The trie updates in BlockExecutionDataRoot contain the _mutated_ registers only, which is
//     helpful for clients to truslessly replicate the state.
//   - In comparison, the chunk data packs contains all the register values at the chunk's starting
//     state that were _touched_ (written and/or read). This is necessary for Verification Nodes to
//     re-run the chunk the computation.
type BlockExecutionDataRoot struct {
	// BlockID is the ID of the block, whose result this execution data is for.
	BlockID Identifier

	// ChunkExecutionDataIDs is a list of the root CIDs for each serialized execution_data.ChunkExecutionData
	// associated with this block.
	ChunkExecutionDataIDs []cid.Cid
}

// EncodeRLP defines an RLP encoding BlockExecutionDataRoot. We need to define a custom RLP encoding since [cid.Cid] doesn't have one. Without it we can't produce a collision-resistant hash.
// No errors are expected during normal operations.
func (b BlockExecutionDataRoot) EncodeRLP(w io.Writer) error {
	encodingCanonicalForm := struct {
		BlockID               Identifier
		ChunkExecutionDataIDs []string
	}{
		BlockID:               b.BlockID,
		ChunkExecutionDataIDs: cidsToStrings(b.ChunkExecutionDataIDs),
	}

	return rlp.Encode(w, encodingCanonicalForm)
}

// MarshalMsgpack implements the msgpack.Marshaler interface
func (b BlockExecutionDataRoot) MarshalMsgpack() ([]byte, error) {
	return msgpack.Marshal(struct {
		BlockID               Identifier
		ChunkExecutionDataIDs []string
	}{
		BlockID:               b.BlockID,
		ChunkExecutionDataIDs: cidsToStrings(b.ChunkExecutionDataIDs),
	})
}

// UnmarshalMsgpack implements the msgpack.Unmarshaler interface
func (b *BlockExecutionDataRoot) UnmarshalMsgpack(data []byte) error {
	var temp struct {
		BlockID               Identifier
		ChunkExecutionDataIDs []string
	}

	if err := msgpack.Unmarshal(data, &temp); err != nil {
		return err
	}

	b.BlockID = temp.BlockID
	cids, err := stringsToCids(temp.ChunkExecutionDataIDs)

	if err != nil {
		return fmt.Errorf("failed to decode chunk execution data ids: %w", err)
	}

	b.ChunkExecutionDataIDs = cids

	return nil
}

// ChunkDataRequest represents a request for the chunk data pack
// which is specified by a chunk ID.
type ChunkDataRequest struct {
	ChunkID Identifier
	Nonce   uint64
}

// ChunkDataResponse is the structurally validated response to a chunk data pack request.
// It contains the chunk data pack of the interest.
type ChunkDataResponse struct {
	ChunkDataPack ChunkDataPack
	Nonce         uint64
}

// Helper function to convert a slice of cid.Cid to a slice of strings
func cidsToStrings(cids []cid.Cid) []string {
	if cids == nil {
		return nil
	}
	strs := make([]string, len(cids))
	for i, c := range cids {
		strs[i] = c.String()
	}
	return strs
}

// Helper function to convert a slice of strings to a slice of cid.Cid
func stringsToCids(strs []string) ([]cid.Cid, error) {
	if strs == nil {
		return nil, nil
	}
	cids := make([]cid.Cid, len(strs))
	for i, s := range strs {
		c, err := cid.Decode(s)
		if err != nil {
			return nil, fmt.Errorf("failed to decode cid %v: %w", s, err)
		}
		cids[i] = c
	}
	return cids, nil
}

// Equals returns true if and only if receiver BlockExecutionDataRoot is equal to the `other`.
func (b BlockExecutionDataRoot) Equals(other BlockExecutionDataRoot) bool {
	if b.BlockID != other.BlockID {
		return false
	}

	if len(b.ChunkExecutionDataIDs) != len(other.ChunkExecutionDataIDs) {
		return false
	}
	for i, cid := range b.ChunkExecutionDataIDs {
		if !cid.Equals(other.ChunkExecutionDataIDs[i]) {
			return false
		}
	}

	return true
}