1 files changed, 32 insertions, 3 deletions

diff --git a/src/s2wasm.h b/src/s2wasm.h
index 85d9f6c9e..a1b5a7ef9 100644
--- a/src/s2wasm.h
+++ b/src/s2wasm.h
@@ -21,7 +21,10 @@ class S2WasmBuilder {
   char *s;
 
 public:
-  S2WasmBuilder(AllocatingModule& wasm, char *s) : wasm(wasm), allocator(wasm.allocator), s(s) {
+  S2WasmBuilder(AllocatingModule& wasm, char *input) : wasm(wasm), allocator(wasm.allocator) {
+    s = input;
+    scan();
+    s = input;
     process();
     fix();
   }
@@ -34,6 +37,8 @@ private:
   typedef std::pair<Const*, Name> Addressing;
   std::vector<Addressing> addressings; // we fix these up
 
+  std::map<Name, WasmType> functionResults; // function name => result type. we scan this early, then use it during processing.
+
   // utilities
 
   void skipWhitespace() {
@@ -219,6 +224,27 @@ private:
 
   // processors
 
+  void scan() {
+    while (*s) {
+      s = strstr(s, "\n	.type	");
+      if (!s) break;
+      mustMatch("\n	.type	");
+      Name name = getCommaSeparated();
+      skipComma();
+      if (!match("@function")) continue;
+      mustMatch(name.str);
+      mustMatch(":");
+      while (1) {
+        skipWhitespace();
+        if (match(".param")) s = strchr(s, '\n');
+        else if (match(".result")) {
+          functionResults[name] = getType();
+          break;
+        } else break;
+      }
+    }
+  }
+
   void process() {
     while (*s) {
       skipWhitespace();
@@ -590,11 +616,14 @@ private:
         }
         Name assign = getAssign();
         if (curr->is<Call>()) {
-          curr->dyn_cast<Call>()->target = getCommaSeparated();
+          Name name = curr->dyn_cast<Call>()->target = getCommaSeparated();
+          curr->type = functionResults[name];
         } else if (curr->is<CallImport>()) {
-          curr->dyn_cast<CallImport>()->target = getCommaSeparated();
+          Name name = curr->dyn_cast<CallImport>()->target = getCommaSeparated();
+          // XXX import definitions would help here, but still undecided in .s format
         } else {
           curr->dyn_cast<CallIndirect>()->target = getInput();
+          // XXX no way to know return type, https://github.com/WebAssembly/experimental/issues/53
         }
         while (1) {
           if (!skipComma()) break;