add apis for accessing module elements

16 files changed, 78 insertions, 64 deletions

diff --git a/src/asm2wasm.h b/src/asm2wasm.h
index 7d9b9652b..e4bfdc341 100644
--- a/src/asm2wasm.h
+++ b/src/asm2wasm.h
@@ -655,10 +655,10 @@ void Asm2WasmBuilder::processAsm(Ref ast) {
         IString value = pair[1][1]->getIString();
         if (key == Name("_emscripten_replace_memory")) {
           // asm.js memory growth provides this special non-asm function, which we don't need (we use grow_memory)
-          assert(wasm.functionsMap.find(value) == wasm.functionsMap.end());
+          assert(!wasm.checkFunction(value));
           continue;
         }
-        assert(wasm.functionsMap.find(value) != wasm.functionsMap.end());
+        assert(wasm.checkFunction(value));
         auto export_ = allocator.alloc<Export>();
         export_->name = key;
         export_->value = value;
@@ -671,18 +671,17 @@ void Asm2WasmBuilder::processAsm(Ref ast) {
 
   std::vector<IString> toErase;
 
-  for (auto& pair : wasm.importsMap) {
-    IString name = pair.first;
-    Import& import = *pair.second;
+  for (auto* import : wasm.imports) {
+    IString name = import->name;
     if (importedFunctionTypes.find(name) != importedFunctionTypes.end()) {
       // special math builtins
-      FunctionType* builtin = getBuiltinFunctionType(import.module, import.base);
+      FunctionType* builtin = getBuiltinFunctionType(import->module, import->base);
       if (builtin) {
-        import.type = builtin;
+        import->type = builtin;
         continue;
       }
-      import.type = ensureFunctionType(getSig(&importedFunctionTypes[name]), &wasm, allocator);
-    } else if (import.module != ASM2WASM) { // special-case the special module
+      import->type = ensureFunctionType(getSig(&importedFunctionTypes[name]), &wasm, allocator);
+    } else if (import->module != ASM2WASM) { // special-case the special module
       // never actually used
       toErase.push_back(name);
     }
@@ -1186,7 +1185,7 @@ Function* Asm2WasmBuilder::processFunction(Ref ast) {
           }
         }
         Call* ret;
-        if (wasm.importsMap.find(name) != wasm.importsMap.end()) {
+        if (wasm.checkImport(name)) {
           Ref parent = astStackHelper.getParent();
           WasmType type = !!parent ? detectWasmType(parent, &asmData) : none;
           auto specific = allocator.alloc<CallImport>();
diff --git a/src/asm_v_wasm.h b/src/asm_v_wasm.h
index 394c7a321..29344e9bb 100644
--- a/src/asm_v_wasm.h
+++ b/src/asm_v_wasm.h
@@ -113,8 +113,8 @@ FunctionType sigToFunctionType(std::string sig) {
 
 FunctionType* ensureFunctionType(std::string sig, Module* wasm, MixedArena& allocator) {
   cashew::IString name(("FUNCSIG$" + sig).c_str(), false);
-  if (wasm->functionTypesMap.find(name) != wasm->functionTypesMap.end()) {
-    return wasm->functionTypesMap[name];
+  if (wasm->checkFunctionType(name)) {
+    return wasm->getFunctionType(name);
   }
   // add new type
   auto type = allocator.alloc<FunctionType>();
diff --git a/src/binaryen-shell.cpp b/src/binaryen-shell.cpp
index c2b6fcaa7..8fc067562 100644
--- a/src/binaryen-shell.cpp
+++ b/src/binaryen-shell.cpp
@@ -249,7 +249,7 @@ static void run_asserts(size_t* i, bool* checked, AllocatingModule* wasm,
     interface = new ShellExternalInterface();
     instance = new ModuleInstance(*wasm, interface);
     if (entry.is()) {
-      Function* function = wasm->functionsMap[entry];
+      Function* function = wasm->getFunction(entry);
       if (!function) {
         std::cerr << "Unknown entry " << entry << std::endl;
       } else {
diff --git a/src/passes/RemoveImports.cpp b/src/passes/RemoveImports.cpp
index cd741180f..6463df10c 100644
--- a/src/passes/RemoveImports.cpp
+++ b/src/passes/RemoveImports.cpp
@@ -29,15 +29,15 @@ namespace wasm {
 
 struct RemoveImports : public WalkerPass<WasmWalker<RemoveImports>> {
   MixedArena* allocator;
-  std::map<Name, Import*> importsMap;
+  Module* module;
 
-  void prepare(PassRunner* runner, Module *module) override {
+  void prepare(PassRunner* runner, Module *module_) override {
     allocator = runner->allocator;
-    importsMap = module->importsMap;
+    module = module_;
   }
 
   void visitCallImport(CallImport *curr) {
-    WasmType type = importsMap[curr->target]->type->result;
+    WasmType type = module->getImport(curr->target)->type->result;
     if (type == none) {
       replaceCurrent(allocator->alloc<Nop>());
     } else {
@@ -48,8 +48,10 @@ struct RemoveImports : public WalkerPass<WasmWalker<RemoveImports>> {
   }
 
   void visitModule(Module *curr) {
-    curr->importsMap.clear();
-    curr->imports.clear();
+    auto imports = curr->imports; // copy
+    for (auto* import : imports) {
+      curr->removeImport(import->name);
+    }
   }
 };
 
diff --git a/src/s2wasm.h b/src/s2wasm.h
index 1fa624dd0..4188003fd 100644
--- a/src/s2wasm.h
+++ b/src/s2wasm.h
@@ -752,7 +752,7 @@ class S2WasmBuilder {
           indirect->operands.push_back(inputs[i]);
         }
         setOutput(indirect, assign);
-        indirect->fullType = wasm.functionTypesMap[ensureFunctionType(getSig(indirect), &wasm, allocator)->name];
+        indirect->fullType = wasm.getFunctionType(ensureFunctionType(getSig(indirect), &wasm, allocator)->name);
       } else {
         // non-indirect call
         CallBase* curr;
@@ -782,7 +782,7 @@ class S2WasmBuilder {
         setOutput(curr, assign);
         if (curr->is<CallImport>()) {
           auto target = curr->cast<CallImport>()->target;
-          if (wasm.importsMap.count(target) == 0) {
+          if (!wasm.checkImport(target)) {
             auto import = allocator.alloc<Import>();
             import->name = import->base = target;
             import->module = ENV;
@@ -1232,7 +1232,7 @@ class S2WasmBuilder {
 
     // XXX For now, export all functions marked .globl.
     for (Name name : globls) {
-      if (wasm.functionsMap.count(name)) {
+      if (wasm.checkFunction(name)) {
         auto exp = allocator.alloc<Export>();
         exp->name = exp->value = name;
         wasm.addExport(exp);
@@ -1258,7 +1258,7 @@ class S2WasmBuilder {
         if (debug) std::cerr << "  ==> " << *(relocation.data) << '\n';
       } else {
         // must be a function address
-        if (wasm.functionsMap.count(name) == 0) {
+        if (!wasm.checkFunction(name)) {
           std::cerr << "Unknown symbol: " << name << '\n';
           if (!ignoreUnknownSymbols) abort();
           *(relocation.data) = 0;
@@ -1273,7 +1273,7 @@ class S2WasmBuilder {
         std::cerr << "Unknown start function: `" << startFunction << "`\n";
         abort();
       }
-      const auto *target = wasm.functionsMap[startFunction];
+      const auto *target = wasm.getFunction(startFunction);
       Name start("_start");
       if (implementedFunctions.count(start) != 0) {
         std::cerr << "Start function already present: `" << start << "`\n";
@@ -1308,7 +1308,7 @@ class S2WasmBuilder {
 
     // ensure an explicit function type for indirect call targets
     for (auto& name : wasm.table.names) {
-      auto* func = wasm.functionsMap[name];
+      auto* func = wasm.getFunction(name);
       func->type = ensureFunctionType(getSig(func), &wasm, allocator)->name;
     }
   }
diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index edc904787..5695cffcc 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -1709,8 +1709,7 @@ public:
   void visitCallImport(CallImport *curr) {
     if (debug) std::cerr << "zz node: CallImport" << std::endl;
     curr->target = wasm.imports[getU32LEB()]->name;
-    assert(wasm.importsMap.find(curr->target) != wasm.importsMap.end());
-    auto type = wasm.importsMap[curr->target]->type;
+    auto type = wasm.getImport(curr->target)->type;
     assert(type);
     auto num = type->params.size();
     if (debug) std::cerr << "zz node: CallImport " << curr->target << " with type " << type->name << " and " << num << " params\n";
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index 4af410de1..2be386f95 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -110,7 +110,7 @@ public:
   }
 
   Literal callExport(IString name, LiteralList& arguments) {
-    Export *export_ = wasm.exportsMap[name];
+    Export *export_ = wasm.getExport(name);
     if (!export_) externalInterface->trap("callExport not found");
     return callFunction(export_->value, arguments);
   }
@@ -337,7 +337,7 @@ private:
         LiteralList arguments;
         Flow flow = generateArguments(curr->operands, arguments);
         if (flow.breaking()) return flow;
-        return instance.externalInterface->callImport(instance.wasm.importsMap[curr->target], arguments);
+        return instance.externalInterface->callImport(instance.wasm.getImport(curr->target), arguments);
       }
       Flow visitCallIndirect(CallIndirect *curr) {
         NOTE_ENTER("CallIndirect");
@@ -346,7 +346,7 @@ private:
         size_t index = target.value.geti32();
         if (index >= instance.wasm.table.names.size()) trap("callIndirect: overflow");
         Name name = instance.wasm.table.names[index];
-        Function *func = instance.wasm.functionsMap[name];
+        Function *func = instance.wasm.getFunction(name);
         if (func->type.is() && func->type != curr->fullType->name) trap("callIndirect: bad type");
         LiteralList arguments;
         Flow flow = generateArguments(curr->operands, arguments);
@@ -669,7 +669,7 @@ private:
     callDepth++;
     functionStack.push_back(name);
 
-    Function *function = wasm.functionsMap[name];
+    Function *function = wasm.getFunction(name);
     assert(function);
     FunctionScope scope(function, arguments);
 
diff --git a/src/wasm-js.cpp b/src/wasm-js.cpp
index 61eac6e6c..6612e1e00 100644
--- a/src/wasm-js.cpp
+++ b/src/wasm-js.cpp
@@ -110,7 +110,7 @@ void finalizeModule() {
     exit(EXIT_FAILURE);
   }
   module->memory.initial = providedMemory / Memory::kPageSize;
-  module->memory.max = (module->exportsMap.find(GROW_WASM_MEMORY) != module->exportsMap.end()) ? -1 : module->memory.initial;
+  module->memory.max = module->checkExport(GROW_WASM_MEMORY) ? -1 : module->memory.initial;
 
   // global mapping is done in js in post.js
 }
@@ -165,8 +165,7 @@ extern "C" void EMSCRIPTEN_KEEPALIVE instantiate() {
   EM_ASM({
     Module['asmExports'] = {};
   });
-  for (auto& pair : module->exportsMap) {
-    auto& curr = pair.second;
+  for (auto* curr : module->exports) {
     EM_ASM_({
       var name = Pointer_stringify($0);
       Module['asmExports'][name] = function() {
@@ -178,15 +177,13 @@ extern "C" void EMSCRIPTEN_KEEPALIVE instantiate() {
   }
 
   // verify imports are provided
-  for (auto& pair : module->importsMap) {
-    auto& name = pair.first;
-    auto* import = pair.second;
+  for (auto* import : module->imports) {
     EM_ASM_({
       var mod = Pointer_stringify($0);
       var base = Pointer_stringify($1);
       var name = Pointer_stringify($2);
       assert(Module['lookupImport'](mod, base), 'checking import ' + name + ' = ' + mod + '.' + base);
-    }, import->module.str, import->base.str, name.str);
+    }, import->module.str, import->base.str, import->name.str);
   }
 
   if (wasmJSDebug) std::cerr << "creating instance...\n";
@@ -432,8 +429,8 @@ extern "C" void EMSCRIPTEN_KEEPALIVE call_from_js(const char *target) {
   if (wasmJSDebug) std::cout << "call_from_js " << target << '\n';
 
   IString exportName(target);
-  IString functionName = instance->wasm.exportsMap[exportName]->value;
-  Function *function = instance->wasm.functionsMap[functionName];
+  IString functionName = instance->wasm.getExport(exportName)->value;
+  Function *function = instance->wasm.getFunction(functionName);
   assert(function);
   size_t seen = EM_ASM_INT_V({ return Module['tempArguments'].length });
   size_t actual = function->params.size();
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index 95f8a280d..c13cd3512 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -281,8 +281,8 @@ private:
         return;
       } else if (id == TYPE) {
         Name typeName = curr[1]->str();
-        if (wasm.functionTypesMap.find(typeName) == wasm.functionTypesMap.end()) onError();
-        FunctionType* type = wasm.functionTypesMap[typeName];
+        if (!wasm.checkFunctionType(typeName)) onError();
+        FunctionType* type = wasm.getFunctionType(typeName);
         functionTypes[name] = type->result;
         return;
       }
@@ -384,8 +384,8 @@ private:
       } else if (id == TYPE) {
         Name name = curr[1]->str();
         func->type = name;
-        if (wasm.functionTypesMap.find(name) == wasm.functionTypesMap.end()) onError();
-        FunctionType* type = wasm.functionTypesMap[name];
+        if (!wasm.checkFunctionType(name)) onError();
+        FunctionType* type = wasm.getFunctionType(name);
         func->result = type->result;
         for (size_t j = 0; j < type->params.size(); j++) {
           IString name = Name::fromInt(j);
@@ -960,7 +960,7 @@ private:
   Expression* makeCallImport(Element& s) {
     auto ret = allocator.alloc<CallImport>();
     ret->target = s[1]->str();
-    Import* import = wasm.importsMap[ret->target];
+    Import* import = wasm.getImport(ret->target);
     ret->type = import->type->result;
     parseCallOperands(s, 2, ret);
     return ret;
@@ -969,8 +969,8 @@ private:
   Expression* makeCallIndirect(Element& s) {
     auto ret = allocator.alloc<CallIndirect>();
     IString type = s[1]->str();
-    assert(wasm.functionTypesMap.find(type) != wasm.functionTypesMap.end());
-    ret->fullType = wasm.functionTypesMap[type];
+    ret->fullType = wasm.getFunctionType(type);
+    assert(ret->fullType);
     ret->type = ret->fullType->result;
     ret->target = parseExpression(s[2]);
     parseCallOperands(s, 3, ret);
@@ -1129,8 +1129,8 @@ private:
         type.result = stringToWasmType(params[1]->str());
       } else if (id == TYPE) {
         IString name = params[1]->str();
-        if (wasm.functionTypesMap.find(name) == wasm.functionTypesMap.end()) onError();
-        type = *wasm.functionTypesMap[name];
+        if (!wasm.checkFunctionType(name)) onError();
+        type = *wasm.getFunctionType(name);
       } else {
         onError();
       }
diff --git a/src/wasm-validator.h b/src/wasm-validator.h
index 2918c918d..77f475990 100644
--- a/src/wasm-validator.h
+++ b/src/wasm-validator.h
@@ -108,9 +108,8 @@ public:
     }
     // start
     if (curr->start.is()) {
-      auto iter = curr->functionsMap.find(curr->start);
-      if (shouldBeTrue(iter != curr->functionsMap.end())) {
-        auto func = iter->second;
+      auto func = curr->checkFunction(curr->start);
+      if (shouldBeTrue(func)) {
         shouldBeTrue(func->params.size() == 0); // must be nullary
       }
     }
diff --git a/src/wasm.h b/src/wasm.h
index af767b688..a07ab3079 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -1114,23 +1114,41 @@ public:
 
 class Module {
 public:
-  // internal wasm contents (don't access these from outside; use add*() and the *Map objects)
+  // wasm contents (generally you shouldn't access these from outside, except maybe for iterating; use add*() and the get() functions)
   std::vector<FunctionType*> functionTypes;
   std::vector<Import*> imports;
   std::vector<Export*> exports;
   std::vector<Function*> functions;
 
-  // publicly-accessible content
+  Table table;
+  Memory memory;
+  Name start;
+
+private:
+  // TODO: add a build option where Names are just indices, and then these methods are not needed
   std::map<Name, FunctionType*> functionTypesMap;
   std::map<Name, Import*> importsMap;
   std::map<Name, Export*> exportsMap;
   std::map<Name, Function*> functionsMap;
-  Table table;
-  Memory memory;
-  Name start;
 
+public:
   Module() : functionTypeIndex(0), importIndex(0), exportIndex(0), functionIndex(0) {}
 
+  FunctionType* getFunctionType(size_t i) { assert(i < functionTypes.size());return functionTypes[i]; }
+  Import* getImport(size_t i) { assert(i < imports.size()); return imports[i]; }
+  Export* getExport(size_t i) { assert(i < exports.size()); return exports[i]; }
+  Function* getFunction(size_t i) { assert(i < functions.size()); return functions[i]; }
+
+  FunctionType* getFunctionType(Name name) { assert(functionTypesMap[name]); return functionTypesMap[name]; }
+  Import* getImport(Name name) { assert(importsMap[name]); return importsMap[name]; }
+  Export* getExport(Name name) { assert(exportsMap[name]); return exportsMap[name]; }
+  Function* getFunction(Name name) { assert(functionsMap[name]); return functionsMap[name]; }
+
+  FunctionType* checkFunctionType(Name name) { if (functionTypesMap.find(name) == functionTypesMap.end()) return nullptr; return functionTypesMap[name]; }
+  Import* checkImport(Name name) { if (importsMap.find(name) == importsMap.end()) return nullptr; return importsMap[name]; }
+  Export* checkExport(Name name) { if (exportsMap.find(name) == exportsMap.end()) return nullptr; return exportsMap[name]; }
+  Function* checkFunction(Name name) { if (functionsMap.find(name) == functionsMap.end()) return nullptr; return functionsMap[name]; }
+
   void addFunctionType(FunctionType* curr) {
     Name numericName = Name::fromInt(functionTypeIndex); // TODO: remove all these, assert on names already existing, do numeric stuff in wasm-s-parser etc.
     if (curr->name.isNull()) {
diff --git a/src/wasm2asm.h b/src/wasm2asm.h
index 310ba4224..0d8422ab1 100644
--- a/src/wasm2asm.h
+++ b/src/wasm2asm.h
@@ -287,7 +287,7 @@ void Wasm2AsmBuilder::addTables(Ref ast, Module *wasm) {
   std::map<std::string, std::vector<IString>> tables; // asm.js tables, sig => contents of table
   for (size_t i = 0; i < wasm->table.names.size(); i++) {
     Name name = wasm->table.names[i];
-    auto func = wasm->functionsMap[name];
+    auto func = wasm->getFunction(name);
     std::string sig = getSig(func);
     auto& table = tables[sig];
     if (table.size() == 0) {
diff --git a/test/emcc_O2_hello_world.fromasm b/test/emcc_O2_hello_world.fromasm
index 749f7d005..f0a74eb1e 100644
--- a/test/emcc_O2_hello_world.fromasm
+++ b/test/emcc_O2_hello_world.fromasm
@@ -4,10 +4,10 @@
   (type $FUNCSIG$iiii (func (param i32 i32 i32) (result i32)))
   (type $FUNCSIG$ii (func (param i32) (result i32)))
   (type $FUNCSIG$vi (func (param i32)))
+  (type $FUNCSIG$i (func (result i32)))
   (type $FUNCSIG$iii (func (param i32 i32) (result i32)))
   (type $FUNCSIG$v (func))
   (type $FUNCSIG$vii (func (param i32 i32)))
-  (type $FUNCSIG$i (func (result i32)))
   (import $abort "env" "abort" (param i32))
   (import $_pthread_cleanup_pop "env" "_pthread_cleanup_pop" (param i32))
   (import $_pthread_self "env" "_pthread_self" (result i32))
diff --git a/test/emcc_O2_hello_world.fromasm.imprecise b/test/emcc_O2_hello_world.fromasm.imprecise
index 749f7d005..f0a74eb1e 100644
--- a/test/emcc_O2_hello_world.fromasm.imprecise
+++ b/test/emcc_O2_hello_world.fromasm.imprecise
@@ -4,10 +4,10 @@
   (type $FUNCSIG$iiii (func (param i32 i32 i32) (result i32)))
   (type $FUNCSIG$ii (func (param i32) (result i32)))
   (type $FUNCSIG$vi (func (param i32)))
+  (type $FUNCSIG$i (func (result i32)))
   (type $FUNCSIG$iii (func (param i32 i32) (result i32)))
   (type $FUNCSIG$v (func))
   (type $FUNCSIG$vii (func (param i32 i32)))
-  (type $FUNCSIG$i (func (result i32)))
   (import $abort "env" "abort" (param i32))
   (import $_pthread_cleanup_pop "env" "_pthread_cleanup_pop" (param i32))
   (import $_pthread_self "env" "_pthread_self" (result i32))
diff --git a/test/emcc_hello_world.fromasm b/test/emcc_hello_world.fromasm
index 08f8106b7..a08e45d05 100644
--- a/test/emcc_hello_world.fromasm
+++ b/test/emcc_hello_world.fromasm
@@ -5,10 +5,10 @@
   (type $FUNCSIG$id (func (param f64) (result i32)))
   (type $FUNCSIG$ii (func (param i32) (result i32)))
   (type $FUNCSIG$vi (func (param i32)))
-  (type $FUNCSIG$iii (func (param i32 i32) (result i32)))
   (type $FUNCSIG$v (func))
-  (type $FUNCSIG$vii (func (param i32 i32)))
   (type $FUNCSIG$i (func (result i32)))
+  (type $FUNCSIG$iii (func (param i32 i32) (result i32)))
+  (type $FUNCSIG$vii (func (param i32 i32)))
   (import $abort "env" "abort")
   (import $nullFunc_ii "env" "nullFunc_ii" (param i32))
   (import $nullFunc_iiii "env" "nullFunc_iiii" (param i32))
diff --git a/test/emcc_hello_world.fromasm.imprecise b/test/emcc_hello_world.fromasm.imprecise
index 262353709..26bd96236 100644
--- a/test/emcc_hello_world.fromasm.imprecise
+++ b/test/emcc_hello_world.fromasm.imprecise
@@ -4,10 +4,10 @@
   (type $FUNCSIG$iiii (func (param i32 i32 i32) (result i32)))
   (type $FUNCSIG$ii (func (param i32) (result i32)))
   (type $FUNCSIG$vi (func (param i32)))
-  (type $FUNCSIG$iii (func (param i32 i32) (result i32)))
   (type $FUNCSIG$v (func))
-  (type $FUNCSIG$vii (func (param i32 i32)))
   (type $FUNCSIG$i (func (result i32)))
+  (type $FUNCSIG$iii (func (param i32 i32) (result i32)))
+  (type $FUNCSIG$vii (func (param i32 i32)))
   (import $abort "env" "abort")
   (import $nullFunc_ii "env" "nullFunc_ii" (param i32))
   (import $nullFunc_iiii "env" "nullFunc_iiii" (param i32))