wip. zktrie part2: add zktrie; allow switch trie type by config; storage witness generator

.gitattributes

@@ -1,3 +1,4 @@
 # Auto detect text files and perform LF normalization
 * text=auto
 *.sol linguist-language=Solidity
+*.go text eol=lf

cmd/geth/genesis_test.go

@@ -121,7 +121,7 @@ func TestCustomGenesis(t *testing.T) {
 		runGeth(t, "--datadir", datadir, "init", json).WaitExit()
 
 		// Query the custom genesis block
-		geth := runGeth(t, "--networkid", "1337", "--syncmode=full", "--cache", "16",
+		geth := runGeth(t, "--networkid", "1337", "--syncmode=full", "--snapshot=false", "--cache", "16",
 			"--datadir", datadir, "--maxpeers", "0", "--port", "0",
 			"--nodiscover", "--nat", "none", "--ipcdisable",
 			"--exec", tt.query, "console")

common/types.go

@@ -229,6 +229,12 @@ func IsHexAddress(s string) bool {
 // Bytes gets the string representation of the underlying address.
 func (a Address) Bytes() []byte { return a[:] }
 
+func (a Address) Bytes32() []byte {
+	ret := make([]byte, 32)
+	copy(ret, a.Bytes())
+	return ret
+}
+
 // Hash converts an address to a hash by left-padding it with zeros.
 func (a Address) Hash() Hash { return BytesToHash(a[:]) }
 

core/blockchain.go

@@ -246,6 +246,7 @@ func NewBlockChain(db ethdb.Database, cacheConfig *CacheConfig, chainConfig *par
 			Cache:     cacheConfig.TrieCleanLimit,
 			Journal:   cacheConfig.TrieCleanJournal,
 			Preimages: cacheConfig.Preimages,
+			Zktrie:    chainConfig.Zktrie,
 		}),
 		quit:             make(chan struct{}),
 		chainmu:          syncx.NewClosableMutex(),
@@ -1193,7 +1194,7 @@ func (bc *BlockChain) writeKnownBlock(block *types.Block) error {
 }
 
 // WriteBlockWithState writes the block and all associated state to the database.
-func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.Receipt, logs []*types.Log, evmTraces []*types.ExecutionResult, state *state.StateDB, emitHeadEvent bool) (status WriteStatus, err error) {
+func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.Receipt, logs []*types.Log, evmTraces *types.EvmTxTraces, state *state.StateDB, emitHeadEvent bool) (status WriteStatus, err error) {
 	if !bc.chainmu.TryLock() {
 		return NonStatTy, errInsertionInterrupted
 	}
@@ -1203,7 +1204,7 @@ func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.
 
 // writeBlockWithState writes the block and all associated state to the database,
 // but is expects the chain mutex to be held.
-func (bc *BlockChain) writeBlockWithState(block *types.Block, receipts []*types.Receipt, logs []*types.Log, evmTraces []*types.ExecutionResult, state *state.StateDB, emitHeadEvent bool) (status WriteStatus, err error) {
+func (bc *BlockChain) writeBlockWithState(block *types.Block, receipts []*types.Receipt, logs []*types.Log, evmTraces *types.EvmTxTraces, state *state.StateDB, emitHeadEvent bool) (status WriteStatus, err error) {
 	if bc.insertStopped() {
 		return NonStatTy, errInsertionInterrupted
 	}
@@ -1351,57 +1352,48 @@ func (bc *BlockChain) writeBlockWithState(block *types.Block, receipts []*types.
 }
 
 // Fill blockResult content
-func (bc *BlockChain) writeBlockResult(state *state.StateDB, block *types.Block, evmTraces []*types.ExecutionResult) *types.BlockResult {
+func (bc *BlockChain) writeBlockResult(state *state.StateDB, block *types.Block, evmTraces *types.EvmTxTraces) *types.BlockResult {
 	blockResult := &types.BlockResult{
-		ExecutionResults: evmTraces,
+		ExecutionResults: evmTraces.TxResults,
+		StorageTrace:     evmTraces.Storage,
 	}
-	coinbase := types.AccountProofWrapper{
+	coinbase := types.AccountWrapper{
 		Address:  block.Coinbase(),
 		Nonce:    state.GetNonce(block.Coinbase()),
 		Balance:  (*hexutil.Big)(state.GetBalance(block.Coinbase())),
 		CodeHash: state.GetCodeHash(block.Coinbase()),
 	}
-	// Get coinbase address's account proof.
-	proof, err := state.GetProof(block.Coinbase())
-	if err != nil {
-		log.Error("Failed to get proof", "blockNumber", block.NumberU64(), "address", block.Coinbase().String(), "err", err)
-	} else {
-		coinbase.Proof = make([]string, len(proof))
-		for i := range proof {
-			coinbase.Proof[i] = hexutil.Encode(proof[i])
-		}
-	}
 
 	blockResult.BlockTrace = types.NewTraceBlock(bc.chainConfig, block, &coinbase)
+	blockResult.StorageTrace.RootAfter = state.GetRootHash()
 	for i, tx := range block.Transactions() {
 		evmTrace := blockResult.ExecutionResults[i]
-
-		from := evmTrace.From.Address
-		// Get proof
-		proof, err := state.GetProof(from)
-		if err != nil {
-			log.Error("Failed to get proof", "blockNumber", block.NumberU64(), "address", from.String(), "err", err)
-		} else {
-			evmTrace.From.Proof = make([]string, len(proof))
-			for i := range proof {
-				evmTrace.From.Proof[i] = hexutil.Encode(proof[i])
-			}
-		}
-
-		if evmTrace.To != nil {
-			to := evmTrace.To.Address
+		/*
+			from := evmTrace.From.Address
 			// Get proof
-			proof, err = state.GetProof(to)
+			proof, err := state.GetProof(from)
 			if err != nil {
-				log.Error("Failed to get proof", "blockNumber", block.NumberU64(), "address", to.String(), "err", err)
+				log.Error("Failed to get proof", "blockNumber", block.NumberU64(), "address", from.String(), "err", err)
 			} else {
-				evmTrace.To.Proof = make([]string, len(proof))
+				evmTrace.From.Proof = make([]string, len(proof))
 				for i := range proof {
-					evmTrace.To.Proof[i] = hexutil.Encode(proof[i])
+					evmTrace.From.Proof[i] = hexutil.Encode(proof[i])
 				}
 			}
-		}
 
+			if evmTrace.To != nil {
+				to := evmTrace.To.Address
+				// Get proof
+				proof, err = state.GetProof(to)
+				if err != nil {
+					log.Error("Failed to get proof", "blockNumber", block.NumberU64(), "address", to.String(), "err", err)
+				} else {
+					evmTrace.To.Proof = make([]string, len(proof))
+					for i := range proof {
+						evmTrace.To.Proof[i] = hexutil.Encode(proof[i])
+					}
+				}
+		*/
 		// Contract is called
 		if len(tx.Data()) != 0 && tx.To() != nil {
 			evmTrace.ByteCode = hexutil.Encode(state.GetCode(*tx.To()))
@@ -1412,6 +1404,7 @@ func (bc *BlockChain) writeBlockResult(state *state.StateDB, block *types.Block,
 			evmTrace.ByteCode = hexutil.Encode(tx.Data())
 		}
 	}
+
 	return blockResult
 }
 

core/genesis.go

@@ -180,7 +180,18 @@ func SetupGenesisBlockWithOverride(db ethdb.Database, genesis *Genesis, override
 	// We have the genesis block in database(perhaps in ancient database)
 	// but the corresponding state is missing.
 	header := rawdb.ReadHeader(db, stored, 0)
-	if _, err := state.New(header.Root, state.NewDatabaseWithConfig(db, nil), nil); err != nil {
+
+	var trieCfg *trie.Config
+
+	if genesis == nil {
+		storedcfg := rawdb.ReadChainConfig(db, stored)
+		if storedcfg == nil {
+			panic("this should never be reached: if genesis is nil, the config is already present or 'geth init' is being called which created it (in the code above, which means genesis != nil)")
+		}
+		trieCfg = &trie.Config{Zktrie: storedcfg.Zktrie}
+	}
+
+	if _, err := state.New(header.Root, state.NewDatabaseWithConfig(db, trieCfg), nil); err != nil {
 		if genesis == nil {
 			genesis = DefaultGenesisBlock()
 		}
@@ -261,7 +272,11 @@ func (g *Genesis) ToBlock(db ethdb.Database) *types.Block {
 	if db == nil {
 		db = rawdb.NewMemoryDatabase()
 	}
-	statedb, err := state.New(common.Hash{}, state.NewDatabase(db), nil)
+	var trieCfg *trie.Config
+	if g.Config != nil {
+		trieCfg = &trie.Config{Zktrie: g.Config.Zktrie}
+	}
+	statedb, err := state.New(common.Hash{}, state.NewDatabaseWithConfig(db, trieCfg), nil)
 	if err != nil {
 		panic(err)
 	}
@@ -301,8 +316,14 @@ func (g *Genesis) ToBlock(db ethdb.Database) *types.Block {
 			head.BaseFee = new(big.Int).SetUint64(params.InitialBaseFee)
 		}
 	}
-	statedb.Commit(false)
-	statedb.Database().TrieDB().Commit(root, true, nil)
+	_, err = statedb.Commit(false)
+	if err != nil {
+		panic(err)
+	}
+	err = statedb.Database().TrieDB().Commit(root, true, nil)
+	if err != nil {
+		panic(err)
+	}
 
 	return types.NewBlock(head, nil, nil, nil, trie.NewStackTrie(nil))
 }
@@ -441,7 +462,8 @@ func DeveloperGenesisBlock(period uint64, gasLimit uint64, faucet common.Address
 			common.BytesToAddress([]byte{7}): {Balance: big.NewInt(1)}, // ECScalarMul
 			common.BytesToAddress([]byte{8}): {Balance: big.NewInt(1)}, // ECPairing
 			common.BytesToAddress([]byte{9}): {Balance: big.NewInt(1)}, // BLAKE2b
-			faucet:                           {Balance: new(big.Int).Sub(new(big.Int).Lsh(big.NewInt(1), 256), big.NewInt(9))},
+			// LSH 250 due to finite field limitation
+			faucet: {Balance: new(big.Int).Sub(new(big.Int).Lsh(big.NewInt(1), 250), big.NewInt(9))},
 		},
 	}
 }

core/state/database.go

@@ -120,6 +120,7 @@ func NewDatabase(db ethdb.Database) Database {
 func NewDatabaseWithConfig(db ethdb.Database, config *trie.Config) Database {
 	csc, _ := lru.New(codeSizeCacheSize)
 	return &cachingDB{
+		zktrie:        config != nil && config.Zktrie,
 		db:            trie.NewDatabaseWithConfig(db, config),
 		codeSizeCache: csc,
 		codeCache:     fastcache.New(codeCacheSize),
@@ -130,10 +131,18 @@ type cachingDB struct {
 	db            *trie.Database
 	codeSizeCache *lru.Cache
 	codeCache     *fastcache.Cache
+	zktrie        bool
 }
 
 // OpenTrie opens the main account trie at a specific root hash.
 func (db *cachingDB) OpenTrie(root common.Hash) (Trie, error) {
+	if db.zktrie {
+		tr, err := trie.NewZkTrie(root, trie.NewZktrieDatabaseFromTriedb(db.db))
+		if err != nil {
+			return nil, err
+		}
+		return tr, nil
+	}
 	tr, err := trie.NewSecure(root, db.db)
 	if err != nil {
 		return nil, err
@@ -143,6 +152,13 @@ func (db *cachingDB) OpenTrie(root common.Hash) (Trie, error) {
 
 // OpenStorageTrie opens the storage trie of an account.
 func (db *cachingDB) OpenStorageTrie(addrHash, root common.Hash) (Trie, error) {
+	if db.zktrie {
+		tr, err := trie.NewZkTrie(root, trie.NewZktrieDatabaseFromTriedb(db.db))
+		if err != nil {
+			return nil, err
+		}
+		return tr, nil
+	}
 	tr, err := trie.NewSecure(root, db.db)
 	if err != nil {
 		return nil, err
@@ -155,6 +171,8 @@ func (db *cachingDB) CopyTrie(t Trie) Trie {
 	switch t := t.(type) {
 	case *trie.SecureTrie:
 		return t.Copy()
+	case *trie.ZkTrie:
+		return t.Copy()
 	default:
 		panic(fmt.Errorf("unknown trie type %T", t))
 	}

core/state/snapshot/snapshot.go

@@ -185,6 +185,9 @@ type Tree struct {
 //   a background thread.
 func New(diskdb ethdb.KeyValueStore, triedb *trie.Database, cache int, root common.Hash, async bool, rebuild bool, recovery bool) (*Tree, error) {
 	// Create a new, empty snapshot tree
+	if triedb.Zktrie {
+		panic("zktrie does not support snapshot yet")
+	}
 	snap := &Tree{
 		diskdb: diskdb,
 		triedb: triedb,

core/state/state_object.go

@@ -107,7 +107,7 @@ func newObject(db *StateDB, address common.Address, data types.StateAccount) *st
 		data.CodeHash = emptyCodeHash
 	}
 	if data.Root == (common.Hash{}) {
-		data.Root = emptyRoot
+		data.Root = db.db.TrieDB().EmptyRoot()
 	}
 	return &stateObject{
 		db:             db,
@@ -151,7 +151,7 @@ func (s *stateObject) getTrie(db Database) Trie {
 	if s.trie == nil {
 		// Try fetching from prefetcher first
 		// We don't prefetch empty tries
-		if s.data.Root != emptyRoot && s.db.prefetcher != nil {
+		if s.data.Root != s.db.db.TrieDB().EmptyRoot() && s.db.prefetcher != nil {
 			// When the miner is creating the pending state, there is no
 			// prefetcher
 			s.trie = s.db.prefetcher.trie(s.data.Root)
@@ -245,12 +245,16 @@ func (s *stateObject) GetCommittedState(db Database, key common.Hash) common.Has
 		}
 	}
 	var value common.Hash
-	if len(enc) > 0 {
-		_, content, _, err := rlp.Split(enc)
-		if err != nil {
-			s.setError(err)
+	if db.TrieDB().Zktrie {
+		value = common.BytesToHash(enc)
+	} else {
+		if len(enc) > 0 {
+			_, content, _, err := rlp.Split(enc)
+			if err != nil {
+				s.setError(err)
+			}
+			value.SetBytes(content)
 		}
-		value.SetBytes(content)
 	}
 	s.originStorage[key] = value
 	return value
@@ -275,6 +279,25 @@ func (s *stateObject) SetState(db Database, key, value common.Hash) {
 		prevalue: prev,
 	})
 	s.setState(key, value)
+	/*
+		// always update trie
+		tr := s.getTrie(db)
+		// copied from "updateTrie"
+		var v []byte
+		if (value == common.Hash{}) {
+			s.setError(tr.TryDelete(key[:]))
+			//s.db.StorageDeleted += 1
+		} else {
+			if db.TrieDB().Zktrie {
+				v = common.CopyBytes(value[:])
+			} else {
+				// Encoding []byte cannot fail, ok to ignore the error.
+				v, _ = rlp.EncodeToBytes(common.TrimLeftZeroes(value[:]))
+			}
+			s.setError(tr.TryUpdate(key[:], v))
+			//s.db.StorageUpdated += 1
+		}
+	*/
 }
 
 // SetStorage replaces the entire state storage with the given one.
@@ -309,7 +332,7 @@ func (s *stateObject) finalise(prefetch bool) {
 			slotsToPrefetch = append(slotsToPrefetch, common.CopyBytes(key[:])) // Copy needed for closure
 		}
 	}
-	if s.db.prefetcher != nil && prefetch && len(slotsToPrefetch) > 0 && s.data.Root != emptyRoot {
+	if s.db.prefetcher != nil && prefetch && len(slotsToPrefetch) > 0 && s.data.Root != s.db.db.TrieDB().EmptyRoot() {
 		s.db.prefetcher.prefetch(s.data.Root, slotsToPrefetch)
 	}
 	if len(s.dirtyStorage) > 0 {
@@ -348,8 +371,12 @@ func (s *stateObject) updateTrie(db Database) Trie {
 			s.setError(tr.TryDelete(key[:]))
 			s.db.StorageDeleted += 1
 		} else {
-			// Encoding []byte cannot fail, ok to ignore the error.
-			v, _ = rlp.EncodeToBytes(common.TrimLeftZeroes(value[:]))
+			if db.TrieDB().Zktrie {
+				v = common.CopyBytes(value[:])
+			} else {
+				// Encoding []byte cannot fail, ok to ignore the error.
+				v, _ = rlp.EncodeToBytes(common.TrimLeftZeroes(value[:]))
+			}
 			s.setError(tr.TryUpdate(key[:], v))
 			s.db.StorageUpdated += 1
 		}