call_import changes: no more call_import, shared index space with functions

3 files changed, 81 insertions, 54 deletions

diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
index 43ddc954c..e0fab9901 100644
--- a/src/passes/Print.cpp
+++ b/src/passes/Print.cpp
@@ -224,7 +224,7 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
     printCallBody(curr);
   }
   void visitCallImport(CallImport *curr) {
-    printOpening(o, "call_import ");
+    printOpening(o, "call ");
     printCallBody(curr);
   }
   void visitCallIndirect(CallIndirect *curr) {
diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index 561a310b0..9b8451b98 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -411,7 +411,6 @@ enum ASTNodes {
   SetLocal = 0x15,
   CallFunction = 0x16,
   CallIndirect = 0x17,
-  CallImport = 0x18,
   TeeLocal = 0x19,
   GetGlobal = 0x1a,
   SetGlobal = 0x1b,
@@ -727,27 +726,21 @@ public:
     }
     finishSection(start);
   }
-
-  std::map<Name, uint32_t> mappedImports; // name of the Import => index
-  uint32_t getImportIndex(Name name) {
-    if (!mappedImports.size()) {
-      // Create name => index mapping.
-      for (size_t i = 0; i < wasm->imports.size(); i++) {
-        assert(mappedImports.count(wasm->imports[i]->name) == 0);
-        mappedImports[wasm->imports[i]->name] = i;
-      }
-    }
-    assert(mappedImports.count(name));
-    return mappedImports[name];
-  }
   
-  std::map<Name, uint32_t> mappedFunctions; // name of the Function => index
+  std::map<Name, Index> mappedFunctions; // name of the Function => index. first imports, then internals
   uint32_t getFunctionIndex(Name name) {
     if (!mappedFunctions.size()) {
       // Create name => index mapping.
+      for (auto& import : wasm->imports) {
+        if (import->kind != Import::Function) continue;
+        assert(mappedFunctions.count(import->name) == 0);
+        auto index = mappedFunctions.size();
+        mappedFunctions[import->name] = index;
+      }
       for (size_t i = 0; i < wasm->functions.size(); i++) {
         assert(mappedFunctions.count(wasm->functions[i]->name) == 0);
-        mappedFunctions[wasm->functions[i]->name] = i;
+        auto index = mappedFunctions.size();
+        mappedFunctions[wasm->functions[i]->name] = index;
       }
     }
     assert(mappedFunctions.count(name));
@@ -957,7 +950,7 @@ public:
     for (auto* operand : curr->operands) {
       recurse(operand);
     }
-    o << int8_t(BinaryConsts::CallImport) << U32LEB(curr->operands.size()) << U32LEB(getImportIndex(curr->target));
+    o << int8_t(BinaryConsts::CallFunction) << U32LEB(curr->operands.size()) << U32LEB(getFunctionIndex(curr->target));
   }
   void visitCallIndirect(CallIndirect *curr) {
     if (debug) std::cerr << "zz node: CallIndirect" << std::endl;
@@ -1501,6 +1494,20 @@ public:
     }
   }
 
+  std::vector<Name> functionImportIndexes; // index in function index space => name of function import
+
+  // gets a name in the combined function import+defined function space
+  Name getFunctionIndexName(Index i) {
+    if (i < functionImportIndexes.size()) {
+      auto* import = wasm.getImport(functionImportIndexes[i]);
+      assert(import->kind == Import::Function);
+      return import->name;
+    } else {
+      i -= functionImportIndexes.size();
+      return wasm.functions.at(i)->name;
+    }
+  }
+
   void readImports() {
     if (debug) std::cerr << "== readImports" << std::endl;
     size_t num = getU32LEB();
@@ -1511,16 +1518,17 @@ public:
       curr->name = Name(std::string("import$") + std::to_string(i));
       curr->kind = (Import::Kind)getU32LEB();
       switch (curr->kind) {
-        case Export::Function: {
+        case Import::Function: {
           auto index = getU32LEB();
           assert(index < wasm.functionTypes.size());
           curr->functionType = wasm.getFunctionType(index);
           assert(curr->functionType->name.is());
+          functionImportIndexes.push_back(curr->name);
           break;
         }
-        case Export::Table: break;
-        case Export::Memory: break;
-        case Export::Global: curr->globalType = getWasmType(); break;
+        case Import::Table: break;
+        case Import::Memory: break;
+        case Import::Global: curr->globalType = getWasmType(); break;
         default: WASM_UNREACHABLE();
       }
       curr->module = getInlineString();
@@ -1529,7 +1537,7 @@ public:
     }
   }
 
-  std::vector<FunctionType*> functionTypes;
+  std::vector<FunctionType*> functionTypes; // types of defined functions
 
   void readFunctionSignatures() {
     if (debug) std::cerr << "== readFunctionSignatures" << std::endl;
@@ -1551,7 +1559,7 @@ public:
   // We read functions before we know their names, so we need to backpatch the names later
 
   std::vector<Function*> functions; // we store functions here before wasm.addFunction after we know their names
-  std::map<size_t, std::vector<Call*>> functionCalls; // at index i we have all calls to i
+  std::map<Index, std::vector<Call*>> functionCalls; // at index i we have all calls to the defined function i
   Function* currFunction = nullptr;
   size_t endOfFunction;
 
@@ -1611,7 +1619,7 @@ public:
     }
   }
 
-  std::map<Export*, size_t> exportIndexes;
+  std::map<Export*, Index> exportIndexes;
 
   void readExports() {
     if (debug) std::cerr << "== readExports" << std::endl;
@@ -1705,7 +1713,10 @@ public:
     for (auto& iter : exportIndexes) {
       Export* curr = iter.first;
       switch (curr->kind) {
-        case Export::Function: curr->value = wasm.functions[iter.second]->name; break;
+        case Export::Function: {
+          curr->value = getFunctionIndexName(iter.second);
+          break;
+        }
         case Export::Table: curr->value = Name::fromInt(0); break;
         case Export::Memory: curr->value = Name::fromInt(0); break;
         case Export::Global: curr->value = getGlobalName(iter.second); break;
@@ -1726,7 +1737,7 @@ public:
       auto i = pair.first;
       auto& indexes = pair.second;
       for (auto j : indexes) {
-        wasm.table.segments[i].data.push_back(wasm.functions[j]->name);
+        wasm.table.segments[i].data.push_back(getFunctionIndexName(j));
       }
     }
   }
@@ -1795,8 +1806,7 @@ public:
       case BinaryConsts::Br:
       case BinaryConsts::BrIf:         visitBreak((curr = allocator.alloc<Break>())->cast<Break>(), code); break; // code distinguishes br from br_if
       case BinaryConsts::TableSwitch:  visitSwitch((curr = allocator.alloc<Switch>())->cast<Switch>()); break;
-      case BinaryConsts::CallFunction: visitCall((curr = allocator.alloc<Call>())->cast<Call>()); break;
-      case BinaryConsts::CallImport:   visitCallImport((curr = allocator.alloc<CallImport>())->cast<CallImport>()); break;
+      case BinaryConsts::CallFunction: curr = visitCall(); break; // we don't know if it's a call or call_import yet
       case BinaryConsts::CallIndirect: visitCallIndirect((curr = allocator.alloc<CallIndirect>())->cast<CallIndirect>()); break;
       case BinaryConsts::GetLocal:     visitGetLocal((curr = allocator.alloc<GetLocal>())->cast<GetLocal>()); break;
       case BinaryConsts::TeeLocal:
@@ -1938,38 +1948,45 @@ public:
     }
     curr->default_ = getBreakName(getInt32());
   }
-  void visitCall(Call *curr) {
-    if (debug) std::cerr << "zz node: Call" << std::endl;
-    auto arity = getU32LEB();
-    WASM_UNUSED(arity);
-    auto index = getU32LEB();
-    assert(index < functionTypes.size());
-    auto type = functionTypes[index];
+
+  template<typename T>
+  void fillCall(T* call, FunctionType* type, Index arity) {
+    assert(type);
     auto num = type->params.size();
     assert(num == arity);
-    curr->operands.resize(num);
+    call->operands.resize(num);
     for (size_t i = 0; i < num; i++) {
-      curr->operands[num - i - 1] = popExpression();
+      call->operands[num - i - 1] = popExpression();
     }
-    curr->type = type->result;
-    functionCalls[index].push_back(curr);
+    call->type = type->result;
   }
-  void visitCallImport(CallImport *curr) {
-    if (debug) std::cerr << "zz node: CallImport" << std::endl;
+
+  Expression* visitCall() {
+    if (debug) std::cerr << "zz node: Call" << std::endl;
     auto arity = getU32LEB();
     WASM_UNUSED(arity);
-    auto import = wasm.getImport(getU32LEB());
-    curr->target = import->name;
-    auto type = import->functionType;
-    assert(type);
-    auto num = type->params.size();
-    assert(num == arity);
-    if (debug) std::cerr << "zz node: CallImport " << curr->target << " with type " << type->name << " and " << num << " params\n";
-    curr->operands.resize(num);
-    for (size_t i = 0; i < num; i++) {
-      curr->operands[num - i - 1] = popExpression();
+    auto index = getU32LEB();
+    FunctionType* type;
+    Expression* ret;
+    if (index < functionImportIndexes.size()) {
+      // this is a call of an imported function
+      auto* call = allocator.alloc<CallImport>();
+      auto* import = wasm.getImport(functionImportIndexes[index]);
+      call->target = import->name;
+      type = import->functionType;
+      fillCall(call, type, arity);
+      ret = call;
+    } else {
+      // this is a call of a defined function
+      auto* call = allocator.alloc<Call>();
+      auto adjustedIndex = index - functionImportIndexes.size();
+      assert(adjustedIndex < functionTypes.size());
+      type = functionTypes[adjustedIndex];
+      fillCall(call, type, arity);
+      functionCalls[adjustedIndex].push_back(call); // we don't know function names yet
+      ret = call;
     }
-    curr->type = type->result;
+    return ret;
   }
   void visitCallIndirect(CallIndirect *curr) {
     if (debug) std::cerr << "zz node: CallIndirect" << std::endl;
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index 9abdea744..1a9e71bd4 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -1207,8 +1207,18 @@ private:
   }
 
   Expression* makeCall(Element& s) {
+    auto target = s[1]->str();
+    auto* import = wasm.checkImport(target);
+    if (import && import->kind == Import::Function) {
+      auto ret = allocator.alloc<CallImport>();
+      ret->target = target;
+      Import* import = wasm.getImport(ret->target);
+      ret->type = import->functionType->result;
+      parseCallOperands(s, 2, s.size(), ret);
+      return ret;
+    }
     auto ret = allocator.alloc<Call>();
-    ret->target = s[1]->str();
+    ret->target = target;
     ret->type = functionTypes[ret->target];
     parseCallOperands(s, 2, s.size(), ret);
     return ret;