2 files changed, 25 insertions, 27 deletions

diff --git a/src/tools/wasm-shell.cpp b/src/tools/wasm-shell.cpp
index a2d35f444..c6815c92d 100644
--- a/src/tools/wasm-shell.cpp
+++ b/src/tools/wasm-shell.cpp
@@ -81,9 +81,9 @@ struct ShellOptions : public Options {
 
 class Shell {
 protected:
-  std::map<Name, std::unique_ptr<Module>> modules;
-  std::map<Name, std::unique_ptr<SExpressionWasmBuilder>> builders;
-  std::map<Name, std::unique_ptr<ShellExternalInterface>> interfaces;
+  std::map<Name, std::shared_ptr<Module>> modules;
+  std::map<Name, std::shared_ptr<SExpressionWasmBuilder>> builders;
+  std::map<Name, std::shared_ptr<ShellExternalInterface>> interfaces;
   std::map<Name, std::shared_ptr<ModuleInstance>> instances;
   // used for imports
   std::map<Name, std::shared_ptr<ModuleInstance>> linkedInstances;
@@ -92,7 +92,7 @@ protected:
 
   void instantiate(Module* wasm) {
     auto tempInterface =
-      wasm::make_unique<ShellExternalInterface>(linkedInstances);
+      std::make_shared<ShellExternalInterface>(linkedInstances);
     auto tempInstance = std::make_shared<ModuleInstance>(
       *wasm, tempInterface.get(), linkedInstances);
     interfaces[wasm->name].swap(tempInterface);
@@ -123,12 +123,12 @@ protected:
     Colors::green(std::cerr);
     std::cerr << "BUILDING MODULE [line: " << s.line << "]\n";
     Colors::normal(std::cerr);
-    auto module = wasm::make_unique<Module>();
+    auto module = std::make_shared<Module>();
     auto builder =
-      wasm::make_unique<SExpressionWasmBuilder>(*module, s, IRProfile::Normal);
+      std::make_shared<SExpressionWasmBuilder>(*module, s, IRProfile::Normal);
     auto moduleName = module->name;
     lastModule = module->name;
-    builders[moduleName].swap(builder);
+    builders[moduleName] = builder;
     modules[moduleName].swap(module);
     modules[moduleName]->features = FeatureSet::All;
     bool valid = WasmValidator().validate(*modules[moduleName]);
@@ -150,15 +150,12 @@ protected:
     auto name = s[1]->str();
     linkedInstances[name] = instance;
 
-    // swap to the new name in all maps
-    modules[name].swap(modules[lastModule]);
-    modules.erase(lastModule);
-    builders[name].swap(builders[lastModule]);
-    builders.erase(lastModule);
-    interfaces[name].swap(interfaces[lastModule]);
-    interfaces.erase(lastModule);
-    instances[name].swap(instances[lastModule]);
-    instances.erase(lastModule);
+    // we copy pointers as a registered module's name might still be used
+    // in an assertion or invoke command.
+    modules[name] = modules[lastModule];
+    builders[name] = builders[lastModule];
+    interfaces[name] = interfaces[lastModule];
+    instances[name] = instances[lastModule];
 
     Colors::green(std::cerr);
     std::cerr << "REGISTER MODULE INSTANCE AS \"" << name.c_str()
@@ -316,7 +313,7 @@ protected:
   // TODO: spectest module is considered deprecated by the spec. Remove when
   // is actually removed from the spec test.
   void buildSpectestModule() {
-    auto spectest = std::make_unique<Module>();
+    auto spectest = std::make_shared<Module>();
     spectest->name = "spectest";
     Builder builder(*spectest);
 
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index 8c772a96e..844716f95 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -2422,23 +2422,24 @@ private:
   protected:
     // Returns the instance that defines the memory used by this one.
     SubType* getMemoryInstance() {
-      if (instance.wasm.memory.imported()) {
-        return instance.linkedInstances.at(instance.wasm.memory.module).get();
-      } else {
-        return static_cast<SubType*>(&instance);
+      auto* inst = instance.self();
+      while (inst->wasm.memory.imported()) {
+        inst = inst->linkedInstances.at(inst->wasm.memory.module).get();
       }
+      return inst;
     }
 
     // Returns a reference to the current value of a potentially imported global
     Literals& getGlobal(Name name) {
-      auto* global = instance.wasm.getGlobal(name);
-      if (global->imported()) {
-        auto inst = instance.linkedInstances.at(global->module);
+      auto* inst = instance.self();
+      auto* global = inst->wasm.getGlobal(name);
+      while (global->imported()) {
+        inst = inst->linkedInstances.at(global->module).get();
         Export* globalExport = inst->wasm.getExport(global->base);
-        return inst->globals[globalExport->value];
-      } else {
-        return instance.globals[name];
+        global = inst->wasm.getGlobal(globalExport->value);
       }
+
+      return inst->globals[global->name];
     }
 
   public: