Support start

As spec'd in: https://github.com/WebAssembly/design/pull/495
And discussed in: https://github.com/WebAssembly/spec/issues/231

This will make it simpler and more uniform to add a start entry point.

s2wasm is the right place to add start because it'll eventually need to
do other basic setup, e.g. put code in start to setup the stack, as
dschuff is doing in:
  https://github.com/WebAssembly/binaryen/pull/179
Or rather, the linker is the right place and s2wasm happens to act as
our linker right now.

6 files changed, 68 insertions, 9 deletions

diff --git a/src/emscripten-optimizer/istring.h b/src/emscripten-optimizer/istring.h
index d114540e2..abc1e327b 100644
--- a/src/emscripten-optimizer/istring.h
+++ b/src/emscripten-optimizer/istring.h
@@ -114,8 +114,8 @@ struct IString {
   const char *c_str() const { return str; }
   bool equals(const char *other) const { return !strcmp(str, other); }
 
-  bool is()     { return str != nullptr; }
-  bool isNull() { return str == nullptr; }
+  bool is() const     { return str != nullptr; }
+  bool isNull() const { return str == nullptr; }
 };
 
 } // namespace cashew
diff --git a/src/s2wasm-main.cpp b/src/s2wasm-main.cpp
index e75147b75..03321a608 100644
--- a/src/s2wasm-main.cpp
+++ b/src/s2wasm-main.cpp
@@ -28,6 +28,7 @@ using namespace wasm;
 
 int main(int argc, const char *argv[]) {
   bool ignoreUnknownSymbols = false;
+  std::string startFunction;
   Options options("s2wasm", "Link .s file into .wast");
   options
       .add("--output", "-o", "Output file (stdout if not specified)",
@@ -41,6 +42,11 @@ int main(int argc, const char *argv[]) {
            [&ignoreUnknownSymbols](Options *, const std::string &) {
              ignoreUnknownSymbols = true;
            })
+      .add("--start", "", "Generate the start method (default: main)",
+           Options::Arguments::Optional,
+           [&startFunction](Options *, const std::string &argument) {
+             startFunction = argument.size() ? argument : "main";
+           })
       .add("--global-base", "-g", "Where to start to place globals",
            Options::Arguments::One,
            [](Options *o, const std::string &argument) {
@@ -70,7 +76,7 @@ int main(int argc, const char *argv[]) {
           : 0;
   if (options.debug) std::cerr << "Global base " << globalBase << '\n';
   S2WasmBuilder s2wasm(wasm, input.c_str(), options.debug, globalBase,
-                       stackAllocation, ignoreUnknownSymbols);
+                       stackAllocation, ignoreUnknownSymbols, startFunction);
 
   if (options.debug) std::cerr << "Emscripten gluing..." << std::endl;
   std::stringstream meta;
diff --git a/src/s2wasm.h b/src/s2wasm.h
index 1e0a28b43..5f8ba850f 100644
--- a/src/s2wasm.h
+++ b/src/s2wasm.h
@@ -39,15 +39,17 @@ class S2WasmBuilder {
   const char* s;
   bool debug;
   bool ignoreUnknownSymbols;
+  const std::string &startFunction;
 
  public:
   S2WasmBuilder(AllocatingModule& wasm, const char* input, bool debug,
                 size_t globalBase, size_t stackAllocation,
-                bool ignoreUnknownSymbols)
+                bool ignoreUnknownSymbols, const std::string& startFunction)
       : wasm(wasm),
         allocator(wasm.allocator),
         debug(debug),
         ignoreUnknownSymbols(ignoreUnknownSymbols),
+        startFunction(startFunction),
         globalBase(globalBase),
         nextStatic(globalBase) {
     s = input;
@@ -78,7 +80,7 @@ class S2WasmBuilder {
   };
   std::vector<Relocation> relocations;
 
-  std::set<Name> implementedFunctions;
+  std::map<Name, Function*> implementedFunctions;
   std::map<Name, Name> aliasedFunctions;
 
   std::map<size_t, size_t> addressSegments; // address => segment index
@@ -365,7 +367,7 @@ class S2WasmBuilder {
       if (match(".hidden")) mustMatch(name.str);
       mustMatch(name.str);
       if (match(":")) {
-        implementedFunctions.insert(name);
+        implementedFunctions.insert({name, nullptr});
       } else if (match("=")) {
         Name alias = getAtSeparated();
         mustMatch("@FUNCTION");
@@ -485,6 +487,7 @@ class S2WasmBuilder {
     };
 
     auto func = allocator.alloc<Function>();
+    implementedFunctions[name] = func;
     func->name = name;
     std::map<Name, WasmType> localTypes;
     // params and result
@@ -722,7 +725,7 @@ class S2WasmBuilder {
         Name target = cleanFunction(getCommaSeparated());
         auto aliased = aliasedFunctions.find(target);
         if (aliased != aliasedFunctions.end()) target = aliased->second;
-        if (implementedFunctions.count(target) > 0) {
+        if (implementedFunctions.count(target) != 0) {
           auto specific = allocator.alloc<Call>();
           specific->target = target;
           curr = specific;
@@ -1171,13 +1174,15 @@ class S2WasmBuilder {
       if (functionIndexes.count(name) == 0) {
         functionIndexes[name] = wasm.table.names.size();
         wasm.table.names.push_back(name);
-        if (debug) std::cerr << "function index: " << name << ": " << functionIndexes[name] << '\n';
+        if (debug)
+          std::cerr << "function index: " << name << ": "
+                    << functionIndexes[name] << '\n';
       }
     };
     for (auto& relocation : relocations) {
       Name name = relocation.value;
       if (debug) std::cerr << "fix relocation " << name << '\n';
-      const auto &symbolAddress = staticAddresses.find(name);
+      const auto& symbolAddress = staticAddresses.find(name);
       if (symbolAddress != staticAddresses.end()) {
         *(relocation.data) = symbolAddress->second + relocation.offset;
         if (debug) std::cerr << "  ==> " << *(relocation.data) << '\n';
@@ -1193,6 +1198,42 @@ class S2WasmBuilder {
         }
       }
     }
+    if (startFunction.size()) {
+      if (implementedFunctions.count(startFunction) == 0) {
+        std::cerr << "Unknown start function: `" << startFunction << "`\n";
+        abort();
+      }
+      const auto *target = implementedFunctions[startFunction];
+      Name start("_start");
+      if (implementedFunctions.count(start) != 0) {
+        std::cerr << "Start function already present: `" << start << "`\n";
+        abort();
+      }
+      auto* func = allocator.alloc<Function>();
+      func->name = start;
+      wasm.addFunction(func);
+      auto* exp = allocator.alloc<Export>();
+      exp->name = exp->value = start;
+      wasm.addExport(exp);
+      wasm.addStart(start);
+      auto* block = allocator.alloc<Block>();
+      func->body = block;
+      {  // Create the call, matching its parameters.
+        // TODO allow calling with non-default values.
+        auto* call = allocator.alloc<Call>();
+        call->target = startFunction;
+        size_t paramNum = 0;
+        for (const NameType& nt : target->params) {
+          Name name = Name::fromInt(paramNum++);
+          func->locals.emplace_back(name, nt.type);
+          auto* param = allocator.alloc<GetLocal>();
+          param->name = name;
+          param->type = nt.type;
+          call->operands.push_back(param);
+        }
+        block->list.push_back(call);
+      }
+    }
   }
 
   template<class C>
diff --git a/src/shared-constants.h b/src/shared-constants.h
index a928756da..6662c50b2 100644
--- a/src/shared-constants.h
+++ b/src/shared-constants.h
@@ -52,6 +52,7 @@ cashew::IString GLOBAL("global"),
                 GROW_WASM_MEMORY("__growWasmMemory"),
                 NEW_SIZE("newSize"),
                 MODULE("module"),
+                START("start"),
                 FUNC("func"),
                 PARAM("param"),
                 RESULT("result"),
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index 1c704ac74..e2a880e51 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -283,6 +283,7 @@ private:
 
   void parseModuleElement(Element& curr) {
     IString id = curr[0]->str();
+    if (id == START) return parseStart(curr);
     if (id == FUNC) return parseFunction(curr);
     if (id == MEMORY) return parseMemory(curr);
     if (id == EXPORT) return parseExport(curr);
@@ -304,6 +305,8 @@ private:
     return IString((prefix + std::to_string(otherIndex++)).c_str(), false);
   }
 
+  void parseStart(Element& s) { wasm.addStart(s[1]->str()); }
+
   void parseFunction(Element& s) {
     auto func = currFunction = allocator.alloc<Function>();
     size_t i = 1;
diff --git a/src/wasm.h b/src/wasm.h
index c030d9d60..b9d0d6822 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -1125,6 +1125,7 @@ public:
   std::map<Name, Function*> functionsMap;
   Table table;
   Memory memory;
+  Name start;
 
   Module() : functionTypeIndex(0), importIndex(0), exportIndex(0), functionIndex(0) {}
 
@@ -1168,6 +1169,9 @@ public:
     functionsMap[numericName] = curr;
     functionIndex++;
   }
+  void addStart(const Name &s) {
+    start = s;
+  }
 
   void removeImport(Name name) {
     for (size_t i = 0; i < imports.size(); i++) {
@@ -1211,6 +1215,10 @@ public:
       o << "\")";
     }
     o << (module.memory.segments.size() > 0 ? "\n  " : "") << ")\n";
+    if (module.start.is()) {
+      doIndent(o, indent);
+      printOpening(o, "start") << " " << module.start << ")\n";
+    }
     for (auto& curr : module.functionTypes) {
       doIndent(o, indent);
       curr->print(o, indent, true);