diff --git a/compiler/bin-wasm_of_ocaml/compile.ml b/compiler/bin-wasm_of_ocaml/compile.ml
index 4f542258b..ddf240248 100644
--- a/compiler/bin-wasm_of_ocaml/compile.ml
+++ b/compiler/bin-wasm_of_ocaml/compile.ml
@@ -148,7 +148,9 @@ let generate_prelude ~out_file =
     Driver.f ~target:Wasm (Parse_bytecode.Debug.create ~include_cmis:false false) code
   in
   let context = Wa_generate.start () in
-  let _ = Wa_generate.f ~context ~unit_name:(Some "prelude") ~live_vars ~in_cps p in
+  let _ =
+    Wa_generate.f ~context ~unit_name:(Some "prelude") ~live_vars ~in_cps ~debug p
+  in
   Wa_generate.output ch ~context ~debug;
   uinfo.provides
 
@@ -283,7 +285,7 @@ let run
     in
     let context = Wa_generate.start () in
     let toplevel_name, generated_js =
-      Wa_generate.f ~context ~unit_name ~live_vars ~in_cps p
+      Wa_generate.f ~context ~unit_name ~live_vars ~in_cps ~debug p
     in
     if standalone then Wa_generate.add_start_function ~context toplevel_name;
     Wa_generate.output ch ~context ~debug;
diff --git a/compiler/lib/wasm/wa_code_generation.ml b/compiler/lib/wasm/wa_code_generation.ml
index 293de76d0..5040b2fc3 100644
--- a/compiler/lib/wasm/wa_code_generation.ml
+++ b/compiler/lib/wasm/wa_code_generation.ml
@@ -291,15 +291,16 @@ let blk l st =
   List.rev st.instrs, { st with instrs }
 
 let with_location loc instrs st =
-  let current_instrs = st.instrs in
-  let (), st = instrs { st with instrs = [] } in
-  let[@tail_mod_cons] rec add_loc loc = function
-    | [] -> current_instrs
-    | W.Nop :: rem -> W.Nop :: add_loc loc rem
-    | Location _ :: _ as l -> l @ current_instrs (* Stop on the first location *)
-    | i :: rem -> W.Location (loc, i) :: add_loc loc rem
-  in
-  (), { st with instrs = add_loc loc st.instrs }
+  let (), st = instrs st in
+  ( ()
+  , { st with
+      instrs =
+        (match st.instrs with
+        | [] -> []
+        | Location _ :: _ when Poly.equal loc No -> st.instrs
+        | Location (_, i) :: rem -> Location (loc, i) :: rem
+        | i :: rem -> Location (loc, i) :: rem)
+    } )
 
 let cast ?(nullable = false) typ e =
   let* e = e in
@@ -469,6 +470,13 @@ let get_i31_value x st =
       let x = Var.fresh () in
       let x, st = add_var ~typ:I32 x st in
       Some x, { st with instrs = LocalSet (x', RefI31 (LocalTee (x, e))) :: rem }
+  | Location (loc, LocalSet (x', RefI31 e)) :: rem when Code.Var.equal x x' && is_smi e ->
+      let x = Var.fresh () in
+      let x, st = add_var ~typ:I32 x st in
+      ( Some x
+      , { st with
+          instrs = Location (loc, LocalSet (x', RefI31 (LocalTee (x, e)))) :: rem
+        } )
   | _ -> None, st
 
 let load x =
diff --git a/compiler/lib/wasm/wa_generate.ml b/compiler/lib/wasm/wa_generate.ml
index 6eb088ad7..fedb49d50 100644
--- a/compiler/lib/wasm/wa_generate.ml
+++ b/compiler/lib/wasm/wa_generate.ml
@@ -37,6 +37,7 @@ module Generate (Target : Wa_target_sig.S) = struct
     ; blocks : block Addr.Map.t
     ; closures : Wa_closure_conversion.closure Var.Map.t
     ; global_context : Wa_code_generation.context
+    ; debug : Parse_bytecode.Debug.t
     }
 
   let func_type n =
@@ -818,6 +819,20 @@ module Generate (Target : Wa_target_sig.S) = struct
       params
       ((pc, _) as cont)
       acc =
+    let ctx =
+      let loc = Before pc in
+      match Parse_bytecode.Debug.find_loc ctx.debug loc with
+      | Some _ ->
+          let block = Addr.Map.find pc ctx.blocks in
+          let block =
+            match block.body with
+            | (i, _) :: rem -> { block with body = (i, loc) :: rem }
+            | [] -> { block with branch = fst block.branch, loc }
+          in
+          let blocks = Addr.Map.add pc block ctx.blocks in
+          { ctx with blocks }
+      | None -> ctx
+    in
     let stack_info =
       Stack.generate_spilling_information
         p
@@ -1107,13 +1122,16 @@ module Generate (Target : Wa_target_sig.S) = struct
       ~in_cps (*
     ~should_export
     ~warn_on_unhandled_effect
-      _debug *) =
+*)
+      ~debug =
     global_context.unit_name <- unit_name;
     let p, closures = Wa_closure_conversion.f p in
     (*
   Code.Print.program (fun _ _ -> "") p;
 *)
-    let ctx = { live = live_vars; in_cps; blocks = p.blocks; closures; global_context } in
+    let ctx =
+      { live = live_vars; in_cps; blocks = p.blocks; closures; global_context; debug }
+    in
     let toplevel_name = Var.fresh_n "toplevel" in
     let functions =
       Code.fold_closures_outermost_first
@@ -1223,15 +1241,15 @@ let start () =
       | `Core -> Wa_core_target.Value.value
       | `GC -> Wa_gc_target.Value.value)
 
-let f ~context ~unit_name p ~live_vars ~in_cps =
+let f ~context ~unit_name p ~live_vars ~in_cps ~debug =
   let p = if Config.Flag.effects () then fix_switch_branches p else p in
   match target with
   | `Core ->
       let module G = Generate (Wa_core_target) in
-      G.f ~context ~unit_name ~live_vars ~in_cps p
+      G.f ~context ~unit_name ~live_vars ~in_cps ~debug p
   | `GC ->
       let module G = Generate (Wa_gc_target) in
-      G.f ~context ~unit_name ~live_vars ~in_cps p
+      G.f ~context ~unit_name ~live_vars ~in_cps ~debug p
 
 let add_start_function =
   match target with
diff --git a/compiler/lib/wasm/wa_generate.mli b/compiler/lib/wasm/wa_generate.mli
index 83f49d662..d7e2e8662 100644
--- a/compiler/lib/wasm/wa_generate.mli
+++ b/compiler/lib/wasm/wa_generate.mli
@@ -26,6 +26,7 @@ val f :
   -> Code.program
   -> live_vars:int array
   -> in_cps:Effects.in_cps
+  -> debug:Parse_bytecode.Debug.t
   -> Wa_ast.var * (string list * (string * Javascript.expression) list)
 
 val add_start_function : context:Wa_code_generation.context -> Wa_ast.var -> unit