Fix regression in jiterpreter trace generation

When I migrated the jiterpreter to a new branch target scanning pass I failed to filter the branch target tables it generates properly. This combined with an existing bug in the cfg code meant that traces in large methods ended up having big branch dispatch tables at the top where most of the branch targets were never hit, eating up a valuable chunk of the trace's 4KB limit.

This PR adds appropriate filtering in both the scanning pass and the cfg so that the dispatch table is as small as possible.

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

@@ -1148,7 +1148,7 @@ class Cfg {
     blockStack: Array<MintOpcodePtr> = [];
     backDispatchOffsets: Array<MintOpcodePtr> = [];
     dispatchTable = new Map<MintOpcodePtr, number>();
-    observedBranchTargets = new Set<MintOpcodePtr>();
+    observedBackBranchTargets = new Set<MintOpcodePtr>();
     trace = 0;
 
     constructor(builder: WasmBuilder) {
@@ -1165,7 +1165,7 @@ class Cfg {
         this.lastSegmentEnd = 0;
         this.overheadBytes = 10; // epilogue
         this.dispatchTable.clear();
-        this.observedBranchTargets.clear();
+        this.observedBackBranchTargets.clear();
         this.trace = trace;
         this.backDispatchOffsets.length = 0;
     }
@@ -1212,7 +1212,9 @@ class Cfg {
     }
 
     branch(target: MintOpcodePtr, isBackward: boolean, branchType: CfgBranchType) {
-        this.observedBranchTargets.add(target);
+        if (isBackward)
+            this.observedBackBranchTargets.add(target);
+
         this.appendBlob();
         this.segments.push({
             type: "branch",
@@ -1224,7 +1226,10 @@ class Cfg {
         // some branches will generate bailouts instead so we allocate 4 bytes per branch
         //  to try and balance this out and avoid underestimating too much
         this.overheadBytes += 4; // forward branches are a constant br + depth (optimally 2 bytes)
+
         if (isBackward) {
+            // TODO: Make this smaller by setting the flag inside the dispatcher when disp != 0,
+            //  this will save space for any trace with more than one back-branch
             // get_local <cinfo>
             // i32_const 1
             // i32_store 0 0
@@ -1298,7 +1303,7 @@ class Cfg {
                 const breakDepth = this.blockStack.indexOf(offset);
                 if (breakDepth < 0)
                     continue;
-                if (!this.observedBranchTargets.has(offset))
+                if (!this.observedBackBranchTargets.has(offset))
                     continue;
 
                 this.dispatchTable.set(offset, this.backDispatchOffsets.length + 1);

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

@@ -142,6 +142,8 @@ export function generateBackwardBranchTable(
     // IP of the start of the trace in U16s, relative to startOfBody.
     const rbase16 = (<any>ip - <any>startOfBody) / 2;
 
+    // FIXME: This will potentially scan the entire method and record branches that won't
+    //  ever run since the trace compilation will end before we reach them.
     while (ip < endOfBody) {
         // IP of the current opcode in U16s, relative to startOfBody. This is what the back branch table uses
         const rip16 = (<any>ip - <any>startOfBody) / 2;
@@ -166,16 +168,23 @@ export function generateBackwardBranchTable(
             break;
         }
 
-        const rtarget16 = rip16 + (displacement);
-        if (rtarget16 < 0) {
-            mono_log_info(`opcode @${ip}'s displacement of ${displacement} goes before body: ${rtarget16}. aborting backbranch table generation`);
-            break;
-        }
+        // Only record *backward* branches
+        // We will filter this down further in the Cfg because it takes note of which branches it sees,
+        //  but it is also beneficial to have a null table (further down) due to seeing no potential
+        //  back branch targets at all, as it allows the Cfg to skip additional code generation entirely
+        //  if it knows there will never be any backwards branches in a given trace
+        if (displacement < 0) {
+            const rtarget16 = rip16 + (displacement);
+            if (rtarget16 < 0) {
+                mono_log_info(`opcode @${ip}'s displacement of ${displacement} goes before body: ${rtarget16}. aborting backbranch table generation`);
+                break;
+            }
 
-        // If the relative target is before the start of the trace, don't record it.
-        // The trace will be unable to successfully branch to it so it would just make the table bigger.
-        if (rtarget16 >= rbase16)
-            table.push(rtarget16);
+            // If the relative target is before the start of the trace, don't record it.
+            // The trace will be unable to successfully branch to it so it would just make the table bigger.
+            if (rtarget16 >= rbase16)
+                table.push(rtarget16);
+        }
 
         switch (opcode) {
             case MintOpcode.MINT_CALL_HANDLER: