node/vms/proposervm/tree/tree.go

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

package tree

import (
	"context"
	"maps"
	"slices"
	"sync"

	"github.com/luxfi/ids"
	chain "github.com/luxfi/vm/chain"
)

// Tree handles the propagation of block acceptance and rejection to inner
// blocks.
//
// The Tree is needed because:
// 1. The consensus engine guarantees that for each verified block, either
// Accept() or Reject() are eventually called, and they are called only once.
// The proposervm must maintain these invariants for the wrapped VM.
// 2. A given inner block may be wrapped into multiple different proposervm
// blocks (e.g. same inner block generated by two validators).
//
// The Tree prevents Accept() and Reject() from being called multiple times on
// the same inner block by:
// 1. tracking inner blocks in a tree-like structure, to be able to easily spot
// siblings
// 2. rejecting an inner block only when one of the siblings is accepted.
// Rejection of a proposervm block does not imply its inner block rejection
// (it may be held by a different proposervm block).
type Tree interface {
	// Add places the block in the tree
	Add(chain.Block)

	// Get returns the block that was added to this tree whose parent and ID
	// match the provided block. If non-exists, then false will be returned.
	Get(chain.Block) (chain.Block, bool)

	// Accept marks the provided block as accepted and rejects every conflicting
	// block.
	Accept(context.Context, chain.Block) error
}

type tree struct {
	// lock guards [nodes]. The proposervm calls Add/Get (during Verify) and
	// Accept concurrently from the consensus engine's handler goroutines, so
	// every access to [nodes] must hold this lock. Accept makes block callouts
	// (Accept/Reject on the inner VM) only AFTER releasing the lock, so this
	// lock is a leaf lock and cannot deadlock against the inner VM.
	lock sync.RWMutex
	// parentID -> childID -> childBlock
	nodes map[ids.ID]map[ids.ID]chain.Block
}

func New() Tree {
	return &tree{
		nodes: make(map[ids.ID]map[ids.ID]chain.Block),
	}
}

func (t *tree) Add(blk chain.Block) {
	t.lock.Lock()
	defer t.lock.Unlock()

	parentID := blk.Parent()
	children, exists := t.nodes[parentID]
	if !exists {
		children = make(map[ids.ID]chain.Block)
		t.nodes[parentID] = children
	}
	children[blk.ID()] = blk
}

func (t *tree) Get(blk chain.Block) (chain.Block, bool) {
	t.lock.RLock()
	defer t.lock.RUnlock()

	parentID := blk.Parent()
	children := t.nodes[parentID]
	originalBlk, exists := children[blk.ID()]
	return originalBlk, exists
}

func (t *tree) Accept(ctx context.Context, blk chain.Block) error {
	// accept the provided block. This callout is made before any map mutation,
	// preserving the original semantics: if Accept fails the tree is unchanged.
	if err := blk.Accept(ctx); err != nil {
		return err
	}

	// Phase 1 (locked, no callouts): detach the accepted block's subtree from
	// the node map and collect every conflicting block that must be rejected.
	// Collecting the full reject set under the lock — rather than interleaving
	// map reads with Reject callouts — is what keeps this lock a leaf lock.
	t.lock.Lock()
	parentID := blk.Parent()
	children := t.nodes[parentID]
	delete(children, blk.ID())
	delete(t.nodes, parentID)

	// frontier holds blocks still to be expanded; rejected accumulates the full
	// set of blocks to reject (each exactly once).
	frontier := slices.Collect(maps.Values(children))
	rejected := make([]chain.Block, 0, len(frontier))
	for len(frontier) > 0 {
		i := len(frontier) - 1
		child := frontier[i]
		frontier = frontier[:i]

		rejected = append(rejected, child)

		// mark the progeny of this block as being rejectable
		childID := child.ID()
		frontier = append(frontier, slices.Collect(maps.Values(t.nodes[childID]))...)
		delete(t.nodes, childID)
	}
	t.lock.Unlock()

	// Phase 2 (unlocked): reject all conflicting blocks via inner-VM callouts.
	for _, child := range rejected {
		if err := child.Reject(ctx); err != nil {
			return err
		}
	}
	return nil
}