simplify ordered maps to simple maps. sort keys on commit

interpreter/storage.go

@@ -20,6 +20,7 @@ package interpreter
 
 import (
 	"bytes"
+	"cmp"
 	"io"
 	"math"
 	"strings"
@@ -106,6 +107,19 @@ type StorageDomainKey struct {
 	Address common.Address
 }
 
+func (k StorageDomainKey) Compare(o StorageDomainKey) int {
+	switch bytes.Compare(k.Address[:], o.Address[:]) {
+	case -1:
+		return -1
+	case 0:
+		return cmp.Compare(k.Domain, o.Domain)
+	case 1:
+		return 1
+	default:
+		panic(errors.NewUnreachableError())
+	}
+}
+
 func NewStorageDomainKey(
 	memoryGauge common.MemoryGauge,
 	address common.Address,

runtime/account_storage_v1.go

@@ -19,10 +19,11 @@
 package runtime
 
 import (
+	"sort"
+
 	"github.com/onflow/atree"
 
 	"github.com/onflow/cadence/common"
-	"github.com/onflow/cadence/common/orderedmap"
 	"github.com/onflow/cadence/errors"
 	"github.com/onflow/cadence/interpreter"
 )
@@ -35,7 +36,7 @@ type AccountStorageV1 struct {
 	// newDomainStorageMapSlabIndices contains root slab indices of new domain storage maps.
 	// The indices are saved using Ledger.SetValue() during commit().
 	// Key is StorageDomainKey{common.StorageDomain, Address} and value is 8-byte slab index.
-	newDomainStorageMapSlabIndices *orderedmap.OrderedMap[interpreter.StorageDomainKey, atree.SlabIndex]
+	newDomainStorageMapSlabIndices map[interpreter.StorageDomainKey]atree.SlabIndex
 }
 
 func NewAccountStorageV1(
@@ -91,9 +92,9 @@ func (s *AccountStorageV1) storeNewDomainStorageMap(
 	storageKey := interpreter.NewStorageDomainKey(s.memoryGauge, address, domain)
 
 	if s.newDomainStorageMapSlabIndices == nil {
-		s.newDomainStorageMapSlabIndices = &orderedmap.OrderedMap[interpreter.StorageDomainKey, atree.SlabIndex]{}
+		s.newDomainStorageMapSlabIndices = map[interpreter.StorageDomainKey]atree.SlabIndex{}
 	}
-	s.newDomainStorageMapSlabIndices.Set(storageKey, slabIndex)
+	s.newDomainStorageMapSlabIndices[storageKey] = slabIndex
 
 	return domainStorageMap
 }
@@ -103,20 +104,50 @@ func (s *AccountStorageV1) commit() error {
 		return nil
 	}
 
-	for pair := s.newDomainStorageMapSlabIndices.Oldest(); pair != nil; pair = pair.Next() {
+	type domainStorageMapSlabIndex struct {
+		StorageDomainKey interpreter.StorageDomainKey
+		SlabIndex        atree.SlabIndex
+	}
+
+	slabIndices := make([]domainStorageMapSlabIndex, 0, len(s.newDomainStorageMapSlabIndices))
+	for storageDomainKey, slabIndex := range s.newDomainStorageMapSlabIndices { //nolint:maprange
+		slabIndices = append(
+			slabIndices,
+			domainStorageMapSlabIndex{
+				StorageDomainKey: storageDomainKey,
+				SlabIndex:        slabIndex,
+			},
+		)
+	}
+	sort.Slice(
+		slabIndices,
+		func(i, j int) bool {
+			slabIndex1 := slabIndices[i]
+			slabIndex2 := slabIndices[j]
+			domainKey1 := slabIndex1.StorageDomainKey
+			domainKey2 := slabIndex2.StorageDomainKey
+			return domainKey1.Compare(domainKey2) < 0
+		},
+	)
+
+	for _, slabIndex := range slabIndices {
+		storageDomainKey := slabIndex.StorageDomainKey
+
 		var err error
 		errors.WrapPanic(func() {
 			err = s.ledger.SetValue(
-				pair.Key.Address[:],
-				[]byte(pair.Key.Domain.Identifier()),
-				pair.Value[:],
+				storageDomainKey.Address[:],
+				[]byte(storageDomainKey.Domain.Identifier()),
+				slabIndex.SlabIndex[:],
 			)
 		})
 		if err != nil {
 			return interpreter.WrappedExternalError(err)
 		}
 	}
 
+	s.newDomainStorageMapSlabIndices = nil
+
 	return nil
 }
 

runtime/account_storage_v2.go

@@ -19,10 +19,11 @@
 package runtime
 
 import (
+	"sort"
+
 	"github.com/onflow/atree"
 
 	"github.com/onflow/cadence/common"
-	"github.com/onflow/cadence/common/orderedmap"
 	"github.com/onflow/cadence/errors"
 	"github.com/onflow/cadence/interpreter"
 )
@@ -35,10 +36,9 @@ type AccountStorageV2 struct {
 	// cachedAccountStorageMaps is a cache of account storage maps.
 	cachedAccountStorageMaps map[common.Address]*interpreter.AccountStorageMap
 
-	// newAccountStorageMapSlabIndices contains root slab index of new account storage maps.
-	// The indices are saved using Ledger.SetValue() during Commit().
-	// Key is StorageKey{address, accountStorageKey} and value is 8-byte slab index.
-	newAccountStorageMapSlabIndices *orderedmap.OrderedMap[interpreter.StorageKey, atree.SlabIndex]
+	// newAccountStorageMapSlabIndices contains root slab indices of new account storage maps.
+	// The indices are saved using Ledger.SetValue() during commit().
+	newAccountStorageMapSlabIndices map[common.Address]atree.SlabIndex
 }
 
 func NewAccountStorageV2(
@@ -53,14 +53,6 @@ func NewAccountStorageV2(
 	}
 }
 
-func (s *AccountStorageV2) accountStorageKey(address common.Address) interpreter.StorageKey {
-	return interpreter.NewStorageKey(
-		s.memoryGauge,
-		address,
-		AccountStorageKey,
-	)
-}
-
 func (s *AccountStorageV2) GetDomainStorageMap(
 	inter *interpreter.Interpreter,
 	address common.Address,
@@ -148,9 +140,10 @@ func (s *AccountStorageV2) storeNewAccountStorageMap(
 
 	slabIndex := accountStorageMap.SlabID().Index()
 
-	accountStorageKey := s.accountStorageKey(address)
-
-	s.SetNewAccountStorageMapSlabIndex(accountStorageKey, slabIndex)
+	s.SetNewAccountStorageMapSlabIndex(
+		address,
+		slabIndex,
+	)
 
 	s.cacheAccountStorageMap(
 		address,
@@ -161,34 +154,65 @@ func (s *AccountStorageV2) storeNewAccountStorageMap(
 }
 
 func (s *AccountStorageV2) SetNewAccountStorageMapSlabIndex(
-	accountStorageKey interpreter.StorageKey,
+	address common.Address,
 	slabIndex atree.SlabIndex,
 ) {
 	if s.newAccountStorageMapSlabIndices == nil {
-		s.newAccountStorageMapSlabIndices = &orderedmap.OrderedMap[interpreter.StorageKey, atree.SlabIndex]{}
+		s.newAccountStorageMapSlabIndices = map[common.Address]atree.SlabIndex{}
 	}
-	s.newAccountStorageMapSlabIndices.Set(accountStorageKey, slabIndex)
+	s.newAccountStorageMapSlabIndices[address] = slabIndex
 }
 
 func (s *AccountStorageV2) commit() error {
 	if s.newAccountStorageMapSlabIndices == nil {
 		return nil
 	}
 
-	for pair := s.newAccountStorageMapSlabIndices.Oldest(); pair != nil; pair = pair.Next() {
+	type accountStorageMapSlabIndex struct {
+		Address   common.Address
+		SlabIndex atree.SlabIndex
+	}
+
+	slabIndices := make([]accountStorageMapSlabIndex, 0, len(s.newAccountStorageMapSlabIndices))
+	for address, slabIndex := range s.newAccountStorageMapSlabIndices { //nolint:maprange
+		slabIndices = append(
+			slabIndices,
+			accountStorageMapSlabIndex{
+				Address:   address,
+				SlabIndex: slabIndex,
+			},
+		)
+	}
+	sort.Slice(
+		slabIndices,
+		func(i, j int) bool {
+			slabIndex1 := slabIndices[i]
+			slabIndex2 := slabIndices[j]
+			address1 := slabIndex1.Address
+			address2 := slabIndex2.Address
+			return address1.Compare(address2) < 0
+		},
+	)
+
+	storageKey := []byte(AccountStorageKey)
+
+	for _, slabIndex := range slabIndices {
+
 		var err error
 		errors.WrapPanic(func() {
 			err = s.ledger.SetValue(
-				pair.Key.Address[:],
-				[]byte(pair.Key.Key),
-				pair.Value[:],
+				slabIndex.Address[:],
+				storageKey,
+				slabIndex.SlabIndex[:],
 			)
 		})
 		if err != nil {
 			return interpreter.WrappedExternalError(err)
 		}
 	}
 
+	s.newAccountStorageMapSlabIndices = nil
+
 	return nil
 }
 

runtime/storage_test.go

@@ -68,15 +68,10 @@ func withWritesToStorage(
 		var address common.Address
 		random.Read(address[:])
 
-		storageKey := interpreter.StorageKey{
-			Address: address,
-			Key:     AccountStorageKey,
-		}
-
 		var slabIndex atree.SlabIndex
 		binary.BigEndian.PutUint32(slabIndex[:], randomIndex)
 
-		storage.AccountStorageV2.SetNewAccountStorageMapSlabIndex(storageKey, slabIndex)
+		storage.AccountStorageV2.SetNewAccountStorageMapSlabIndex(address, slabIndex)
 	}
 
 	handler(storage, inter)