Move regs logic to jit_regs with maps backend

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

Author: pguyot

libs/jit/src/CMakeLists.txt

@@ -25,6 +25,7 @@ include(BuildErlang)
 set(ERLANG_MODULES
     jit
     jit_precompile
+    jit_regs
     jit_stream_binary
     jit_stream_flash
     jit_stream_mmap

libs/jit/src/jit_aarch64.erl

@@ -136,7 +136,7 @@
     branches :: [{non_neg_integer(), non_neg_integer(), non_neg_integer()}],
     jump_table_start :: non_neg_integer(),
     available_regs :: [aarch64_register()],
-    used_regs :: [aarch64_register()],
+    used_regs :: jit_regs:regs(),
     labels :: [{integer() | reference(), integer()}],
     variant :: non_neg_integer()
 }).
@@ -238,7 +238,7 @@ new(Variant, StreamModule, Stream) ->
         jump_table_start = 0,
         offset = StreamModule:offset(Stream),
         available_regs = ?AVAILABLE_REGS,
-        used_regs = [],
+        used_regs = jit_regs:new(),
         labels = [],
         variant = Variant
     }.
@@ -293,7 +293,7 @@ debugger(#state{stream_module = StreamModule, stream = Stream0} = State) ->
 %% @return The list of used registers
 %%-----------------------------------------------------------------------------
 -spec used_regs(state()) -> [aarch64_register()].
-used_regs(#state{used_regs = Used}) -> Used.
+used_regs(#state{used_regs = Used}) -> lists:sort(jit_regs:to_list(Used)).
 
 %%-----------------------------------------------------------------------------
 %% @doc Return the list of currently available native scratch registers. This
@@ -343,7 +343,7 @@ free_native_register(State, _Other) ->
 %%-----------------------------------------------------------------------------
 -spec assert_all_native_free(state()) -> ok.
 assert_all_native_free(#state{
-    available_regs = ?AVAILABLE_REGS, used_regs = []
+    available_regs = ?AVAILABLE_REGS, used_regs = #{}
 }) ->
     ok.
 
@@ -508,7 +508,7 @@ call_primitive_last(
     ArgsRegs = args_regs(Args),
     ScratchRegs = ?AVAILABLE_REGS -- ArgsRegs -- ParamRegs,
     [Temp | AvailableRegs1] = ScratchRegs,
-    UsedRegs = ?AVAILABLE_REGS -- AvailableRegs1,
+    UsedRegs = jit_regs:from_list(?AVAILABLE_REGS -- AvailableRegs1),
     PrepCall =
         case Primitive of
             0 ->
@@ -526,7 +526,7 @@ call_primitive_last(
     #state{stream = Stream2} = State1,
     Call = jit_aarch64_asm:br(Temp),
     Stream3 = StreamModule:append(Stream2, Call),
-    State1#state{stream = Stream3, available_regs = ?AVAILABLE_REGS, used_regs = []}.
+    State1#state{stream = Stream3, available_regs = ?AVAILABLE_REGS, used_regs = jit_regs:new()}.
 
 %%-----------------------------------------------------------------------------
 %% @doc Emit a return of a value if it's not equal to ctx.
@@ -634,7 +634,7 @@ jump_to_continuation(
     Code = <<I1/binary, I2/binary, I3/binary>>,
     Stream1 = StreamModule:append(Stream0, Code),
     % Free all registers since this is a tail jump
-    State#state{stream = Stream1, available_regs = ?AVAILABLE_REGS, used_regs = []}.
+    State#state{stream = Stream1, available_regs = ?AVAILABLE_REGS, used_regs = jit_regs:new()}.
 
 %% @private
 -spec rewrite_branch_instruction(
@@ -1064,22 +1064,24 @@ if_block_free_reg(Reg, State0) when ?IS_GPR(Reg) ->
     State0.
 
 %% @private
--spec merge_used_regs(state(), [aarch64_register()]) -> state().
-merge_used_regs(#state{used_regs = UR0, available_regs = AvR0} = State, [
-    Reg | T
-]) ->
-    case lists:member(Reg, UR0) of
-        true ->
-            merge_used_regs(State, T);
-        false ->
-            AvR1 = lists:delete(Reg, AvR0),
-            UR1 = [Reg | UR0],
-            merge_used_regs(
-                State#state{used_regs = UR1, available_regs = AvR1}, T
-            )
-    end;
-merge_used_regs(State, []) ->
-    State.
+-spec merge_used_regs(state(), jit_regs:regs()) -> state().
+merge_used_regs(State, OtherRegs) ->
+    maps:fold(
+        fun(Reg, _, AccState) ->
+            #state{used_regs = UR, available_regs = AvR} = AccState,
+            case jit_regs:is_used(UR, Reg) of
+                true ->
+                    AccState;
+                false ->
+                    AccState#state{
+                        used_regs = jit_regs:add(UR, Reg),
+                        available_regs = lists:delete(Reg, AvR)
+                    }
+            end
+        end,
+        State,
+        OtherRegs
+    ).
 
 %%-----------------------------------------------------------------------------
 %% @doc Emit a shift register right by a fixed number of bits, effectively
@@ -1111,7 +1113,10 @@ shift_right(
 ->
     I = jit_aarch64_asm:lsr(ResultReg, Reg, Shift),
     Stream1 = StreamModule:append(Stream0, I),
-    {State#state{stream = Stream1, available_regs = T, used_regs = [ResultReg | UR]}, ResultReg}.
+    {
+        State#state{stream = Stream1, available_regs = T, used_regs = jit_regs:add(UR, ResultReg)},
+        ResultReg
+    }.
 
 %%-----------------------------------------------------------------------------
 %% @doc Emit a shift register left by a fixed number of bits, effectively
@@ -1159,8 +1164,9 @@ call_func_ptr(
         end,
         [FuncPtrTuple | Args]
     ),
-    UsedRegs1 = UsedRegs0 -- FreeRegs,
-    SavedRegs = [?LR_REG, ?CTX_REG, ?JITSTATE_REG, ?NATIVE_INTERFACE_REG | UsedRegs1],
+    UsedRegs1 = jit_regs:without(UsedRegs0, FreeRegs),
+    UsedRegsList = [R || R <- lists:reverse(?AVAILABLE_REGS), jit_regs:is_used(UsedRegs1, R)],
+    SavedRegs = [?LR_REG, ?CTX_REG, ?JITSTATE_REG, ?NATIVE_INTERFACE_REG | UsedRegsList],
     {SavedRegsOdd, Stream1} = push_registers(SavedRegs, StreamModule, Stream0),
 
     % Set up arguments following AArch64 calling convention
@@ -1201,9 +1207,8 @@ call_func_ptr(
 
     Stream6 = pop_registers(SavedRegsOdd, lists:reverse(SavedRegs), StreamModule, Stream5),
 
-    AvailableRegs2 = lists:delete(ResultReg, AvailableRegs1),
-    AvailableRegs3 = ?AVAILABLE_REGS -- (?AVAILABLE_REGS -- AvailableRegs2),
-    UsedRegs2 = [ResultReg | UsedRegs1],
+    UsedRegs2 = jit_regs:add(UsedRegs1, ResultReg),
+    AvailableRegs3 = jit_regs:available(?AVAILABLE_REGS, UsedRegs2),
     {
         State1#state{
             stream = Stream6,
@@ -1244,7 +1249,13 @@ set_args(
     ParamRegs = parameter_regs(Args),
     ArgsRegs = args_regs(Args),
     AvailableScratchGP =
-        ?SCRATCH_REGS -- ParamRegs -- ArgsRegs -- UsedRegs,
+        [
+            R
+         || R <- ?SCRATCH_REGS,
+            not lists:member(R, ParamRegs),
+            not lists:member(R, ArgsRegs),
+            not jit_regs:is_used(UsedRegs, R)
+        ],
     Offset = StreamModule:offset(Stream0),
     Args1 = [
         case Arg of
@@ -1257,17 +1268,22 @@ set_args(
     Stream1 = StreamModule:append(Stream0, SetArgsCode),
     NewUsedRegs = lists:foldl(
         fun
-            ({free, {ptr, Reg}}, AccUsed) -> lists:delete(Reg, AccUsed);
-            ({free, Reg}, AccUsed) -> lists:delete(Reg, AccUsed);
+            ({free, {ptr, Reg}}, AccUsed) -> jit_regs:remove(AccUsed, Reg);
+            ({free, Reg}, AccUsed) -> jit_regs:remove(AccUsed, Reg);
             (_, AccUsed) -> AccUsed
         end,
         UsedRegs,
         Args
     ),
     State0#state{
         stream = Stream1,
-        available_regs = ?AVAILABLE_REGS -- ParamRegs -- NewUsedRegs,
-        used_regs = ParamRegs ++ (NewUsedRegs -- ParamRegs)
+        available_regs = [
+            R
+         || R <- ?AVAILABLE_REGS,
+            not lists:member(R, ParamRegs),
+            not jit_regs:is_used(NewUsedRegs, R)
+        ],
+        used_regs = jit_regs:merge(NewUsedRegs, ParamRegs)
     }.
 
 %% @private
@@ -1628,7 +1644,9 @@ get_array_element(
     Stream1 = StreamModule:append(Stream0, <<I1/binary>>),
     {
         State#state{
-            stream = Stream1, available_regs = AvailableT, used_regs = [ElemReg | UsedRegs0]
+            stream = Stream1,
+            available_regs = AvailableT,
+            used_regs = jit_regs:add(UsedRegs0, ElemReg)
         },
         ElemReg
     }.
@@ -1743,7 +1761,10 @@ move_to_native_register(
 ) ->
     I1 = jit_aarch64_asm:ldr(Reg, ?CP),
     Stream1 = StreamModule:append(Stream0, I1),
-    {State#state{stream = Stream1, used_regs = [Reg | Used], available_regs = AvailT}, Reg};
+    {
+        State#state{stream = Stream1, used_regs = jit_regs:add(Used, Reg), available_regs = AvailT},
+        Reg
+    };
 move_to_native_register(State, Reg) when is_atom(Reg) ->
     {State, Reg};
 move_to_native_register(
@@ -1765,7 +1786,10 @@ move_to_native_register(
 ->
     I1 = jit_aarch64_asm:mov(Reg, Imm),
     Stream1 = StreamModule:append(Stream0, I1),
-    {State#state{stream = Stream1, used_regs = [Reg | Used], available_regs = AvailT}, Reg};
+    {
+        State#state{stream = Stream1, used_regs = jit_regs:add(Used, Reg), available_regs = AvailT},
+        Reg
+    };
 move_to_native_register(
     #state{
         stream_module = StreamModule,
@@ -1777,7 +1801,10 @@ move_to_native_register(
 ) ->
     I1 = jit_aarch64_asm:ldr(Reg, ?X_REG(?MAX_REG)),
     Stream1 = StreamModule:append(Stream0, I1),
-    {State#state{stream = Stream1, used_regs = [Reg | Used], available_regs = AvailT}, Reg};
+    {
+        State#state{stream = Stream1, used_regs = jit_regs:add(Used, Reg), available_regs = AvailT},
+        Reg
+    };
 move_to_native_register(
     #state{
         stream_module = StreamModule,
@@ -1791,7 +1818,10 @@ move_to_native_register(
 ->
     I1 = jit_aarch64_asm:ldr(Reg, ?X_REG(X)),
     Stream1 = StreamModule:append(Stream0, I1),
-    {State#state{stream = Stream1, used_regs = [Reg | Used], available_regs = AvailT}, Reg};
+    {
+        State#state{stream = Stream1, used_regs = jit_regs:add(Used, Reg), available_regs = AvailT},
+        Reg
+    };
 move_to_native_register(
     #state{
         stream_module = StreamModule,
@@ -1805,7 +1835,10 @@ move_to_native_register(
     I2 = jit_aarch64_asm:ldr(Reg, {Reg, Y * ?WORD_SIZE}),
     Code = <<I1/binary, I2/binary>>,
     Stream1 = StreamModule:append(Stream0, Code),
-    {State#state{stream = Stream1, available_regs = AvailT, used_regs = [Reg | Used]}, Reg}.
+    {
+        State#state{stream = Stream1, available_regs = AvailT, used_regs = jit_regs:add(Used, Reg)},
+        Reg
+    }.
 
 %%-----------------------------------------------------------------------------
 %% @doc Move a value (integer, vm register, pointer or native register) to a
@@ -1873,7 +1906,12 @@ copy_to_native_register(
 ) when is_atom(Reg) ->
     I1 = jit_aarch64_asm:mov(SaveReg, Reg),
     Stream1 = StreamModule:append(Stream0, I1),
-    {State#state{stream = Stream1, available_regs = AvailT, used_regs = [SaveReg | Used]}, SaveReg};
+    {
+        State#state{
+            stream = Stream1, available_regs = AvailT, used_regs = jit_regs:add(Used, SaveReg)
+        },
+        SaveReg
+    };
 copy_to_native_register(
     #state{
         stream_module = StreamModule,
@@ -1885,7 +1923,12 @@ copy_to_native_register(
 ) when is_atom(Reg) ->
     I1 = jit_aarch64_asm:ldr(SaveReg, {Reg, 0}),
     Stream1 = StreamModule:append(Stream0, I1),
-    {State#state{stream = Stream1, available_regs = AvailT, used_regs = [SaveReg | Used]}, SaveReg};
+    {
+        State#state{
+            stream = Stream1, available_regs = AvailT, used_regs = jit_regs:add(Used, SaveReg)
+        },
+        SaveReg
+    };
 copy_to_native_register(State, Reg) ->
     move_to_native_register(State, Reg).
 
@@ -2024,7 +2067,9 @@ get_module_index(
     Code = <<I1/binary, I2/binary>>,
     Stream1 = StreamModule:append(Stream0, Code),
     {
-        State#state{stream = Stream1, available_regs = AvailableT, used_regs = [Reg | UsedRegs0]},
+        State#state{
+            stream = Stream1, available_regs = AvailableT, used_regs = jit_regs:add(UsedRegs0, Reg)
+        },
         Reg
     }.
 
@@ -2073,7 +2118,11 @@ and_(
     Val
 ) ->
     NewState = op_imm(
-        State#state{available_regs = T, used_regs = [ResultReg | UR]}, and_, ResultReg, Reg, Val
+        State#state{available_regs = T, used_regs = jit_regs:add(UR, ResultReg)},
+        and_,
+        ResultReg,
+        Reg,
+        Val
     ),
     {NewState, ResultReg}.
 
@@ -2395,22 +2444,12 @@ return_labels_and_lines(
 %% @private
 -spec free_reg([aarch64_register()], [aarch64_register()], aarch64_register()) ->
     {[aarch64_register()], [aarch64_register()]}.
-free_reg(AvailableRegs0, UsedRegs0, Reg) when ?IS_GPR(Reg) ->
-    AvailableRegs1 = free_reg0(?AVAILABLE_REGS, AvailableRegs0, Reg, []),
-    true = lists:member(Reg, UsedRegs0),
-    UsedRegs1 = lists:delete(Reg, UsedRegs0),
+free_reg(_AvailableRegs0, UsedRegs0, Reg) when ?IS_GPR(Reg) ->
+    true = jit_regs:is_used(UsedRegs0, Reg),
+    UsedRegs1 = jit_regs:remove(UsedRegs0, Reg),
+    AvailableRegs1 = jit_regs:available(?AVAILABLE_REGS, UsedRegs1),
     {AvailableRegs1, UsedRegs1}.
 
-%% @private
--spec free_reg0([aarch64_register()], [aarch64_register()], aarch64_register(), [aarch64_register()]) ->
-    [aarch64_register()].
-free_reg0([Reg | _SortedT], PrevRegs0, Reg, Acc) ->
-    lists:reverse(Acc, [Reg | PrevRegs0]);
-free_reg0([PrevReg | SortedT], [PrevReg | PrevT], Reg, Acc) ->
-    free_reg0(SortedT, PrevT, Reg, [PrevReg | Acc]);
-free_reg0([_Other | SortedT], PrevRegs, Reg, Acc) ->
-    free_reg0(SortedT, PrevRegs, Reg, Acc).
-
 %% @private
 -spec args_regs([arg()]) -> [aarch64_register() | imm].
 args_regs(Args) ->