Unify imported and non-imported things (#1678)

Fixes #1649

This moves us to a single object for functions, which can be imported or nor, and likewise for globals (as a result, GetGlobals do not need to check if the global is imported or not, etc.). All imported things now inherit from Importable, which has the module and base of the import, and if they are set then it is an import.

For convenient iteration, there are a few helpers like

ModuleUtils::iterDefinedGlobals(wasm, [&](Global* global) {
  .. use global ..
});

as often iteration only cares about imported or defined (non-imported) things.

1 files changed, 90 insertions, 78 deletions

diff --git a/src/tools/wasm-merge.cpp b/src/tools/wasm-merge.cpp
index e8910860e..70611e269 100644
--- a/src/tools/wasm-merge.cpp
+++ b/src/tools/wasm-merge.cpp
@@ -34,20 +34,10 @@
 #include "wasm-binary.h"
 #include "wasm-builder.h"
 #include "wasm-validator.h"
+#include "ir/module-utils.h"
 
 using namespace wasm;
 
-// Calls note() on every import that has form "env".(base)
-static void findImportsByBase(Module& wasm, Name base, std::function<void (Name)> note) {
-  for (auto& curr : wasm.imports) {
-    if (curr->module == ENV) {
-      if (curr->base == base) {
-        note(curr->name);
-      }
-    }
-  }
-}
-
 // Ensure a memory or table is of at least a size
 template<typename T>
 static void ensureSize(T& what, Index size) {
@@ -119,31 +109,33 @@ struct Mergeable {
   }
 
   void findImports() {
-    findImportsByBase(wasm, MEMORY_BASE, [&](Name name) {
-      memoryBaseGlobals.insert(name);
+    ModuleUtils::iterImportedGlobals(wasm, [&](Global* import) {
+      if (import->module == ENV && import->base == MEMORY_BASE) {
+        memoryBaseGlobals.insert(import->name);
+      }
     });
     if (memoryBaseGlobals.size() == 0) {
       // add one
-      auto* import = new Import;
+      auto* import = new Global;
       import->name = MEMORY_BASE;
       import->module = ENV;
       import->base = MEMORY_BASE;
-      import->kind = ExternalKind::Global;
-      import->globalType = i32;
-      wasm.addImport(import);
+      import->type = i32;
+      wasm.addGlobal(import);
       memoryBaseGlobals.insert(import->name);
     }
-    findImportsByBase(wasm, TABLE_BASE, [&](Name name) {
-      tableBaseGlobals.insert(name);
+    ModuleUtils::iterImportedGlobals(wasm, [&](Global* import) {
+      if (import->module == ENV && import->base == TABLE_BASE) {
+        tableBaseGlobals.insert(import->name);
+      }
     });
     if (tableBaseGlobals.size() == 0) {
-      auto* import = new Import;
+      auto* import = new Global;
       import->name = TABLE_BASE;
       import->module = ENV;
       import->base = TABLE_BASE;
-      import->kind = ExternalKind::Global;
-      import->globalType = i32;
-      wasm.addImport(import);
+      import->type = i32;
+      wasm.addGlobal(import);
       tableBaseGlobals.insert(import->name);
     }
   }
@@ -244,7 +236,7 @@ struct Mergeable {
 struct OutputMergeable : public PostWalker<OutputMergeable, Visitor<OutputMergeable>>, public Mergeable {
   OutputMergeable(Module& wasm) : Mergeable(wasm) {}
 
-  void visitCallImport(CallImport* curr) {
+  void visitCall(Call* curr) {
     auto iter = implementedFunctionImports.find(curr->target);
     if (iter != implementedFunctionImports.end()) {
       // this import is now in the module - call it
@@ -264,10 +256,10 @@ struct OutputMergeable : public PostWalker<OutputMergeable, Visitor<OutputMergea
   void visitModule(Module* curr) {
     // remove imports that are being implemented
     for (auto& pair : implementedFunctionImports) {
-      curr->removeImport(pair.first);
+      curr->removeFunction(pair.first);
     }
     for (auto& pair : implementedGlobalImports) {
-      curr->removeImport(pair.first);
+      curr->removeGlobal(pair.first);
     }
   }
 };
@@ -288,11 +280,6 @@ struct InputMergeable : public ExpressionStackWalker<InputMergeable, Visitor<Inp
   std::map<Name, Name> gNames; // globals
 
   void visitCall(Call* curr) {
-    curr->target = fNames[curr->target];
-    assert(curr->target.is());
-  }
-
-  void visitCallImport(CallImport* curr) {
     auto iter = implementedFunctionImports.find(curr->target);
     if (iter != implementedFunctionImports.end()) {
       // this import is now in the module - call it
@@ -333,29 +320,38 @@ struct InputMergeable : public ExpressionStackWalker<InputMergeable, Visitor<Inp
   void merge() {
     // find function imports in us that are implemented in the output
     // TODO make maps, avoid N^2
-    for (auto& imp : wasm.imports) {
+    ModuleUtils::iterImportedFunctions(wasm, [&](Function* import) {
       // per wasm dynamic library rules, we expect to see exports on 'env'
-      if ((imp->kind == ExternalKind::Function || imp->kind == ExternalKind::Global) && imp->module == ENV) {
+      if (import->module == ENV) {
         // seek an export on the other side that matches
         for (auto& exp : outputMergeable.wasm.exports) {
-          if (exp->kind == imp->kind && exp->name == imp->base) {
+          if (exp->name == import->base) {
             // fits!
-            if (imp->kind == ExternalKind::Function) {
-              implementedFunctionImports[imp->name] = exp->value;
-            } else {
-              implementedGlobalImports[imp->name] = exp->value;
-            }
+            implementedFunctionImports[import->name] = exp->value;
             break;
           }
         }
       }
-    }
+    });
+    ModuleUtils::iterImportedGlobals(wasm, [&](Global* import) {
+      // per wasm dynamic library rules, we expect to see exports on 'env'
+      if (import->module == ENV) {
+        // seek an export on the other side that matches
+        for (auto& exp : outputMergeable.wasm.exports) {
+          if (exp->name == import->base) {
+            // fits!
+            implementedGlobalImports[import->name] = exp->value;
+            break;
+          }
+        }
+      }
+    });
     // remove the unneeded ones
     for (auto& pair : implementedFunctionImports) {
-      wasm.removeImport(pair.first);
+      wasm.removeFunction(pair.first);
     }
     for (auto& pair : implementedGlobalImports) {
-      wasm.removeImport(pair.first);
+      wasm.removeGlobal(pair.first);
     }
 
     // find new names
@@ -364,27 +360,26 @@ struct InputMergeable : public ExpressionStackWalker<InputMergeable, Visitor<Inp
         return outputMergeable.wasm.getFunctionTypeOrNull(name);
       });
     }
-    for (auto& curr : wasm.imports) {
-      if (curr->kind == ExternalKind::Function) {
-        curr->name = fNames[curr->name] = getNonColliding(curr->name, [&](Name name) -> bool {
-          return !!outputMergeable.wasm.getImportOrNull(name) || !!outputMergeable.wasm.getFunctionOrNull(name);
-        });
-      } else if (curr->kind == ExternalKind::Global) {
-        curr->name = gNames[curr->name] = getNonColliding(curr->name, [&](Name name) -> bool {
-          return !!outputMergeable.wasm.getImportOrNull(name) || !!outputMergeable.wasm.getGlobalOrNull(name);
-        });
-      }
-    }
-    for (auto& curr : wasm.functions) {
+    ModuleUtils::iterImportedFunctions(wasm, [&](Function* curr) {
+      curr->name = fNames[curr->name] = getNonColliding(curr->name, [&](Name name) -> bool {
+        return !!outputMergeable.wasm.getFunctionOrNull(name);
+      });
+    });
+    ModuleUtils::iterImportedGlobals(wasm, [&](Global* curr) {
+      curr->name = gNames[curr->name] = getNonColliding(curr->name, [&](Name name) -> bool {
+        return !!outputMergeable.wasm.getGlobalOrNull(name);
+      });
+    });
+    ModuleUtils::iterDefinedFunctions(wasm, [&](Function* curr) {
       curr->name = fNames[curr->name] = getNonColliding(curr->name, [&](Name name) -> bool {
         return outputMergeable.wasm.getFunctionOrNull(name);
       });
-    }
-    for (auto& curr : wasm.globals) {
+    });
+    ModuleUtils::iterDefinedGlobals(wasm, [&](Global* curr) {
       curr->name = gNames[curr->name] = getNonColliding(curr->name, [&](Name name) -> bool {
         return outputMergeable.wasm.getGlobalOrNull(name);
       });
-    }
+    });
 
     // update global names in input
     {
@@ -403,20 +398,26 @@ struct InputMergeable : public ExpressionStackWalker<InputMergeable, Visitor<Inp
     }
 
     // find function imports in output that are implemented in the input
-    for (auto& imp : outputMergeable.wasm.imports) {
-      if ((imp->kind == ExternalKind::Function || imp->kind == ExternalKind::Global) && imp->module == ENV) {
+    ModuleUtils::iterImportedFunctions(outputMergeable.wasm, [&](Function* import) {
+      if (import->module == ENV) {
         for (auto& exp : wasm.exports) {
-          if (exp->kind == imp->kind && exp->name == imp->base) {
-            if (imp->kind == ExternalKind::Function) {
-              outputMergeable.implementedFunctionImports[imp->name] = fNames[exp->value];
-            } else {
-              outputMergeable.implementedGlobalImports[imp->name] = gNames[exp->value];
-            }
+          if (exp->name == import->base) {
+            outputMergeable.implementedFunctionImports[import->name] = fNames[exp->value];
             break;
           }
         }
       }
-    }
+    });
+    ModuleUtils::iterImportedGlobals(outputMergeable.wasm, [&](Global* import) {
+      if (import->module == ENV) {
+        for (auto& exp : wasm.exports) {
+          if (exp->name == import->base) {
+            outputMergeable.implementedGlobalImports[import->name] = gNames[exp->value];
+            break;
+          }
+        }
+      }
+    });
 
     // update the output before bringing anything in. avoid doing so when possible, as in the
     // common case the output module is very large.
@@ -448,16 +449,19 @@ struct InputMergeable : public ExpressionStackWalker<InputMergeable, Visitor<Inp
     for (auto& curr : wasm.functionTypes) {
       outputMergeable.wasm.addFunctionType(curr.release());
     }
-    for (auto& curr : wasm.imports) {
-      if (curr->kind == ExternalKind::Memory || curr->kind == ExternalKind::Table) {
-        continue; // wasm has just 1 of each, they must match
+    for (auto& curr : wasm.globals) {
+      if (curr->imported()) {
+        outputMergeable.wasm.addGlobal(curr.release());
       }
-      // update and add
-      if (curr->functionType.is()) {
-        curr->functionType = ftNames[curr->functionType];
-        assert(curr->functionType.is());
+    }
+    for (auto& curr : wasm.functions) {
+      if (curr->imported()) {
+        if (curr->type.is()) {
+          curr->type = ftNames[curr->type];
+          assert(curr->type.is());
+        }
+        outputMergeable.wasm.addFunction(curr.release());
       }
-      outputMergeable.wasm.addImport(curr.release());
     }
     for (auto& curr : wasm.exports) {
       if (curr->kind == ExternalKind::Memory || curr->kind == ExternalKind::Table) {
@@ -477,13 +481,21 @@ struct InputMergeable : public ExpressionStackWalker<InputMergeable, Visitor<Inp
         }
       }
     }
+    // Copy over the remaining non-imports (we have already transferred
+    // the imports, and they are nullptrs).
     for (auto& curr : wasm.functions) {
-      curr->type = ftNames[curr->type];
-      assert(curr->type.is());
-      outputMergeable.wasm.addFunction(curr.release());
+      if (curr) {
+        assert(!curr->imported());
+        curr->type = ftNames[curr->type];
+        assert(curr->type.is());
+        outputMergeable.wasm.addFunction(curr.release());
+      }
     }
     for (auto& curr : wasm.globals) {
-      outputMergeable.wasm.addGlobal(curr.release());
+      if (curr) {
+        assert(!curr->imported());
+        outputMergeable.wasm.addGlobal(curr.release());
+      }
     }
   }