[wasm] Re-enable null check optimization for mid-method traces

src/mono/mono/mini/interp/interp-internals.h

@@ -182,6 +182,7 @@ struct InterpMethod {
 	unsigned int contains_traces : 1;
 	guint16 *backward_branch_offsets;
 	unsigned int backward_branch_offsets_count;
+	MonoBitSet *address_taken_bits;
 #endif
 #if PROFILE_INTERP
 	long calls;

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

@@ -942,6 +942,30 @@ trace_info_alloc () {
 	return index;
 }
 
+static void
+build_address_taken_bitset (TransformData *td, InterpBasicBlock *bb, guint32 bitset_size)
+{
+	for (InterpInst *ins = bb->first_ins; ins != NULL; ins = ins->next) {
+		if (ins->opcode == MINT_LDLOCA_S) {
+			InterpMethod *imethod = td->rtm;
+			InterpLocal *loc = &td->locals[ins->sregs[0]];
+
+			// Allocate on demand so if a method contains no ldlocas we don't allocate the bitset
+			if (!imethod->address_taken_bits)
+				imethod->address_taken_bits = mono_bitset_new (bitset_size, 0);
+
+			// Ensure that every bit in the set corresponding to space occupied by this local
+			//  is set, so that large locals (structs etc) being ldloca'd properly sets the
+			//  whole range covered by the struct as a no-go for optimization.
+			// FIXME: Do this per slot instead of per byte.
+			for (int j = 0; j < loc->size; j++) {
+				guint32 b = (loc->offset + j) / MINT_STACK_SLOT_SIZE;
+				mono_bitset_set (imethod->address_taken_bits, b);
+			}
+		}
+	}
+}
+
 /*
  * Insert jiterpreter entry points at the correct candidate locations:
  * The first basic block of the function,
@@ -959,13 +983,15 @@ jiterp_insert_entry_points (void *_imethod, void *_td)
 	TransformData *td = (TransformData *)_td;
 	// Insert an entry opcode for the next basic block (call resume and first bb)
 	// FIXME: Should we do this based on relationships between BBs instead of insn sequence?
-	gboolean enter_at_next = TRUE;
+	gboolean enter_at_next = TRUE, table_full = FALSE;
 
 	if (!mono_opt_jiterpreter_traces_enabled)
 		return;
 
 	// Start with a high instruction counter so the distance check will pass
 	int instruction_count = mono_opt_jiterpreter_minimum_distance_between_traces;
+	// Pre-calculate how big the address-taken-locals bitset needs to be
+	guint32 bitset_size = td->total_locals_size / MINT_STACK_SLOT_SIZE;
 
 	for (InterpBasicBlock *bb = td->entry_bb; bb != NULL; bb = bb->next_bb) {
 		// Enter trace at top of functions
@@ -977,7 +1003,7 @@ jiterp_insert_entry_points (void *_imethod, void *_td)
 		if (mono_opt_jiterpreter_backward_branch_entries_enabled && bb->backwards_branch_target)
 			is_backwards_branch = TRUE;
 
-		gboolean enabled = (is_backwards_branch || is_resume_or_first);
+		gboolean enabled = (is_backwards_branch || is_resume_or_first) && !table_full;
 		// FIXME: This scan will likely proceed forward all the way out of the current block,
 		//  which means that for large methods we will sometimes scan the same instruction
 		//  multiple times and waste some work. At present this is unavoidable because
@@ -1004,7 +1030,7 @@ jiterp_insert_entry_points (void *_imethod, void *_td)
 			gint32 trace_index = trace_info_alloc ();
 			if (trace_index < 0) {
 				// We're out of space in the TraceInfo table.
-				return;
+				table_full = TRUE;
 			} else {
 				td->cbb = bb;
 				imethod->contains_traces = TRUE;
@@ -1032,6 +1058,14 @@ jiterp_insert_entry_points (void *_imethod, void *_td)
 		//  the new instruction counter will be the number of instructions in the block, so if
 		//  it's big enough we'll be able to insert another entry point right away.
 		instruction_count += bb->in_count;
+
+		build_address_taken_bitset (td, bb, bitset_size);
+	}
+
+	// If we didn't insert any entry points and we allocated the bitset, free it.
+	if (!imethod->contains_traces && imethod->address_taken_bits) {
+		mono_bitset_free (imethod->address_taken_bits);
+		imethod->address_taken_bits = NULL;
 	}
 }
 
@@ -1446,6 +1480,16 @@ mono_jiterp_boost_back_branch_target (guint16 *ip) {
 	*/
 }
 
+EMSCRIPTEN_KEEPALIVE int
+mono_jiterp_is_imethod_var_address_taken (InterpMethod *imethod, int offset) {
+	g_assert (imethod);
+	g_assert (offset >= 0);
+	if (!imethod->address_taken_bits)
+		return FALSE;
+
+	return mono_bitset_test (imethod->address_taken_bits, offset / MINT_STACK_SLOT_SIZE);
+}
+
 // HACK: fix C4206
 EMSCRIPTEN_KEEPALIVE
 #endif // HOST_BROWSER

src/mono/wasm/runtime/cwraps.ts

@@ -123,6 +123,7 @@ const fn_signatures: SigLine[] = [
     [true, "mono_jiterp_get_polling_required_address", "number", []],
     [true, "mono_jiterp_get_rejected_trace_count", "number", []],
     [true, "mono_jiterp_boost_back_branch_target", "void", ["number"]],
+    [true, "mono_jiterp_is_imethod_var_address_taken", "number", ["number", "number"]],
     ...legacy_interop_cwraps
 ];
 
@@ -242,6 +243,7 @@ export interface t_Cwraps {
     mono_jiterp_write_number_unaligned(destination: VoidPtr, value: number, mode: number): void;
     mono_jiterp_get_rejected_trace_count(): number;
     mono_jiterp_boost_back_branch_target(destination: number): void;
+    mono_jiterp_is_imethod_var_address_taken(imethod: VoidPtr, offsetBytes: number): number;
 }
 
 const wrapped_c_functions: t_Cwraps = <any>{};

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

@@ -151,6 +151,7 @@ export class WasmBuilder {
     argumentCount!: number;
     activeBlocks!: number;
     base!: MintOpcodePtr;
+    frame: NativePointer = <any>0;
     traceBuf: Array<string> = [];
     branchTargets = new Set<MintOpcodePtr>();
     options!: JiterpreterOptions;