Handle importing globals in s2wasm (#843)

* Handle importing globals in s2wasm

* Make importedGlobals a set of Names, make Names hashable

* Revert "Make importedGlobals a set of Names, make Names hashable"

This reverts commit 1d0ca7a5e3839b15ca60593330979864c9c3ed60.

* Refactor relocation parsing to handle expressions directly

* PR Feedback

- Move comment where it belongs
- Add comment about ownership to addRelocation
- Remove do-nothing parseImportGlobal

* Reword "imported globals" to "imported objects"

- Flip isObjectImported to isObjectImplemented, for consistency

* Add tests for s2wasm globals.

Also implement import relocation expression handling

* Simplify globals.s test

* Fix memory leak of relocation

* Use unique_ptr instead of delete in getRelocatableExpression

6 files changed, 127 insertions, 36 deletions

diff --git a/src/s2wasm.h b/src/s2wasm.h
index c5b502c6d..faa31ae92 100644
--- a/src/s2wasm.h
+++ b/src/s2wasm.h
@@ -229,29 +229,68 @@ class S2WasmBuilder {
   // gets a constant, which may be a relocation for later.
   // returns whether this is a relocation
   // TODO: Clean up this and the way relocs are created from parsed objects
-  bool getRelocatableConst(uint32_t* target) {
+  LinkerObject::Relocation* getRelocatableConst(uint32_t* target) {
     if (isdigit(*s) || *s == '-') {
       int32_t val = getInt();
       memcpy(target, &val, sizeof(val));
-      return false;
-    } else {
-      // a global constant, we need to fix it up later
-      Name name = getStrToSep();
-      LinkerObject::Relocation::Kind kind = isFunctionName(name) ?
-          LinkerObject::Relocation::kFunction :
-          LinkerObject::Relocation::kData;
-      int offset = 0;
-      if (*s == '+') {
-        s++;
-        offset = getInt();
-      } else if (*s == '-') {
-        s++;
-        offset = -getInt();
-      }
-      linkerObj->addRelocation(kind, target,
-                               fixEmLongjmp(cleanFunction(name)), offset);
-      return true;
+      return nullptr;
+    }
+
+    // a global constant, we need to fix it up later
+    Name name = getStrToSep();
+    LinkerObject::Relocation::Kind kind = isFunctionName(name) ?
+        LinkerObject::Relocation::kFunction :
+        LinkerObject::Relocation::kData;
+    int offset = 0;
+    if (*s == '+') {
+      s++;
+      offset = getInt();
+    } else if (*s == '-') {
+      s++;
+      offset = -getInt();
+    }
+    return new LinkerObject::Relocation(
+        kind, target, fixEmLongjmp(cleanFunction(name)), offset);
+  }
+  Expression* relocationToGetGlobal(LinkerObject::Relocation* relocation) {
+    if (!relocation) {
+      return nullptr;
+    }
+
+    auto name = relocation->symbol;
+    auto g = allocator->alloc<GetGlobal>();
+    g->name = name;
+    g->type = i32;
+
+    // Optimization: store any nonnegative addends in their natural place.
+    // Only do this for positive addends because load/store offsets cannot be
+    // negative.
+    if (relocation->addend >= 0) {
+      *relocation->data = relocation->addend;
+      return g;
+    }
+
+    auto c = allocator->alloc<Const>();
+    c->type = i32;
+    c->value = Literal(relocation->addend);
+
+    auto add = allocator->alloc<Binary>();
+    add->type = i32;
+    add->op = AddInt32;
+    add->left = c;
+    add->right = g;
+    return add;
+  }
+  Expression* getRelocatableExpression(uint32_t* target) {
+    auto relocation = std::unique_ptr<LinkerObject::Relocation>(getRelocatableConst(target));
+    if (!relocation) {
+      return nullptr;
+    }
+    if (linkerObj->isObjectImplemented(relocation->symbol)) {
+      linkerObj->addRelocation(relocation.release());
+      return nullptr;
     }
+    return relocationToGetGlobal(relocation.get());
   }
 
   int64_t getInt64() {
@@ -415,8 +454,12 @@ class S2WasmBuilder {
         } else {
           abort_on("unknown directive");
         }
-      // add data aliases
+      } else if (match(".import_global")) {
+        Name name = getStr();
+        info->importedObjects.insert(name);
+        s = strchr(s, '\n');
       } else {
+        // add data aliases
         Name lhs = getStrToSep();
         // When the current line contains only one word, e.g.".text"
         if (match("\n"))
@@ -466,7 +509,8 @@ class S2WasmBuilder {
       else if (match("data")) {}
       else if (match("ident")) skipToEOL();
       else if (match("section")) parseToplevelSection();
-      else if (match("align") || match("p2align")) skipToEOL();
+      else if (match("align") || match("p2align") || match("import_global"))
+        skipToEOL();
       else if (match("globl")) parseGlobl();
       else if (match("functype")) parseFuncType();
       else skipObjectAlias(true);
@@ -834,6 +878,24 @@ class S2WasmBuilder {
       curr->finalize();
       setOutput(curr, assign);
     };
+    auto useRelocationExpression = [&](Expression *expr, Expression *reloc) {
+      if (!reloc) {
+        return expr;
+      }
+      // Optimization: if the given expr is (i32.const 0), ignore it
+      if (expr->_id == Expression::ConstId &&
+          ((Const*)expr)->value.getInteger() == 0) {
+        return reloc;
+      }
+
+      // Otherwise, need to add relocation expr to given expr
+      auto add = allocator->alloc<Binary>();
+      add->type = i32;
+      add->op = AddInt32;
+      add->left = expr;
+      add->right = reloc;
+      return (Expression*)add;
+    };
     auto makeLoad = [&](WasmType type) {
       skipComma();
       auto curr = allocator->alloc<Load>();
@@ -843,10 +905,10 @@ class S2WasmBuilder {
       curr->signed_ = match("_s");
       match("_u");
       Name assign = getAssign();
-      getRelocatableConst(&curr->offset.addr);
+      auto relocation = getRelocatableExpression(&curr->offset.addr);
       mustMatch("(");
       auto attributes = getAttributes(1);
-      curr->ptr = getInput();
+      curr->ptr = useRelocationExpression(getInput(), relocation);
       curr->align = curr->bytes;
       if (attributes[0]) {
         assert(strncmp(attributes[0], "p2align=", 8) == 0);
@@ -864,11 +926,11 @@ class S2WasmBuilder {
         curr->bytes = getWasmTypeSize(type);
       }
       skipWhitespace();
-      getRelocatableConst(&curr->offset.addr);
+      auto relocation = getRelocatableExpression(&curr->offset.addr);
       mustMatch("(");
       auto attributes = getAttributes(2);
       auto inputs = getInputs(2);
-      curr->ptr = inputs[0];
+      curr->ptr = useRelocationExpression(inputs[0], relocation);
       curr->align = curr->bytes;
       if (attributes[0]) {
         assert(strncmp(attributes[0], "p2align=", 8) == 0);
@@ -946,8 +1008,9 @@ class S2WasmBuilder {
               // may be a relocation
               auto curr = allocator->alloc<Const>();
               curr->type = curr->value.type = i32;
-              getRelocatableConst((uint32_t*)curr->value.geti32Ptr());
-              setOutput(curr, assign);
+              auto relocation = getRelocatableExpression((uint32_t*)curr->value.geti32Ptr());
+              auto expr = useRelocationExpression(curr, relocation);
+              setOutput(expr, assign);
             } else {
               cashew::IString str = getStr();
               setOutput(parseConst(str, type, *allocator), assign);
@@ -1307,8 +1370,13 @@ class S2WasmBuilder {
       } else if (match(".int32")) {
         Address size = raw.size();
         raw.resize(size + 4);
-        if (getRelocatableConst((uint32_t*)&raw[size])) { // just the size, as we may reallocate; we must fix this later, if it's a relocation
-          currRelocations.emplace_back(linkerObj->getCurrentRelocation(), size);
+        auto relocation = getRelocatableConst((uint32_t*)&raw[size]); // just the size, as we may reallocate; we must fix this later, if it's a relocation
+        if (relocation) {
+          if (!linkerObj->isObjectImplemented(relocation->symbol)) {
+            abort_on("s2wasm is currently unable to model imported globals in data segment initializers");
+          }
+          linkerObj->addRelocation(relocation);
+          currRelocations.emplace_back(relocation, size);
         }
         zero = false;
       } else if (match(".int64")) {
diff --git a/src/wasm-emscripten.cpp b/src/wasm-emscripten.cpp
index 5fa27b4f1..4be5e6b43 100644
--- a/src/wasm-emscripten.cpp
+++ b/src/wasm-emscripten.cpp
@@ -86,6 +86,17 @@ std::vector<Function*> makeDynCallThunks(Module& wasm, std::vector<Name> const&
   return generatedFunctions;
 }
 
+void addObjectImports(Module& wasm, std::unordered_set<cashew::IString> const& objectImports) {
+  for (Name name : objectImports) {
+    auto import = new Import;
+    import->name = import->base = name;
+    import->module = ENV;
+    import->kind = ExternalKind::Global;
+    import->globalType = i32;
+    wasm.addImport(import);
+  }
+}
+
 struct AsmConstWalker : public PostWalker<AsmConstWalker, Visitor<AsmConstWalker>> {
   Module& wasm;
   std::unordered_map<Address, Address> segmentsByAddress; // address => segment index
diff --git a/src/wasm-emscripten.h b/src/wasm-emscripten.h
index 07c470898..f1aa15129 100644
--- a/src/wasm-emscripten.h
+++ b/src/wasm-emscripten.h
@@ -29,6 +29,8 @@ void generateMemoryGrowthFunction(Module&);
 // signature in the indirect function table.
 std::vector<Function*> makeDynCallThunks(Module& wasm, std::vector<Name> const& tableSegmentData);
 
+void addObjectImports(Module& wasm, std::unordered_set<cashew::IString> const& objectImports);
+
 void generateEmscriptenMetadata(std::ostream& o,
                                 Module& wasm,
                                 std::unordered_map<Address, Address> segmentsByAddress,
diff --git a/src/wasm-linker.cpp b/src/wasm-linker.cpp
index 56ec03851..76bf8942d 100644
--- a/src/wasm-linker.cpp
+++ b/src/wasm-linker.cpp
@@ -40,7 +40,9 @@ void Linker::placeStackPointer(Address stackAllocation) {
     // stack pointer to point to one past-the-end of the stack allocation.
     std::vector<char> raw;
     raw.resize(pointerSize);
-    out.addRelocation(LinkerObject::Relocation::kData, (uint32_t*)&raw[0], ".stack", stackAllocation);
+    auto relocation = new LinkerObject::Relocation(
+      LinkerObject::Relocation::kData, (uint32_t*)&raw[0], ".stack", stackAllocation);
+    out.addRelocation(relocation);
     assert(out.wasm.memory.segments.empty());
     out.addSegment("__stack_pointer", raw);
   }
@@ -327,6 +329,8 @@ void Linker::emscriptenGlue(std::ostream& o) {
     exportFunction(f->name, true);
   }
 
+  emscripten::addObjectImports(out.wasm, out.symbolInfo.importedObjects);
+
   auto staticBump = nextStatic - globalBase;
   emscripten::generateEmscriptenMetadata(o, out.wasm, segmentsByAddress, staticBump, out.initializerFunctions);
 }
diff --git a/src/wasm-linker.h b/src/wasm-linker.h
index 3660b3038..76ea62d7d 100644
--- a/src/wasm-linker.h
+++ b/src/wasm-linker.h
@@ -58,6 +58,7 @@ class LinkerObject {
   struct SymbolInfo {
     std::unordered_set<cashew::IString> implementedFunctions;
     std::unordered_set<cashew::IString> undefinedFunctions;
+    std::unordered_set<cashew::IString> importedObjects;
     // TODO: it's not clear that this really belongs here.
     std::unordered_map<cashew::IString, SymbolAlias> aliasedSymbols;
 
@@ -70,6 +71,8 @@ class LinkerObject {
       }
       implementedFunctions.insert(other.implementedFunctions.begin(),
                                   other.implementedFunctions.end());
+      importedObjects.insert(other.importedObjects.begin(),
+                             other.importedObjects.end());
       aliasedSymbols.insert(other.aliasedSymbols.begin(),
                             other.aliasedSymbols.end());
     }
@@ -86,18 +89,20 @@ class LinkerObject {
     globls.push_back(name);
   }
 
-  void addRelocation(Relocation::Kind kind, uint32_t* target, Name name, int addend) {
-    relocations.emplace_back(new Relocation(kind, target, name, addend));
-  }
-
-  Relocation* getCurrentRelocation() {
-    return relocations.back().get();
+  // This takes ownership of the added Relocation
+  void addRelocation(Relocation* relocation) {
+    relocations.emplace_back(relocation);
   }
 
   bool isFunctionImplemented(Name name) {
     return symbolInfo.implementedFunctions.count(name) != 0;
   }
 
+  // An object is considered implemented if it is not imported
+  bool isObjectImplemented(Name name) {
+    return symbolInfo.importedObjects.count(name) == 0;
+  }
+
   // If name is an alias, return what it points to. Otherwise return name.
   Name resolveAlias(Name name, Relocation::Kind kind) {
     auto aliased = symbolInfo.aliasedSymbols.find(name);
diff --git a/src/wasm-validator.h b/src/wasm-validator.h
index ad2b7e1a3..28e575ca4 100644
--- a/src/wasm-validator.h
+++ b/src/wasm-validator.h
@@ -407,6 +407,7 @@ public:
 
   void visitGlobal(Global* curr) {
     if (!validateGlobally) return;
+    shouldBeTrue(curr->init != nullptr, curr->name, "global init must be non-null");
     shouldBeTrue(curr->init->is<Const>() || curr->init->is<GetGlobal>(), curr->name, "global init must be valid");
     shouldBeEqual(curr->type, curr->init->type, nullptr, "global init must have correct type");
   }