Add an option to allocate space for user stack at link time

Currently s2wasm pre-allocates space for the special __stack_pointer
variable but not for the actual stack, and __stack_pointer is null. On
real systems the stack is allocated by the dynamic linker, thread
runtime, or other toolchain/system-provided startup code. Since we don't
have any of that in wasm yet, just add an option to allocate the stack
in the linker for now, so we can continue to run single-file libc-free tests.

2 files changed, 40 insertions, 5 deletions

diff --git a/src/s2wasm-main.cpp b/src/s2wasm-main.cpp
index c3f3d0a2f..e75147b75 100644
--- a/src/s2wasm-main.cpp
+++ b/src/s2wasm-main.cpp
@@ -46,6 +46,11 @@ int main(int argc, const char *argv[]) {
            [](Options *o, const std::string &argument) {
              o->extra["global-base"] = argument;
            })
+      .add("--allocate-stack", "-s", "Size of the user stack in linear memory",
+           Options::Arguments::One,
+           [](Options *o, const std::string &argument) {
+             o->extra["stack-allocation"] = argument;
+           })
       .add_positional("INFILE", Options::Arguments::One,
                       [](Options *o, const std::string &argument) {
                         o->extra["infile"] = argument;
@@ -59,9 +64,13 @@ int main(int argc, const char *argv[]) {
   size_t globalBase = options.extra.find("global-base") != options.extra.end()
                           ? std::stoull(options.extra["global-base"])
                           : 1;
+  size_t stackAllocation =
+      options.extra.find("stack-allocation") != options.extra.end()
+          ? std::stoull(options.extra["stack-allocation"])
+          : 0;
   if (options.debug) std::cerr << "Global base " << globalBase << '\n';
   S2WasmBuilder s2wasm(wasm, input.c_str(), options.debug, globalBase,
-                       ignoreUnknownSymbols);
+                       stackAllocation, ignoreUnknownSymbols);
 
   if (options.debug) std::cerr << "Emscripten gluing..." << std::endl;
   std::stringstream meta;
diff --git a/src/s2wasm.h b/src/s2wasm.h
index 2a2a36028..1e0a28b43 100644
--- a/src/s2wasm.h
+++ b/src/s2wasm.h
@@ -42,7 +42,8 @@ class S2WasmBuilder {
 
  public:
   S2WasmBuilder(AllocatingModule& wasm, const char* input, bool debug,
-                size_t globalBase, bool ignoreUnknownSymbols)
+                size_t globalBase, size_t stackAllocation,
+                bool ignoreUnknownSymbols)
       : wasm(wasm),
         allocator(wasm.allocator),
         debug(debug),
@@ -52,8 +53,13 @@ class S2WasmBuilder {
     s = input;
     scan();
     s = input;
-    prepare();
+    // Place the stack pointer at the bottom of the linear memory, to keep its
+    // address small (and thus with a small encoding).
+    placeStackPointer(stackAllocation);
     process();
+    // Place the stack after the user's static data, to keep those addresses
+    // small.
+    if (stackAllocation) placeStack(stackAllocation);
     fix();
   }
 
@@ -370,10 +376,30 @@ class S2WasmBuilder {
     }
   }
 
-  void prepare() {
+  void placeStackPointer(size_t stackAllocation) {
     assert(nextStatic == globalBase); // we are the first allocation
+    // Allocate space for the stack pointer
     staticAddresses["__stack_pointer"] = nextStatic;
-    nextStatic += 4;
+    const size_t pointerSize = 4;
+    if (stackAllocation) {
+      // If we are also allocating the stack, initialize the stack pointer to
+      // point to one past-the-end of the stack allocation.
+      auto* raw = new uint32_t;
+      relocations.emplace_back(raw, ".stack", stackAllocation);
+      assert(wasm.memory.segments.size() == 0);
+      addressSegments[nextStatic] = wasm.memory.segments.size();
+      wasm.memory.segments.emplace_back(
+          nextStatic, reinterpret_cast<char*>(raw), pointerSize);
+      wasm.memory.initial = nextStatic + pointerSize;
+    }
+    nextStatic += pointerSize;
+  }
+
+  void placeStack(size_t stackAllocation) {
+    // Allocate space for a user stack. It starts zeroed-out so needs no segment
+    staticAddresses[".stack"] = nextStatic;
+    nextStatic += stackAllocation;
+    wasm.memory.initial = nextStatic;
   }
 
   void process() {