github.com/ava-labs/avalanchego@v1.11.11/vms/proposervm/tree/tree.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package tree
     5  
     6  import (
     7  	"context"
     8  
     9  	"golang.org/x/exp/maps"
    10  
    11  	"github.com/ava-labs/avalanchego/ids"
    12  	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
    13  )
    14  
    15  // Tree handles the propagation of block acceptance and rejection to inner
    16  // blocks.
    17  //
    18  // The Tree is needed because:
    19  // 1. The consensus engine guarantees that for each verified block, either
    20  // Accept() or Reject() are eventually called, and they are called only once.
    21  // The proposervm must maintain these invariants for the wrapped VM.
    22  // 2. A given inner block may be wrapped into multiple different proposervm
    23  // blocks (e.g. same inner block generated by two validators).
    24  //
    25  // The Tree prevents Accept() and Reject() from being called multiple times on
    26  // the same inner block by:
    27  // 1. tracking inner blocks in a tree-like structure, to be able to easily spot
    28  // siblings
    29  // 2. rejecting an inner block only when one of the siblings is accepted.
    30  // Rejection of a proposervm block does not imply its inner block rejection
    31  // (it may be held by a different proposervm block).
    32  type Tree interface {
    33  	// Add places the block in the tree
    34  	Add(snowman.Block)
    35  
    36  	// Get returns the block that was added to this tree whose parent and ID
    37  	// match the provided block. If non-exists, then false will be returned.
    38  	Get(snowman.Block) (snowman.Block, bool)
    39  
    40  	// Accept marks the provided block as accepted and rejects every conflicting
    41  	// block.
    42  	Accept(context.Context, snowman.Block) error
    43  }
    44  
    45  type tree struct {
    46  	// parentID -> childID -> childBlock
    47  	nodes map[ids.ID]map[ids.ID]snowman.Block
    48  }
    49  
    50  func New() Tree {
    51  	return &tree{
    52  		nodes: make(map[ids.ID]map[ids.ID]snowman.Block),
    53  	}
    54  }
    55  
    56  func (t *tree) Add(blk snowman.Block) {
    57  	parentID := blk.Parent()
    58  	children, exists := t.nodes[parentID]
    59  	if !exists {
    60  		children = make(map[ids.ID]snowman.Block)
    61  		t.nodes[parentID] = children
    62  	}
    63  	blkID := blk.ID()
    64  	children[blkID] = blk
    65  }
    66  
    67  func (t *tree) Get(blk snowman.Block) (snowman.Block, bool) {
    68  	parentID := blk.Parent()
    69  	children := t.nodes[parentID]
    70  	blkID := blk.ID()
    71  	originalBlk, exists := children[blkID]
    72  	return originalBlk, exists
    73  }
    74  
    75  func (t *tree) Accept(ctx context.Context, blk snowman.Block) error {
    76  	// accept the provided block
    77  	if err := blk.Accept(ctx); err != nil {
    78  		return err
    79  	}
    80  
    81  	// get the siblings of the block
    82  	parentID := blk.Parent()
    83  	children := t.nodes[parentID]
    84  	blkID := blk.ID()
    85  	delete(children, blkID)
    86  	delete(t.nodes, parentID)
    87  
    88  	// mark the siblings of the accepted block as rejectable
    89  	childrenToReject := maps.Values(children)
    90  
    91  	// reject all the rejectable blocks
    92  	for len(childrenToReject) > 0 {
    93  		i := len(childrenToReject) - 1
    94  		child := childrenToReject[i]
    95  		childrenToReject = childrenToReject[:i]
    96  
    97  		// reject the block
    98  		if err := child.Reject(ctx); err != nil {
    99  			return err
   100  		}
   101  
   102  		// mark the progeny of this block as being rejectable
   103  		blkID := child.ID()
   104  		children := t.nodes[blkID]
   105  		childrenToReject = append(childrenToReject, maps.Values(children)...)
   106  		delete(t.nodes, blkID)
   107  	}
   108  	return nil
   109  }