Refactor state and contract modules

blockchain/encoder_initializer.go

@@ -20,6 +20,7 @@ func RegisterCoreTypesToEncoder() {
 			reflect.TypeOf(core.DeployAccountTransaction{}),
 			reflect.TypeOf(core.Cairo0Class{}),
 			reflect.TypeOf(core.Cairo1Class{}),
+			reflect.TypeOf(core.StateContract{}),
 		}
 
 		for _, t := range types {

clients/feeder/feeder.go

@@ -92,7 +92,7 @@ func NopBackoff(d time.Duration) time.Duration {
 }
 
 // NewTestClient returns a client and a function to close a test server.
-func NewTestClient(t *testing.T, network *utils.Network) *Client {
+func NewTestClient(t testing.TB, network *utils.Network) *Client {
 	srv := newTestServer(t, network)
 	t.Cleanup(srv.Close)
 	ua := "Juno/v0.0.1-test Starknet Implementation"
@@ -117,7 +117,7 @@ func NewTestClient(t *testing.T, network *utils.Network) *Client {
 	return c
 }
 
-func newTestServer(t *testing.T, network *utils.Network) *httptest.Server {
+func newTestServer(t testing.TB, network *utils.Network) *httptest.Server {
 	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		queryMap, err := url.ParseQuery(r.URL.RawQuery)
 		if err != nil {

cmd/juno/dbcmd.go

@@ -206,7 +206,7 @@ func dbSize(cmd *cobra.Command, args []string) error {
 		totalSize += bucketItem.Size
 		totalCount += bucketItem.Count
 
-		if utils.AnyOf(b, db.StateTrie, db.ContractStorage, db.Class, db.ContractNonce, db.ContractDeploymentHeight) {
+		if utils.AnyOf(b, db.StateTrie, db.ContractStorage, db.Class, db.Contract) {
 			withoutHistorySize += bucketItem.Size
 			withHistorySize += bucketItem.Size
 

core/contract.go

@@ -7,6 +7,7 @@
 	"github.com/NethermindEth/juno/core/felt"
 	"github.com/NethermindEth/juno/core/trie"
 	"github.com/NethermindEth/juno/db"
+	"github.com/NethermindEth/juno/encoder"
 )
 
 // contract storage has fixed height at 251
@@ -17,187 +18,211 @@
 	ErrContractAlreadyDeployed = errors.New("contract already deployed")
 )
 
-// NewContractUpdater creates an updater for the contract instance at the given address.
-// Deploy should be called for contracts that were just deployed to the network.
-func NewContractUpdater(addr *felt.Felt, txn db.Transaction) (*ContractUpdater, error) {
-	contractDeployed, err := deployed(addr, txn)
-	if err != nil {
-		return nil, err
-	}
+type OnValueChanged = func(location, oldValue *felt.Felt) error
 
-	if !contractDeployed {
-		return nil, ErrContractNotDeployed
-	}
+// StateContract represents a contract instance.
+// The usage of a `StateContract` is as follows:
+// 1. Create or obtain `StateContract` instance from the database.
+// 2. Update the contract fields
+// 3. Commit the contract to the database
+type StateContract struct {
+	// ClassHash is the hash of the contract's class
+	ClassHash *felt.Felt
+	// Nonce is the contract's nonce
+	Nonce *felt.Felt
+	// DeployHeight is the height at which the contract is deployed
+	DeployHeight uint64
+	// Address that this contract instance is deployed to
+	Address *felt.Felt `cbor:"-"`
+	// dirtyStorage is a map of storage locations that have been updated
+	dirtyStorage map[felt.Felt]*felt.Felt `cbor:"-"`
+}
 
-	return &ContractUpdater{
-		Address: addr,
-		txn:     txn,
-	}, nil
+// NewStateContract creates a new contract instance.
+func NewStateContract(
+	addr *felt.Felt,
+	classHash *felt.Felt,
+	nonce *felt.Felt,
+	deployHeight uint64,
+) *StateContract {
+	sc := &StateContract{
+		Address:      addr,
+		ClassHash:    classHash,
+		Nonce:        nonce,
+		DeployHeight: deployHeight,
+		dirtyStorage: make(map[felt.Felt]*felt.Felt),
+	}
+
+	return sc
 }
 
-// DeployContract sets up the database for a new contract.
-func DeployContract(addr, classHash *felt.Felt, txn db.Transaction) (*ContractUpdater, error) {
-	contractDeployed, err := deployed(addr, txn)
+// StorageRoot returns the root of the contract's storage trie.
+func (c *StateContract) StorageRoot(txn db.Transaction) (*felt.Felt, error) {
+	storageTrie, err := storage(c.Address, txn)
 	if err != nil {
 		return nil, err
 	}
 
-	if contractDeployed {
-		return nil, ErrContractAlreadyDeployed
-	}
+	return storageTrie.Root()
+}
 
-	err = setClassHash(txn, addr, classHash)
-	if err != nil {
-		return nil, err
+// UpdateStorage updates the storage of a contract.
+// Note that this does not modify the storage trie, which must be committed separately.
+func (c *StateContract) UpdateStorage(key, value *felt.Felt) {
+	if c.dirtyStorage == nil {
+		c.dirtyStorage = make(map[felt.Felt]*felt.Felt)
 	}
 
-	c, err := NewContractUpdater(addr, txn)
-	if err != nil {
-		return nil, err
+	c.dirtyStorage[*key] = value
+}
+
+// GetStorage retrieves the value of a storage location from the contract's storage
+func (c *StateContract) GetStorage(key *felt.Felt, txn db.Transaction) (*felt.Felt, error) {
+	if c.dirtyStorage != nil {
+		if val, ok := c.dirtyStorage[*key]; ok {
+			return val, nil
+		}
 	}
 
-	err = c.UpdateNonce(&felt.Zero)
+	// get from db
+	storage, err := storage(c.Address, txn)
 	if err != nil {
 		return nil, err
 	}
 
-	return c, nil
+	return storage.Get(key)
 }
 
-// ContractAddress computes the address of a Starknet contract.
-func ContractAddress(callerAddress, classHash, salt *felt.Felt, constructorCallData []*felt.Felt) *felt.Felt {
-	prefix := new(felt.Felt).SetBytes([]byte("STARKNET_CONTRACT_ADDRESS"))
-	callDataHash := crypto.PedersenArray(constructorCallData...)
-
-	// https://docs.starknet.io/architecture-and-concepts/smart-contracts/contract-address/
-	return crypto.PedersenArray(
-		prefix,
-		callerAddress,
-		salt,
-		classHash,
-		callDataHash,
-	)
+// logOldValue is a helper function to record the history of a contract's value
+func (c *StateContract) logOldValue(key []byte, oldValue *felt.Felt, height uint64, txn db.Transaction) error {
+	return txn.Set(logDBKey(key, height), oldValue.Marshal())
 }
 
-func deployed(addr *felt.Felt, txn db.Transaction) (bool, error) {
-	_, err := ContractClassHash(addr, txn)
-	if errors.Is(err, db.ErrKeyNotFound) {
-		return false, nil
-	}
-	if err != nil {
-		return false, err
-	}
-	return true, nil
+// LogStorage records the history of the contract's storage
+func (c *StateContract) LogStorage(location, oldVal *felt.Felt, height uint64, txn db.Transaction) error {
+	key := storageLogKey(c.Address, location)
+	return c.logOldValue(key, oldVal, height, txn)
 }
 
-// ContractUpdater is a helper to update an existing contract instance.
-type ContractUpdater struct {
-	// Address that this contract instance is deployed to
-	Address *felt.Felt
-	// txn to access the database
-	txn db.Transaction
+// LogNonce records the history of the contract's nonce
+func (c *StateContract) LogNonce(height uint64, txn db.Transaction) error {
+	key := nonceLogKey(c.Address)
+	return c.logOldValue(key, c.Nonce, height, txn)
 }
 
-// Purge eliminates the contract instance, deleting all associated data from storage
-// assumes storage is cleared in revert process
-func (c *ContractUpdater) Purge() error {
-	addrBytes := c.Address.Marshal()
-	buckets := []db.Bucket{db.ContractNonce, db.ContractClassHash}
-
-	for _, bucket := range buckets {
-		if err := c.txn.Delete(bucket.Key(addrBytes)); err != nil {
-			return err
-		}
-	}
-
-	return nil
+// LogClassHash records the history of the contract's class hash
+func (c *StateContract) LogClassHash(height uint64, txn db.Transaction) error {
+	key := classHashLogKey(c.Address)
+	return c.logOldValue(key, c.ClassHash, height, txn)
 }
 
-// ContractNonce returns the amount transactions sent from this contract.
-// Only account contracts can have a non-zero nonce.
-func ContractNonce(addr *felt.Felt, txn db.Transaction) (*felt.Felt, error) {
-	key := db.ContractNonce.Key(addr.Marshal())
-	var nonce *felt.Felt
-	if err := txn.Get(key, func(val []byte) error {
-		nonce = new(felt.Felt)
-		nonce.SetBytes(val)
-		return nil
-	}); err != nil {
-		return nil, err
-	}
-	return nonce, nil
-}
+// BufferedCommit creates a buffered transaction and commits the contract to the database
+func (c *StateContract) BufferedCommit(txn db.Transaction, logChanges bool, blockNum uint64) (*db.BufferedTransaction, error) {
+	bufferedTxn := db.NewBufferedTransaction(txn)
 
-// UpdateNonce updates the nonce value in the database.
-func (c *ContractUpdater) UpdateNonce(nonce *felt.Felt) error {
-	nonceKey := db.ContractNonce.Key(c.Address.Marshal())
-	return c.txn.Set(nonceKey, nonce.Marshal())
-}
-
-// ContractRoot returns the root of the contract storage.
-func ContractRoot(addr *felt.Felt, txn db.Transaction) (*felt.Felt, error) {
-	cStorage, err := storage(addr, txn)
-	if err != nil {
+	if err := c.Commit(bufferedTxn, logChanges, blockNum); err != nil {
 		return nil, err
 	}
-	return cStorage.Root()
-}
 
-type OnValueChanged = func(location, oldValue *felt.Felt) error
+	return bufferedTxn, nil
+}
 
-// UpdateStorage applies a change-set to the contract storage.
-func (c *ContractUpdater) UpdateStorage(diff map[felt.Felt]*felt.Felt, cb OnValueChanged) error {
-	cStorage, err := storage(c.Address, c.txn)
+func (c *StateContract) Commit(txn db.Transaction, logChanges bool, blockNum uint64) error {
+	storageTrie, err := storage(c.Address, txn)
 	if err != nil {
 		return err
 	}
-	// apply the diff
-	for key, value := range diff {
-		oldValue, pErr := cStorage.Put(&key, value)
-		if pErr != nil {
-			return pErr
+
+	for key, value := range c.dirtyStorage {
+		oldVal, err := storageTrie.Put(&key, value)
+		if err != nil {
+			return err
 		}
 
-		if oldValue != nil {
-			if err = cb(&key, oldValue); err != nil {
+		if oldVal != nil && logChanges {
+			if err = c.LogStorage(&key, oldVal, blockNum, txn); err != nil {
 				return err
 			}
 		}
 	}
 
-	return cStorage.Commit()
+	if err := storageTrie.Commit(); err != nil {
+		return err
+	}
+
+	contractBytes, err := encoder.Marshal(c)
+	if err != nil {
+		return err
+	}
+
+	return txn.Set(db.Contract.Key(c.Address.Marshal()), contractBytes)
 }
 
-func ContractStorage(addr, key *felt.Felt, txn db.Transaction) (*felt.Felt, error) {
-	cStorage, err := storage(addr, txn)
+// Purge eliminates the contract instance, deleting all associated data from database
+// assumes storage is cleared in revert process
+func (c *StateContract) Purge(txn db.Transaction) error {
+	addrBytes := c.Address.Marshal()
+
+	return txn.Delete(db.Contract.Key(addrBytes))
+}
+
+func storageLogKey(contractAddress, storageLocation *felt.Felt) []byte {
+	return db.ContractStorageHistory.Key(contractAddress.Marshal(), storageLocation.Marshal())
+}
+
+func nonceLogKey(contractAddress *felt.Felt) []byte {
+	return db.ContractNonceHistory.Key(contractAddress.Marshal())
+}
+
+func classHashLogKey(contractAddress *felt.Felt) []byte {
+	return db.ContractClassHashHistory.Key(contractAddress.Marshal())
+}
+
+// GetContract is a wrapper around getContract which checks if a contract is deployed
+func GetContract(addr *felt.Felt, txn db.Transaction) (*StateContract, error) {
+	contract, err := getContract(addr, txn)
 	if err != nil {
+		if errors.Is(err, db.ErrKeyNotFound) {
+			return nil, ErrContractNotDeployed
+		}
 		return nil, err
 	}
-	return cStorage.Get(key)
+
+	return contract, nil
 }
 
-// ContractClassHash returns hash of the class that the contract at the given address instantiates.
-func ContractClassHash(addr *felt.Felt, txn db.Transaction) (*felt.Felt, error) {
-	key := db.ContractClassHash.Key(addr.Marshal())
-	var classHash *felt.Felt
+// getContract gets a contract instance from the database.
+func getContract(addr *felt.Felt, txn db.Transaction) (*StateContract, error) {
+	key := db.Contract.Key(addr.Marshal())
+	var contract StateContract
 	if err := txn.Get(key, func(val []byte) error {
-		classHash = new(felt.Felt)
-		classHash.SetBytes(val)
+		if err := encoder.Unmarshal(val, &contract); err != nil {
+			return err
+		}
+
+		contract.Address = addr
+		contract.dirtyStorage = make(map[felt.Felt]*felt.Felt)
+
 		return nil
 	}); err != nil {
 		return nil, err
 	}
-	return classHash, nil
+	return &contract, nil
 }
 
-func setClassHash(txn db.Transaction, addr, classHash *felt.Felt) error {
-	classHashKey := db.ContractClassHash.Key(addr.Marshal())
-	return txn.Set(classHashKey, classHash.Marshal())
-}
+// ContractAddress computes the address of a Starknet contract.
+func ContractAddress(callerAddress, classHash, salt *felt.Felt, constructorCallData []*felt.Felt) *felt.Felt {
+	prefix := new(felt.Felt).SetBytes([]byte("STARKNET_CONTRACT_ADDRESS"))
+	callDataHash := crypto.PedersenArray(constructorCallData...)
 
-// Replace replaces the class that the contract instantiates
-func (c *ContractUpdater) Replace(classHash *felt.Felt) error {
-	return setClassHash(c.txn, c.Address, classHash)
+	// https://docs.starknet.io/architecture-and-concepts/smart-contracts/contract-address/
+	return crypto.PedersenArray(
+		prefix,
+		callerAddress,
+		salt,
+		classHash,
+		callDataHash,
+	)
 }
 
 // storage returns the [core.Trie] that represents the