support expressions in segment offsets

14 files changed, 66 insertions, 40 deletions

diff --git a/src/binaryen-c.cpp b/src/binaryen-c.cpp
index b7fc4347a..b0e241ab3 100644
--- a/src/binaryen-c.cpp
+++ b/src/binaryen-c.cpp
@@ -736,7 +736,7 @@ void BinaryenSetFunctionTable(BinaryenModuleRef module, BinaryenFunctionRef* fun
 
 // Memory. One per module
 
-void BinaryenSetMemory(BinaryenModuleRef module, BinaryenIndex initial, BinaryenIndex maximum, const char* exportName, const char **segments, BinaryenIndex* segmentOffsets, BinaryenIndex* segmentSizes, BinaryenIndex numSegments) {
+void BinaryenSetMemory(BinaryenModuleRef module, BinaryenIndex initial, BinaryenIndex maximum, const char* exportName, const char **segments, BinaryenExpressionRef* segmentOffsets, BinaryenIndex* segmentSizes, BinaryenIndex numSegments) {
   if (tracing) {
     std::cout << "  {\n";
     for (BinaryenIndex i = 0; i < numSegments; i++) {
@@ -754,10 +754,10 @@ void BinaryenSetMemory(BinaryenModuleRef module, BinaryenIndex initial, Binaryen
     }
     if (numSegments == 0) std::cout << "0"; // ensure the array is not empty, otherwise a compiler error on VS
     std::cout << " };\n";
-    std::cout << "    BinaryenIndex segmentOffsets[] = { ";
+    std::cout << "    BinaryenExpressionRef segmentOffsets[] = { ";
     for (BinaryenIndex i = 0; i < numSegments; i++) {
       if (i > 0) std::cout << ", ";
-      std::cout << segmentOffsets[i];
+      std::cout << "expressions[" << expressions[segmentOffsets[i]] << "]";
     }
     if (numSegments == 0) std::cout << "0"; // ensure the array is not empty, otherwise a compiler error on VS
     std::cout << " };\n";
@@ -779,7 +779,7 @@ void BinaryenSetMemory(BinaryenModuleRef module, BinaryenIndex initial, Binaryen
   wasm->memory.max = maximum;
   if (exportName) wasm->memory.exportName = exportName;
   for (BinaryenIndex i = 0; i < numSegments; i++) {
-    wasm->memory.segments.emplace_back(segmentOffsets[i], segments[i], segmentSizes[i]);
+    wasm->memory.segments.emplace_back((Expression*)segmentOffsets[i], segments[i], segmentSizes[i]);
   }
 }
 
diff --git a/src/binaryen-c.h b/src/binaryen-c.h
index b32fa566a..e2086072b 100644
--- a/src/binaryen-c.h
+++ b/src/binaryen-c.h
@@ -346,7 +346,7 @@ void BinaryenSetFunctionTable(BinaryenModuleRef module, BinaryenFunctionRef* fun
 
 // Each segment has data in segments, a start offset in segmentOffsets, and a size in segmentSizes.
 // exportName can be NULL
-void BinaryenSetMemory(BinaryenModuleRef module, BinaryenIndex initial, BinaryenIndex maximum, const char* exportName, const char **segments, BinaryenIndex* segmentOffsets, BinaryenIndex* segmentSizes, BinaryenIndex numSegments);
+void BinaryenSetMemory(BinaryenModuleRef module, BinaryenIndex initial, BinaryenIndex maximum, const char* exportName, const char **segments, BinaryenExpressionRef* segmentOffsets, BinaryenIndex* segmentSizes, BinaryenIndex numSegments);
 
 // Start function. One per module
 
diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
index 47e503fa4..eb97eead5 100644
--- a/src/passes/Print.cpp
+++ b/src/passes/Print.cpp
@@ -589,9 +589,12 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
     doIndent(o, indent);
     printOpening(o, "memory") << ' ' << curr->memory.initial;
     if (curr->memory.max && curr->memory.max != Memory::kMaxSize) o << ' ' << curr->memory.max;
+    o << ")\n";
     for (auto segment : curr->memory.segments) {
-      o << maybeNewLine;
-      o << (minify ? "" : "    ") << "(segment " << segment.offset << " \"";
+      doIndent(o, indent);
+      printOpening(o, "data ", true);
+      visit(segment.offset);
+      o << " \"";
       for (size_t i = 0; i < segment.data.size(); i++) {
         unsigned char c = segment.data[i];
         switch (c) {
@@ -612,10 +615,8 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
           }
         }
       }
-      o << "\")";
+      o << "\")\n";
     }
-    o << ((curr->memory.segments.size() > 0 && !minify) ? "\n  " : "") << ')';
-    o << maybeNewLine;
     if (curr->memory.exportName.is()) {
       doIndent(o, indent);
       printOpening(o, "export ");
diff --git a/src/shared-constants.h b/src/shared-constants.h
index fb94ff11b..8599d8144 100644
--- a/src/shared-constants.h
+++ b/src/shared-constants.h
@@ -24,6 +24,7 @@ extern Name GROW_WASM_MEMORY,
             PARAM,
             RESULT,
             MEMORY,
+            DATA,
             SEGMENT,
             EXPORT,
             IMPORT,
diff --git a/src/shell-interface.h b/src/shell-interface.h
index f6c12e8ca..55f9b2a0b 100644
--- a/src/shell-interface.h
+++ b/src/shell-interface.h
@@ -92,9 +92,10 @@ struct ShellExternalInterface : ModuleInstance::ExternalInterface {
     memory.resize(wasm.memory.initial * wasm::Memory::kPageSize);
     // apply memory segments
     for (auto& segment : wasm.memory.segments) {
-      assert(segment.offset + segment.data.size() <= wasm.memory.initial * wasm::Memory::kPageSize);
+      Address offset = ConstantExpressionRunner().visit(segment.offset).value.geti32();
+      assert(offset + segment.data.size() <= wasm.memory.initial * wasm::Memory::kPageSize);
       for (size_t i = 0; i != segment.data.size(); ++i) {
-        memory.set(segment.offset + i, segment.data[i]);
+        memory.set(offset + i, segment.data[i]);
       }
     }
   }
diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index aad142d44..b67a79c3d 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -706,7 +706,8 @@ public:
     o << U32LEB(num);
     for (auto& segment : wasm->memory.segments) {
       if (segment.data.size() == 0) continue;
-      o << U32LEB(segment.offset);
+      writeExpression(segment.offset);
+      o << int8_t(BinaryConsts::End);
       writeInlineBuffer(&segment.data[0], segment.data.size());
     }
     finishSection(start);
@@ -1572,6 +1573,15 @@ public:
     }
   }
 
+  Expression* readExpression() {
+    assert(depth == 0);
+    processExpressions();
+    assert(expressionStack.size() == 1);
+    auto* ret = popExpression();
+    assert(depth == 0);
+    return ret;
+  }
+
   void readGlobals() {
     if (debug) std::cerr << "== readGlobals" << std::endl;
     size_t num = getU32LEB();
@@ -1580,11 +1590,7 @@ public:
       if (debug) std::cerr << "read one" << std::endl;
       auto curr = new Global;
       curr->type = getWasmType();
-      assert(depth == 0);
-      processExpressions();
-      assert(expressionStack.size() == 1);
-      curr->init = popExpression();
-      assert(depth == 0);
+      curr->init = readExpression();
       wasm.addGlobal(curr);
     }
   }
@@ -1649,7 +1655,7 @@ public:
     auto num = getU32LEB();
     for (size_t i = 0; i < num; i++) {
       Memory::Segment curr;
-      auto offset = getU32LEB();
+      auto offset = readExpression();
       auto size = getU32LEB();
       std::vector<char> buffer;
       buffer.resize(size);
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index fcf03c591..0be36a677 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -497,8 +497,8 @@ public:
   }
 };
 
-// Execute an expression in global init
-class GlobalInitRunner : public ExpressionRunner<GlobalInitRunner> {
+// Execute an constant expression in a global init or memory offset
+class ConstantExpressionRunner : public ExpressionRunner<ConstantExpressionRunner> {
 public:
   Flow visitLoop(Loop* curr) { WASM_UNREACHABLE(); }
   Flow visitCall(Call* curr) { WASM_UNREACHABLE(); }
@@ -547,7 +547,7 @@ public:
   ModuleInstance(Module& wasm, ExternalInterface* externalInterface) : wasm(wasm), externalInterface(externalInterface) {
     memorySize = wasm.memory.initial;
     for (Index i = 0; i < wasm.globals.size(); i++) {
-      globals.push_back(GlobalInitRunner().visit(wasm.globals[i]->init).value);
+      globals.push_back(ConstantExpressionRunner().visit(wasm.globals[i]->init).value);
     }
     externalInterface->init(wasm);
     if (wasm.start.is()) {
diff --git a/src/wasm-js.cpp b/src/wasm-js.cpp
index 33b0faa55..2aacd94d4 100644
--- a/src/wasm-js.cpp
+++ b/src/wasm-js.cpp
@@ -193,7 +193,7 @@ extern "C" void EMSCRIPTEN_KEEPALIVE instantiate() {
           var source = Module['HEAP8'].subarray($1, $1 + $2);
           var target = new Int8Array(Module['outside']['newBuffer']);
           target.set(source, $0);
-        }, (uint32_t)segment.offset, &segment.data[0], segment.data.size());
+        }, ConstantExpressionRunner().visit(segment.offset).value.geti32(), &segment.data[0], segment.data.size());
       }
     }
 
diff --git a/src/wasm-linker.cpp b/src/wasm-linker.cpp
index d8bae7558..1f06c6ec2 100644
--- a/src/wasm-linker.cpp
+++ b/src/wasm-linker.cpp
@@ -83,7 +83,7 @@ void Linker::layout() {
   // Update the segments with their addresses now that they have been allocated.
   for (const auto& seg : out.segments) {
     Address address = staticAddresses[seg.first];
-    out.wasm.memory.segments[seg.second].offset = address;
+    out.wasm.memory.segments[seg.second].offset = out.wasm.allocator.alloc<Const>()->set(Literal(uint32_t(address)));
     segmentsByAddress[address] = seg.second;
   }
 
diff --git a/src/wasm-linker.h b/src/wasm-linker.h
index 3f1e8c7ac..bf76a8f94 100644
--- a/src/wasm-linker.h
+++ b/src/wasm-linker.h
@@ -114,12 +114,12 @@ class LinkerObject {
   // Add an initializer segment for the named static variable.
   void addSegment(Name name, const char* data, Address size) {
     segments[name] = wasm.memory.segments.size();
-    wasm.memory.segments.emplace_back(0, data, size);
+    wasm.memory.segments.emplace_back(wasm.allocator.alloc<Const>()->set(Literal(uint32_t(0))), data, size);
   }
 
   void addSegment(Name name, std::vector<char>& data) {
     segments[name] = wasm.memory.segments.size();
-    wasm.memory.segments.emplace_back(0, data);
+    wasm.memory.segments.emplace_back(wasm.allocator.alloc<Const>()->set(Literal(uint32_t(0))), data);
   }
 
   void addInitializerFunction(Name name) {
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index 750d485fb..87217cd20 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -367,13 +367,14 @@ private:
     if (id == START) return parseStart(curr);
     if (id == FUNC) return parseFunction(curr);
     if (id == MEMORY) return parseMemory(curr);
+    if (id == DATA) return parseData(curr);
     if (id == EXPORT) return parseExport(curr);
     if (id == IMPORT) return; // already done
     if (id == GLOBAL) return parseGlobal(curr);
     if (id == TABLE) return parseTable(curr);
     if (id == TYPE) return; // already done
     std::cerr << "bad module element " << id.str << '\n';
-    throw ParseException("unknown module element");
+    throw ParseException("unknown module element", curr.line, curr.col);
   }
 
   // function parsing state
@@ -1321,19 +1322,40 @@ private:
     }
     while (i < s.size()) {
       Element& curr = *s[i];
-      assert(curr[0]->str() == SEGMENT);
-      const char *input = curr[2]->c_str();
+      size_t j = 1;
+      Address offsetValue;
+      if (curr[0]->str() == DATA) {
+        offsetValue = 0;
+      } else {
+        offsetValue = atoi(curr[j++]->c_str());
+      }
+      const char *input = curr[j]->c_str();
+      auto* offset = allocator.alloc<Const>();
+      offset->type = i32;
+      offset->value = Literal(int32_t(offsetValue));
       if (auto size = strlen(input)) {
         std::vector<char> data;
         stringToBinary(input, size, data);
-        wasm.memory.segments.emplace_back(atoi(curr[1]->c_str()), data.data(), data.size());
+        wasm.memory.segments.emplace_back(offset, data.data(), data.size());
       } else {
-        wasm.memory.segments.emplace_back(atoi(curr[1]->c_str()), "", 0);
+        wasm.memory.segments.emplace_back(offset, "", 0);
       }
       i++;
     }
   }
 
+  void parseData(Element& s) {
+    auto* offset = parseExpression(s[1]);
+    const char *input = s[2]->c_str();
+    if (auto size = strlen(input)) {
+      std::vector<char> data;
+      stringToBinary(input, size, data);
+      wasm.memory.segments.emplace_back(offset, data.data(), data.size());
+    } else {
+      wasm.memory.segments.emplace_back(offset, "", 0);
+    }
+  }
+
   void parseExport(Element& s) {
     if (!s[2]->dollared() && !std::isdigit(s[2]->str()[0])) {
       assert(s[2]->str() == MEMORY);
diff --git a/src/wasm-validator.h b/src/wasm-validator.h
index 90930d8f4..3d9a0e1c3 100644
--- a/src/wasm-validator.h
+++ b/src/wasm-validator.h
@@ -281,12 +281,6 @@ public:
   void visitMemory(Memory *curr) {
     shouldBeFalse(curr->initial > curr->max, "memory", "memory max >= initial");
     shouldBeTrue(curr->max <= Memory::kMaxSize, "memory", "max memory must be <= 4GB");
-    size_t top = 0;
-    for (auto& segment : curr->segments) {
-      shouldBeFalse(segment.offset < top, "memory", "segment offset is small enough");
-      top = segment.offset + segment.data.size();
-    }
-    shouldBeFalse(top > curr->initial * Memory::kPageSize, "memory", "total segments must be small enough");
   }
   void visitModule(Module *curr) {
     // exports
diff --git a/src/wasm.cpp b/src/wasm.cpp
index dbbd209b4..2138ec34d 100644
--- a/src/wasm.cpp
+++ b/src/wasm.cpp
@@ -49,6 +49,7 @@ Name GROW_WASM_MEMORY("__growWasmMemory"),
      PARAM("param"),
      RESULT("result"),
      MEMORY("memory"),
+     DATA("data"),
      SEGMENT("segment"),
      EXPORT("export"),
      IMPORT("import"),
diff --git a/src/wasm.h b/src/wasm.h
index 38a030e1a..c0129cf23 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -1441,14 +1441,14 @@ public:
   static const Address::address_t kMaxSize = ~Address::address_t(0) / kPageSize;
   static const Address::address_t kPageMask = ~(kPageSize - 1);
   struct Segment {
-    Address offset;
+    Expression* offset;
     std::vector<char> data; // TODO: optimize
     Segment() {}
-    Segment(Address offset, const char *init, Address size) : offset(offset) {
+    Segment(Expression* offset, const char *init, Address size) : offset(offset) {
       data.resize(size);
       std::copy_n(init, size, data.begin());
     }
-    Segment(Address offset, std::vector<char>& init) : offset(offset) {
+    Segment(Expression* offset, std::vector<char>& init) : offset(offset) {
       data.swap(init);
     }
   };