From 77032fcfcd1f3072c64438aa6234331d2aecf651 Mon Sep 17 00:00:00 2001
From: deelawn <dboltz03@gmail.com>
Date: Fri, 29 Mar 2024 18:15:28 -0700
Subject: [PATCH] save package mempackage after init

---
 gnovm/pkg/gnolang/machine.go | 16 ++++++++++++++--
 gnovm/pkg/gnolang/realm.go   |  9 ++++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/gnovm/pkg/gnolang/machine.go b/gnovm/pkg/gnolang/machine.go
index 80354ecd1fd..f767cc10415 100644
--- a/gnovm/pkg/gnolang/machine.go
+++ b/gnovm/pkg/gnolang/machine.go
@@ -596,10 +596,10 @@ func (m *Machine) runFiles(memPkg *std.MemPackage, fns ...*FileNode) {
 		}
 	}
 
-	// Save the mempackage if provided. We do this here so that it occurs before
+	// Save the realm mempackage if provided. We do this here so that it occurs before
 	// the init functions are run. This avoids any realm object ownership issues that may arise
 	// when passing around unpersisted realm object pointers to other realms during initialization.
-	if memPkg != nil {
+	if memPkg != nil && IsRealmPath(memPkg.Path) {
 		// store package values and types
 		m.savePackageValuesAndTypes()
 		// store mempackage
@@ -625,6 +625,18 @@ func (m *Machine) runFiles(memPkg *std.MemPackage, fns ...*FileNode) {
 			}
 		}
 	}
+
+	// If this is a package, we can save it after the init function because init is the
+	// only way stateful realm variables can be dynamically set in a package. We don't need to
+	// worry about pointer ownership issues for packages because they (eventually) will only be
+	// able to import other packages and those packages won't be able to persist any pointers
+	// to their own state that they receive from other packages.
+	if memPkg != nil && IsPackagePath(memPkg.Path) {
+		// store package values and types
+		m.savePackageValuesAndTypes()
+		// store mempackage
+		m.Store.AddMemPackage(memPkg)
+	}
 }
 
 // Save the machine's package using realm finalization deep crawl.
diff --git a/gnovm/pkg/gnolang/realm.go b/gnovm/pkg/gnolang/realm.go
index 4a8fe504301..17d82e8c553 100644
--- a/gnovm/pkg/gnolang/realm.go
+++ b/gnovm/pkg/gnolang/realm.go
@@ -1521,7 +1521,10 @@ func isUnsaved(oo Object) bool {
 
 // realmPathPrefix is the prefix used to identify pkgpaths which are meant to
 // be realms and as such to have their state persisted. This is used by [IsRealmPath].
-const realmPathPrefix = "gno.land/r/"
+const (
+	realmPathPrefix   = "gno.land/r/"
+	packagePathPrefix = "gno.land/p"
+)
 
 // IsRealmPath determines whether the given pkgpath is for a realm, and as such
 // should persist the global state.
@@ -1529,6 +1532,10 @@ func IsRealmPath(pkgPath string) bool {
 	return strings.HasPrefix(pkgPath, realmPathPrefix)
 }
 
+func IsPackagePath(pkgPath string) bool {
+	return strings.HasPrefix(pkgPath, packagePathPrefix)
+}
+
 func prettyJSON(jstr []byte) []byte {
 	var c interface{}
 	err := json.Unmarshal(jstr, &c)