Simplify statics, segments, and relocations (#380)

Also defer address assignment until layout time in preparation for
separating linker objects out from Linker

4 files changed, 78 insertions, 45 deletions

diff --git a/src/s2wasm-main.cpp b/src/s2wasm-main.cpp
index a0fb5a1b5..c6a77c583 100644
--- a/src/s2wasm-main.cpp
+++ b/src/s2wasm-main.cpp
@@ -102,7 +102,8 @@ int main(int argc, const char *argv[]) {
           : 0;
   if (options.debug) std::cerr << "Global base " << globalBase << '\n';
   Linker lm(wasm, globalBase, stackAllocation, initialMem, maxMem,
-                  ignoreUnknownSymbols, startFunction, options.debug);
+            ignoreUnknownSymbols, startFunction, options.debug);
+
   S2WasmBuilder s2wasm(wasm, input.c_str(), options.debug, lm);
 
   std::stringstream meta;
diff --git a/src/s2wasm.h b/src/s2wasm.h
index 1aa132c6c..2e979f589 100644
--- a/src/s2wasm.h
+++ b/src/s2wasm.h
@@ -200,7 +200,8 @@ class S2WasmBuilder {
       return false;
     } else {
       // a global constant, we need to fix it up later
-      Name name = cleanFunction(getStrToSep());
+      Name name = getStrToSep();
+      Linker::Relocation::Kind kind = isFunctionName(name) ? Linker::Relocation::kFunction : Linker::Relocation::kData;
       int offset = 0;
       if (*s == '+') {
         s++;
@@ -209,7 +210,7 @@ class S2WasmBuilder {
         s++;
         offset = -getInt();
       }
-      linker.addRelocation(target, name, offset);
+      linker.addRelocation(kind, target, cleanFunction(name), offset);
       return true;
     }
   }
@@ -318,7 +319,11 @@ class S2WasmBuilder {
     abort_on("getType");
   }
 
-  // The LLVM backend emits function names as name@FUNCTION. We can drop the @ and after it.
+  // The LLVM backend emits function names as name@FUNCTION.
+  bool isFunctionName(Name name) {
+    return strstr(name.str, "@FUNCTION");
+  }
+  // Drop the @ and after it.
   Name cleanFunction(Name name) {
     if (!strchr(name.str, '@')) return name;
     char *temp = strdup(name.str);
@@ -1105,9 +1110,9 @@ class S2WasmBuilder {
       r->data = (uint32_t*)&(*raw)[i];
     }
     // assign the address, add to memory
-    size_t address = linker.allocateStatic(size, align, name);
+    linker.addStatic(size, align, name);
     if (!zero) {
-      linker.addAddressSegment(address, (const char*)&(*raw)[0], size);
+      linker.addSegment(name, (const char*)&(*raw)[0], size);
     }
   }
 
@@ -1119,7 +1124,7 @@ class S2WasmBuilder {
       skipComma();
       getInt();
     }
-    linker.allocateStatic(size, align, name);
+    linker.addStatic(size, align, name);
   }
 
   void skipImports() {
diff --git a/src/wasm-linker.cpp b/src/wasm-linker.cpp
index 90ef4c2b9..2c5e0003d 100644
--- a/src/wasm-linker.cpp
+++ b/src/wasm-linker.cpp
@@ -37,21 +37,29 @@ void Linker::placeStackPointer(size_t stackAllocation) {
   const size_t pointerSize = 4;
   // Unconditionally allocate space for the stack pointer. Emscripten
   // allocates the stack itself, and initializes the stack pointer itself.
-  size_t address = allocateStatic(pointerSize, pointerSize, "__stack_pointer");
+  addStatic(pointerSize, pointerSize, "__stack_pointer");
   if (stackAllocation) {
     // If we are allocating the stack, set up a relocation to initialize the
     // stack pointer to point to one past-the-end of the stack allocation.
     auto* raw = new uint32_t;
-    relocations.emplace_back(
-        make_unique<Relocation>(raw, ".stack", stackAllocation));
+    addRelocation(Relocation::kData, raw, ".stack", stackAllocation);
     assert(wasm.memory.segments.size() == 0);
-    addressSegments[address] = wasm.memory.segments.size();
-    wasm.memory.segments.emplace_back(
-        address, reinterpret_cast<char*>(raw), pointerSize);
+    addSegment("__stack_pointer", reinterpret_cast<char*>(raw), pointerSize);
   }
 }
 
 void Linker::layout() {
+  // Allocate all user statics
+  for (const auto& obj : staticObjects) {
+    allocateStatic(obj.allocSize, obj.alignment, obj.name);
+  }
+  // Update the segments with their addresses now that they have been allocated.
+  for (auto& seg : segments) {
+    size_t address = staticAddresses[seg.first];
+    wasm.memory.segments[seg.second].offset = address;
+    segmentsByAddress[address] = seg.second;
+  }
+
   // Place the stack after the user's static data, to keep those addresses
   // small.
   if (stackAllocation) allocateStatic(stackAllocation, 16, ".stack");
@@ -88,36 +96,35 @@ void Linker::layout() {
     }
   };
   for (auto& relocation : relocations) {
-    Name name = relocation->value;
+    Name name = relocation->symbol;
     if (debug) std::cerr << "fix relocation " << name << '\n';
-    const auto& symbolAddress = staticAddresses.find(name);
-    if (symbolAddress != staticAddresses.end()) {
-      *(relocation->data) = symbolAddress->second + relocation->offset;
+
+    if (relocation->kind == Relocation::kData) {
+      const auto& symbolAddress = staticAddresses.find(name);
+      assert(symbolAddress != staticAddresses.end());
+      *(relocation->data) = symbolAddress->second + relocation->addend;
       if (debug) std::cerr << "  ==> " << *(relocation->data) << '\n';
     } else {
-      // must be a function address
-      auto aliased = aliasedFunctions.find(name);
-      if (aliased != aliasedFunctions.end()) name = aliased->second;
+      // function address
+      name = resolveAlias(name);
       if (!wasm.checkFunction(name)) {
         std::cerr << "Unknown symbol: " << name << '\n';
-        if (!ignoreUnknownSymbols) abort();
+        if (!ignoreUnknownSymbols) Fatal() << "undefined reference\n";
         *(relocation->data) = 0;
       } else {
         ensureFunctionIndex(name);
-        *(relocation->data) = functionIndexes[name] + relocation->offset;
+        *(relocation->data) = functionIndexes[name] + relocation->addend;
       }
     }
   }
   if (!!startFunction) {
     if (implementedFunctions.count(startFunction) == 0) {
-      std::cerr << "Unknown start function: `" << startFunction << "`\n";
-      abort();
+      Fatal() << "Unknown start function: `" << startFunction << "`\n";
     }
     const auto *target = wasm.getFunction(startFunction);
     Name start("_start");
     if (implementedFunctions.count(start) != 0) {
-      std::cerr << "Start function already present: `" << start << "`\n";
-      abort();
+      Fatal() << "Start function already present: `" << start << "`\n";
     }
     auto* func = wasm.allocator.alloc<Function>();
     func->name = start;
@@ -173,7 +180,7 @@ void Linker::emscriptenGlue(std::ostream& o) {
     void visitCallImport(CallImport* curr) {
       if (curr->target == EMSCRIPTEN_ASM_CONST) {
         auto arg = curr->operands[0]->cast<Const>();
-        size_t segmentIndex = parent->addressSegments[arg->value.geti32()];
+        size_t segmentIndex = parent->segmentsByAddress[arg->value.geti32()];
         std::string code = escape(parent->wasm.memory.segments[segmentIndex].data);
         int32_t id;
         if (ids.count(code) == 0) {
diff --git a/src/wasm-linker.h b/src/wasm-linker.h
index 9929a6f70..9065c4b3e 100644
--- a/src/wasm-linker.h
+++ b/src/wasm-linker.h
@@ -29,15 +29,18 @@
 
 namespace wasm {
 
-// Wasm module linking/layout information
 class Linker {
  public:
   struct Relocation {
+    enum Kind { kData, kFunction };
+    Kind kind; // Whether the symbol refers to data or a function.
+    // Instead of section offsets as relocation targets, for now this is just
+    // a pointer to the memory to rewrite.
     uint32_t* data;
-    Name value;
-    int offset;
-    Relocation(uint32_t* data, Name value, int offset) :
-        data(data), value(value), offset(offset) {}
+    Name symbol; // Like the symbol index in ELF r_info field
+    int addend; // Like the ELF r_addend field
+    Relocation(Kind kind, uint32_t* data, Name symbol, int addend) :
+        kind(kind), data(data), symbol(symbol), addend(addend) {}
   };
 
   Linker(Module& wasm, size_t globalBase, size_t stackAllocation,
@@ -74,7 +77,7 @@ class Linker {
     // wasm modules can't import data objects. Its value is 0 for the main
     // executable, which is all we have with static linking. In the future this
     // can go in a crtbegin or similar file.
-    allocateStatic(4, 4, "__dso_handle");
+    addStatic(4, 4, "__dso_handle");
   }
 
   // Allocate a static variable and return its address in linear memory
@@ -85,15 +88,17 @@ class Linker {
     return address;
   }
 
-  // Allocate space for a stack pointer and (if stackAllocation > 0) set up a
-  // relocation for it to point to the top of the stack.
-  void placeStackPointer(size_t stackAllocation);
+  // Allocate a static object
+  void addStatic(size_t allocSize, size_t alignment, Name name) {
+    staticObjects.emplace_back(allocSize, alignment, name);
+  }
+
   void addGlobal(Name name) {
     globls.push_back(name);
   }
 
-  void addRelocation(uint32_t* target, Name name, int offset) {
-    relocations.emplace_back(make_unique<Relocation>(target, name, offset));
+  void addRelocation(Relocation::Kind kind, uint32_t* target, Name name, int addend) {
+    relocations.emplace_back(make_unique<Relocation>(kind, target, name, addend));
   }
   Relocation* getCurrentRelocation() {
     return relocations.back().get();
@@ -116,9 +121,10 @@ class Linker {
     return name;
   }
 
-  void addAddressSegment(size_t address, const char* data, size_t size) {
-    addressSegments[address] = wasm.memory.segments.size();
-    wasm.memory.segments.emplace_back(address, data, size);
+  // Add an initializer segment for the named static variable.
+  void addSegment(Name name, const char* data, size_t size) {
+    segments[name] = wasm.memory.segments.size();
+    wasm.memory.segments.emplace_back(0, data, size);
   }
 
   void addInitializerFunction(Name name) {
@@ -136,6 +142,18 @@ class Linker {
   void emscriptenGlue(std::ostream& o);
 
  private:
+  struct StaticObject {
+    size_t allocSize;
+    size_t alignment;
+    Name name;
+    StaticObject(size_t allocSize, size_t alignment, Name name) :
+        allocSize(allocSize), alignment(alignment), name(name) {}
+  };
+
+  // Allocate space for a stack pointer and (if stackAllocation > 0) set up a
+  // relocation for it to point to the top of the stack.
+  void placeStackPointer(size_t stackAllocation);
+
   template<class C>
   void printSet(std::ostream& o, C& c) {
     o << "[";
@@ -182,16 +200,18 @@ class Linker {
   size_t stackAllocation;
   bool debug;
 
-  std::map<Name, int32_t> staticAddresses; // name => address
+  std::vector<StaticObject> staticObjects;
+  std::unordered_map<cashew::IString, int32_t> staticAddresses; // name => address
 
   std::vector<std::unique_ptr<Relocation>> relocations;
 
   std::set<Name> implementedFunctions;
-  std::map<Name, Name> aliasedFunctions;
+  std::unordered_map<cashew::IString, Name> aliasedFunctions;
 
-  std::map<size_t, size_t> addressSegments; // address => segment index
+  std::map<Name, size_t> segments; // name => segment index (in wasm module)
+  std::unordered_map<size_t, size_t> segmentsByAddress; // address => segment index
 
-  std::map<Name, size_t> functionIndexes;
+  std::unordered_map<cashew::IString, size_t> functionIndexes;
 
   std::vector<Name> initializerFunctions;