[wasm] Introduce jiterpreter control flow pass (#83247)

* Checkpoint CFG

Generate fallthrough in CFG

Generate branch block header partially in cfg

Emit branches in CFG

Checkpoint: Emit loop and exit return in CFG (broken)

Fix CFG emitting function header in the wrong place

Improve accuracy of cfg size estimation

Remove log messages

Checkpoint: Forward branches partially working

Fix non-conditional branches not being added to target table
Remove fallthrough

* Implement backward branches via a dispatch table

* Cleanup

* Cleanup

* Remove use of DataView since it has a hazard around heap growth
Use copyWithin to implement appendBytes where possible
More accurate overhead calculation

src/mono/mono/mini/interp/jiterpreter.c

@@ -1281,6 +1281,31 @@ mono_jiterp_get_member_offset (int member) {
 	}
 }
 
+#define JITERP_NUMBER_MODE_U32 0
+#define JITERP_NUMBER_MODE_I32 1
+#define JITERP_NUMBER_MODE_F32 2
+#define JITERP_NUMBER_MODE_F64 3
+
+EMSCRIPTEN_KEEPALIVE void
+mono_jiterp_write_number_unaligned (void *dest, double value, int mode) {
+	switch (mode) {
+		case JITERP_NUMBER_MODE_U32:
+			*((uint32_t *)dest) = (uint32_t)value;
+			return;
+		case JITERP_NUMBER_MODE_I32:
+			*((int32_t *)dest) = (int32_t)value;
+			return;
+		case JITERP_NUMBER_MODE_F32:
+			*((float *)dest) = (float)value;
+			return;
+		case JITERP_NUMBER_MODE_F64:
+			*((double *)dest) = value;
+			return;
+		default:
+			g_assert_not_reached();
+	}
+}
+
 // HACK: fix C4206
 EMSCRIPTEN_KEEPALIVE
 #endif // HOST_BROWSER

src/mono/wasm/runtime/cwraps.ts

@@ -100,6 +100,7 @@ const fn_signatures: SigLine[] = [
     [false, "mono_jiterp_encode_leb52", "number", ["number", "number", "number"]],
     [false, "mono_jiterp_encode_leb64_ref", "number", ["number", "number", "number"]],
     [false, "mono_jiterp_encode_leb_signed_boundary", "number", ["number", "number", "number"]],
+    [false, "mono_jiterp_write_number_unaligned", "void", ["number", "number", "number"]],
     [true, "mono_jiterp_type_is_byref", "number", ["number"]],
     [true, "mono_jiterp_get_size_of_stackval", "number", []],
     [true, "mono_jiterp_parse_option", "number", ["string"]],
@@ -234,6 +235,7 @@ export interface t_Cwraps {
     mono_jiterp_debug_count(): number;
     mono_jiterp_get_trace_hit_count(traceIndex: number): number;
     mono_jiterp_get_polling_required_address(): Int32Ptr;
+    mono_jiterp_write_number_unaligned(destination: VoidPtr, value: number, mode: number): void;
 }
 
 const wrapped_c_functions: t_Cwraps = <any>{};

src/mono/wasm/runtime/jiterpreter-support.ts

@@ -17,6 +17,67 @@ export declare interface MintOpcodePtr extends NativePointer {
     __brand: "MintOpcodePtr"
 }
 
+export const enum JiterpNumberMode {
+    U32 = 0,
+    I32 = 1,
+    F32 = 2,
+    F64 = 3
+}
+
+export const enum BailoutReason {
+    Unknown,
+    InterpreterTiering,
+    NullCheck,
+    VtableNotInitialized,
+    Branch,
+    BackwardBranch,
+    ConditionalBranch,
+    ConditionalBackwardBranch,
+    ComplexBranch,
+    ArrayLoadFailed,
+    ArrayStoreFailed,
+    StringOperationFailed,
+    DivideByZero,
+    Overflow,
+    Return,
+    Call,
+    Throw,
+    AllocFailed,
+    SpanOperationFailed,
+    CastFailed,
+    SafepointBranchTaken,
+    UnboxFailed,
+    CallDelegate,
+    Debugging
+}
+
+export const BailoutReasonNames = [
+    "Unknown",
+    "InterpreterTiering",
+    "NullCheck",
+    "VtableNotInitialized",
+    "Branch",
+    "BackwardBranch",
+    "ConditionalBranch",
+    "ConditionalBackwardBranch",
+    "ComplexBranch",
+    "ArrayLoadFailed",
+    "ArrayStoreFailed",
+    "StringOperationFailed",
+    "DivideByZero",
+    "Overflow",
+    "Return",
+    "Call",
+    "Throw",
+    "AllocFailed",
+    "SpanOperationFailed",
+    "CastFailed",
+    "SafepointBranchTaken",
+    "UnboxFailed",
+    "CallDelegate",
+    "Debugging"
+];
+
 type FunctionType = [
     index: FunctionTypeIndex,
     parameters: { [name: string]: WasmValtype },
@@ -52,6 +113,7 @@ type ImportedFunctionInfo = {
 }
 
 export class WasmBuilder {
+    cfg: Cfg;
     stack: Array<BlobBuilder>;
     stackSize!: number;
     inSection!: boolean;
@@ -89,6 +151,7 @@ export class WasmBuilder {
     constructor (constantSlotCount: number) {
         this.stack = [new BlobBuilder()];
         this.clear(constantSlotCount);
+        this.cfg = new Cfg(this);
     }
 
     clear (constantSlotCount: number) {
@@ -144,7 +207,7 @@ export class WasmBuilder {
             this.appendBytes(av);
             return null;
         } else
-            return av.slice();
+            return av.slice(0, current.size);
     }
 
     // HACK: Approximate amount of space we need to generate the full module at present
@@ -375,7 +438,7 @@ export class WasmBuilder {
             type: string,
             name: string,
             export: boolean,
-            locals: { [name: string]: WasmValtype }
+            locals: { [name: string]: WasmValtype },
         }, generator: Function
     ) {
         const rec : FunctionInfo = {
@@ -387,7 +450,7 @@ export class WasmBuilder {
             locals: options.locals,
             generator,
             error: null,
-            blob: null
+            blob: null,
         };
         this.functions.push(rec);
         if (rec.export)
@@ -403,15 +466,9 @@ export class WasmBuilder {
                 exportCount++;
 
             this.beginFunction(func.typeName, func.locals);
-            try {
-                func.generator();
-            /*
-            } catch (exc) {
-                func.error = <any>exc;
-            */
-            } finally {
+            func.blob = func.generator();
+            if (!func.blob)
                 func.blob = this.endFunction(false);
-            }
         }
 
         this._generateImportSection();
@@ -663,7 +720,6 @@ export class WasmBuilder {
 
 export class BlobBuilder {
     buffer: number;
-    view!: DataView;
     size: number;
     capacity: number;
     encoder?: TextEncoder;
@@ -684,7 +740,6 @@ export class BlobBuilder {
         // FIXME: This should not be necessary
         Module.HEAPU8.fill(0, this.buffer, this.buffer + this.size);
         this.size = 0;
-        this.view = new DataView(Module.HEAPU8.buffer, this.buffer, this.capacity);
     }
 
     appendU8 (value: number | WasmOpcode) {
@@ -696,44 +751,30 @@ export class BlobBuilder {
         return result;
     }
 
-    appendU16 (value: number) {
-        const result = this.size;
-        this.view.setUint16(this.size, value, true);
-        this.size += 2;
-        return result;
-    }
-
-    appendI16 (value: number) {
-        const result = this.size;
-        this.view.setInt16(this.size, value, true);
-        this.size += 2;
-        return result;
-    }
-
     appendU32 (value: number) {
         const result = this.size;
-        this.view.setUint32(this.size, value, true);
+        cwraps.mono_jiterp_write_number_unaligned(<any>this.buffer + this.size, value, JiterpNumberMode.U32);
         this.size += 4;
         return result;
     }
 
     appendI32 (value: number) {
         const result = this.size;
-        this.view.setInt32(this.size, value, true);
+        cwraps.mono_jiterp_write_number_unaligned(<any>this.buffer + this.size, value, JiterpNumberMode.I32);
         this.size += 4;
         return result;
     }
 
     appendF32 (value: number) {
         const result = this.size;
-        this.view.setFloat32(this.size, value, true);
+        cwraps.mono_jiterp_write_number_unaligned(<any>this.buffer + this.size, value, JiterpNumberMode.F32);
         this.size += 4;
         return result;
     }
 
     appendF64 (value: number) {
         const result = this.size;
-        this.view.setFloat64(this.size, value, true);
+        cwraps.mono_jiterp_write_number_unaligned(<any>this.buffer + this.size, value, JiterpNumberMode.F64);
         this.size += 8;
         return result;
     }
@@ -784,11 +825,18 @@ export class BlobBuilder {
 
     appendBytes (bytes: Uint8Array, count?: number) {
         const result = this.size;
-        if (typeof (count) === "number")
-            bytes = new Uint8Array(bytes.buffer, bytes.byteOffset, count);
-        const av = this.getArrayView(true);
-        av.set(bytes, this.size);
-        this.size += bytes.length;
+        if (bytes.buffer === Module.HEAPU8.buffer) {
+            if (typeof (count) !== "number")
+                count = bytes.length;
+            Module.HEAPU8.copyWithin(this.buffer + result, bytes.byteOffset, bytes.byteOffset + count);
+            this.size += count;
+        } else {
+            if (typeof (count) === "number")
+                bytes = new Uint8Array(bytes.buffer, bytes.byteOffset, count);
+            const av = this.getArrayView(true);
+            av.set(bytes, this.size);
+            this.size += bytes.length;
+        }
         return result;
     }
 
@@ -828,6 +876,267 @@ export class BlobBuilder {
     }
 }
 
+type CfgBlob = {
+    type: "blob";
+    ip: MintOpcodePtr;
+    start: number;
+    length: number;
+}
+
+type CfgBranchBlockHeader = {
+    type: "branch-block-header";
+    ip: MintOpcodePtr;
+    isBackBranchTarget: boolean;
+}
+
+type CfgBranch = {
+    type: "branch";
+    from: MintOpcodePtr;
+    target: MintOpcodePtr;
+    isBackward: boolean; // FIXME: This should be inferred automatically
+    isConditional: boolean;
+}
+
+type CfgSegment = CfgBlob | CfgBranchBlockHeader | CfgBranch;
+
+class Cfg {
+    builder: WasmBuilder;
+    startOfBody!: MintOpcodePtr;
+    segments: Array<CfgSegment> = [];
+    backBranchTargets: Uint16Array | null = null;
+    base!: MintOpcodePtr;
+    ip!: MintOpcodePtr;
+    entryIp!: MintOpcodePtr;
+    exitIp!: MintOpcodePtr;
+    lastSegmentStartIp!: MintOpcodePtr;
+    lastSegmentEnd = 0;
+    overheadBytes = 0;
+    entryBlob!: CfgBlob;
+    blockStack: Array<MintOpcodePtr> = [];
+    dispatchTable = new Map<MintOpcodePtr, number>();
+    trace = false;
+
+    constructor (builder: WasmBuilder) {
+        this.builder = builder;
+    }
+
+    initialize (startOfBody: MintOpcodePtr, backBranchTargets: Uint16Array | null, trace: boolean) {
+        this.segments.length = 0;
+        this.blockStack.length = 0;
+        this.startOfBody = startOfBody;
+        this.backBranchTargets = backBranchTargets;
+        this.base = this.builder.base;
+        this.ip = this.lastSegmentStartIp = this.builder.base;
+        this.lastSegmentEnd = 0;
+        this.overheadBytes = 10; // epilogue
+        this.dispatchTable.clear();
+        this.trace = trace;
+    }
+
+    // We have a header containing the table of locals and we need to preserve it
+    entry (ip: MintOpcodePtr) {
+        this.entryIp = ip;
+        this.appendBlob();
+        mono_assert(this.segments.length === 1, "expected 1 segment");
+        mono_assert(this.segments[0].type === "blob", "expected blob");
+        this.entryBlob = <CfgBlob>this.segments[0];
+        this.segments.length = 0;
+        this.overheadBytes += 9; // entry eip init + block + optional loop
+        if (this.backBranchTargets)
+            this.overheadBytes += 24; // some extra padding for the dispatch br_table
+    }
+
+    appendBlob () {
+        if (this.builder.current.size === this.lastSegmentEnd)
+            return;
+
+        this.segments.push({
+            type: "blob",
+            ip: this.lastSegmentStartIp,
+            start: this.lastSegmentEnd,
+            length: this.builder.current.size - this.lastSegmentEnd,
+        });
+        this.lastSegmentStartIp = this.ip;
+        this.lastSegmentEnd = this.builder.current.size;
+    }
+
+    startBranchBlock (ip: MintOpcodePtr, isBackBranchTarget: boolean) {
+        this.appendBlob();
+        this.segments.push({
+            type: "branch-block-header",
+            ip,
+            isBackBranchTarget,
+        });
+        this.overheadBytes += 3; // each branch block just costs us a block (2 bytes) and an end
+        if (this.backBranchTargets)
+            this.overheadBytes += 3; // size of the br_table entry for this branch target
+    }
+
+    branch (target: MintOpcodePtr, isBackward: boolean, isConditional: boolean) {
+        this.appendBlob();
+        this.segments.push({
+            type: "branch",
+            from: this.ip,
+            target,
+            isBackward,
+            isConditional,
+        });
+        this.overheadBytes += 3; // forward branches are a constant br + depth (optimally 2 bytes)
+        if (isBackward)
+            this.overheadBytes += 4; // back branches are more complex
+    }
+
+    emitBlob (segment: CfgBlob, source: Uint8Array) {
+        // console.log(`segment @${(<any>segment.ip).toString(16)} ${segment.start}-${segment.start + segment.length}`);
+        const view = source.subarray(segment.start, segment.start + segment.length);
+        this.builder.appendBytes(view);
+    }
+
+    generate (): Uint8Array {
+        // HACK: Make sure any remaining bytes are inserted into a trailing segment
+        this.appendBlob();
+
+        // Now finish generating the function body and copy it
+        const source = this.builder.endFunction(false)!;
+
+        // Now reclaim the builder that was being used so we can stitch segments together
+        this.builder._push();
+        // HACK: Make sure ip_const works
+        this.builder.base = this.base;
+
+        // Emit the function header
+        this.emitBlob(this.entryBlob, source);
+
+        // We wrap the entire trace in a loop that starts with a dispatch br_table in order to support
+        //  backwards branches.
+        if (this.backBranchTargets) {
+            this.builder.i32_const(0);
+            this.builder.local("disp", WasmOpcode.set_local);
+            this.builder.block(WasmValtype.void, WasmOpcode.loop);
+        }
+
+        // We create a block for each of our forward branch targets, which can be used to skip forward to it
+        // The block for each target will end *right before* the branch target, so that br <block nesting level>
+        //  will skip every opcode before it
+        for (let i = 0; i < this.segments.length; i++) {
+            const segment = this.segments[i];
+            if (segment.type !== "branch-block-header")
+                continue;
+            this.blockStack.push(segment.ip);
+        }
+
+        this.blockStack.sort((lhs, rhs) => <any>lhs - <any>rhs);
+        for (let i = 0; i < this.blockStack.length; i++)
+            this.builder.block(WasmValtype.void);
+
+        const dispatchIp = <MintOpcodePtr><any>0;
+        if (this.backBranchTargets) {
+            // the loop needs to start with a br_table that performs dispatch based on the current value
+            //  of the dispatch index local
+            // br_table has to be surrounded by a block in order for a depth of 0 to be fallthrough
+            // We wrap it in an additional block so we can have a trap for unexpected disp values
+            this.builder.block(WasmValtype.void);
+            this.builder.block(WasmValtype.void);
+            this.builder.local("disp");
+            this.builder.appendU8(WasmOpcode.br_table);
+            // br_table <number of values starting from 0> <labels for values starting from 0> <default>
+            // we have to assign disp==0 to fallthrough so that we start at the top of the fn body, then
+            //  assign disp values starting from 1 to branch targets
+            this.builder.appendULeb(this.blockStack.length + 1);
+            this.builder.appendULeb(1); // br depth of 1 = skip the unreachable and fall through to the start
+            for (let i = 0; i < this.blockStack.length; i++) {
+                this.dispatchTable.set(this.blockStack[i], i + 1);
+                this.builder.appendULeb(i + 2); // add 2 to the depth because of the double block around it
+            }
+            this.builder.appendULeb(0); // for unrecognized value we br 0, which causes us to trap
+            this.builder.endBlock();
+            this.builder.appendU8(WasmOpcode.unreachable);
+            this.builder.endBlock();
+            // We put a dummy IP at the end of the block stack to represent the dispatch loop
+            // We will use this dummy IP to find the appropriate br depth when restarting the loop later
+            this.blockStack.push(dispatchIp);
+        }
+
+        if (this.trace)
+            console.log(`blockStack=${this.blockStack}`);
+
+        for (let i = 0; i < this.segments.length; i++) {
+            const segment = this.segments[i];
+            switch (segment.type) {
+                case "blob": {
+                    // FIXME: If back branch target, generate a loop and put it on the block stack
+                    this.emitBlob(segment, source);
+                    break;
+                }
+                case "branch-block-header": {
+                    // When we reach a branch target, we pop the current block off the stack, because it is used
+                    //  to jump to this instruction pointer. So the result is that when previous code BRs to the
+                    //  current block, it will skip everything remaining in it and resume from segment.ip
+                    const indexInStack = this.blockStack.indexOf(segment.ip);
+                    mono_assert(indexInStack === 0, () => `expected ${segment.ip} on top of blockStack but found it at index ${indexInStack}, top is ${this.blockStack[0]}`);
+                    this.builder.endBlock();
+                    this.blockStack.shift();
+                    break;
+                }
+                case "branch": {
+                    const lookupTarget = segment.isBackward ? dispatchIp : segment.target;
+                    let indexInStack = this.blockStack.indexOf(lookupTarget);
+
+                    // Back branches will target the dispatcher loop so we need to update the dispatch index
+                    //  which will be used by the loop dispatch br_table to jump to the correct location
+                    if (segment.isBackward && (indexInStack >= 0)) {
+                        if (this.dispatchTable.has(segment.target)) {
+                            const disp = this.dispatchTable.get(segment.target)!;
+                            if (this.trace)
+                                console.log(`backward br from ${segment.from} to ${segment.target}: disp=${disp}`);
+                            this.builder.i32_const(disp);
+                            this.builder.local("disp", WasmOpcode.set_local);
+                        } else {
+                            if (this.trace)
+                                console.log(`br from ${segment.from} to ${segment.target} failed: back branch target not in dispatch table`);
+                            indexInStack = -1;
+                        }
+                    }
+
+                    if (indexInStack >= 0) {
+                        // Conditional branches are nested in an extra block, so the depth is +1
+                        const offset = segment.isConditional ? 1 : 0;
+                        this.builder.appendU8(WasmOpcode.br);
+                        this.builder.appendULeb(offset + indexInStack);
+                        if (this.trace)
+                            console.log(`br from ${segment.from} to ${segment.target} breaking out ${offset + indexInStack + 1} level(s)`);
+                    } else {
+                        if (this.trace)
+                            console.log(`br from ${segment.from} to ${segment.target} failed`);
+                        append_bailout(this.builder, segment.target, BailoutReason.Branch);
+                    }
+                    break;
+                }
+                default:
+                    throw new Error("unreachable");
+            }
+        }
+
+        // Close the dispatch loop
+        if (this.backBranchTargets) {
+            mono_assert(this.blockStack[0] === <any>0, "expected one zero entry on the block stack for the dispatch loop");
+            this.blockStack.shift();
+            this.builder.endBlock();
+        }
+
+        mono_assert(this.blockStack.length === 0, () => `expected block stack to be empty at end of function but it was ${this.blockStack}`);
+
+        // Now we generate a ret at the end of the function body so it's Valid(tm)
+        // We will only hit this if execution falls through every block without hitting a bailout
+        this.builder.ip_const(this.exitIp);
+        this.builder.appendU8(WasmOpcode.return_);
+        this.builder.appendU8(WasmOpcode.end);
+
+        const result = this.builder._pop(false)!;
+        return result;
+    }
+}
+
 export const enum WasmValtype {
     void = 0x40,
     i32 = 0x7F,
@@ -864,6 +1173,15 @@ export const _now = (globalThis.performance && globalThis.performance.now)
 
 let scratchBuffer : NativePointer = <any>0;
 
+export function append_bailout (builder: WasmBuilder, ip: MintOpcodePtr, reason: BailoutReason) {
+    builder.ip_const(ip);
+    if (builder.options.countBailouts) {
+        builder.i32_const(reason);
+        builder.callImport("bailout");
+    }
+    builder.appendU8(WasmOpcode.return_);
+}
+
 export function copyIntoScratchBuffer (src: NativePointer, size: number) : NativePointer {
     if (!scratchBuffer)
         scratchBuffer = Module._malloc(64);

src/mono/wasm/runtime/jiterpreter-trace-generator.ts

@@ -12,9 +12,10 @@ import { MintOpcode, OpcodeInfo } from "./mintops";
 import cwraps from "./cwraps";
 import {
     MintOpcodePtr, WasmValtype, WasmBuilder,
-    append_memset_dest, append_memmove_dest_src,
-    try_append_memset_fast, try_append_memmove_fast,
-    counters, getMemberOffset, JiterpMember
+    append_memset_dest, append_bailout,
+    append_memmove_dest_src, try_append_memset_fast,
+    try_append_memmove_fast, counters,
+    getMemberOffset, JiterpMember, BailoutReason,
 } from "./jiterpreter-support";
 import {
     sizeOfDataItem,
@@ -29,8 +30,6 @@ import {
 
     mostRecentOptions,
 
-    BailoutReason,
-
     record_abort,
 } from "./jiterpreter";
 
@@ -150,29 +149,11 @@ export function generate_wasm_body (
     ip += <any>(OpcodeInfo[MintOpcode.MINT_TIER_ENTER_JITERPRETER][1] * 2);
     let rip = ip;
 
-    // Initialize eip, so that we will never return a 0 displacement
-    // Otherwise we could return 0 in the scenario where none of our blocks executed
-    // (This shouldn't happen though!)
-    builder.ip_const(ip);
-    builder.local("eip", WasmOpcode.set_local);
-
-    // If a method contains backward branches we also need to wrap the whole trace in a loop
-    //  that we can jump to the top of in order to begin executing the trace again
-    // FIXME: It would be much more efficient to use br_table to dispatch to the appropriate
-    //  branch block somehow but the code generation is tough due to WASM's IR
-    if (backwardBranchTable) {
-        builder.block(WasmValtype.void, WasmOpcode.loop);
-    }
-
-    // We wrap all instructions in a 'branch block' that is used
-    //  when performing a branch and will be skipped over if the
-    //  current instruction pointer does not match. This means
-    //  that if ip points to a branch target we don't handle,
-    //  the trace will automatically bail out at the end after
-    //  skipping past all the branch targets
-    builder.block();
+    builder.cfg.entry(ip);
 
     while (ip) {
+        builder.cfg.ip = ip;
+
         if (ip >= endOfBody) {
             record_abort(traceIp, ip, traceName, "end-of-body");
             break;
@@ -181,8 +162,9 @@ export function generate_wasm_body (
         // HACK: Browsers set a limit of 4KB, we lower it slightly since a single opcode
         //  might generate a ton of code and we generate a bit of an epilogue after
         //  we finish
-        const maxModuleSize = 3850;
-        if (builder.size >= maxModuleSize - builder.bytesGeneratedSoFar) {
+        const maxModuleSize = 3850,
+            spaceLeft = maxModuleSize - builder.bytesGeneratedSoFar - builder.cfg.overheadBytes;
+        if (builder.size >= spaceLeft) {
             // console.log(`trace too big, estimated size is ${builder.size + builder.bytesGeneratedSoFar}`);
             record_abort(traceIp, ip, traceName, "trace-too-big");
             break;
@@ -203,11 +185,10 @@ export function generate_wasm_body (
         const isBackBranchTarget = builder.options.noExitBackwardBranches &&
             is_backward_branch_target(ip, startOfBody, backwardBranchTable),
             isForwardBranchTarget = builder.branchTargets.has(ip),
-            needsEipCheck = isBackBranchTarget || isForwardBranchTarget ||
+            startBranchBlock = isBackBranchTarget || isForwardBranchTarget ||
                 // If a method contains backward branches, we also need to check eip at the first insn
                 //  because a backward branch might target a point in the middle of the trace
-                (isFirstInstruction && backwardBranchTable),
-            needsFallthroughEipUpdate = needsEipCheck && !isFirstInstruction;
+                (isFirstInstruction && backwardBranchTable);
         let isLowValueOpcode = false,
             skipDregInvalidation = false;
 
@@ -219,7 +200,7 @@ export function generate_wasm_body (
             builder.backBranchOffsets.push(ip);
         }
 
-        if (needsEipCheck) {
+        if (startBranchBlock) {
             // If execution runs past the end of the current branch block, ensure
             //  that the instruction pointer is updated appropriately. This will
             //  also guarantee that the branch target block's comparison will
@@ -227,11 +208,7 @@ export function generate_wasm_body (
             // We make sure above that this isn't done for the start of the trace,
             //  otherwise loops will run forever and never terminate since after
             //  branching to the top of the loop we would blow away eip
-            if (needsFallthroughEipUpdate) {
-                builder.ip_const(rip);
-                builder.local("eip", WasmOpcode.set_local);
-            }
-            append_branch_target_block(builder, ip);
+            append_branch_target_block(builder, ip, isBackBranchTarget);
             inBranchBlock = true;
             firstOpcodeInBlock = true;
             eraseInferredState();
@@ -1218,24 +1195,14 @@ export function generate_wasm_body (
     if (emitPadding)
         builder.appendU8(WasmOpcode.nop);
 
-    // Ensure that if execution runs past the end of our last branch block, we
-    //  update eip appropriately so that we will return the right ip
-    builder.ip_const(rip);
-    builder.local("eip", WasmOpcode.set_local);
-
     // We need to close any open blocks before generating our closing ret,
     //  because wasm would allow branching past the ret otherwise
     while (builder.activeBlocks > 0)
         builder.endBlock();
 
-    // Now we generate a ret at the end of the function body so it's Valid(tm)
-    // When branching is enabled, we will have potentially updated eip due to a
-    //  branch and then executed forward without ever finding it, so we want to
-    //  return the branch target and ensure that the interpreter starts running
-    //  from there.
-    builder.local("eip");
-    builder.appendU8(WasmOpcode.return_);
-    builder.appendU8(WasmOpcode.end);
+    builder.cfg.exitIp = rip;
+
+    // console.log(`estimated size: ${builder.size + builder.cfg.overheadBytes + builder.bytesGeneratedSoFar}`);
 
     return result;
 }
@@ -1260,17 +1227,8 @@ function invalidate_local_range (start: number, bytes: number) {
         invalidate_local(start + i);
 }
 
-function append_branch_target_block (builder: WasmBuilder, ip: MintOpcodePtr) {
-    builder.endBlock();
-    // Create a new branch block that conditionally executes depending on the eip local
-    // FIXME: For methods containing backward branches, we will have one of these compares
-    //  at the top of the trace and pay the cost of it on every entry even though it will
-    //  always pass. If we never generate any backward branches during compilation, we should
-    //  patch it out
-    builder.local("eip");
-    builder.ip_const(ip);
-    builder.appendU8(WasmOpcode.i32_eq);
-    builder.block(WasmValtype.void, WasmOpcode.if_);
+function append_branch_target_block (builder: WasmBuilder, ip: MintOpcodePtr, isBackBranchTarget: boolean) {
+    builder.cfg.startBranchBlock(ip, isBackBranchTarget);
 }
 
 function append_ldloc (builder: WasmBuilder, offset: number, opcode: WasmOpcode) {
@@ -2378,10 +2336,7 @@ function emit_branch (
                     // append_safepoint(builder, ip);
                     if (traceBackBranches)
                         console.log(`performing backward branch to 0x${destination.toString(16)}`);
-                    builder.ip_const(destination);
-                    builder.local("eip", WasmOpcode.set_local);
-                    builder.appendU8(WasmOpcode.br);
-                    builder.appendULeb(1);
+                    builder.cfg.branch(destination, true, false);
                     counters.backBranchesEmitted++;
                     return true;
                 } else {
@@ -2397,10 +2352,7 @@ function emit_branch (
                 //  don't need to wrap things in a block here, we can just exit
                 //  the current branch block after updating eip
                 builder.branchTargets.add(destination);
-                builder.ip_const(destination);
-                builder.local("eip", WasmOpcode.set_local);
-                builder.appendU8(WasmOpcode.br);
-                builder.appendULeb(0);
+                builder.cfg.branch(destination, false, false);
                 return true;
             }
         }
@@ -2470,15 +2422,7 @@ function emit_branch (
             //  we update eip and branch out to the top of the loop block
             if (traceBackBranches)
                 console.log(`performing conditional backward branch to 0x${destination.toString(16)}`);
-            builder.ip_const(destination);
-            builder.local("eip", WasmOpcode.set_local);
-            builder.appendU8(WasmOpcode.br);
-            // break out 3 levels, because the current stack layout is
-            // loop {
-            //   branch target block {
-            //     branch dispatch block {
-            // and we want to target the loop in order to jump to the top of it
-            builder.appendULeb(2);
+            builder.cfg.branch(destination, true, true);
             counters.backBranchesEmitted++;
         } else {
             if (traceBackBranches)
@@ -2493,12 +2437,7 @@ function emit_branch (
             append_safepoint(builder, ip);
         // Branching is enabled, so set eip and exit the current branch block
         builder.branchTargets.add(destination);
-        builder.ip_const(destination);
-        builder.local("eip", WasmOpcode.set_local);
-        builder.appendU8(WasmOpcode.br);
-        // The branch block encloses this tiny branch dispatch block, so break
-        //  out two levels
-        builder.appendULeb(1);
+        builder.cfg.branch(destination, false, true);
     }
 
     builder.endBlock();
@@ -3052,15 +2991,6 @@ function emit_arrayop (builder: WasmBuilder, frame: NativePointer, ip: MintOpcod
     return true;
 }
 
-function append_bailout (builder: WasmBuilder, ip: MintOpcodePtr, reason: BailoutReason) {
-    builder.ip_const(ip);
-    if (builder.options.countBailouts) {
-        builder.i32_const(reason);
-        builder.callImport("bailout");
-    }
-    builder.appendU8(WasmOpcode.return_);
-}
-
 function append_safepoint (builder: WasmBuilder, ip: MintOpcodePtr) {
     // Check whether a safepoint is required
     builder.ptr_const(cwraps.mono_jiterp_get_polling_required_address());

src/mono/wasm/runtime/jiterpreter.ts

@@ -15,7 +15,8 @@ import {
     _now, elapsedTimes, shortNameBase,
     counters, getRawCwrap, importDef,
     JiterpreterOptions, getOptions, recordFailure,
-    JiterpMember, getMemberOffset
+    JiterpMember, getMemberOffset,
+    BailoutReasonNames
 } from "./jiterpreter-support";
 import {
     generate_wasm_body
@@ -165,60 +166,6 @@ struct InterpMethod {
        void **data_items;
 */
 
-export const enum BailoutReason {
-    Unknown,
-    InterpreterTiering,
-    NullCheck,
-    VtableNotInitialized,
-    Branch,
-    BackwardBranch,
-    ConditionalBranch,
-    ConditionalBackwardBranch,
-    ComplexBranch,
-    ArrayLoadFailed,
-    ArrayStoreFailed,
-    StringOperationFailed,
-    DivideByZero,
-    Overflow,
-    Return,
-    Call,
-    Throw,
-    AllocFailed,
-    SpanOperationFailed,
-    CastFailed,
-    SafepointBranchTaken,
-    UnboxFailed,
-    CallDelegate,
-    Debugging
-}
-
-export const BailoutReasonNames = [
-    "Unknown",
-    "InterpreterTiering",
-    "NullCheck",
-    "VtableNotInitialized",
-    "Branch",
-    "BackwardBranch",
-    "ConditionalBranch",
-    "ConditionalBackwardBranch",
-    "ComplexBranch",
-    "ArrayLoadFailed",
-    "ArrayStoreFailed",
-    "StringOperationFailed",
-    "DivideByZero",
-    "Overflow",
-    "Return",
-    "Call",
-    "Throw",
-    "AllocFailed",
-    "SpanOperationFailed",
-    "CastFailed",
-    "SafepointBranchTaken",
-    "UnboxFailed",
-    "CallDelegate",
-    "Debugging"
-];
-
 export let traceBuilder : WasmBuilder;
 export let traceImports : Array<[string, string, Function]> | undefined;
 
@@ -723,7 +670,7 @@ function generate_wasm (
                 name: traceName,
                 export: true,
                 locals: {
-                    "eip": WasmValtype.i32,
+                    "disp": WasmValtype.i32,
                     "temp_ptr": WasmValtype.i32,
                     "cknull_ptr": WasmValtype.i32,
                     "math_lhs32": WasmValtype.i32,
@@ -743,6 +690,8 @@ function generate_wasm (
                 if (getU16(ip) !== MintOpcode.MINT_TIER_PREPARE_JITERPRETER)
                     throw new Error(`Expected *ip to be MINT_TIER_PREPARE_JITERPRETER but was ${getU16(ip)}`);
 
+                builder.cfg.initialize(startOfBody, backwardBranchTable, !!instrument);
+
                 // TODO: Call generate_wasm_body before generating any of the sections and headers.
                 // This will allow us to do things like dynamically vary the number of locals, in addition
                 //  to using global constants and figuring out how many constant slots we need in advance
@@ -751,7 +700,10 @@ function generate_wasm (
                     frame, traceName, ip, startOfBody, endOfBody,
                     builder, instrumentedTraceId, backwardBranchTable
                 );
+
                 keep = (opcodesProcessed >= mostRecentOptions!.minimumTraceLength);
+
+                return builder.cfg.generate();
             }
         );
 
@@ -769,6 +721,8 @@ function generate_wasm (
 
         compileStarted = _now();
         const buffer = builder.getArrayView();
+        // console.log(`bytes generated: ${buffer.length}`);
+
         if (trace > 0)
             console.log(`${(<any>(builder.base)).toString(16)} ${methodFullName || traceName} generated ${buffer.length} byte(s) of wasm`);
         counters.bytesGenerated += buffer.length;