atree register migration fix

onflow · Aug 25, 2023 · 6767f31 · 6767f31
1 parent 48e955b
commit 6767f31
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 9 deletions.
diff --git a/cmd/util/ledger/migrations/atree_register_migration.go b/cmd/util/ledger/migrations/atree_register_migration.go
@@ -260,7 +260,6 @@ func (m *AtreeRegisterMigrator) newMigratorRuntime() (
 	// due to this being reentered, only the top stack is interesting.
 	getOrLoadDepth := 0
 	sampler := m.log.Sample(util2.NewTimedSampler(10 * time.Second))
-	wrappedFunc := ri.GetOrLoadProgramFunc
 	ri.GetOrLoadProgramFunc = func(location runtime.Location, load func() (*interpreter.Program, error)) (*interpreter.Program, error) {
 		if getOrLoadDepth == 0 {
 			sampler.Info().Func(
@@ -272,7 +271,7 @@ func (m *AtreeRegisterMigrator) newMigratorRuntime() (
 		getOrLoadDepth += 1
 		defer func() { getOrLoadDepth -= 1 }()
 
-		return wrappedFunc(location, load)
+		return ri.Programs.GetOrLoadProgram(location, load)
 	}
 
 	env.Configure(

diff --git a/cmd/util/ledger/util/migration_runtime_interface.go b/cmd/util/ledger/util/migration_runtime_interface.go
@@ -34,8 +34,6 @@ func newMigrationRuntimeInterface(
 		Programs: Programs,
 	}
 
-	mri.GetOrLoadProgramFunc = mri.defaultGetOrLoadProgram
-
 	return mri
 }
 
@@ -120,12 +118,12 @@ func (m MigrationRuntimeInterface) GetAccountContractCode(
 	return m.Accounts.GetContract(l.Name, flow.Address(l.Address))
 }
 
-func (m MigrationRuntimeInterface) defaultGetOrLoadProgram(location runtime.Location, load func() (*interpreter.Program, error)) (*interpreter.Program, error) {
-	return m.Programs.GetOrLoadProgram(location, load)
-}
-
 func (m MigrationRuntimeInterface) GetOrLoadProgram(location runtime.Location, load func() (*interpreter.Program, error)) (*interpreter.Program, error) {
-	return m.GetOrLoadProgramFunc(location, load)
+	if m.GetOrLoadProgramFunc != nil {
+		return m.GetOrLoadProgramFunc(location, load)
+	}
+
+	return m.Programs.GetOrLoadProgram(location, load)
 }
 
 func (m MigrationRuntimeInterface) MeterMemory(_ common.MemoryUsage) error {