Add import resolver and use it for globals

Author: stevenfontanella

src/shell-interface.h

@@ -129,18 +129,6 @@ struct ShellExternalInterface : ModuleRunner::ExternalInterface {
       wasm, [&](Table* table) { tables[table->name].resize(table->initial); });
   }
 
-  void importGlobals(std::map<Name, Literals>& globals, Module& wasm) override {
-    ModuleUtils::iterImportedGlobals(wasm, [&](Global* import) {
-      auto inst = getImportInstance(import);
-      auto* exportedGlobal = inst->wasm.getExportOrNull(import->base);
-      if (!exportedGlobal || exportedGlobal->kind != ExternalKind::Global) {
-        Fatal() << "importGlobals: unknown import: " << import->module.str
-                << "." << import->name.str;
-      }
-      globals[import->name] = inst->globals[*exportedGlobal->getInternalName()];
-    });
-  }
-
   Literal getImportedFunction(Function* import) override {
     // TODO: We should perhaps restrict the types with which the well-known
     // functions can be imported.

src/tools/wasm-ctor-eval.cpp

@@ -67,13 +67,40 @@ bool isNullableAndMutable(Expression* ref, Index fieldIndex) {
 // the output.
 #define RECOMMENDATION "\n       recommendation: "
 
+class EvallingModuleRunner;
+
+class EvallingImportResolver : public ImportResolver {
+public:
+  EvallingImportResolver(
+    std::map<Name, std::shared_ptr<EvallingModuleRunner>> linkedInstances,
+    ModuleRunnerBase<EvallingModuleRunner>::ExternalInterface*
+      externalInterface)
+    : externalInterface(externalInterface) {}
+
+  std::optional<Literals*> getGlobal(QualifiedName name,
+                                     Type type) const override {
+    externalInterface->trap((std::stringstream()
+                             << "EvallingImportResolver: unexpected getGlobal "
+                             << name)
+                              .str());
+    return std::nullopt;
+  }
+
+private:
+  ModuleRunnerBase<EvallingModuleRunner>::ExternalInterface* externalInterface;
+};
+
 class EvallingModuleRunner : public ModuleRunnerBase<EvallingModuleRunner> {
 public:
   EvallingModuleRunner(
     Module& wasm,
     ExternalInterface* externalInterface,
     std::map<Name, std::shared_ptr<EvallingModuleRunner>> linkedInstances_ = {})
-    : ModuleRunnerBase(wasm, externalInterface, linkedInstances_) {}
+    : ModuleRunnerBase(wasm,
+                       externalInterface,
+                       std::make_shared<EvallingImportResolver>(
+                         linkedInstances_, externalInterface),
+                       linkedInstances_) {}
 
   Flow visitGlobalGet(GlobalGet* curr) {
     // Error on reads of imported globals.
@@ -147,19 +174,6 @@ std::unique_ptr<Module> buildEnvModule(Module& wasm) {
     }
   });
 
-  ModuleUtils::iterImportedGlobals(wasm, [&](Global* global) {
-    if (global->module == env->name) {
-      auto* copied = ModuleUtils::copyGlobal(global, *env);
-      copied->module = Name();
-      copied->base = Name();
-
-      Builder builder(*env);
-      copied->init = builder.makeConst(Literal::makeZero(global->type));
-      env->addExport(
-        builder.makeExport(global->base, copied->name, ExternalKind::Global));
-    }
-  });
-
   // create an exported memory with the same initial and max size
   ModuleUtils::iterImportedMemories(wasm, [&](Memory* memory) {
     if (memory->module == env->name) {
@@ -231,26 +245,6 @@ struct CtorEvalExternalInterface : EvallingModuleRunner::ExternalInterface {
     }
   }
 
-  void importGlobals(GlobalValueSet& globals, Module& wasm_) override {
-    ModuleUtils::iterImportedGlobals(wasm_, [&](Global* global) {
-      auto it = linkedInstances.find(global->module);
-      if (it != linkedInstances.end()) {
-        auto* inst = it->second.get();
-        auto* globalExport = inst->wasm.getExportOrNull(global->base);
-        if (!globalExport || globalExport->kind != ExternalKind::Global) {
-          throw FailToEvalException(std::string("importGlobals: ") +
-                                    global->module.toString() + "." +
-                                    global->base.toString());
-        }
-        globals[global->name] = inst->globals[*globalExport->getInternalName()];
-      } else {
-        throw FailToEvalException(std::string("importGlobals: ") +
-                                  global->module.toString() + "." +
-                                  global->base.toString());
-      }
-    });
-  }
-
   Literal getImportedFunction(Function* import) override {
     auto f = [import, this](const Literals& arguments) -> Flow {
       Name WASI("wasi_snapshot_preview1");
@@ -558,8 +552,8 @@ struct CtorEvalExternalInterface : EvallingModuleRunner::ExternalInterface {
   void applyGlobalsToModule() {
     if (!wasm->features.hasGC()) {
       // Without GC, we can simply serialize the globals in place as they are.
-      for (const auto& [name, values] : instance->globals) {
-        wasm->getGlobal(name)->init = getSerialization(values);
+      for (const auto& [name, values] : instance->allGlobals) {
+        wasm->getGlobal(name)->init = getSerialization(*values);
       }
       return;
     }
@@ -590,9 +584,9 @@ struct CtorEvalExternalInterface : EvallingModuleRunner::ExternalInterface {
       // for it. (If there is no value, then this is a new global we've added
       // during execution, for whom we've already set up a proper serialized
       // value when we created it.)
-      auto iter = instance->globals.find(oldGlobal->name);
-      if (iter != instance->globals.end()) {
-        oldGlobal->init = getSerialization(iter->second, name);
+      auto iter = instance->allGlobals.find(oldGlobal->name);
+      if (iter != instance->allGlobals.end()) {
+        oldGlobal->init = getSerialization(*iter->second, name);
       }
 
       // Add the global back to the module.

src/tools/wasm-shell.cpp

@@ -278,7 +278,7 @@ struct Shell {
       }
       auto& instance = it->second;
       try {
-        return instance->getExportedGlobal(get->name);
+        return instance->getExportedGlobalOrTrap(get->name);
       } catch (TrapException&) {
         return TrapResult{};
       } catch (...) {

src/wasm-interpreter.h

@@ -29,6 +29,7 @@
 #define wasm_wasm_interpreter_h
 
 #include <cmath>
+#include <iomanip>
 #include <limits.h>
 #include <sstream>
 #include <variant>
@@ -46,6 +47,7 @@
 #include "wasm-limits.h"
 #include "wasm-traversal.h"
 #include "wasm.h"
+#include "wasm/import-resolver.h"
 
 #if __has_feature(leak_sanitizer) || __has_feature(address_sanitizer)
 #include <sanitizer/lsan_interface.h>
@@ -2940,8 +2942,6 @@ class ConstantExpressionRunner : public ExpressionRunner<SubType> {
   }
 };
 
-using GlobalValueSet = std::map<Name, Literals>;
-
 //
 // A runner for a module. Each runner contains the information to execute the
 // module, such as the state of globals, and so forth, so it basically
@@ -2969,7 +2969,6 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
       std::map<Name, std::shared_ptr<SubType>> linkedInstances = {}) {}
     virtual ~ExternalInterface() = default;
     virtual void init(Module& wasm, SubType& instance) {}
-    virtual void importGlobals(GlobalValueSet& globals, Module& wasm) = 0;
     virtual Literal getImportedFunction(Function* import) = 0;
     virtual bool growMemory(Name name, Address oldSize, Address newSize) = 0;
     virtual bool growTable(Name name,
@@ -3174,18 +3173,22 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
   // TODO: this duplicates module in ExpressionRunner, and can be removed
   Module& wasm;
 
-  // Values of globals
-  GlobalValueSet globals;
-
   // Multivalue ABI support (see push/pop).
   std::vector<Literals> multiValues;
 
+  // keyed by internal name
+  std::map<Name, Literals> definedGlobals;
+  std::map<Name, Literals*> allGlobals;
+
   ModuleRunnerBase(
     Module& wasm,
     ExternalInterface* externalInterface,
+    std::shared_ptr<ImportResolver> importResolver,
     std::map<Name, std::shared_ptr<SubType>> linkedInstances_ = {})
     : ExpressionRunner<SubType>(&wasm), wasm(wasm),
-      externalInterface(externalInterface), linkedInstances(linkedInstances_) {
+      externalInterface(externalInterface),
+      linkedInstances(std::move(linkedInstances_)),
+      importResolver(std::move(importResolver)) {
     // Set up a single shared CurrContinuations for all these linked instances,
     // reusing one if it exists.
     std::shared_ptr<ContinuationStore> shared;
@@ -3208,16 +3211,11 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
   // (This is separate from the constructor so that it does not occur
   // synchronously, which makes some code patterns harder to write.)
   void instantiate() {
-    // import globals from the outside
-    externalInterface->importGlobals(globals, wasm);
-    // generate internal (non-imported) globals
-    ModuleUtils::iterDefinedGlobals(wasm, [&](Global* global) {
-      globals[global->name] = self()->visit(global->init).values;
-    });
-
     // initialize the rest of the external interface
     externalInterface->init(wasm, *self());
 
+    initializeGlobals();
+
     initializeTableContents();
     initializeMemoryContents();
 
@@ -3254,20 +3252,30 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
                    func->type);
   }
 
-  // get an exported global
-  Literals getExportedGlobal(Name name) {
+  std::optional<Literals*> getExportedGlobal(Name name) {
     Export* export_ = wasm.getExportOrNull(name);
     if (!export_ || export_->kind != ExternalKind::Global) {
-      externalInterface->trap("getExport external not found");
+      return std::nullopt;
     }
     Name internalName = *export_->getInternalName();
-    auto iter = globals.find(internalName);
-    if (iter == globals.end()) {
-      externalInterface->trap("getExport internal not found");
+    auto iter = allGlobals.find(internalName);
+    if (iter == allGlobals.end()) {
+      return std::nullopt;
     }
     return iter->second;
   }
 
+  Literals& getExportedGlobalOrTrap(Name name) {
+    auto global = getExportedGlobal(name);
+    if (!global.has_value()) {
+      externalInterface->trap((std::stringstream()
+                               << "getExportedGlobal: export " << name
+                               << " not found.")
+                                .str());
+    }
+    return **global;
+  }
+
   Tag* getExportedTag(Name name) {
     Export* export_ = wasm.getExportOrNull(name);
     if (!export_ || export_->kind != ExternalKind::Tag) {
@@ -3321,6 +3329,34 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     return TableInstanceInfo{self(), name};
   }
 
+  void initializeGlobals() {
+    for (auto& global : wasm.globals) {
+      if (global->imported()) {
+        QualifiedName name{global->module, global->base};
+        auto importedGlobal = importResolver->getGlobal(name, global->type);
+        if (!importedGlobal) {
+          externalInterface->trap(
+            (std::stringstream() << "Imported global " << name << " not found.")
+              .str());
+        }
+        allGlobals[global->name] = *importedGlobal;
+      } else {
+        Literals init = self()->visit(global->init).values;
+        auto [it, inserted] = definedGlobals.emplace(global->name, init);
+
+        // This was likely checked during parsing or validation
+        if (!inserted) {
+          externalInterface->trap(
+            (std::stringstream()
+             << "Global: " << std::quoted(global->name.toString())
+             << " was defined twice.")
+              .str());
+        }
+        allGlobals[global->name] = &it->second;
+      }
+    }
+  }
+
   void initializeTableContents() {
     for (auto& table : wasm.tables) {
       if (table->type.isNullable()) {
@@ -3511,20 +3547,8 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
   SmallVector<std::pair<WasmException, Name>, 4> exceptionStack;
 
 protected:
-  // Returns a reference to the current value of a potentially imported global.
-  Literals& getGlobal(Name name) {
-    auto* inst = self();
-    auto* global = inst->wasm.getGlobal(name);
-    while (global->imported()) {
-      inst = inst->linkedInstances.at(global->module).get();
-      Export* globalExport = inst->wasm.getExport(global->base);
-      global = inst->wasm.getGlobal(*globalExport->getInternalName());
-    }
-
-    return inst->globals[global->name];
-  }
-
-  // As above, but for a function.
+  // Returns a reference to the current value of a potentially imported
+  // function.
   Literal getFunction(Name name) {
     auto* inst = self();
     auto* func = inst->wasm.getFunction(name);
@@ -3846,13 +3870,13 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
 
   Flow visitGlobalGet(GlobalGet* curr) {
     auto name = curr->name;
-    return getGlobal(name);
+    return *allGlobals.at(name);
   }
   Flow visitGlobalSet(GlobalSet* curr) {
     auto name = curr->name;
     VISIT(flow, curr->value)
 
-    getGlobal(name) = flow.values;
+    *allGlobals.at(name) = flow.values;
     return Flow();
   }
 
@@ -5055,6 +5079,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
 
   ExternalInterface* externalInterface;
   std::map<Name, std::shared_ptr<SubType>> linkedInstances;
+  std::shared_ptr<ImportResolver> importResolver;
 };
 
 class ModuleRunner : public ModuleRunnerBase<ModuleRunner> {
@@ -5063,7 +5088,12 @@ class ModuleRunner : public ModuleRunnerBase<ModuleRunner> {
     Module& wasm,
     ExternalInterface* externalInterface,
     std::map<Name, std::shared_ptr<ModuleRunner>> linkedInstances = {})
-    : ModuleRunnerBase(wasm, externalInterface, linkedInstances) {}
+    : ModuleRunnerBase(
+        wasm,
+        externalInterface,
+        std::make_shared<LinkedInstancesImportResolver<ModuleRunner>>(
+          linkedInstances),
+        linkedInstances) {}
 
   Literal makeFuncData(Name name, Type type) {
     // As the super's |makeFuncData|, but here we also provide a way to