squash

fix unit tests

cleanups

Integrate common interfaces into rawdb

fixes after rebase

minor linter changes

clean-up blockchain.go, fix linter

unit tests

linter

minor cleanups

linter

linter

fixes after rebase

Author: MaksymMalicki

Makefile

@@ -73,6 +73,9 @@ clean-testcache: ## Clean Go test cache
 test: clean-testcache rustdeps ## Run tests
 	go test $(GO_TAGS) ./...
 
+test-new-state: clean-testcache rustdeps ## Run tests with new state
+	USE_NEW_STATE=true go test $(GO_TAGS) ./...
+
 test-cached: rustdeps ## Run cached tests
 	go test $(GO_TAGS) ./...
 

blockchain/blockchain.go

@@ -6,8 +6,10 @@ import (
 
 	"github.com/NethermindEth/juno/core"
 	"github.com/NethermindEth/juno/core/felt"
+	"github.com/NethermindEth/juno/core/state"
+	"github.com/NethermindEth/juno/core/state/statefactory"
+	"github.com/NethermindEth/juno/core/trie2/triedb"
 	"github.com/NethermindEth/juno/db"
-	"github.com/NethermindEth/juno/db/memory"
 	"github.com/NethermindEth/juno/feed"
 	"github.com/NethermindEth/juno/utils"
 	"github.com/ethereum/go-ethereum/common"
@@ -19,6 +21,7 @@ type L1HeadSubscription struct {
 
 //go:generate mockgen -destination=../mocks/mock_blockchain.go -package=mocks github.com/NethermindEth/juno/blockchain Reader
 type Reader interface {
+	StateProvider
 	Height() (height uint64, err error)
 
 	Head() (head *core.Block, err error)
@@ -51,10 +54,6 @@ type Reader interface {
 	StateUpdateByHash(hash *felt.Felt) (update *core.StateUpdate, err error)
 	L1HandlerTxnHash(msgHash *common.Hash) (l1HandlerTxnHash felt.Felt, err error)
 
-	HeadState() (core.StateReader, StateCloser, error)
-	StateAtBlockHash(blockHash *felt.Felt) (core.StateReader, StateCloser, error)
-	StateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error)
-
 	BlockCommitmentsByNumber(blockNumber uint64) (*core.BlockCommitments, error)
 
 	EventFilter(
@@ -66,6 +65,12 @@ type Reader interface {
 	Network() *utils.Network
 }
 
+type StateProvider interface {
+	HeadState() (core.StateReader, StateCloser, error)
+	StateAtBlockHash(blockHash *felt.Felt) (core.StateReader, StateCloser, error)
+	StateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error)
+}
+
 var ErrParentDoesNotMatchHead = errors.New("block's parent hash does not match head block hash")
 
 var _ Reader = (*Blockchain)(nil)
@@ -74,14 +79,24 @@ var _ Reader = (*Blockchain)(nil)
 type Blockchain struct {
 	network           *utils.Network
 	database          db.KeyValueStore
+	stateDB           *state.StateDB
 	listener          EventListener
 	l1HeadFeed        *feed.Feed[*core.L1Head]
 	cachedFilters     *AggregatedBloomFilterCache
 	runningFilter     *core.RunningEventFilter
 	transactionLayout core.TransactionLayout
+	StateFactory      *statefactory.StateFactory
 }
 
-func New(database db.KeyValueStore, network *utils.Network) *Blockchain {
+func New(database db.KeyValueStore, network *utils.Network, stateVersion bool) *Blockchain {
+	trieDB, err := triedb.New(database, nil) // TODO: handle hashdb and pathdb
+	if err != nil {
+		panic(err)
+	}
+	stateDB := state.NewStateDB(database, trieDB)
+
+	stateFactory := statefactory.NewStateFactory(stateVersion, trieDB, stateDB)
+
 	cachedFilters := NewAggregatedBloomCache(AggregatedBloomFilterCacheSize)
 	fallback := func(key EventFiltersCacheKey) (core.AggregatedBloomFilter, error) {
 		return core.GetAggregatedBloomFilter(database, key.fromBlock, key.toBlock)
@@ -92,12 +107,14 @@ func New(database db.KeyValueStore, network *utils.Network) *Blockchain {
 
 	return &Blockchain{
 		database:          database,
+		stateDB:           stateDB,
 		network:           network,
 		listener:          &SelectiveListener{},
 		l1HeadFeed:        feed.New[*core.L1Head](),
 		cachedFilters:     &cachedFilters,
 		runningFilter:     runningFilter,
 		transactionLayout: core.TransactionLayoutPerTx, // default to per-tx for backward compatibility
+		StateFactory:      stateFactory,
 	}
 }
 
@@ -295,8 +312,7 @@ func (b *Blockchain) SubscribeL1Head() L1HeadSubscription {
 
 func (b *Blockchain) L1Head() (core.L1Head, error) {
 	b.listener.OnRead("L1Head")
-	l1Head, err := core.GetL1Head(b.database)
-	return l1Head, err
+	return core.GetL1Head(b.database)
 }
 
 func (b *Blockchain) SetL1Head(update *core.L1Head) error {
@@ -310,6 +326,21 @@ func (b *Blockchain) Store(
 	blockCommitments *core.BlockCommitments,
 	stateUpdate *core.StateUpdate,
 	newClasses map[felt.Felt]core.ClassDefinition,
+) error {
+	// old state
+	// TODO(maksymmalick): remove this once we have a new state implementation
+	if !b.StateFactory.UseNewState() {
+		return b.deprecatedStore(block, blockCommitments, stateUpdate, newClasses)
+	}
+
+	return b.store(block, blockCommitments, stateUpdate, newClasses)
+}
+
+func (b *Blockchain) deprecatedStore(
+	block *core.Block,
+	blockCommitments *core.BlockCommitments,
+	stateUpdate *core.StateUpdate,
+	newClasses map[felt.Felt]core.ClassDefinition,
 ) error {
 	err := b.database.Update(func(txn db.IndexedBatch) error {
 		if err := verifyBlock(txn, block); err != nil {
@@ -348,6 +379,7 @@ func (b *Blockchain) Store(
 		}
 
 		err = storeCasmHashMetadata(
+			b.database,
 			txn,
 			block.Number,
 			block.ProtocolVersion,
@@ -432,7 +464,8 @@ func (b *Blockchain) store(
 // storeCasmHashMetadata stores CASM hash metadata for declared and migrated classes.
 // See [core.ClassCasmHashMetadata]
 func storeCasmHashMetadata(
-	txn db.IndexedBatch,
+	reader db.KeyValueReader,
+	writer db.KeyValueWriter,
 	blockNumber uint64,
 	protocolVersion string,
 	stateUpdate *core.StateUpdate,
@@ -446,16 +479,17 @@ func storeCasmHashMetadata(
 	isV2Protocol := ver.GreaterThanEqual(core.Ver0_14_1)
 
 	if isV2Protocol {
-		return storeCasmHashMetadataV2(txn, blockNumber, stateUpdate)
+		return storeCasmHashMetadataV2(reader, writer, blockNumber, stateUpdate)
 	}
 
-	return storeCasmHashMetadataV1(txn, blockNumber, stateUpdate, newClasses)
+	return storeCasmHashMetadataV1(writer, blockNumber, stateUpdate, newClasses)
 }
 
 // storeCasmHashMetadataV2 stores metadata for classes declared with casm hash v2 or
 // migrated from v1. casm hash v2 is after protocol version >= 0.14.1.
 func storeCasmHashMetadataV2(
-	txn db.IndexedBatch,
+	reader db.KeyValueReader,
+	writer db.KeyValueWriter,
 	blockNumber uint64,
 	stateUpdate *core.StateUpdate,
 ) error {
@@ -465,7 +499,7 @@ func storeCasmHashMetadataV2(
 			(*felt.CasmClassHash)(casmHash),
 		)
 		err := core.WriteClassCasmHashMetadata(
-			txn,
+			writer,
 			(*felt.SierraClassHash)(&sierraClassHash),
 			&metadata,
 		)
@@ -475,7 +509,7 @@ func storeCasmHashMetadataV2(
 	}
 
 	for sierraClassHash := range stateUpdate.StateDiff.MigratedClasses {
-		metadata, err := core.GetClassCasmHashMetadata(txn, &sierraClassHash)
+		metadata, err := core.GetClassCasmHashMetadata(reader, &sierraClassHash)
 		if err != nil {
 			return fmt.Errorf("cannot migrate class %s: metadata not found",
 				sierraClassHash.String(),
@@ -490,7 +524,7 @@ func storeCasmHashMetadataV2(
 			)
 		}
 
-		err = core.WriteClassCasmHashMetadata(txn, &sierraClassHash, &metadata)
+		err = core.WriteClassCasmHashMetadata(writer, &sierraClassHash, &metadata)
 		if err != nil {
 			return err
 		}
@@ -501,7 +535,7 @@ func storeCasmHashMetadataV2(
 // storeDeclaredV1Classes stores metadata for classes declared with V1 hash (protocol < 0.14.1).
 // It computes the V2 hash from the class definition.
 func storeCasmHashMetadataV1(
-	txn db.IndexedBatch,
+	writer db.KeyValueWriter,
 	blockNumber uint64,
 	stateUpdate *core.StateUpdate,
 	newClasses map[felt.Felt]core.ClassDefinition,
@@ -530,7 +564,7 @@ func storeCasmHashMetadataV1(
 
 		metadata := core.NewCasmHashMetadataDeclaredV1(blockNumber, casmHashV1, &casmHashV2)
 		err := core.WriteClassCasmHashMetadata(
-			txn,
+			writer,
 			(*felt.SierraClassHash)(&sierraClassHash),
 			&metadata,
 		)
@@ -604,7 +638,7 @@ func (b *Blockchain) HeadState() (core.CommonStateReader, StateCloser, error) {
 	b.listener.OnRead("HeadState")
 	txn := b.database.NewIndexedBatch()
 
-	_, err := core.GetChainHeight(txn)
+	height, err := core.GetChainHeight(txn)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -633,7 +667,25 @@ func (b *Blockchain) StateAtBlockNumber(
 		return nil, nil, err
 	}
 
-	return core.NewDeprecatedStateHistory(core.NewState(txn), blockNumber), noopStateCloser, nil
+	if !b.StateFactory.UseNewState() {
+		return core.NewDeprecatedStateHistory(core.NewState(txn), blockNumber), noopStateCloser, nil
+	}
+
+	height, err := core.GetChainHeight(txn)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	header, err := core.GetBlockHeaderByNumber(txn, height)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	history, err := state.NewStateHistory(blockNumber, header.GlobalStateRoot, b.stateDB)
+	if err != nil {
+		return nil, nil, err
+	}
+	return &history, noopStateCloser, nil
 }
 
 // StateAtBlockHash returns a StateReader that provides
@@ -643,19 +695,34 @@ func (b *Blockchain) StateAtBlockHash(
 ) (core.StateReader, StateCloser, error) {
 	b.listener.OnRead("StateAtBlockHash")
 	if blockHash.IsZero() {
-		memDB := memory.New()
-		txn := memDB.NewIndexedBatch()
-		emptyState := core.NewState(txn)
-		return emptyState, noopStateCloser, nil
+		emptyState, err := b.StateFactory.EmptyState()
+		return emptyState, noopStateCloser, err
 	}
 
 	txn := b.database.NewIndexedBatch()
 	header, err := core.GetBlockHeaderByHash(txn, blockHash)
 	if err != nil {
 		return nil, nil, err
 	}
+	if !b.StateFactory.UseNewState() {
+		return core.NewDeprecatedStateHistory(core.NewState(txn), header.Number), noopStateCloser, nil
+	}
+
+	height, err := core.GetChainHeight(txn)
+	if err != nil {
+		return nil, nil, err
+	}
 
-	return core.NewDeprecatedStateHistory(core.NewState(txn), header.Number), noopStateCloser, nil
+	headHeader, err := core.GetBlockHeaderByNumber(txn, height)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	history, err := state.NewStateHistory(header.Number, headHeader.GlobalStateRoot, b.stateDB)
+	if err != nil {
+		return nil, nil, err
+	}
+	return &history, noopStateCloser, nil
 }
 
 // EventFilter returns an EventFilter object that is tied to a snapshot of the blockchain
@@ -685,10 +752,22 @@ func (b *Blockchain) EventFilter(
 
 // RevertHead reverts the head block
 func (b *Blockchain) RevertHead() error {
-	return b.database.Update(b.revertHead)
+	if !b.StateFactory.UseNewState() {
+		return b.database.Update(b.deprecatedRevertHead)
+	}
+	return b.database.Write(b.revertHead)
 }
 
 func (b *Blockchain) GetReverseStateDiff() (core.StateDiff, error) {
+	if !b.StateFactory.UseNewState() {
+		return b.deprecatedGetReverseStateDiff()
+	}
+
+	return b.getReverseStateDiff()
+}
+
+// TODO(maksymmalick): remove this once we have a new state integrated
+func (b *Blockchain) deprecatedGetReverseStateDiff() (core.StateDiff, error) {
 	txn := b.database.NewIndexedBatch()
 	blockNum, err := core.GetChainHeight(txn)
 	if err != nil {
@@ -753,7 +832,7 @@ func (b *Blockchain) revertHead(txn db.IndexedBatch) error {
 		}
 	}
 
-	if err := b.transactionLayout.DeleteTxsAndReceipts(txn, blockNumber); err != nil {
+	if err := b.transactionLayout.DeleteTxsAndReceipts(b.database, txn, blockNumber); err != nil {
 		return err
 	}
 
@@ -979,6 +1058,7 @@ func (b *Blockchain) updateStateRoots(
 	block *core.Block,
 	stateUpdate *core.StateUpdate,
 	newClasses map[felt.Felt]core.ClassDefinition,
+	flushChanges bool,
 ) error {
 	var height uint64
 	var err error
@@ -1011,7 +1091,7 @@ func (b *Blockchain) updateStateRoots(
 	stateUpdate.OldRoot = &oldStateRoot
 
 	// Apply state update
-	if err = state.Update(block.Number, stateUpdate, newClasses, true); err != nil {
+	if err = state.Update(block.Number, stateUpdate, newClasses, true, flushChanges); err != nil {
 		return err
 	}
 
@@ -1065,7 +1145,7 @@ func (b *Blockchain) signBlock(
 
 // storeBlockData persists all block-related data to the database
 func (b *Blockchain) storeBlockData(
-	txn db.IndexedBatch,
+	txn db.KeyValueWriter,
 	block *core.Block,
 	stateUpdate *core.StateUpdate,
 	commitments *core.BlockCommitments,