container/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"

	chainblock "github.com/luxfi/consensus/engine/chain/block"
	"github.com/luxfi/ids"
)

// 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(chainblock.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(chainblock.Block) (chainblock.Block, bool)

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

type tree struct {
	// parentID -> childID -> childBlock
	nodes map[ids.ID]map[ids.ID]chainblock.Block
}

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

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

func (t *tree) Get(blk chainblock.Block) (chainblock.Block, bool) {
	parentID := blk.Parent()
	children := t.nodes[parentID]
	originalBlk, exists := children[blk.ID()]
	return originalBlk, exists
}

func (t *tree) Accept(ctx context.Context, blk chainblock.Block) error {
	// accept the provided block
	if err := blk.Accept(ctx); err != nil {
		return err
	}

	// get the siblings of the block
	parentID := blk.Parent()
	children := t.nodes[parentID]
	delete(children, blk.ID())
	delete(t.nodes, parentID)

	// mark the siblings of the accepted block as rejectable
	childrenToReject := slices.Collect(maps.Values(children))

	// reject all the rejectable blocks
	for len(childrenToReject) > 0 {
		i := len(childrenToReject) - 1
		child := childrenToReject[i]
		childrenToReject = childrenToReject[:i]

		// reject the block
		if err := child.Reject(ctx); err != nil {
			return err
		}

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