square.go

// Package square implements the logic to construct the original data square
// based on a list of transactions.
package square

import (
	"bytes"
	"crypto/sha256"
	"fmt"
	"math"

	"github.com/celestiaorg/go-square/v4/inclusion"
	"github.com/celestiaorg/go-square/v4/share"
	"golang.org/x/exp/constraints"
)

// validateTxOrdering validates that transactions are ordered correctly:
//  1. Normal transactions
//  2. Pay for blob transactions
//  3. Pay for fibre transactions
//
// Returns an error if the ordering is invalid.
func validateTxOrdering(txs []ClassifiedTx) error {
	seenBlobTx := false
	seenFibreTx := false

	for idx, classified := range txs {
		if err := classified.Validate(); err != nil {
			return fmt.Errorf("classified tx at index %d: %w", idx, err)
		}

		if classified.FibreTx != nil {
			seenFibreTx = true
			continue
		}

		_, isBlobTx, err := classified.blobTx()
		if err != nil && isBlobTx {
			return fmt.Errorf("unmarshalling blob tx at index %d: %w", idx, err)
		}

		if isBlobTx {
			if seenFibreTx {
				return fmt.Errorf("blob tx at index %d cannot be appended after pay-for-fibre tx", idx)
			}
			seenBlobTx = true
			continue
		}

		if seenBlobTx {
			return fmt.Errorf("normal tx at index %d cannot be appended after blob tx", idx)
		}
		if seenFibreTx {
			return fmt.Errorf("normal tx at index %d cannot be appended after pay-for-fibre tx", idx)
		}
	}

	return nil
}

// Construct takes the exact list of ordered, pre-classified transactions and
// constructs a square, validating that
//   - each transaction's classification is self-consistent,
//   - transactions are ordered: normal txs, pay for blob txs, pay for fibre txs,
//   - the transactions don't collectively exceed the maxSquareSize.
//
// Note that this function does not check the underlying validity of the
// transactions. Callers classify transactions themselves; see ClassifiedTx.
func Construct(txs []ClassifiedTx, maxSquareSize, subtreeRootThreshold int) (Square, error) {
	if err := validateTxOrdering(txs); err != nil {
		return nil, err
	}

	builder, err := NewBuilder(maxSquareSize, subtreeRootThreshold, txs...)
	if err != nil {
		return nil, err
	}

	return builder.Export()
}

// populateBuilder appends each pre-classified transaction to the builder. Fibre
// txs are appended atomically with the system blob the caller synthesized, blob
// txs are taken from the caller's pre-decoded form or unmarshalled from
// go-square's own BlobTx format, and everything else is appended as a normal
// tx.
func populateBuilder(builder *Builder, txs []ClassifiedTx) error {
	for idx, classified := range txs {
		fibreTx, blobTx, err := classified.decode()
		if err != nil {
			return fmt.Errorf("classified tx at index %d: %w", idx, err)
		}

		if fibreTx != nil {
			added, err := builder.AppendFibreTx(fibreTx)
			if err != nil {
				return fmt.Errorf("appending fibre tx at index %d: %w", idx, err)
			}
			if !added {
				return fmt.Errorf("not enough space to append fibre tx at index %d", idx)
			}
			continue
		}

		if blobTx != nil {
			added, err := builder.AppendBlobTx(blobTx)
			if err != nil {
				return fmt.Errorf("appending blob tx at index %d: %w", idx, err)
			}
			if !added {
				return fmt.Errorf("not enough space to append blob tx at index %d", idx)
			}
			continue
		}

		if !builder.AppendTx(classified.Bytes) {
			return fmt.Errorf("not enough space to append tx at index %d", idx)
		}
	}
	return nil
}

// TxShareRange returns the range of share indexes that the tx, specified by txIndex, occupies.
// The range is end exclusive.
func TxShareRange(txs []ClassifiedTx, txIndex, maxSquareSize, subtreeRootThreshold int) (share.Range, error) {
	if err := validateTxOrdering(txs); err != nil {
		return share.Range{}, err
	}

	builder, err := NewBuilder(maxSquareSize, subtreeRootThreshold, txs...)
	if err != nil {
		return share.Range{}, err
	}

	return builder.FindTxShareRange(txIndex)
}

// BlobShareRange returns the range of share indexes that the blob, identified by txIndex and blobIndex, occupies.
// The range is end exclusive.
//
// Note: this function only supports PFB blobs. It does not support FibreTx
// system blobs because their share indexes are not tracked via IndexWrapper.
// Callers needing system blob positions can use GetShareRangeForNamespace on
// the constructed square instead.
func BlobShareRange(txs []ClassifiedTx, txIndex, blobIndex, maxSquareSize, subtreeRootThreshold int) (share.Range, error) {
	if err := validateTxOrdering(txs); err != nil {
		return share.Range{}, err
	}

	builder, err := NewBuilder(maxSquareSize, subtreeRootThreshold, txs...)
	if err != nil {
		return share.Range{}, err
	}

	start, err := builder.FindBlobStartingIndex(txIndex, blobIndex)
	if err != nil {
		return share.Range{}, err
	}

	blobLen, err := builder.BlobShareLength(txIndex, blobIndex)
	if err != nil {
		return share.Range{}, err
	}
	end := start + blobLen

	return share.NewRange(start, end), nil
}

// Square is a 2D square of shares with symmetrical sides that are always a power of 2.
type Square []share.Share

// Size returns the size of the sides of a square
func (s Square) Size() (int, error) {
	return Size(len(s))
}

// Size returns the size of the row or column in shares of a square. This
// function is currently a wrapper around the da packages equivalent function to
// avoid breaking the api. In future versions there will not be a copy of this
// code here.
func Size(length int) (int, error) {
	return RoundUpPowerOfTwo(int(math.Ceil(math.Sqrt(float64(length)))))
}

// Deprecated: Use inclusion.RoundUpPowerOfTwo directly.
func RoundUpPowerOfTwo[I constraints.Integer](input I) (I, error) {
	return inclusion.RoundUpPowerOfTwo(input)
}

// Equals returns true if two squares are equal
func (s Square) Equals(other Square) bool {
	if len(s) != len(other) {
		return false
	}
	for i := range s {
		if !bytes.Equal(s[i].ToBytes(), other[i].ToBytes()) {
			return false
		}
	}
	return true
}

// WrappedPFBs returns the wrapped PFBs in a square
func (s Square) WrappedPFBs() ([][]byte, error) {
	wpfbShareRange := share.GetShareRangeForNamespace(s, share.PayForBlobNamespace)
	if wpfbShareRange.IsEmpty() {
		return [][]byte{}, nil
	}
	return share.ParseTxs(s[wpfbShareRange.Start:wpfbShareRange.End])
}

func (s Square) IsEmpty() bool {
	return s.Equals(EmptySquare())
}

// ToBytes returns all the shares in the square flattened into a single byte
// slice.
func (s Square) ToBytes() []byte {
	result := make([]byte, 0, len(s)*share.ShareSize)
	for _, share := range s {
		result = append(result, share.ToBytes()...)
	}
	return result
}

// Hash returns a hash based on the shares in the square.
func (s Square) Hash() [32]byte {
	return sha256.Sum256(s.ToBytes())
}

// EmptySquare returns a 1x1 square with a single tail padding share
func EmptySquare() Square {
	return share.TailPaddingShares(share.MinShareCount)
}

func WriteSquare(
	txWriter *share.CompactShareSplitter,
	pfbWriter *share.CompactShareSplitter,
	payForFibreWriter *share.CompactShareSplitter,
	blobWriter *share.SparseShareSplitter,
	nonReservedStart int,
	squareSize int,
) (Square, error) {
	totalShares := squareSize * squareSize
	pfbStartIndex := txWriter.Count()
	payForFibreStartIndex := pfbStartIndex + pfbWriter.Count()
	paddingStartIndex := payForFibreStartIndex + payForFibreWriter.Count()
	if nonReservedStart < paddingStartIndex {
		return nil, fmt.Errorf("nonReservedStart %d is too small to fit all txs, PayForBlob txs, and PayForFibre txs", nonReservedStart)
	}
	padding := share.ReservedPaddingShares(nonReservedStart - paddingStartIndex)
	endOfLastBlob := nonReservedStart + blobWriter.Count()
	if totalShares < endOfLastBlob {
		return nil, fmt.Errorf("square size %d is too small to fit all blobs", totalShares)
	}

	txShares, err := txWriter.Export()
	if err != nil {
		return nil, fmt.Errorf("failed to export tx shares: %w", err)
	}

	pfbShares, err := pfbWriter.Export()
	if err != nil {
		return nil, fmt.Errorf("failed to export pfb shares: %w", err)
	}

	payForFibreShares, err := payForFibreWriter.Export()
	if err != nil {
		return nil, fmt.Errorf("failed to export pay-for-fibre shares: %w", err)
	}

	square := make([]share.Share, totalShares)
	copy(square, txShares)
	copy(square[pfbStartIndex:], pfbShares)
	copy(square[payForFibreStartIndex:], payForFibreShares)
	copy(square[paddingStartIndex:], padding)
	if blobWriter.Count() > 0 {
		copy(square[nonReservedStart:], blobWriter.Export())
	}
	if totalShares > endOfLastBlob {
		copy(square[endOfLastBlob:], share.TailPaddingShares(totalShares-endOfLastBlob))
	}

	return square, nil
}

type PFBDecoder func(txBytes []byte) ([]uint32, error)