From 00d9d885ffc74c15221f152504d77af7db695a7e Mon Sep 17 00:00:00 2001
From: Clifton Kaznocha <ckaznocha@gmail.com>
Date: Sun, 16 Apr 2023 17:32:08 -0700
Subject: [PATCH] Cleanup aliased modules when the main module is deleted
 (#1365)

Signed-off-by: Clifton Kaznocha <ckaznocha@users.noreply.github.com>
---
 internal/wasm/store.go                  |  5 +++++
 internal/wasm/store_module_list.go      | 16 +++++++++++++++-
 internal/wasm/store_module_list_test.go | 19 +++++++++++++++----
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 83bfe5af..bc6bc627 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -119,6 +119,11 @@ type (
 		s *Store
 		// prev and next hold the nodes in the linked list of ModuleInstance held by Store.
 		prev, next *ModuleInstance
+		// aliases holds the module names that are aliases of this module registered in the store.
+		// Access to this field must be guarded by s.mux.
+		//
+		// Note: This is currently only used for spectests and will be nil in most cases.
+		aliases []string
 		// Definitions is derived from *Module, and is constructed during compilation phrase.
 		Definitions []FunctionDefinition
 	}
diff --git a/internal/wasm/store_module_list.go b/internal/wasm/store_module_list.go
index 02e2842c..bdfbbdec 100644
--- a/internal/wasm/store_module_list.go
+++ b/internal/wasm/store_module_list.go
@@ -30,6 +30,12 @@ func (s *Store) deleteModule(m *ModuleInstance) error {
 	if m.ModuleName != "" {
 		delete(s.nameToModule, m.ModuleName)
 
+		// Under normal circumstances, m.aliases will be nil this loop will not
+		// be entered unless aliases have been created. See `*store.AliasModule`
+		for _, alias := range m.aliases {
+			delete(s.nameToModule, alias)
+		}
+
 		// Shrink the map if it's allocated more than twice the size of the list
 		newCap := len(s.nameToModule)
 		if newCap < nameToModuleShrinkThreshold {
@@ -93,7 +99,15 @@ func (s *Store) registerModule(m *ModuleInstance) error {
 func (s *Store) AliasModule(src, dst string) error {
 	s.mux.Lock()
 	defer s.mux.Unlock()
-	s.nameToModule[dst] = s.nameToModule[src]
+	if _, ok := s.nameToModule[dst]; ok {
+		return nil
+	}
+	m, ok := s.nameToModule[src]
+	if !ok {
+		return nil
+	}
+	m.aliases = append(m.aliases, dst)
+	s.nameToModule[dst] = m
 	if len(s.nameToModule) > s.nameToModuleCap {
 		s.nameToModuleCap = len(s.nameToModule)
 	}
diff --git a/internal/wasm/store_module_list_test.go b/internal/wasm/store_module_list_test.go
index 75bdbf71..462c6880 100644
--- a/internal/wasm/store_module_list_test.go
+++ b/internal/wasm/store_module_list_test.go
@@ -108,17 +108,28 @@ func TestStore_module(t *testing.T) {
 }
 
 func TestStore_AliasModule(t *testing.T) {
-	s := newStore()
-	m1 := &ModuleInstance{ModuleName: "m1"}
-	s.nameToModule[m1.ModuleName] = m1
-
 	t.Run("alias module", func(t *testing.T) {
+		s := newStore()
+		m1 := &ModuleInstance{ModuleName: "m1"}
+		s.nameToModule[m1.ModuleName] = m1
+
 		require.NoError(t, s.AliasModule("m1", "m2"))
 		require.Equal(t, map[string]*ModuleInstance{"m1": m1, "m2": m1}, s.nameToModule)
 		// Doesn't affect module names
 		require.Nil(t, s.moduleList)
 		require.Equal(t, nameToModuleShrinkThreshold, s.nameToModuleCap)
 	})
+
+	t.Run("delete aliased module", func(t *testing.T) {
+		s := newStore()
+		m1 := &ModuleInstance{ModuleName: "m1"}
+		s.nameToModule[m1.ModuleName] = m1
+
+		require.NoError(t, s.AliasModule("m1", "m2"))
+		require.NoError(t, s.deleteModule(m1))
+		_, ok := s.nameToModule["m2"]
+		require.False(t, ok)
+	})
 }
 
 func TestStore_nameToModuleCap(t *testing.T) {