Extend cp for more than 255 modules

CHANGELOG.md

@@ -109,6 +109,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed `term_is_uint32` accepting big integers whose low 64 bits are within range on 32-bit
   builds, which made `erlang:crc32/2`, `erlang:crc32_combine/3` and `crypto:pbkdf2_hmac/5`
   silently truncate huge integer arguments instead of raising `badarg`
+- Fixed a bug where AtomVM could only have 256 loaded modules
 
 ## [0.7.0-alpha.1] - 2026-04-06
 

libs/jit/src/jit.erl

@@ -272,7 +272,8 @@ first_pass(<<?OP_CALL_LAST, Rest0/binary>>, MMod, MSt0, #state{tail_cache = TC}
         false ->
             Offset0 = MMod:offset(MSt0),
             MSt1 = MMod:move_to_cp(MSt0, {y_reg, NWords}),
-            MSt2 = MMod:increment_sp(MSt1, NWords + 1),
+            % The saved cp occupies one stack slot on 64-bit, two on 32-bit.
+            MSt2 = MMod:increment_sp(MSt1, NWords + (8 div MMod:word_size())),
             TailCacheKey1 = {op_call_only, Label},
             case tail_cache_find(TailCacheKey1, TC) of
                 false ->
@@ -443,25 +444,48 @@ first_pass(<<?OP_DEALLOCATE, Rest0/binary>>, MMod, MSt0, State0) ->
 first_pass(<<?OP_RETURN, Rest/binary>>, MMod, MSt0, #state{tail_cache = TC} = State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
     ?TRACE("OP_RETURN\n", []),
-    % Optimized return: check if returning within same module
-    {MSt1, CpReg0} = MMod:move_to_native_register(MSt0, cp),
-    {MSt2, ModuleIndexReg} = MMod:get_module_index(MSt1),
-    % Extract module index from cp (upper 8 bits: cp >> 24)
-    {MSt3, CpReg1} = MMod:shift_right(MSt2, CpReg0, 24),
-    % Compare extracted module index with current module index
-    MSt4 = MMod:if_block(
-        MSt3,
-        {{free, CpReg1}, '==', {free, ModuleIndexReg}},
-        % Same module: fast intra-module return
-        fun(BSt0) ->
-            % Mask to get lower 24 bits and shift right by 2 for offset
-            {BSt1, CpReg0} = MMod:and_(BSt0, {free, CpReg0}, 16#FFFFFF),
-            {BSt3, CPReg1} = MMod:shift_right(BSt1, {free, CpReg0}, 2),
-            % Jump to continuation (this is a tail call)
-            MMod:jump_to_continuation(BSt3, {free, CPReg1})
-        end
-    ),
-    MSt5 = MMod:free_native_registers(MSt4, [CpReg0]),
+    % Optimized return: check if returning within the same module, in which case
+    % we jump directly to the continuation rather than going through PRIM_RETURN.
+    MSt5 =
+        case MMod:word_size() of
+            8 ->
+                % 64-bit: cp packs (module_index << 24) | (offset << 2) in one word.
+                {MSt1, CpReg0} = MMod:move_to_native_register(MSt0, cp),
+                {MSt2, ModuleIndexReg} = MMod:get_module_index(MSt1),
+                % Extract module index from cp (upper 8 bits: cp >> 24)
+                {MSt3, CpReg1} = MMod:shift_right(MSt2, CpReg0, 24),
+                % Compare extracted module index with current module index
+                MSt4 = MMod:if_block(
+                    MSt3,
+                    {{free, CpReg1}, '==', {free, ModuleIndexReg}},
+                    % Same module: fast intra-module return
+                    fun(BSt0) ->
+                        % Mask to get lower 24 bits and shift right by 2 for offset
+                        {BSt1, CpReg0} = MMod:and_(BSt0, {free, CpReg0}, 16#FFFFFF),
+                        {BSt3, CPReg1} = MMod:shift_right(BSt1, {free, CpReg0}, 2),
+                        % Jump to continuation (this is a tail call)
+                        MMod:jump_to_continuation(BSt3, {free, CPReg1})
+                    end
+                ),
+                MMod:free_native_registers(MSt4, [CpReg0]);
+            4 ->
+                % 32-bit: cp spans two words, the Module pointer (?CP_MODULE) and
+                % the offset << 2 (?CP_OFFSET). Compare the saved Module pointer
+                % with the current module pointer (jit_state->module).
+                {MSt1, CpModReg} = MMod:get_cp_module(MSt0),
+                {MSt2, CurModReg} = MMod:get_module(MSt1),
+                MMod:if_block(
+                    MSt2,
+                    {{free, CpModReg}, '==', {free, CurModReg}},
+                    % Same module: fast intra-module return
+                    fun(BSt0) ->
+                        {BSt1, OffReg} = MMod:get_cp_offset(BSt0),
+                        {BSt2, OffReg2} = MMod:shift_right(BSt1, {free, OffReg}, 2),
+                        % Jump to continuation (this is a tail call)
+                        MMod:jump_to_continuation(BSt2, {free, OffReg2})
+                    end
+                )
+        end,
     % Different module: use existing slow path
     TailCacheKey = {call_primitive_last, ?PRIM_RETURN},
     case tail_cache_find(TailCacheKey, TC) of
@@ -1241,7 +1265,8 @@ first_pass(<<?OP_APPLY_LAST, Rest0/binary>>, MMod, MSt0, State0) ->
     MSt4 = verify_is_atom(Module, 0, MMod, MSt3),
     MSt5 = verify_is_atom(Function, 0, MMod, MSt4),
     MSt6 = MMod:move_to_cp(MSt5, {y_reg, NWords}),
-    MSt7 = MMod:increment_sp(MSt6, NWords + 1),
+    % The saved cp occupies one stack slot on 64-bit, two on 32-bit.
+    MSt7 = MMod:increment_sp(MSt6, NWords + (8 div MMod:word_size())),
     MSt8 = MMod:call_primitive_last(MSt7, ?PRIM_APPLY, [
         ctx, jit_state, offset, {free, Module}, {free, Function}, Arity
     ]),
@@ -4185,10 +4210,11 @@ term_alloc_bin_match_state(Live, Src, Dest, MMod, MSt0) ->
     MSt6 = MMod:free_native_registers(MSt5, [AllocMatchStateReg]),
     {MSt6, NewSrc}.
 
-term_from_catch_label(Dest, Label, MMod, MSt1) ->
-    {MSt2, Reg} = MMod:get_module_index(MSt1),
-    MSt3 = MMod:shift_left(MSt2, Reg, 24),
-    MSt4 = MMod:or_(MSt3, Reg, (Label bsl ?TERM_IMMED2_TAG_SIZE) bor ?TERM_IMMED2_CATCH),
+term_from_catch_label(Dest, Label, MMod, MSt0) ->
+    {MSt1, Reg} = MMod:get_module_catch_labels_base(MSt0),
+    MSt2 = MMod:add(MSt1, Reg, Label),
+    MSt3 = MMod:shift_left(MSt2, Reg, ?TERM_IMMED2_TAG_SIZE),
+    MSt4 = MMod:or_(MSt3, Reg, ?TERM_IMMED2_CATCH),
     MSt5 = MMod:move_to_vm_register(MSt4, Reg, Dest),
     MMod:free_native_registers(MSt5, [Reg, Dest]).
 

libs/jit/src/jit_aarch64.erl

@@ -59,6 +59,7 @@
     set_continuation_to_offset/1,
     continuation_entry_point/1,
     get_module_index/1,
+    get_module_catch_labels_base/1,
     and_/3,
     or_/3,
     add/3,
@@ -208,6 +209,7 @@
 -define(JITSTATE_REDUCTIONCOUNT, {?JITSTATE_REG, 16#10}).
 -define(PRIMITIVE(N), {?NATIVE_INTERFACE_REG, N * ?WORD_SIZE}).
 -define(MODULE_INDEX(ModuleReg), {ModuleReg, 0}).
+-define(MODULE_CATCH_LABELS_BASE(ModuleReg), {ModuleReg, 4}).
 
 % aarch64 ABI specific
 -define(LR_REG, r30).
@@ -2376,6 +2378,37 @@ get_module_index(
         Reg
     }.
 
+%%-----------------------------------------------------------------------------
+%% @doc Load the catch id of the current module's label 0 into a native
+%% register.
+%% @end
+%% @param State current backend state
+%% @return Tuple of {Updated backend state, Native register containing the base}
+%%-----------------------------------------------------------------------------
+-spec get_module_catch_labels_base(state()) -> {state(), aarch64_register()}.
+get_module_catch_labels_base(
+    #state{
+        stream_module = StreamModule,
+        stream = Stream0,
+        regs = Regs0
+    } = State
+) ->
+    Avail = jit_regs:available_regs(Regs0),
+    Reg = first_avail(Avail),
+    Bit = reg_bit(Reg),
+    I1 = jit_aarch64_asm:ldr(Reg, ?JITSTATE_MODULE),
+    I2 = jit_aarch64_asm:ldr_w(Reg, ?MODULE_CATCH_LABELS_BASE(Reg)),
+    Code = <<I1/binary, I2/binary>>,
+    Stream1 = StreamModule:append(Stream0, Code),
+    Regs1 = jit_regs:set_contents(Regs0, Reg, catch_labels_base),
+    {
+        State#state{
+            stream = Stream1,
+            regs = jit_regs:alloc_reg(Regs1, Bit)
+        },
+        Reg
+    }.
+
 %% @private
 -spec op_imm(state(), atom(), aarch64_register(), aarch64_register(), integer()) -> state().
 op_imm(#state{stream_module = StreamModule, stream = Stream0} = State, Op, Reg, Reg, Val) ->

libs/jit/src/jit_arm32.erl

@@ -60,6 +60,10 @@
     set_continuation_to_offset/1,
     continuation_entry_point/1,
     get_module_index/1,
+    get_module_catch_labels_base/1,
+    get_module/1,
+    get_cp_module/1,
+    get_cp_offset/1,
     and_/3,
     or_/3,
     add/3,
@@ -179,17 +183,21 @@
 -define(NATIVE_INTERFACE_REG, r2).
 -define(Y_REGS, {?CTX_REG, 16#14}).
 -define(X_REG(N), {?CTX_REG, 16#18 + (N * 4)}).
+% ctx->cp is a 64-bit cp_t occupying two slots (little-endian targets):
+% ?CP holds the low word (offset << 2), ?CP_MODULE holds the high word (Module*).
 -define(CP, {?CTX_REG, 16#5C}).
--define(FP_REGS, {?CTX_REG, 16#60}).
--define(BS, {?CTX_REG, 16#64}).
--define(BS_OFFSET, {?CTX_REG, 16#68}).
+-define(CP_MODULE, {?CTX_REG, 16#60}).
+-define(FP_REGS, {?CTX_REG, 16#64}).
+-define(BS, {?CTX_REG, 16#68}).
+-define(BS_OFFSET, {?CTX_REG, 16#6C}).
 % JITSTATE is on stack, accessed via stack offset
 % These macros now expect a register that contains the jit_state pointer
 -define(JITSTATE_MODULE(Reg), {Reg, 0}).
 -define(JITSTATE_CONTINUATION(Reg), {Reg, 16#4}).
 -define(JITSTATE_REDUCTIONCOUNT(Reg), {Reg, 16#8}).
 -define(PRIMITIVE(N), {?NATIVE_INTERFACE_REG, N * 4}).
 -define(MODULE_INDEX(ModuleReg), {ModuleReg, 0}).
+-define(MODULE_CATCH_LABELS_BASE(ModuleReg), {ModuleReg, 4}).
 
 -define(JUMP_TABLE_ENTRY_SIZE, 8).
 
@@ -2835,10 +2843,14 @@ move_to_cp(
     Avail = jit_regs:available_regs(Regs0),
     Reg = first_avail(Avail),
     AvailT = Avail band (bnot reg_bit(Reg)),
+    % The saved cp spans two slots: y[Y] = offset word (-> ?CP), y[Y+1] = Module*
+    % (-> ?CP_MODULE). Copy both into ctx->cp.
     State1 = ldr_y_reg(State, Reg, Y, AvailT),
-    I2 = jit_arm32_asm:str(al, Reg, ?CP),
-    Stream1 = (State1#state.stream_module):append(State1#state.stream, I2),
-    State1#state{stream = Stream1}.
+    SM = State1#state.stream_module,
+    Stream1 = SM:append(State1#state.stream, jit_arm32_asm:str(al, Reg, ?CP)),
+    State2 = ldr_y_reg(State1#state{stream = Stream1}, Reg, Y + 1, AvailT),
+    Stream2 = SM:append(State2#state.stream, jit_arm32_asm:str(al, Reg, ?CP_MODULE)),
+    State2#state{stream = Stream2}.
 
 increment_sp(
     #state{stream_module = StreamModule, stream = Stream0, regs = Regs0} =
@@ -2967,6 +2979,78 @@ get_module_index(
         Reg
     }.
 
+%% @doc Load the catch id of the current module's label 0 into a native register.
+get_module_catch_labels_base(
+    #state{
+        stream_module = StreamModule,
+        stream = Stream0,
+        regs = Regs0
+    } = State
+) ->
+    Avail = jit_regs:available_regs(Regs0),
+    Reg = first_avail(Avail),
+    RegBit = reg_bit(Reg),
+    Avail1 = Avail band (bnot RegBit),
+    TempJitState = first_avail(Avail1),
+    % Load jit_state pointer from stack, then load module
+    I1a = jit_arm32_asm:ldr(al, TempJitState, {sp, ?STACK_OFFSET_JITSTATE}),
+    I1b = jit_arm32_asm:ldr(al, Reg, ?JITSTATE_MODULE(TempJitState)),
+    I2 = jit_arm32_asm:ldr(al, Reg, ?MODULE_CATCH_LABELS_BASE(Reg)),
+    Code = <<I1a/binary, I1b/binary, I2/binary>>,
+    Stream1 = StreamModule:append(Stream0, Code),
+    Regs1 = jit_regs:invalidate_reg(Regs0, TempJitState),
+    Regs2 = jit_regs:set_contents(Regs1, Reg, catch_labels_base),
+    Regs3 = jit_regs:alloc_reg(Regs2, RegBit),
+    {
+        State#state{
+            stream = Stream1,
+            regs = Regs3
+        },
+        Reg
+    }.
+
+%% @doc Load the current module pointer (jit_state->module) into a register.
+get_module(
+    #state{stream_module = StreamModule, stream = Stream0, regs = Regs0} = State
+) ->
+    Avail = jit_regs:available_regs(Regs0),
+    Reg = first_avail(Avail),
+    RegBit = reg_bit(Reg),
+    Avail1 = Avail band (bnot RegBit),
+    TempJitState = first_avail(Avail1),
+    I1a = jit_arm32_asm:ldr(al, TempJitState, {sp, ?STACK_OFFSET_JITSTATE}),
+    I1b = jit_arm32_asm:ldr(al, Reg, ?JITSTATE_MODULE(TempJitState)),
+    Code = <<I1a/binary, I1b/binary>>,
+    Stream1 = StreamModule:append(Stream0, Code),
+    Regs1 = jit_regs:invalidate_reg(jit_regs:invalidate_reg(Regs0, TempJitState), Reg),
+    Regs2 = jit_regs:alloc_reg(Regs1, RegBit),
+    {State#state{stream = Stream1, regs = Regs2}, Reg}.
+
+%% @doc Load the Module pointer stored in ctx->cp (?CP_MODULE) into a register.
+get_cp_module(
+    #state{stream_module = StreamModule, stream = Stream0, regs = Regs0} = State
+) ->
+    Avail = jit_regs:available_regs(Regs0),
+    Reg = first_avail(Avail),
+    RegBit = reg_bit(Reg),
+    I = jit_arm32_asm:ldr(al, Reg, ?CP_MODULE),
+    Stream1 = StreamModule:append(Stream0, I),
+    Regs1 = jit_regs:invalidate_reg(Regs0, Reg),
+    Regs2 = jit_regs:alloc_reg(Regs1, RegBit),
+    {State#state{stream = Stream1, regs = Regs2}, Reg}.
+
+%% @doc Load the offset word (offset << 2) stored in ctx->cp (?CP) into a register.
+get_cp_offset(
+    #state{stream_module = StreamModule, stream = Stream0, regs = Regs0} = State
+) ->
+    Avail = jit_regs:available_regs(Regs0),
+    Reg = first_avail(Avail),
+    RegBit = reg_bit(Reg),
+    I = jit_arm32_asm:ldr(al, Reg, ?CP),
+    Stream1 = StreamModule:append(Stream0, I),
+    Regs1 = jit_regs:invalidate_reg(Regs0, Reg),
+    Regs2 = jit_regs:alloc_reg(Regs1, RegBit),
+    {State#state{stream = Stream1, regs = Regs2}, Reg}.
 %% @doc Perform an AND of a register with an immediate.
 %% JIT currently calls this with two values: ?TERM_PRIMARY_CLEAR_MASK (-4) to
 %% clear bits and ?TERM_BOXED_TAG_MASK (0x3F). We can avoid any literal pool
@@ -3518,30 +3602,24 @@ call_primitive_with_cp(State0, Primitive, Args) ->
 
 -spec set_cp(state()) -> {state(), non_neg_integer(), arm32_register()}.
 set_cp(State0) ->
-    % get module index (dynamically)
-    {
-        #state{stream_module = StreamModule, stream = Stream0, regs = Regs1} = State1,
-        Reg
-    } = get_module_index(
-        State0
-    ),
-    AvailRegs = jit_regs:available_regs(Regs1),
-    % Get a temporary register from available registers
+    % cp is two words: store the Module pointer (jit_state->module) at ?CP_MODULE,
+    % and the return offset << 2 at ?CP (patched by rewrite_cp_offset below).
+    {#state{stream_module = StreamModule, stream = Stream0} = State1, ModReg} =
+        get_module(State0),
+    IModStore = jit_arm32_asm:str(al, ModReg, ?CP_MODULE),
+    Stream1 = StreamModule:append(Stream0, IModStore),
+    State2 = free_native_register(State1#state{stream = Stream1}, ModReg),
+    AvailRegs = jit_regs:available_regs(State2#state.regs),
+    % Get a temporary register to hold the offset value
     TempReg = first_avail(AvailRegs),
-
-    Offset = StreamModule:offset(Stream0),
-    % build cp with module_index << 24
-    I1 = jit_arm32_asm:lsl(al, Reg, Reg, 24),
-    % Placeholder for offset load instruction
-    I2 = <<16#FFFFFFFF:32>>,
-    MOVOffset = Offset + byte_size(I1),
-    % OR the module index with the offset (loaded in temp register)
-    I3 = jit_arm32_asm:orr(al, Reg, Reg, TempReg),
-    I4 = jit_arm32_asm:str(al, Reg, ?CP),
-    Code = <<I1/binary, I2/binary, I3/binary, I4/binary>>,
-    Stream1 = StreamModule:append(Stream0, Code),
-    State2 = State1#state{stream = Stream1},
-    State3 = free_native_register(State2, Reg),
+    Offset = StreamModule:offset(Stream1),
+    % Placeholder for the offset load instruction (patched by rewrite_cp_offset)
+    I1 = <<16#FFFFFFFF:32>>,
+    MOVOffset = Offset,
+    I2 = jit_arm32_asm:str(al, TempReg, ?CP),
+    Code = <<I1/binary, I2/binary>>,
+    Stream2 = StreamModule:append(Stream1, Code),
+    State3 = State2#state{stream = Stream2},
     {State3, MOVOffset, TempReg}.
 
 -spec rewrite_cp_offset(state(), non_neg_integer(), arm32_register()) -> state().