zktrie part2: add zktrie; allow switch trie type by config; (#113)

* induce zktrie

* refactoring zktrie

* fix crash issue in logger

* renaming JSON field

* unify hash scheme

* goimport and mod lint

* backward compatible with go 1.17

* lints

* add option on genesis file

* corrections according to the reviews

* trivial fixes: ValueKey entry, key in prove nodes

* fixing for the proof fix ...

* avoiding panic before loading stateDb in genesis setup

* revert ExtraData.StateList json annotation for compatibility

* fix goimports lint

* fix goimports lint

* better encoding for leaf node

* fix proof's printing issue, add handling on coinbase

* update genesis, and rule out snapshot in zktrie mode

* update readme and lint

* fix an issue

Co-authored-by: HAOYUatHZ <[email protected]>
Co-authored-by: HAOYUatHZ <[email protected]>

README.md

@@ -13,6 +13,15 @@ https://camo.githubusercontent.com/915b7be44ada53c290eb157634330494ebe3e30a/6874
 
 ZK-Rollup adapts the Go Ethereum to run as Layer 2 Sequencer. The codebase is based on v1.10.13.
 
+### ZKTrie Storage
+
+Another implement for storage trie, base on patricia merkle tree, has been induced. It is feasible to zk proving in the storage part. It is specified as a flag
+in gensis, set `config.zktrie` to true for enabling it. Using `genesis_zktrie.json` as an example to create a L2 chain with zktrie sotrage:
+
+> geth init \<repo root\>/genesis_zktrie.json
+
+Notice current the snapshot would be disabled by the zktrie implement.
+
 ## Building the source
 
 For prerequisites and detailed build instructions please read the [Installation Instructions](https://geth.ethereum.org/docs/install-and-build/installing-geth).

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

@@ -236,6 +236,11 @@ func NewBlockChain(db ethdb.Database, cacheConfig *CacheConfig, chainConfig *par
 	if cacheConfig.TraceCacheLimit != 0 {
 		blockResultCache, _ = lru.New(cacheConfig.TraceCacheLimit)
 	}
+	// override snapshot setting
+	if chainConfig.Zktrie && cacheConfig.SnapshotLimit > 0 {
+		log.Warn("snapshot has been disabled by zktrie")
+		cacheConfig.SnapshotLimit = 0
+	}
 
 	bc := &BlockChain{
 		chainConfig: chainConfig,
@@ -246,6 +251,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(),

core/genesis.go

@@ -180,7 +180,21 @@ 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 {
+			log.Warn("Found genesis block without chain config")
+		} else {
+			trieCfg = &trie.Config{Zktrie: storedcfg.Zktrie}
+		}
+	} else {
+		trieCfg = &trie.Config{Zktrie: genesis.Config.Zktrie}
+	}
+
+	if _, err := state.New(header.Root, state.NewDatabaseWithConfig(db, trieCfg), nil); err != nil {
 		if genesis == nil {
 			genesis = DefaultGenesisBlock()
 		}
@@ -261,7 +275,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)
 	}
@@ -441,7 +459,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
@@ -309,7 +313,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 +352,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
 		}

core/state/statedb.go

@@ -40,11 +40,6 @@ type revision struct {
 	journalIndex int
 }
 
-var (
-	// emptyRoot is the known root hash of an empty trie.
-	emptyRoot = common.HexToHash("56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421")
-)
-
 type proofList [][]byte
 
 func (n *proofList) Put(key []byte, value []byte) error {
@@ -187,6 +182,10 @@ func (s *StateDB) Error() error {
 	return s.dbErr
 }
 
+func (s *StateDB) IsZktrie() bool {
+	return s.db.TrieDB().Zktrie
+}
+
 func (s *StateDB) AddLog(log *types.Log) {
 	s.journal.append(addLogChange{txhash: s.thash})
 
@@ -315,11 +314,19 @@ func (s *StateDB) GetState(addr common.Address, hash common.Hash) common.Hash {
 
 // GetProof returns the Merkle proof for a given account.
 func (s *StateDB) GetProof(addr common.Address) ([][]byte, error) {
+	if s.IsZktrie() {
+		var proof proofList
+		err := s.trie.Prove(addr.Bytes32(), 0, &proof)
+		return proof, err
+	}
 	return s.GetProofByHash(crypto.Keccak256Hash(addr.Bytes()))
 }
 
 // GetProofByHash returns the Merkle proof for a given account.
 func (s *StateDB) GetProofByHash(addrHash common.Hash) ([][]byte, error) {
+	if s.IsZktrie() {
+		panic("unimplemented")
+	}
 	var proof proofList
 	err := s.trie.Prove(addrHash[:], 0, &proof)
 	return proof, err
@@ -339,6 +346,7 @@ func (s *StateDB) GetRootHash() common.Hash {
 
 // StorageTrieProof is not in Db interface and used explictily for reading proof in storage trie (not the dirty value)
 func (s *StateDB) GetStorageTrieProof(a common.Address, key common.Hash) ([][]byte, error) {
+
 	// try the trie in stateObject first, else we would create one
 	stateObject := s.getStateObject(a)
 	if stateObject == nil {
@@ -356,7 +364,11 @@ func (s *StateDB) GetStorageTrieProof(a common.Address, key common.Hash) ([][]by
 	}
 
 	var proof proofList
-	err = trie.Prove(crypto.Keccak256(key.Bytes()), 0, &proof)
+	if s.IsZktrie() {
+		err = trie.Prove(key.Bytes(), 0, &proof)
+	} else {
+		err = trie.Prove(crypto.Keccak256(key.Bytes()), 0, &proof)
+	}
 	return proof, err
 }
 
@@ -367,7 +379,12 @@ func (s *StateDB) GetStorageProof(a common.Address, key common.Hash) ([][]byte,
 	if trie == nil {
 		return proof, errors.New("storage trie for requested address does not exist")
 	}
-	err := trie.Prove(crypto.Keccak256(key.Bytes()), 0, &proof)
+	var err error
+	if s.IsZktrie() {
+		err = trie.Prove(key.Bytes(), 0, &proof)
+	} else {
+		err = trie.Prove(crypto.Keccak256(key.Bytes()), 0, &proof)
+	}
 	return proof, err
 }
 
@@ -564,7 +581,7 @@ func (s *StateDB) getDeletedStateObject(addr common.Address) *stateObject {
 				data.CodeHash = emptyCodeHash
 			}
 			if data.Root == (common.Hash{}) {
-				data.Root = emptyRoot
+				data.Root = s.db.TrieDB().EmptyRoot()
 			}
 		}
 	}
@@ -582,7 +599,12 @@ func (s *StateDB) getDeletedStateObject(addr common.Address) *stateObject {
 			return nil
 		}
 		data = new(types.StateAccount)
-		if err := rlp.DecodeBytes(enc, data); err != nil {
+		if s.IsZktrie() {
+			data, err = types.UnmarshalStateAccount(enc)
+		} else {
+			err = rlp.DecodeBytes(enc, data)
+		}
+		if err != nil {
 			log.Error("Failed to decode state object", "addr", addr, "err", err)
 			return nil
 		}
@@ -977,7 +999,7 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 		if err := rlp.DecodeBytes(leaf, &account); err != nil {
 			return nil
 		}
-		if account.Root != emptyRoot {
+		if account.Root != s.db.TrieDB().EmptyRoot() {
 			s.db.TrieDB().Reference(account.Root, parent)
 		}
 		return nil