3 files changed, 21 insertions, 3 deletions

diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index 8a9681cd9..d90c10ca6 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -211,6 +211,7 @@ private:
   void parseFunction(Element& s) {
     auto func = allocator.alloc<Function>();
     func->name = s[1]->str();
+    func->body = nullptr;
     for (unsigned i = 2; i < s.size(); i++) {
       Element& curr = *s[i];
       IString id = curr[0]->str();
@@ -227,7 +228,18 @@ private:
         func->locals.emplace_back(name, type);
         currLocalTypes[name] = type;
       } else {
-        func->body = parseExpression(curr);
+        Expression* ex = parseExpression(curr);
+        if (!func->body) {
+          func->body = ex;
+        } else {
+          auto block = func->body->dyn_cast<Block>();
+          if (!block) {
+            block = allocator.alloc<Block>();
+            block->list.push_back(func->body);
+            func->body = block;
+          }
+          block->list.push_back(ex);
+        }
       }
     }
     wasm.functions.push_back(func);
@@ -620,7 +632,8 @@ private:
     while (i < s.size()) {
       Element& curr = *s[i];
       assert(curr[0]->str() == SEGMENT);
-      wasm.memory.segments.emplace_back(atoi(curr[1]->c_str()), strdup(curr[2]->c_str()));
+      char *data = strdup(curr[2]->c_str()); // TODO: handle non-null-terminated?
+      wasm.memory.segments.emplace_back(atoi(curr[1]->c_str()), data, strlen(data));
       i++;
     }
   }
diff --git a/src/wasm-shell.cpp b/src/wasm-shell.cpp
index 7fe05851d..1ad0a19ad 100644
--- a/src/wasm-shell.cpp
+++ b/src/wasm-shell.cpp
@@ -29,6 +29,10 @@ struct ShellExternalInterface : ModuleInstance::ExternalInterface {
   void init(Module& wasm) override {
     memory = new char[wasm.memory.initial];
     memorySize = wasm.memory.initial;
+    // apply memory segments
+    for (auto segment : wasm.memory.segments) {
+      memcpy(memory + segment.offset, segment.data, segment.size);
+    }
   }
 
   jmp_buf trapState;
diff --git a/src/wasm.h b/src/wasm.h
index 998c0bd50..a91a1bd7f 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -803,7 +803,8 @@ public:
   struct Segment {
     size_t offset;
     const char* data;
-    Segment(size_t offset, const char *data) : offset(offset), data(data) {}
+    size_t size;
+    Segment(size_t offset, const char *data, size_t size) : offset(offset), data(data), size(size) {}
   };
 
   size_t initial, max;