chains/aivm/dag_vertex.go

// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package aivm

import (
	"context"
	"crypto/sha256"
	"encoding/binary"
	"errors"

	"github.com/luxfi/consensus/core/choices"
	"github.com/luxfi/consensus/engine/dag/vertex"
	"github.com/luxfi/ids"

	"github.com/luxfi/ai/pkg/aivm"
)

var _ vertex.DAGVM = (*VM)(nil)

// AIVertex represents a DAG vertex in the AI chain.
// Conflict key: jobID. Independent jobs commute; same job conflicts.
type AIVertex struct {
	id      ids.ID
	bytes   []byte
	height  uint64
	epoch   uint32
	parents []ids.ID
	txIDs   []ids.ID
	status  choices.Status

	tasks   []*aivm.Task
	results []*aivm.TaskResult
	jobIDs  []string
	vm      *VM
}

func (v *AIVertex) ID() ids.ID             { return v.id }
func (v *AIVertex) Bytes() []byte          { return v.bytes }
func (v *AIVertex) Height() uint64         { return v.height }
func (v *AIVertex) Epoch() uint32          { return v.epoch }
func (v *AIVertex) Parents() []ids.ID      { return v.parents }
func (v *AIVertex) Txs() []ids.ID          { return v.txIDs }
func (v *AIVertex) Status() choices.Status { return v.status }

func (v *AIVertex) Verify(ctx context.Context) error {
	for _, task := range v.tasks {
		if task.ID == "" {
			return errors.New("task missing ID")
		}
	}
	return nil
}

// Accept records the vertex. It does NOT move the block chain's tip: a vertex is
// not a block, and writing the vertex id into the tip left the id the VM reported
// and the block the VM held disagreeing — two owners of one fact, and the one a
// restart would read back named no block at all.
//
// The vertex is stored under its own prefix, disjoint from the block keyspace,
// so a block read can never decode a vertex.
func (v *AIVertex) Accept(ctx context.Context) error {
	v.status = choices.Accepted

	v.vm.mu.Lock()
	defer v.vm.mu.Unlock()

	wire, err := marshalVertex(v)
	if err != nil {
		return err
	}
	return v.vm.db.Put(vertexKey(v.id), wire)
}

func (v *AIVertex) Reject(ctx context.Context) error {
	v.status = choices.Rejected
	return nil
}

// jobIDSet returns the set of jobIDs for conflict detection.
func (v *AIVertex) jobIDSet() map[string]struct{} {
	s := make(map[string]struct{}, len(v.jobIDs))
	for _, j := range v.jobIDs {
		s[j] = struct{}{}
	}
	return s
}

// Conflicts returns true if this vertex and other reference the same jobID.
func (v *AIVertex) Conflicts(other *AIVertex) bool {
	ours := v.jobIDSet()
	for _, j := range other.jobIDs {
		if _, ok := ours[j]; ok {
			return true
		}
	}
	return false
}

// ConflictsVertex performs the same check against the vertex.Vertex interface.
func (v *AIVertex) ConflictsVertex(other vertex.Vertex) bool {
	ov, ok := other.(*AIVertex)
	if !ok {
		return false
	}
	return v.Conflicts(ov)
}

// computeID is the vertex's identity, so everything that distinguishes one
// vertex from another has to reach the hash, and no two distinct vertices may
// reach the same one.
//
// Both halves of that were false. The transaction ids were left out entirely,
// so the payload was unauthenticated: two vertices carrying different work
// under the same job names shared an id, and either could be served for the
// other. And the job names were concatenated without separation, so ["ab","c"]
// and ["a","bc"] hashed the same bytes — a second vertex that a peer must treat
// as the first. Each variable-length part is now written with its length ahead
// of it, which is what makes the encoding injective.
//
// The chain id leads, for the reason a block's does: it is not on the wire, so
// every node names a vertex under ITS OWN chain and a vertex cannot be carried
// from one A-Chain to another.
func (v *AIVertex) computeID() ids.ID {
	h := sha256.New()
	chainID := ids.Empty
	if v.vm != nil {
		chainID = v.vm.chainID
	}
	h.Write(chainID[:])
	binary.Write(h, binary.BigEndian, v.height)
	binary.Write(h, binary.BigEndian, v.epoch)

	binary.Write(h, binary.BigEndian, uint32(len(v.parents)))
	for _, p := range v.parents {
		h.Write(p[:]) // fixed width, self-delimiting
	}

	binary.Write(h, binary.BigEndian, uint32(len(v.txIDs)))
	for _, t := range v.txIDs {
		h.Write(t[:]) // fixed width
	}

	binary.Write(h, binary.BigEndian, uint32(len(v.jobIDs)))
	for _, j := range v.jobIDs {
		binary.Write(h, binary.BigEndian, uint32(len(j)))
		h.Write([]byte(j))
	}
	return ids.ID(h.Sum(nil))
}

// BuildVertex creates a vertex from pending tasks/results, batching independent jobs.
func (vm *VM) BuildVertex(ctx context.Context) (vertex.Vertex, error) {
	vm.mu.Lock()
	defer vm.mu.Unlock()

	if !vm.running {
		return nil, ErrNotInitialized
	}

	parent := vm.lastAccepted
	if parent == nil {
		return nil, errors.New("no parent block")
	}

	// Collect pending tasks from core (empty providerID = all pending)
	pendingTasks := vm.core.GetPendingTasks("")
	if len(pendingTasks) == 0 {
		return nil, errors.New("no pending tasks")
	}

	// Greedily batch non-conflicting tasks (unique jobIDs)
	seen := make(map[string]struct{})
	var batch []*aivm.Task
	var jobIDs []string
	for _, task := range pendingTasks {
		if _, dup := seen[task.ID]; dup {
			continue
		}
		seen[task.ID] = struct{}{}
		batch = append(batch, task)
		jobIDs = append(jobIDs, task.ID)
	}

	txIDs := make([]ids.ID, len(batch))
	for i, task := range batch {
		h := sha256.Sum256([]byte(task.ID))
		txIDs[i] = ids.ID(h)
	}

	v := &AIVertex{
		height:  parent.Height_ + 1,
		epoch:   0,
		parents: []ids.ID{parent.ID_},
		txIDs:   txIDs,
		tasks:   batch,
		jobIDs:  jobIDs,
		status:  choices.Processing,
		vm:      vm,
	}
	v.id = v.computeID()
	wire, err := marshalVertex(v)
	if err != nil {
		return nil, err
	}
	v.bytes = wire
	return v, nil
}

// ParseVertex deserializes a vertex from bytes.
func (vm *VM) ParseVertex(ctx context.Context, b []byte) (vertex.Vertex, error) {
	v := &AIVertex{vm: vm}
	if err := parseVertex(b, v); err != nil {
		return nil, err
	}
	v.id = v.computeID()
	v.bytes = b
	return v, nil
}