protocol/tree_node.go

   1 package protocol
   2
   3 import (
   4         "errors"
   5         "fmt"
   6
   7         "github.com/bytom/bytom/protocol/bc"
   8         "github.com/bytom/bytom/protocol/state"
   9 )
  10
  11 type treeNode struct {
  12         checkpoint *state.Checkpoint
  13         children   []*treeNode
  14 }
  15
  16 func makeTree(root *state.Checkpoint, successors []*state.Checkpoint) *treeNode {
  17         parentToSuccessors := make(map[bc.Hash][]*state.Checkpoint)
  18         for _, successor := range successors {
  19                 parentToSuccessors[successor.ParentHash] = append(parentToSuccessors[successor.ParentHash], successor)
  20         }
  21
  22         rootNode := &treeNode{checkpoint: root}
  23         nodes := []*treeNode{rootNode}
  24         for len(nodes) != 0 {
  25                 node := nodes[0]
  26                 for _, successor := range parentToSuccessors[node.checkpoint.Hash] {
  27                         child := &treeNode{checkpoint: successor}
  28                         successor.Parent = node.checkpoint
  29                         node.children = append(node.children, child)
  30                         nodes = append(nodes, child)
  31                 }
  32                 nodes = nodes[1:]
  33         }
  34         return rootNode
  35 }
  36
  37 func (t *treeNode) addChild(child *treeNode) {
  38         for i, n := range t.children {
  39                 if n.checkpoint.Hash == child.checkpoint.Hash {
  40                         t.children[i] = child
  41                         return
  42                 }
  43         }
  44         t.children = append(t.children, child)
  45 }
  46
  47 func (t *treeNode) nodeByHash(blockHash bc.Hash) (*treeNode, error) {
  48         if c := t.findOnlyOne(func(c *state.Checkpoint) bool {
  49                 return c.Hash == blockHash
  50         }); c != nil {
  51                 return c, nil
  52         }
  53
  54         return nil, errors.New(fmt.Sprintf("fail to find checkpoint of hash:%s", blockHash.String()))
  55 }
  56
  57 func (t *treeNode) findOnlyOne(predicate func(*state.Checkpoint) bool) *treeNode {
  58         if predicate(t.checkpoint) {
  59                 return t
  60         }
  61
  62         for _, child := range t.children {
  63                 if node := child.findOnlyOne(predicate); node != nil {
  64                         return node
  65                 }
  66         }
  67         return nil
  68 }