Add support for locals

asm.js

@@ -47,10 +47,14 @@ class Tokenizer {
 }
 
 const State = Object.freeze({
-    TOP:    0,
+    TOP:        0,
-    EXPORT: 1,
+    EXPORT:     1,
-    FUNC:   2,
+    FUNC:       2,
-    RESULT: 3,
+    RESULT:     3,
+    PARAM_NAME: 4,
+    PARAM_TYPE: 5,
+    LOCAL_NAME: 6,
+    LOCAL_TYPE: 7,
 });
 
 const Action = Object.freeze({
@@ -58,6 +62,9 @@ const Action = Object.freeze({
     EXPORT: 1,
     FUNC:   2,
     RESULT: 3,
+    PARAM:  4,
+    SYMBOL: 5,
+    LOCAL:  6,
 });
 
 const types = {
@@ -67,6 +74,9 @@ const types = {
 
 const opcodes = {
     "end":       0x0b,
+    "local.get": 0x20,
+    "local.set": 0x21,
+    "local.tee": 0x22,
     "i32.const": 0x41,
     "i32.mul":   0x6c,
 };
@@ -80,19 +90,39 @@ class Parser {
             ".export": State.EXPORT,
             ".func":   State.FUNC,
             ".result": State.RESULT,
+            ".param":  State.PARAM_NAME,
+            ".local":  State.LOCAL_NAME,
         };
         this.handlers = {
-            [State.TOP]:    (token) => this.token_top(token),
+            [State.TOP]:        (token) => this.token_top(token),
-            [State.EXPORT]: (token) => this.token_export(token),
+            [State.EXPORT]:     (token) => this.token_export(token),
-            [State.FUNC]:   (token) => this.token_func(token),
+            [State.FUNC]:       (token) => this.token_func(token),
-            [State.RESULT]: (token) => this.token_result(token),
+            [State.RESULT]:     (token) => this.token_result(token),
+            [State.PARAM_NAME]: (token) => this.token_param_name(token),
+            [State.PARAM_TYPE]: (token) => this.token_param_type(token),
+            [State.LOCAL_NAME]: (token) => this.token_local_name(token),
+            [State.LOCAL_TYPE]: (token) => this.token_local_type(token),
         };
 
         this.results = [];
+        this.params = {};
+        this.locals = {};
 	}
 
+    integer(token) {
+        let base;
+        switch (token.slice(-1)) {
+        case "b": base = 2; break;
+        case "o": base = 8; break;
+        case "h": base = 16; break;
+        default:  base = 10; break;
+        }
+        const x = parseInt(token, base);
+        return Number.isNaN(x) ? null : x;
+    }
+
     translate_code(token) {
-        return opcodes[token] ?? parseInt(token);
+        return opcodes[token] ?? this.integer(token);
     }
 
     translate_type(token) {
@@ -108,7 +138,10 @@ class Parser {
             return;
         }
         const code = this.translate_code(token);
-        return { type: Action.APPEND, code };
+        if (code)
+            return { type: Action.APPEND, code };
+        else
+            return { type: Action.SYMBOL, symbol: token };
     }
 
     token_export(token) {
@@ -123,7 +156,7 @@ class Parser {
 
     token_result(token) {
         if (token == LINE_END) {
-            const action = { type: Action.RESULT, types: this.results };
+            const action = { type: Action.RESULT, results: this.results };
             this.state = State.TOP;
             this.results = [];
             return action;
@@ -132,6 +165,56 @@ class Parser {
         }
     }
 
+    token_param_name(token) {
+        if (token == LINE_END) {
+            const action = { type: Action.PARAM, params: this.params };
+            this.state = State.TOP;
+            this.params = {};
+            return action;
+        } else {
+            this.current_param = token;
+            this.state = State.PARAM_TYPE;
+        }
+    }
+
+    token_param_type(token) {
+        if (token == LINE_END) {
+            console.error(
+                "ERROR: Unexpected newline in .param: expected type");
+            this.state = State.TOP;
+            this.params = {};
+        } else {
+            this.params[this.current_param] = types[token];
+            this.current_param = undefined;
+            this.state = State.PARAM_NAME;
+        }
+    }
+
+    token_local_name(token) {
+        if (token == LINE_END) {
+            const action = { type: Action.LOCAL, locals: this.locals };
+            this.state = State.TOP;
+            this.locals = {};
+            return action;
+        } else {
+            this.current_local = token;
+            this.state = State.LOCAL_TYPE;
+        }
+    }
+
+    token_local_type(token) {
+        if (token == LINE_END) {
+            console.error(
+                "ERROR: Unexpected newline in .local: expected type");
+            this.state = State.TOP;
+            this.locals = {};
+        } else {
+            this.locals[this.current_local] = types[token];
+            this.current_local = undefined;
+            this.state = State.LOCAL_NAME;
+        }
+    }
+
 	*handle(src) {
         let action;
 		for (const token of this.tokenizer.handle(src)) {
@@ -158,6 +241,9 @@ export class Assembler {
             [Action.EXPORT]: (action) => this.action_export(action),
             [Action.FUNC]:   (action) => this.action_func(action),
             [Action.RESULT]: (action) => this.action_result(action),
+            [Action.PARAM]:  (action) => this.action_param(action),
+            [Action.SYMBOL]: (action) => this.action_symbol(action),
+            [Action.LOCAL]:  (action) => this.action_local(action),
         };
 
         this.exports = [];
@@ -184,7 +270,22 @@ export class Assembler {
     }
 
     action_result(action) {
-        this.funcs[this.current_func].results.push(...action.types);
+        this.funcs[this.current_func].results.push(...action.results);
+    }
+
+    action_param(action) {
+        Object.assign(this.funcs[this.current_func].params, action.params);
+    }
+
+    action_local(action) {
+        Object.assign(this.funcs[this.current_func].locals, action.locals);
+    }
+
+    action_symbol(action) {
+        const func = this.funcs[this.current_func];
+        const index = this.lookup_param(func, action.symbol)
+              ?? this.lookup_local(func, action.symbol);
+        func.body.push(index);
     }
 
 	push(chunk) {
@@ -193,6 +294,17 @@ export class Assembler {
             this.handlers[action.type](action);
 	}
 
+    lookup_param(func, symbol) {
+        const index = Object.keys(func.params).indexOf(symbol);
+        return index == -1 ? null : index;
+    }
+
+    lookup_local(func, symbol) {
+        const param_count = Object.entries(func.params).length;
+        const index = param_count + Object.keys(func.locals).indexOf(symbol);
+        return index == -1 ? null : index;
+    }
+
     wasm_section_type() {
         const funcs = Object.values(this.funcs);
         const contents = funcs.map(({ params, results }) => {
@@ -230,10 +342,13 @@ export class Assembler {
     wasm_section_code() {
         const funcs = Object.values(this.funcs);
         const contents = funcs.map(({ body, locals }) => {
-            const local_count = Object.entries(locals).length;
+            const local_types = Object.values(locals);
+            const local_count = local_types.length;
             return [
-                body.length + 2,
+                body.length + local_count + 3,
                 local_count,
+                local_count,
+                ...local_types,
                 ...body,
                 opcodes["end"]
             ]