JIT: record line info at continuation labels after function splits

Add current_line tracking to the compiler state and
`record_continuation_line/3` to emit line info at continuation points.

Signed-off-by: Paul Guyot <pguyot@kallisys.net>

Author: pguyot

libs/jit/src/jit.erl

@@ -85,6 +85,7 @@
 
 -record(state, {
     line_offsets :: [{integer(), integer()}],
+    current_line :: integer() | undefined,
     labels_count :: pos_integer(),
     atom_resolver :: fun((integer()) -> atom()),
     literal_resolver :: fun((integer()) -> any()),
@@ -136,6 +137,7 @@ compile(
 ) when OpcodeMax =< ?OPCODE_MAX ->
     State0 = #state{
         line_offsets = [],
+        current_line = undefined,
         labels_count = LabelsCount,
         atom_resolver = AtomResolver,
         literal_resolver = LiteralResolver,
@@ -296,11 +298,12 @@ first_pass(<<?OP_CALL_EXT, Rest0/binary>>, MMod, MSt0, State0) ->
     {Index, Rest2} = decode_literal(Rest1),
     ?TRACE("OP_CALL_EXT ~p, ~p\n", [Arity, Index]),
     MSt1 = MMod:decrement_reductions_and_maybe_schedule_next(MSt0),
+    State1 = record_continuation_line(MMod, MSt1, State0),
     MSt2 = MMod:call_primitive_with_cp(MSt1, ?PRIM_CALL_EXT, [
         ctx, jit_state, offset, Arity, Index, ?CALL_EXT_NO_DEALLOC_MFA
     ]),
     ?ASSERT_ALL_NATIVE_FREE(MSt2),
-    first_pass(Rest2, MMod, MSt2, State0);
+    first_pass(Rest2, MMod, MSt2, State1);
 % 8
 first_pass(<<?OP_CALL_EXT_LAST, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
@@ -309,11 +312,12 @@ first_pass(<<?OP_CALL_EXT_LAST, Rest0/binary>>, MMod, MSt0, State0) ->
     {NWords, Rest3} = decode_literal(Rest2),
     ?TRACE("OP_CALL_EXT_LAST ~p, ~p, ~p\n", [Arity, Index, NWords]),
     MSt1 = MMod:decrement_reductions_and_maybe_schedule_next(MSt0),
+    State1 = record_continuation_line(MMod, MSt1, State0),
     MSt2 = MMod:call_primitive_last(MSt1, ?PRIM_CALL_EXT, [
         ctx, jit_state, offset, Arity, Index, NWords
     ]),
     ?ASSERT_ALL_NATIVE_FREE(MSt2),
-    first_pass(Rest3, MMod, MSt2, State0);
+    first_pass(Rest3, MMod, MSt2, State1);
 % 9
 first_pass(<<?OP_BIF0, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
@@ -1033,13 +1037,14 @@ first_pass(<<?OP_CALL_FUN, Rest0/binary>>, MMod, MSt0, State0) ->
     {ArgsCount, Rest1} = decode_literal(Rest0),
     ?TRACE("OP_CALL_FUN ~p\n", [ArgsCount]),
     MSt1 = MMod:decrement_reductions_and_maybe_schedule_next(MSt0),
+    State1a = record_continuation_line(MMod, MSt1, State0),
     {MSt2, FuncReg} = read_any_xreg(ArgsCount, MMod, MSt1),
     {MSt3, Reg} = verify_is_function(FuncReg, MMod, MSt2),
     MSt4 = MMod:call_primitive_with_cp(MSt3, ?PRIM_CALL_FUN, [
         ctx, jit_state, offset, {free, Reg}, ArgsCount
     ]),
     ?ASSERT_ALL_NATIVE_FREE(MSt4),
-    first_pass(Rest1, MMod, MSt4, State0);
+    first_pass(Rest1, MMod, MSt4, State1a);
 % 77
 first_pass(<<?OP_IS_FUNCTION, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
@@ -1056,11 +1061,12 @@ first_pass(<<?OP_CALL_EXT_ONLY, Rest0/binary>>, MMod, MSt0, State0) ->
     {Index, Rest2} = decode_literal(Rest1),
     ?TRACE("OP_CALL_EXT_ONLY ~p, ~p\n", [Arity, Index]),
     MSt1 = MMod:decrement_reductions_and_maybe_schedule_next(MSt0),
+    State1 = record_continuation_line(MMod, MSt1, State0),
     MSt2 = MMod:call_primitive_last(MSt1, ?PRIM_CALL_EXT, [
         ctx, jit_state, offset, Arity, Index, ?CALL_EXT_NO_DEALLOC
     ]),
     ?ASSERT_ALL_NATIVE_FREE(MSt2),
-    first_pass(Rest2, MMod, MSt2, State0);
+    first_pass(Rest2, MMod, MSt2, State1);
 % 96
 first_pass(<<?OP_FMOVE, ?COMPACT_EXTENDED_FP_REGISTER, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
@@ -1188,6 +1194,7 @@ first_pass(<<?OP_APPLY, Rest0/binary>>, MMod, MSt0, State0) ->
     {Arity, Rest1} = decode_literal(Rest0),
     ?TRACE("OP_APPLY ~p\n", [Arity]),
     MSt1 = MMod:decrement_reductions_and_maybe_schedule_next(MSt0),
+    State1a = record_continuation_line(MMod, MSt1, State0),
     {MSt2, Module} = read_any_xreg(Arity, MMod, MSt1),
     {MSt3, Function} = read_any_xreg(Arity + 1, MMod, MSt2),
     MSt4 = verify_is_atom(Module, 0, MMod, MSt3),
@@ -1196,14 +1203,15 @@ first_pass(<<?OP_APPLY, Rest0/binary>>, MMod, MSt0, State0) ->
         ctx, jit_state, offset, {free, Module}, {free, Function}, Arity
     ]),
     ?ASSERT_ALL_NATIVE_FREE(MSt6),
-    first_pass(Rest1, MMod, MSt6, State0);
+    first_pass(Rest1, MMod, MSt6, State1a);
 % 113
 first_pass(<<?OP_APPLY_LAST, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
     {Arity, Rest1} = decode_literal(Rest0),
     {NWords, Rest2} = decode_literal(Rest1),
     ?TRACE("OP_APPLY_LAST ~p, ~p\n", [Arity, NWords]),
     MSt1 = MMod:decrement_reductions_and_maybe_schedule_next(MSt0),
+    State1 = record_continuation_line(MMod, MSt1, State0),
     {MSt2, Module} = read_any_xreg(Arity, MMod, MSt1),
     {MSt3, Function} = read_any_xreg(Arity + 1, MMod, MSt2),
     MSt4 = verify_is_atom(Module, 0, MMod, MSt3),
@@ -1214,7 +1222,7 @@ first_pass(<<?OP_APPLY_LAST, Rest0/binary>>, MMod, MSt0, State0) ->
         ctx, jit_state, offset, {free, Module}, {free, Function}, Arity
     ]),
     ?ASSERT_ALL_NATIVE_FREE(MSt8),
-    first_pass(Rest2, MMod, MSt8, State0);
+    first_pass(Rest2, MMod, MSt8, State1);
 % 114
 first_pass(<<?OP_IS_BOOLEAN, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
@@ -1729,7 +1737,8 @@ first_pass(
     MSt0 = ?DWARF_LINE(MMod, MSt, Line),
     Offset = MMod:offset(MSt0),
     first_pass(Rest1, MMod, MSt0, State0#state{
-        line_offsets = [{Line, Offset} | AccLines]
+        line_offsets = [{Line, Offset} | AccLines],
+        current_line = Line
     });
 % 154
 first_pass(<<?OP_PUT_MAP_ASSOC, Rest0/binary>>, MMod, MSt0, State0) ->
@@ -2446,13 +2455,14 @@ first_pass(<<?OP_CALL_FUN2, Rest0/binary>>, MMod, MSt0, State0) ->
     % We ignore Tag (could be literal 0 or atom unsafe)
     MSt2 = MMod:free_native_registers(MSt1, [Tag]),
     MSt3 = MMod:decrement_reductions_and_maybe_schedule_next(MSt2),
-    {MSt4, Fun, Rest3} = decode_typed_compact_term(Rest2, MMod, MSt3, State0),
+    State1a = record_continuation_line(MMod, MSt3, State0),
+    {MSt4, Fun, Rest3} = decode_typed_compact_term(Rest2, MMod, MSt3, State1a),
     {MSt5, Reg} = verify_is_function(Fun, MMod, MSt4),
     MSt6 = MMod:call_primitive_with_cp(MSt5, ?PRIM_CALL_FUN, [
         ctx, jit_state, offset, {free, Reg}, ArgsCount
     ]),
     ?ASSERT_ALL_NATIVE_FREE(MSt6),
-    first_pass(Rest3, MMod, MSt6, State0);
+    first_pass(Rest3, MMod, MSt6, State1a);
 % 180
 first_pass(<<?OP_BADRECORD, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
@@ -4518,6 +4528,15 @@ memory_ensure_free_with_extra_root(ExtraRoot, Live, Size, MMod, MSt0) when is_tu
     MSt4 = MMod:move_to_vm_register(MSt3, ExtraRootXReg, ExtraRoot),
     {MSt4, ExtraRoot}.
 
+%% Record the current line at the current offset, so that continuation labels
+%% created by function-splitting operations (WASM) have line info.
+%% On native backends this is a no-op as the offset is within the existing range.
+record_continuation_line(_MMod, _MSt, #state{current_line = undefined} = State) ->
+    State;
+record_continuation_line(MMod, MSt, #state{current_line = Line, line_offsets = AccLines} = State) ->
+    Offset = MMod:offset(MSt),
+    State#state{line_offsets = [{Line, Offset} | AccLines]}.
+
 second_pass(MMod, MSt0, #state{line_offsets = Lines}) ->
     ?TRACE("SECOND PASS -- ~B lines\n", [length(Lines)]),
     % Add extra function that returns labels and line information