From ba76c1111ddc6f985953c2d56f969a528abdf905 Mon Sep 17 00:00:00 2001
From: Faye Amacker <33205765+fxamacker@users.noreply.github.com>
Date: Mon, 25 Nov 2024 16:59:00 -0600
Subject: [PATCH] Extract new function NewPersistentSlabStorage()

---
 runtime/storage.go      | 16 ++++++++++++----
 runtime/storage_test.go | 37 ++-----------------------------------
 2 files changed, 14 insertions(+), 39 deletions(-)

diff --git a/runtime/storage.go b/runtime/storage.go
index 8198e69fa..49aaa92b4 100644
--- a/runtime/storage.go
+++ b/runtime/storage.go
@@ -69,11 +69,10 @@ type Storage struct {
 var _ atree.SlabStorage = &Storage{}
 var _ interpreter.Storage = &Storage{}
 
-func NewStorage(
+func NewPersistentSlabStorage(
 	ledger atree.Ledger,
 	memoryGauge common.MemoryGauge,
-	config StorageConfig,
-) *Storage {
+) *atree.PersistentSlabStorage {
 	decodeStorable := func(
 		decoder *cbor.StreamDecoder,
 		slabID atree.SlabID,
@@ -95,13 +94,22 @@ func NewStorage(
 	}
 
 	ledgerStorage := atree.NewLedgerBaseStorage(ledger)
-	persistentSlabStorage := atree.NewPersistentSlabStorage(
+
+	return atree.NewPersistentSlabStorage(
 		ledgerStorage,
 		interpreter.CBOREncMode,
 		interpreter.CBORDecMode,
 		decodeStorable,
 		decodeTypeInfo,
 	)
+}
+
+func NewStorage(
+	ledger atree.Ledger,
+	memoryGauge common.MemoryGauge,
+	config StorageConfig,
+) *Storage {
+	persistentSlabStorage := NewPersistentSlabStorage(ledger, memoryGauge)
 
 	accountStorageV1 := NewAccountStorageV1(
 		ledger,
diff --git a/runtime/storage_test.go b/runtime/storage_test.go
index d1e24cfd9..f0bbaef75 100644
--- a/runtime/storage_test.go
+++ b/runtime/storage_test.go
@@ -30,7 +30,6 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/fxamacker/cbor/v2"
 	"github.com/onflow/atree"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -8363,7 +8362,7 @@ func TestGetDomainStorageMapRegisterReadsForV1Account(t *testing.T) {
 		return func() (storedValues map[string][]byte, StorageIndices map[string]uint64) {
 			ledger := NewTestLedger(nil, nil)
 
-			persistentSlabStorage := newSlabStorage(ledger)
+			persistentSlabStorage := NewPersistentSlabStorage(ledger, nil)
 
 			orderedMap, err := atree.NewMap(
 				persistentSlabStorage,
@@ -8761,7 +8760,7 @@ func TestGetDomainStorageMapRegisterReadsForV2Account(t *testing.T) {
 		return func() (storedValues map[string][]byte, StorageIndices map[string]uint64) {
 			ledger := NewTestLedger(nil, nil)
 
-			persistentSlabStorage := newSlabStorage(ledger)
+			persistentSlabStorage := NewPersistentSlabStorage(ledger, nil)
 
 			accountOrderedMap, err := atree.NewMap(
 				persistentSlabStorage,
@@ -9144,38 +9143,6 @@ func checkAccountStorageMapData(
 	require.Contains(tb, rootSlabIDs, accountSlabID)
 }
 
-func newSlabStorage(ledger atree.Ledger) *atree.PersistentSlabStorage {
-	decodeStorable := func(
-		decoder *cbor.StreamDecoder,
-		slabID atree.SlabID,
-		inlinedExtraData []atree.ExtraData,
-	) (
-		atree.Storable,
-		error,
-	) {
-		return interpreter.DecodeStorable(
-			decoder,
-			slabID,
-			inlinedExtraData,
-			nil,
-		)
-	}
-
-	decodeTypeInfo := func(decoder *cbor.StreamDecoder) (atree.TypeInfo, error) {
-		return interpreter.DecodeTypeInfo(decoder, nil)
-	}
-
-	ledgerStorage := atree.NewLedgerBaseStorage(ledger)
-
-	return atree.NewPersistentSlabStorage(
-		ledgerStorage,
-		interpreter.CBOREncMode,
-		interpreter.CBORDecMode,
-		decodeStorable,
-		decodeTypeInfo,
-	)
-}
-
 func concatRegisterAddressAndKey(
 	address common.Address,
 	key []byte,