1 files changed, 184 insertions, 67 deletions
diff --git a/light/trie.go b/light/trie.go
index 2988a16cf..7502b6e5d 100644
--- a/light/trie.go
+++ b/light/trie.go
@@ -18,99 +18,216 @@ package light
 
 import (
 	"context"
+	"fmt"
 
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
-	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/trie"
 )
 
-// LightTrie is an ODR-capable wrapper around trie.SecureTrie
-type LightTrie struct {
-	trie *trie.SecureTrie
+func NewState(ctx context.Context, head *types.Header, odr OdrBackend) *state.StateDB {
+	state, _ := state.New(head.Root, NewStateDatabase(ctx, head, odr))
+	return state
+}
+
+func NewStateDatabase(ctx context.Context, head *types.Header, odr OdrBackend) state.Database {
+	return &odrDatabase{ctx, StateTrieID(head), odr}
+}
+
+type odrDatabase struct {
+	ctx     context.Context
+	id      *TrieID
+	backend OdrBackend
+}
+
+func (db *odrDatabase) OpenTrie(root common.Hash) (state.Trie, error) {
+	return &odrTrie{db: db, id: db.id}, nil
+}
+
+func (db *odrDatabase) OpenStorageTrie(addrHash, root common.Hash) (state.Trie, error) {
+	return &odrTrie{db: db, id: StorageTrieID(db.id, addrHash, root)}, nil
+}
+
+func (db *odrDatabase) CopyTrie(t state.Trie) state.Trie {
+	switch t := t.(type) {
+	case *odrTrie:
+		cpy := &odrTrie{db: t.db, id: t.id}
+		if t.trie != nil {
+			cpytrie := *t.trie
+			cpy.trie = &cpytrie
+		}
+		return cpy
+	default:
+		panic(fmt.Errorf("unknown trie type %T", t))
+	}
+}
+
+func (db *odrDatabase) ContractCode(addrHash, codeHash common.Hash) ([]byte, error) {
+	if codeHash == sha3_nil {
+		return nil, nil
+	}
+	if code, err := db.backend.Database().Get(codeHash[:]); err == nil {
+		return code, nil
+	}
+	id := *db.id
+	id.AccKey = addrHash[:]
+	req := &CodeRequest{Id: &id, Hash: codeHash}
+	err := db.backend.Retrieve(db.ctx, req)
+	return req.Data, err
+}
+
+func (db *odrDatabase) ContractCodeSize(addrHash, codeHash common.Hash) (int, error) {
+	code, err := db.ContractCode(addrHash, codeHash)
+	return len(code), err
+}
+
+type odrTrie struct {
+	db   *odrDatabase
 	id   *TrieID
-	odr  OdrBackend
-	db   ethdb.Database
-}
-
-// NewLightTrie creates a new LightTrie instance. It doesn't instantly try to
-// access the db or network and retrieve the root node, it only initializes its
-// encapsulated SecureTrie at the first actual operation.
-func NewLightTrie(id *TrieID, odr OdrBackend, useFakeMap bool) *LightTrie {
-	return &LightTrie{
-		// SecureTrie is initialized before first request
-		id:  id,
-		odr: odr,
-		db:  odr.Database(),
+	trie *trie.Trie
+}
+
+func (t *odrTrie) TryGet(key []byte) ([]byte, error) {
+	key = crypto.Keccak256(key)
+	var res []byte
+	err := t.do(key, func() (err error) {
+		res, err = t.trie.TryGet(key)
+		return err
+	})
+	return res, err
+}
+
+func (t *odrTrie) TryUpdate(key, value []byte) error {
+	key = crypto.Keccak256(key)
+	return t.do(key, func() error {
+		return t.trie.TryDelete(key)
+	})
+}
+
+func (t *odrTrie) TryDelete(key []byte) error {
+	key = crypto.Keccak256(key)
+	return t.do(key, func() error {
+		return t.trie.TryDelete(key)
+	})
+}
+
+func (t *odrTrie) CommitTo(db trie.DatabaseWriter) (common.Hash, error) {
+	if t.trie == nil {
+		return t.id.Root, nil
+	}
+	return t.trie.CommitTo(db)
+}
+
+func (t *odrTrie) Hash() common.Hash {
+	if t.trie == nil {
+		return t.id.Root
 	}
+	return t.trie.Hash()
+}
+
+func (t *odrTrie) NodeIterator(startkey []byte) trie.NodeIterator {
+	return newNodeIterator(t, startkey)
 }
 
-// retrieveKey retrieves a single key, returns true and stores nodes in local
-// database if successful
-func (t *LightTrie) retrieveKey(ctx context.Context, key []byte) bool {
-	r := &TrieRequest{Id: t.id, Key: crypto.Keccak256(key)}
-	return t.odr.Retrieve(ctx, r) == nil
+func (t *odrTrie) GetKey(sha []byte) []byte {
+	return nil
 }
 
 // do tries and retries to execute a function until it returns with no error or
 // an error type other than MissingNodeError
-func (t *LightTrie) do(ctx context.Context, key []byte, fn func() error) error {
-	err := fn()
-	for err != nil {
+func (t *odrTrie) do(key []byte, fn func() error) error {
+	for {
+		var err error
+		if t.trie == nil {
+			t.trie, err = trie.New(t.id.Root, t.db.backend.Database())
+		}
+		if err == nil {
+			err = fn()
+		}
 		if _, ok := err.(*trie.MissingNodeError); !ok {
 			return err
 		}
-		if !t.retrieveKey(ctx, key) {
-			break
+		r := &TrieRequest{Id: t.id, Key: key}
+		if err := t.db.backend.Retrieve(t.db.ctx, r); err != nil {
+			return fmt.Errorf("can't fetch trie key %x: %v", key, err)
 		}
-		err = fn()
 	}
-	return err
 }
 
-// Get returns the value for key stored in the trie.
-// The value bytes must not be modified by the caller.
-func (t *LightTrie) Get(ctx context.Context, key []byte) (res []byte, err error) {
-	err = t.do(ctx, key, func() (err error) {
-		if t.trie == nil {
-			t.trie, err = trie.NewSecure(t.id.Root, t.db, 0)
-		}
-		if err == nil {
-			res, err = t.trie.TryGet(key)
-		}
-		return
+type nodeIterator struct {
+	trie.NodeIterator
+	t   *odrTrie
+	err error
+}
+
+func newNodeIterator(t *odrTrie, startkey []byte) trie.NodeIterator {
+	it := &nodeIterator{t: t}
+	// Open the actual non-ODR trie if that hasn't happened yet.
+	if t.trie == nil {
+		it.do(func() error {
+			t, err := trie.New(t.id.Root, t.db.backend.Database())
+			if err == nil {
+				it.t.trie = t
+			}
+			return err
+		})
+	}
+	it.do(func() error {
+		it.NodeIterator = it.t.trie.NodeIterator(startkey)
+		return it.NodeIterator.Error()
 	})
-	return
+	return it
 }
 
-// Update associates key with value in the trie. Subsequent calls to
-// Get will return value. If value has length zero, any existing value
-// is deleted from the trie and calls to Get will return nil.
-//
-// The value bytes must not be modified by the caller while they are
-// stored in the trie.
-func (t *LightTrie) Update(ctx context.Context, key, value []byte) (err error) {
-	err = t.do(ctx, key, func() (err error) {
-		if t.trie == nil {
-			t.trie, err = trie.NewSecure(t.id.Root, t.db, 0)
-		}
-		if err == nil {
-			err = t.trie.TryUpdate(key, value)
-		}
-		return
+func (it *nodeIterator) Next(descend bool) bool {
+	var ok bool
+	it.do(func() error {
+		ok = it.NodeIterator.Next(descend)
+		return it.NodeIterator.Error()
 	})
-	return
+	return ok
 }
 
-// Delete removes any existing value for key from the trie.
-func (t *LightTrie) Delete(ctx context.Context, key []byte) (err error) {
-	err = t.do(ctx, key, func() (err error) {
-		if t.trie == nil {
-			t.trie, err = trie.NewSecure(t.id.Root, t.db, 0)
+// do runs fn and attempts to fill in missing nodes by retrieving.
+func (it *nodeIterator) do(fn func() error) {
+	var lasthash common.Hash
+	for {
+		it.err = fn()
+		missing, ok := it.err.(*trie.MissingNodeError)
+		if !ok {
+			return
 		}
-		if err == nil {
-			err = t.trie.TryDelete(key)
+		if missing.NodeHash == lasthash {
+			it.err = fmt.Errorf("retrieve loop for trie node %x", missing.NodeHash)
+			return
 		}
-		return
-	})
-	return
+		lasthash = missing.NodeHash
+		r := &TrieRequest{Id: it.t.id, Key: nibblesToKey(missing.Path)}
+		if it.err = it.t.db.backend.Retrieve(it.t.db.ctx, r); it.err != nil {
+			return
+		}
+	}
+}
+
+func (it *nodeIterator) Error() error {
+	if it.err != nil {
+		return it.err
+	}
+	return it.NodeIterator.Error()
+}
+
+func nibblesToKey(nib []byte) []byte {
+	if len(nib) > 0 && nib[len(nib)-1] == 0x10 {
+		nib = nib[:len(nib)-1] // drop terminator
+	}
+	if len(nib)&1 == 1 {
+		nib = append(nib, 0) // make even
+	}
+	key := make([]byte, len(nib)/2)
+	for bi, ni := 0, 0; ni < len(nib); bi, ni = bi+1, ni+2 {
+		key[bi] = nib[ni]<<4 | nib[ni+1]
+	}
+	return key
 }