Use unsigned shift in uleb128()

asm.js

@@ -97,6 +97,14 @@ const State = Object.freeze({
     DEF_VALUE:    25,
     BLOCK_NAME:   26,
     BLOCK_TYPE:   27,
+    TYPE_NAME:    28,
+    TYPE_PARAM:   29,
+    TYPE_RESULT:  30,
+    TABLE_NAME:   31,
+    TABLE_SIZE:   32,
+    ELEM_TABLE:   33,
+    ELEM_ELEM:    34,
+    ELEM_LABEL:   35,
 });
 
 const Action = Object.freeze({
@@ -118,11 +126,16 @@ const Action = Object.freeze({
     ENTER:  16,
     EXIT:   17,
     ELSE:   18,
+    TYPE:   19,
+    TABLE:  20,
+    ELEM:   21,
 });
 
 const types = {
     "void":    0x40,
     "func":    0x60,
+    "funcref": 0x70,
+    "f32":     0x7d,
     "i32":     0x7f,
 };
 
@@ -208,6 +221,9 @@ class Parser {
             ".utf8":   State.UTF8,
             ".align":  State.ALIGN,
             ".def":    State.DEF_NAME,
+            ".type":   State.TYPE_NAME,
+            ".table":  State.TABLE_NAME,
+            ".elem":   State.ELEM_TABLE,
         };
         this.blocks = new Set(["block", "loop", "if"]);
         this.handlers = {
@@ -239,6 +255,14 @@ class Parser {
             [State.DEF_VALUE]:    (token) => this.token_def_value(token),
             [State.BLOCK_NAME]:   (token) => this.token_block_name(token),
             [State.BLOCK_TYPE]:   (token) => this.token_block_type(token),
+            [State.TYPE_NAME]:    (token) => this.token_type_name(token),
+            [State.TYPE_PARAM]:   (token) => this.token_type_param(token),
+            [State.TYPE_RESULT]:  (token) => this.token_type_result(token),
+            [State.TABLE_NAME]:   (token) => this.token_table_name(token),
+            [State.TABLE_SIZE]:   (token) => this.token_table_size(token),
+            [State.ELEM_TABLE]:   (token) => this.token_elem_table(token),
+            [State.ELEM_ELEM]:    (token) => this.token_elem_elem(token),
+            [State.ELEM_LABEL]:   (token) => this.token_elem_label(token),
         };
 
         this.results = [];
@@ -293,7 +317,7 @@ class Parser {
         if (opcode)
             return { type: Action.APPEND, opcode };
         const literal = this.integer(token);
-        if (literal)
+        if (literal != null)
             return { type: Action.APPEND, literal };
 
         return { type: Action.SYMBOL, symbol: token };
@@ -508,8 +532,7 @@ class Parser {
             const value = this.integer(token) ?? console.error(
                 `ERROR: Unexpected token ${token} in .global: expected`
                     + " initial value");
-            const const_opcode = const_opcodes[this.global.type];
-            this.global.init = [ const_opcode, value, opcodes["end"] ];
+            this.global.init = value;
             const action = {
                 type: Action.GLOBAL,
                 global: { [this.global_name]: this.global }
@@ -555,7 +578,7 @@ class Parser {
         } else {
             if (value > 0xff)
                 console.error(`WARNING: Value ${token} is truncated`);
-            action.value = [ value & 0xff ];
+            action.value = value;
         }
         return action;
     }
@@ -570,9 +593,9 @@ class Parser {
         if (value == null) {
             action.symbol = token;
         } else {
-            if (value > 0xffff)
+            if (value > 0xffffffff)
                 console.error(`WARNING: Value ${token} is truncated`);
-            action.value = [ value & 0xff, (value >> 8) & 0xff ];
+            action.value = value;
         }
         return action;
     }
@@ -586,8 +609,8 @@ class Parser {
                 `ERROR: Unexpected token ${token}, expected string`);
             return;
         }
-        const value = this.encoder.encode(token.string);
-        const action = { type: Action.DATA, size: value.length, value };
+        const bytes = this.encoder.encode(token.string);
+        const action = { type: Action.DATA, size: bytes.length, bytes };
         return action;
     }
 
@@ -669,6 +692,141 @@ class Parser {
         return action;
     }
 
+    token_type_name(token) {
+        if (token == LINE_END) {
+            console.error(
+                "ERROR: Unexpected end of line in .type, expected name");
+            this.state = State.TOP;
+            return;
+        }
+
+        this.type = { name: token, params: [] };
+        this.state = State.TYPE_PARAM;
+    }
+
+    token_type_param(token) {
+        if (token == LINE_END) {
+            console.error(
+                "ERROR: Unexpected end of line in .type, expected "
+                    + "parameter type");
+            this.type = undefined;
+            this.state = State.TOP;
+            return;
+        }
+
+        if (token == "result") {
+            this.type.results = [];
+            this.state = State.TYPE_RESULT;
+            return;
+        }
+
+        const type = types[token];
+        if (type == undefined) {
+            console.error(
+                `ERROR: Unexpected token ${token} in .type, expected `
+                    + "parameter type");
+            this.type = undefined;
+            this.state = State.TOP;
+            return;
+        }
+
+        this.type.params.push(type);
+    }
+
+    token_type_result(token) {
+        if (token == LINE_END) {
+            const action = { type: Action.TYPE, the_type: this.type };
+            this.type = undefined;
+            this.state = State.TOP;
+            return action;
+        }
+
+        const type = types[token];
+        if (type == undefined) {
+            console.error(
+                `ERROR: Unexpected token ${token} in .type, expected `
+                    + "result type");
+            this.type = undefined;
+            this.state = State.TOP;
+            return;
+        }
+
+        this.type.results.push(type);
+    }
+
+    token_table_name(token) {
+        if (token == LINE_END) {
+            console.error(
+                "ERROR: Unexpected end of line in .table, expected name");
+            this.state = State.TOP;
+            return;
+        }
+
+        this.table = { name: token };
+        this.state = State.TABLE_SIZE;
+    }
+
+    token_table_size(token) {
+        if (token == LINE_END) {
+            console.error(
+                "ERROR: Unexpected end of line in .table, expected size");
+            this.table = undefined;
+            this.state = State.TOP;
+            return;
+        }
+
+        const size = this.integer(token);
+        if (size == null) {
+            console.error(
+                `ERROR: Unexpected token ${token} in .table, expected size`);
+            this.table = undefined;
+            this.state = State.TOP;
+            return;
+        }
+
+        this.table.size = size;
+        const action = { type: Action.TABLE, table: this.table };
+        this.table = undefined;
+        this.state = State.TOP;
+        return action;
+    }
+
+    token_elem_table(token) {
+        if (token == LINE_END) {
+            console.error(
+                "ERROR: Unexpected end of line in .elem, expected "
+                    + "table name");
+            this.state = State.TOP;
+            return;
+        }
+
+        this.elem = { table: token };
+        this.state = State.ELEM_ELEM;
+    }
+
+    token_elem_elem(token) {
+        if (token == LINE_END) {
+            console.error(
+                "ERROR: Unexpected end of line in .elem, expected "
+                    + "element name");
+            this.state = State.TOP;
+            this.elem = undefined;
+            return;
+        }
+
+        this.elem.elem = token;
+        this.state = State.ELEM_LABEL;
+    }
+
+    token_elem_label(token) {
+        if (token != LINE_END)
+            this.elem.label = token;
+        const action = { type: Action.ELEM, elem: this.elem };
+        this.elem = undefined
+        this.state = State.TOP;
+        return action;
+    }
+
     mem_action() {
         const action = {
             type: Action.MEM,
@@ -683,9 +841,15 @@ class Parser {
 	*handle(src) {
         let action;
 		for (const token of this.tokenizer.handle(src)) {
-            if (action = this.handlers[this.state](token)) {
-                yield action;
+            const handler = this.handlers[this.state];
+            if (handler == undefined) {
+                console.error(`ERROR: Unhandled state ${this.state}`);
+                this.state = State.TOP;
+                continue;
             }
+
+            if (action = handler(token))
+                yield action;
         }
 	}
 }
@@ -694,9 +858,11 @@ const Section = Object.freeze({
     TYPE:   0x01,
     IMPORT: 0x02,
     FUNC:   0x03,
+    TABLE:  0x04,
     MEM:    0x05,
     GLOBAL: 0x06,
     EXPORT: 0x07,
+    ELEM:   0x09,
     CODE:   0x0a,
     DATA:   0x0b,
 });
@@ -730,6 +896,9 @@ export class Assembler {
             [Action.ENTER]:  (action) => this.action_enter(action),
             [Action.EXIT]:   (action) => this.action_exit(action),
             [Action.ELSE]:   (action) => this.action_else(action),
+            [Action.TYPE]:   (action) => this.action_type(action),
+            [Action.TABLE]:  (action) => this.action_table(action),
+            [Action.ELEM]:   (action) => this.action_elem(action),
         };
 
         this.exports = [];
@@ -742,6 +911,8 @@ export class Assembler {
         this.defs = {};
         this.blocks = [];
         this.types = [];
+        this.type_bindings = {};
+        this.tables = {};
 	}
 
     action_append(action) {
@@ -784,6 +955,9 @@ export class Assembler {
               ?? this.lookup_param(func, action.symbol)
               ?? this.lookup_local(func, action.symbol)
               ?? this.lookup_global(action.symbol)
+              ?? this.lookup_table(action.symbol)
+              ?? this.lookup_type(action.symbol)
+              ?? this.lookup_func(action.symbol);
         if (value == null) {
             const def_value = this.lookup_def(action.symbol);
             if (def_value == null) {
@@ -793,7 +967,7 @@ export class Assembler {
             }
             func.body.push(...this.leb128(def_value));
         } else {
-            func.body.push(value);
+            func.body.push(...this.leb128(value));
         }
     }
 
@@ -831,10 +1005,26 @@ export class Assembler {
 
     action_data(action) {
         const data = this.data.at(-1).data;
-        const value = action.value != null
-              ? action.value
-              : this.le(this.lookup_def(action.symbol), action.size);
-        data.push(...value);
+		let bytes;
+		if (action.bytes != undefined) {
+			bytes = action.bytes;
+		} else {
+			let value = action.value;
+			if (value == undefined) {
+				if (action.symbol == undefined) {
+					console.error("ERROR: Invalid data action", action);
+					return;
+				}
+				value = this.lookup_def(action.symbol);
+				if (value == undefined) {
+					console.error(
+						`ERROR: Unable to resolve symbol ${action.symbol}`);
+					return;
+				}
+			}
+			bytes = this.le(value, action.size);
+		}
+        data.push(...bytes);
         this.pos.addr += action.size;
     }
 
@@ -874,10 +1064,40 @@ export class Assembler {
         this.blocks.push(undefined);
     }
 
+    action_type(action) {
+        const type = this.func_type(action.the_type);
+        const index = this.ensure_type(type);
+        this.type_bindings[action.the_type.name] = index;
+    }
+
+    action_table(action) {
+        this.tables[action.table.name] = {
+            size: action.table.size,
+            elems: [],
+        };
+    }
+
+    action_elem(action) {
+        const table = this.tables[action.elem.table];
+        const fn = Object.keys(this.funcs).indexOf(action.elem.elem);
+        if (fn == -1) {
+            console.error(`ERROR: ${action.elem.elem}: no such function`);
+            return;
+        }
+        const index = table.elems.push(fn) - 1;
+        if (action.elem.label)
+            this.defs[action.elem.label] = index;
+    }
+
 	push(chunk) {
 		const text = this.decoder.decode(chunk, { stream: true });
-		for (const action of this.parser.handle(text))
-            this.handlers[action.type](action);
+		for (const action of this.parser.handle(text)) {
+            const handler = this.handlers[action.type];
+            if (handler == undefined)
+                console.error("ERROR: Unhandled action", action);
+            else
+                handler(action);
+        }
 	}
 
     lookup_block(symbol) {
@@ -901,15 +1121,29 @@ export class Assembler {
         return index == -1 ? null : index;
     }
 
+    lookup_table(symbol) {
+        const index = Object.keys(this.tables).indexOf(symbol);
+        return index == -1 ? null : index;
+    }
+
+    lookup_type(symbol) {
+        return this.type_bindings[symbol];
+    }
+
+    lookup_func(symbol) {
+        const index = Object.keys(this.funcs).indexOf(symbol);
+        return index == -1 ? null : index;
+    }
+
     lookup_def(symbol) {
         return this.defs[symbol];
     }
 
     le(value, count) {
-        let bytes = []
+        const bytes = []
         while (value != 0) {
             bytes.push(value & 0xff);
-            value >>= 8;
+            value >>>= 8;
         }
         if (count != undefined) {
             while (bytes.length < count)
@@ -919,7 +1153,6 @@ export class Assembler {
     }
 
     leb128(x) {
-        const orig = x;
         const bytes = [];
         while (true) {
             const b = x & 0x7f;
@@ -932,6 +1165,20 @@ export class Assembler {
         }
     }
 
+    uleb128(x) {
+        const bytes = [];
+        while (true) {
+            const b = x & 0x7f;
+            x >>>= 7;
+            if (x == 0) {
+                bytes.push(b);
+                return bytes;
+            }
+            bytes.push(b | 0x80);
+        }
+        return bytes;
+    }
+
     wasm_section_type() {
         if (this.types.length == 0) return null;
         return [ this.types.length ].concat(...this.types);
@@ -962,6 +1209,17 @@ export class Assembler {
         return [ count, ...types ];
     }
 
+    wasm_section_table() {
+        const sizes = Object.values(this.tables).map(({size}) => size);
+        if (sizes.length == 0) return null;
+        const contents = sizes.map((size) => [
+            types.funcref,
+            0x00,
+            this.uleb128(size),
+        ]);
+        return [ sizes.length, contents ].flat(Infinity)
+    }
+
     wasm_section_mem() {
         const mems = Object.values(this.mems).filter(
             ({imported}) => !imported);
@@ -975,8 +1233,13 @@ export class Assembler {
         const globals = Object.values(this.globals);
         if (globals.length == 0)
             return null;
-        const contents = globals.map(
-            ({ type, init }) => [ type, 1, ...init ]);
+        const contents = globals.map(({ type, init }) => [
+            type,
+            1,
+            const_opcodes[type],
+            ...this.leb128(init),
+            opcodes["end"],
+        ]);
         return [ globals.length ].concat(...contents);
     }
 
@@ -995,6 +1258,24 @@ export class Assembler {
         return [ exports.length ].concat(...contents);
     }
 
+    wasm_section_elem() {
+        const table_elems = Object.values(this.tables).map(
+            ({elems}) => elems);
+        const count = table_elems.flat().length;
+        if (count == 0) return null;
+        const contents = table_elems.flatMap((elems, table) =>
+            elems.flatMap((fn, index) => [
+                table == 0 ? 0 : [ 2, table ],
+                opcodes["i32.const"],
+                index,
+                opcodes["end"],
+                1,
+                fn,
+            ])
+        );
+        return [ count, contents ].flat();
+    }
+
     wasm_section_code() {
         const funcs = Object.values(this.funcs);
         if (funcs.length == 0) return null;
@@ -1002,24 +1283,20 @@ export class Assembler {
             const local_types = Object.values(locals);
             const local_count = local_types.length;
             if (local_count == 0) {
-                return [
-                    body.length + 2,
-                    0,
-                    ...body,
-                    opcodes["end"]
-                ]
+                const full_body = [ 0, body, opcodes.end ].flat()
+                return [ full_body.length, full_body ].flat();
             } else {
-                return [
-                    body.length + local_count + 3,
-                    local_count,
-                    local_count,
-                    ...local_types,
-                    ...body,
-                    opcodes["end"]
-                ];
+                const groups = this.group(local_types);
+                const full_body = [
+                    groups.length,
+					...groups.flat(),
+                    body,
+                    opcodes.end,
+                ].flat();
+                return [ full_body.length, full_body ].flat();
             }
         });
-        return [ contents.length ].concat(...contents);
+        return [ contents.length, contents ].flat(Infinity);
     }
 
     wasm_section_data() {
@@ -1034,7 +1311,7 @@ export class Assembler {
                 ...data,
             ]
         });
-        return [ contents.length ].concat(...contents);
+        return [ contents.length, contents ].flat(Infinity);
     }
 
 	wasm() {
@@ -1044,18 +1321,20 @@ export class Assembler {
             [ Section.TYPE,   () => this.wasm_section_type() ],
             [ Section.IMPORT, () => this.wasm_section_import() ],
             [ Section.FUNC,   () => this.wasm_section_func() ],
+            [ Section.TABLE,  () => this.wasm_section_table() ],
             [ Section.MEM,    () => this.wasm_section_mem() ],
             [ Section.GLOBAL, () => this.wasm_section_global() ],
             [ Section.EXPORT, () => this.wasm_section_export() ],
+            [ Section.ELEM,   () => this.wasm_section_elem() ],
             [ Section.CODE,   () => this.wasm_section_code() ],
             [ Section.DATA,   () => this.wasm_section_data() ],
         ];
         const sections = template.map(([ code, generator ]) => {
             const body = generator();
-            return body == null ? [] : [ code, body.length, ...body ];
+            return body == null ? [] : [ code, body.length, body ];
         });
 
-		return new Uint8Array(HEADER.concat(...sections));
+		return new Uint8Array([ HEADER, sections ].flat(Infinity));
 	}
 
     mem_wasm({ flags, init, max }) {
@@ -1066,7 +1345,9 @@ export class Assembler {
     }
 
     func_type({ params, results }) {
-        const param_types = Object.values(params);
+        const param_types = params.length == undefined
+              ? Object.values(params)
+              : params;
         return [
             types["func"],
             param_types.length,
@@ -1089,4 +1370,15 @@ export class Assembler {
         for (const func of Object.values(this.funcs))
             func.type = this.ensure_type(this.func_type(func));
     }
+
+    group(array) {
+        return array.reduce((acc, val) => {
+			const last = acc.at(-1);
+			if (last != undefined && last[1] == val)
+				++last[0]
+			else
+				acc.push([1, val]);
+			return acc;
+		}, []);
+    }
 }