github.com/chain5j/chain5j-pkg@v1.0.7/collection/trees/tree/trie.go (about)

     1  // Copyright 2014 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  // Package trie implements Merkle Patricia Tries.
    18  package tree
    19  
    20  import (
    21  	"bytes"
    22  	"fmt"
    23  
    24  	"github.com/chain5j/chain5j-pkg/crypto/keccak"
    25  	"github.com/chain5j/chain5j-pkg/types"
    26  )
    27  
    28  var (
    29  	// emptyRoot is the known root hash of an empty trie.
    30  	emptyRoot = types.HexToHash("56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421")
    31  
    32  	// emptyState is the known hash of an empty state trie entry.
    33  	emptyState = keccak.Keccak256Hash(nil)
    34  )
    35  
    36  // LeafCallback is a callback type invoked when a trie operation reaches a leaf
    37  // node. It's used by state sync and commit to allow handling external references
    38  // between account and storage tries.
    39  type LeafCallback func(leaf []byte, parent types.Hash) error
    40  
    41  // Trie is a Merkle Patricia Trie.
    42  // The zero value is an empty trie with no database.
    43  // Use New to create a trie that sits on top of a database.
    44  //
    45  // Trie is not safe for concurrent use.
    46  type Trie struct {
    47  	db   *Database
    48  	root node
    49  }
    50  
    51  // newFlag returns the cache flag value for a newly created node.
    52  func (t *Trie) newFlag() nodeFlag {
    53  	return nodeFlag{dirty: true}
    54  }
    55  
    56  // New creates a trie with an existing root node from db.
    57  //
    58  // If root is the zero hash or the sha3 hash of an empty string, the
    59  // trie is initially empty and does not require a database. Otherwise,
    60  // New will panic if db is nil and returns a MissingNodeError if root does
    61  // not exist in the database. Accessing the trie loads nodes from db on demand.
    62  func New(root types.Hash, db *Database) (*Trie, error) {
    63  	return NewWithGenesis(root, db, types.Hash{})
    64  }
    65  func NewWithGenesis(root types.Hash, db *Database, genesisRoot types.Hash) (*Trie, error) {
    66  	if db == nil {
    67  		panic("trie.New called without a database")
    68  	}
    69  	trie := &Trie{
    70  		db: db,
    71  	}
    72  	if root != (types.Hash{}) && root != emptyRoot && root != genesisRoot {
    73  		rootnode, err := trie.resolveHash(root[:], nil)
    74  		if err != nil {
    75  			return nil, err
    76  		}
    77  		trie.root = rootnode
    78  	}
    79  	return trie, nil
    80  }
    81  
    82  // NodeIterator returns an iterator that returns nodes of the trie. Iteration starts at
    83  // the key after the given start key.
    84  func (t *Trie) NodeIterator(start []byte) NodeIterator {
    85  	return newNodeIterator(t, start)
    86  }
    87  
    88  // Get returns the value for key stored in the trie.
    89  // The value bytes must not be modified by the caller.
    90  func (t *Trie) Get(key []byte) []byte {
    91  	res, err := t.TryGet(key)
    92  	if err != nil {
    93  		logger().Error(fmt.Sprintf("Unhandled trie error: %v", err))
    94  	}
    95  	return res
    96  }
    97  
    98  // TryGet returns the value for key stored in the trie.
    99  // The value bytes must not be modified by the caller.
   100  // If a node was not found in the database, a MissingNodeError is returned.
   101  func (t *Trie) TryGet(key []byte) ([]byte, error) {
   102  	key = keybytesToHex(key)
   103  	value, newroot, didResolve, err := t.tryGet(t.root, key, 0)
   104  	if err == nil && didResolve {
   105  		t.root = newroot
   106  	}
   107  	return value, err
   108  }
   109  
   110  func (t *Trie) tryGet(origNode node, key []byte, pos int) (value []byte, newnode node, didResolve bool, err error) {
   111  	switch n := (origNode).(type) {
   112  	case nil:
   113  		return nil, nil, false, nil
   114  	case valueNode:
   115  		return n, n, false, nil
   116  	case *shortNode:
   117  		if len(key)-pos < len(n.Key) || !bytes.Equal(n.Key, key[pos:pos+len(n.Key)]) {
   118  			// key not found in trie
   119  			return nil, n, false, nil
   120  		}
   121  		value, newnode, didResolve, err = t.tryGet(n.Val, key, pos+len(n.Key))
   122  		if err == nil && didResolve {
   123  			n = n.copy()
   124  			n.Val = newnode
   125  		}
   126  		return value, n, didResolve, err
   127  	case *fullNode:
   128  		value, newnode, didResolve, err = t.tryGet(n.Children[key[pos]], key, pos+1)
   129  		if err == nil && didResolve {
   130  			n = n.copy()
   131  			n.Children[key[pos]] = newnode
   132  		}
   133  		return value, n, didResolve, err
   134  	case hashNode:
   135  		child, err := t.resolveHash(n, key[:pos])
   136  		if err != nil {
   137  			return nil, n, true, err
   138  		}
   139  		value, newnode, _, err := t.tryGet(child, key, pos)
   140  		return value, newnode, true, err
   141  	default:
   142  		panic(fmt.Sprintf("%T: invalid node: %v", origNode, origNode))
   143  	}
   144  }
   145  
   146  // Update associates key with value in the trie. Subsequent calls to
   147  // Get will return value. If value has length zero, any existing value
   148  // is deleted from the trie and calls to Get will return nil.
   149  //
   150  // The value bytes must not be modified by the caller while they are
   151  // stored in the trie.
   152  func (t *Trie) Update(key, value []byte) {
   153  	if err := t.TryUpdate(key, value); err != nil {
   154  		logger().Error(fmt.Sprintf("Unhandled trie error: %v", err))
   155  	}
   156  }
   157  
   158  // TryUpdate associates key with value in the trie. Subsequent calls to
   159  // Get will return value. If value has length zero, any existing value
   160  // is deleted from the trie and calls to Get will return nil.
   161  //
   162  // The value bytes must not be modified by the caller while they are
   163  // stored in the trie.
   164  //
   165  // If a node was not found in the database, a MissingNodeError is returned.
   166  func (t *Trie) TryUpdate(key, value []byte) error {
   167  	k := keybytesToHex(key)
   168  	if len(value) != 0 {
   169  		_, n, err := t.insert(t.root, nil, k, valueNode(value))
   170  		if err != nil {
   171  			return err
   172  		}
   173  		t.root = n
   174  	} else {
   175  		_, n, err := t.delete(t.root, nil, k)
   176  		if err != nil {
   177  			return err
   178  		}
   179  		t.root = n
   180  	}
   181  	return nil
   182  }
   183  
   184  func (t *Trie) insert(n node, prefix, key []byte, value node) (bool, node, error) {
   185  	if len(key) == 0 {
   186  		if v, ok := n.(valueNode); ok {
   187  			return !bytes.Equal(v, value.(valueNode)), value, nil
   188  		}
   189  		return true, value, nil
   190  	}
   191  	switch n := n.(type) {
   192  	case *shortNode:
   193  		matchlen := prefixLen(key, n.Key)
   194  		// If the whole key matches, keep this short node as is
   195  		// and only update the value.
   196  		if matchlen == len(n.Key) {
   197  			dirty, nn, err := t.insert(n.Val, append(prefix, key[:matchlen]...), key[matchlen:], value)
   198  			if !dirty || err != nil {
   199  				return false, n, err
   200  			}
   201  			return true, &shortNode{n.Key, nn, t.newFlag()}, nil
   202  		}
   203  		// Otherwise branch out at the index where they differ.
   204  		branch := &fullNode{flags: t.newFlag()}
   205  		var err error
   206  		_, branch.Children[n.Key[matchlen]], err = t.insert(nil, append(prefix, n.Key[:matchlen+1]...), n.Key[matchlen+1:], n.Val)
   207  		if err != nil {
   208  			return false, nil, err
   209  		}
   210  		_, branch.Children[key[matchlen]], err = t.insert(nil, append(prefix, key[:matchlen+1]...), key[matchlen+1:], value)
   211  		if err != nil {
   212  			return false, nil, err
   213  		}
   214  		// Replace this shortNode with the branch if it occurs at index 0.
   215  		if matchlen == 0 {
   216  			return true, branch, nil
   217  		}
   218  		// Otherwise, replace it with a short node leading up to the branch.
   219  		return true, &shortNode{key[:matchlen], branch, t.newFlag()}, nil
   220  
   221  	case *fullNode:
   222  		dirty, nn, err := t.insert(n.Children[key[0]], append(prefix, key[0]), key[1:], value)
   223  		if !dirty || err != nil {
   224  			return false, n, err
   225  		}
   226  		n = n.copy()
   227  		n.flags = t.newFlag()
   228  		n.Children[key[0]] = nn
   229  		return true, n, nil
   230  
   231  	case nil:
   232  		return true, &shortNode{key, value, t.newFlag()}, nil
   233  
   234  	case hashNode:
   235  		// We've hit a part of the trie that isn't loaded yet. Load
   236  		// the node and insert into it. This leaves all child nodes on
   237  		// the path to the value in the trie.
   238  		rn, err := t.resolveHash(n, prefix)
   239  		if err != nil {
   240  			return false, nil, err
   241  		}
   242  		dirty, nn, err := t.insert(rn, prefix, key, value)
   243  		if !dirty || err != nil {
   244  			return false, rn, err
   245  		}
   246  		return true, nn, nil
   247  
   248  	default:
   249  		panic(fmt.Sprintf("%T: invalid node: %v", n, n))
   250  	}
   251  }
   252  
   253  // Delete removes any existing value for key from the trie.
   254  func (t *Trie) Delete(key []byte) {
   255  	if err := t.TryDelete(key); err != nil {
   256  		logger().Error(fmt.Sprintf("Unhandled trie error: %v", err))
   257  	}
   258  }
   259  
   260  // TryDelete removes any existing value for key from the trie.
   261  // If a node was not found in the database, a MissingNodeError is returned.
   262  func (t *Trie) TryDelete(key []byte) error {
   263  	k := keybytesToHex(key)
   264  	_, n, err := t.delete(t.root, nil, k)
   265  	if err != nil {
   266  		return err
   267  	}
   268  	t.root = n
   269  	return nil
   270  }
   271  
   272  // delete returns the new root of the trie with key deleted.
   273  // It reduces the trie to minimal form by simplifying
   274  // nodes on the way up after deleting recursively.
   275  func (t *Trie) delete(n node, prefix, key []byte) (bool, node, error) {
   276  	switch n := n.(type) {
   277  	case *shortNode:
   278  		matchlen := prefixLen(key, n.Key)
   279  		if matchlen < len(n.Key) {
   280  			return false, n, nil // don't replace n on mismatch
   281  		}
   282  		if matchlen == len(key) {
   283  			return true, nil, nil // remove n entirely for whole matches
   284  		}
   285  		// The key is longer than n.Key. Remove the remaining suffix
   286  		// from the subtrie. Child can never be nil here since the
   287  		// subtrie must contain at least two other values with keys
   288  		// longer than n.Key.
   289  		dirty, child, err := t.delete(n.Val, append(prefix, key[:len(n.Key)]...), key[len(n.Key):])
   290  		if !dirty || err != nil {
   291  			return false, n, err
   292  		}
   293  		switch child := child.(type) {
   294  		case *shortNode:
   295  			// Deleting from the subtrie reduced it to another
   296  			// short node. Merge the nodes to avoid creating a
   297  			// shortNode{..., shortNode{...}}. Use concat (which
   298  			// always creates a new slice) instead of append to
   299  			// avoid modifying n.Key since it might be shared with
   300  			// other nodes.
   301  			return true, &shortNode{concat(n.Key, child.Key...), child.Val, t.newFlag()}, nil
   302  		default:
   303  			return true, &shortNode{n.Key, child, t.newFlag()}, nil
   304  		}
   305  
   306  	case *fullNode:
   307  		dirty, nn, err := t.delete(n.Children[key[0]], append(prefix, key[0]), key[1:])
   308  		if !dirty || err != nil {
   309  			return false, n, err
   310  		}
   311  		n = n.copy()
   312  		n.flags = t.newFlag()
   313  		n.Children[key[0]] = nn
   314  
   315  		// Check how many non-nil entries are left after deleting and
   316  		// reduce the full node to a short node if only one entry is
   317  		// left. Since n must've contained at least two children
   318  		// before deletion (otherwise it would not be a full node) n
   319  		// can never be reduced to nil.
   320  		//
   321  		// When the loop is done, pos contains the index of the single
   322  		// value that is left in n or -2 if n contains at least two
   323  		// values.
   324  		pos := -1
   325  		for i, cld := range &n.Children {
   326  			if cld != nil {
   327  				if pos == -1 {
   328  					pos = i
   329  				} else {
   330  					pos = -2
   331  					break
   332  				}
   333  			}
   334  		}
   335  		if pos >= 0 {
   336  			if pos != 16 {
   337  				// If the remaining entry is a short node, it replaces
   338  				// n and its key gets the missing nibble tacked to the
   339  				// front. This avoids creating an invalid
   340  				// shortNode{..., shortNode{...}}.  Since the entry
   341  				// might not be loaded yet, resolve it just for this
   342  				// check.
   343  				cnode, err := t.resolve(n.Children[pos], prefix)
   344  				if err != nil {
   345  					return false, nil, err
   346  				}
   347  				if cnode, ok := cnode.(*shortNode); ok {
   348  					k := append([]byte{byte(pos)}, cnode.Key...)
   349  					return true, &shortNode{k, cnode.Val, t.newFlag()}, nil
   350  				}
   351  			}
   352  			// Otherwise, n is replaced by a one-nibble short node
   353  			// containing the child.
   354  			return true, &shortNode{[]byte{byte(pos)}, n.Children[pos], t.newFlag()}, nil
   355  		}
   356  		// n still contains at least two values and cannot be reduced.
   357  		return true, n, nil
   358  
   359  	case valueNode:
   360  		return true, nil, nil
   361  
   362  	case nil:
   363  		return false, nil, nil
   364  
   365  	case hashNode:
   366  		// We've hit a part of the trie that isn't loaded yet. Load
   367  		// the node and delete from it. This leaves all child nodes on
   368  		// the path to the value in the trie.
   369  		rn, err := t.resolveHash(n, prefix)
   370  		if err != nil {
   371  			return false, nil, err
   372  		}
   373  		dirty, nn, err := t.delete(rn, prefix, key)
   374  		if !dirty || err != nil {
   375  			return false, rn, err
   376  		}
   377  		return true, nn, nil
   378  
   379  	default:
   380  		panic(fmt.Sprintf("%T: invalid node: %v (%v)", n, n, key))
   381  	}
   382  }
   383  
   384  func concat(s1 []byte, s2 ...byte) []byte {
   385  	r := make([]byte, len(s1)+len(s2))
   386  	copy(r, s1)
   387  	copy(r[len(s1):], s2)
   388  	return r
   389  }
   390  
   391  func (t *Trie) resolve(n node, prefix []byte) (node, error) {
   392  	if n, ok := n.(hashNode); ok {
   393  		return t.resolveHash(n, prefix)
   394  	}
   395  	return n, nil
   396  }
   397  
   398  func (t *Trie) resolveHash(n hashNode, prefix []byte) (node, error) {
   399  	hash := types.BytesToHash(n)
   400  	if node := t.db.node(hash); node != nil {
   401  		return node, nil
   402  	}
   403  	return nil, &MissingNodeError{NodeHash: hash, Path: prefix}
   404  }
   405  
   406  // Hash returns the root hash of the trie. It does not write to the
   407  // database and can be used even if the trie doesn't have one.
   408  func (t *Trie) Hash() types.Hash {
   409  	hash, cached, _ := t.hashRoot(nil, nil)
   410  	t.root = cached
   411  	return types.BytesToHash(hash.(hashNode))
   412  }
   413  
   414  // Commit writes all nodes to the trie's memory database, tracking the internal
   415  // and external (for account tries) references.
   416  func (t *Trie) Commit(onleaf LeafCallback) (root types.Hash, err error) {
   417  	if t.db == nil {
   418  		panic("commit called on trie with nil database")
   419  	}
   420  	hash, cached, err := t.hashRoot(t.db, onleaf)
   421  	if err != nil {
   422  		return types.Hash{}, err
   423  	}
   424  	t.root = cached
   425  	return types.BytesToHash(hash.(hashNode)), nil
   426  }
   427  
   428  func (t *Trie) hashRoot(db *Database, onleaf LeafCallback) (node, node, error) {
   429  	if t.root == nil {
   430  		return hashNode(emptyRoot.Bytes()), nil, nil
   431  	}
   432  	h := newHasher(onleaf)
   433  	defer returnHasherToPool(h)
   434  	return h.hash(t.root, db, true)
   435  }