lust/jit/
trace.rs

1use crate::bytecode::value::NativeFn;
2use crate::bytecode::Instruction;
3use crate::bytecode::{Register, Value};
4use crate::LustError;
5use alloc::{
6    format,
7    rc::Rc,
8    string::{String, ToString},
9    vec::Vec,
10};
11use core::fmt;
12use hashbrown::{HashMap, HashSet};
13
14#[derive(Clone)]
15pub struct TracedNativeFn {
16    function: NativeFn,
17}
18
19impl TracedNativeFn {
20    pub fn new(function: NativeFn) -> Self {
21        Self { function }
22    }
23
24    pub fn pointer(&self) -> *const () {
25        Rc::as_ptr(&self.function) as *const ()
26    }
27}
28
29impl fmt::Debug for TracedNativeFn {
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        write!(f, "NativeFn({:p})", Rc::as_ptr(&self.function))
32    }
33}
34
35#[derive(Debug, Clone)]
36pub struct Trace {
37    pub function_idx: usize,
38    pub start_ip: usize,
39    /// Operations executed once at trace entry (unboxing, guards, etc.)
40    pub preamble: Vec<TraceOp>,
41    /// Operations in the trace loop body
42    pub ops: Vec<TraceOp>,
43    /// Operations executed once at trace exit (reboxing to restore state)
44    pub postamble: Vec<TraceOp>,
45    pub inputs: Vec<Register>,
46    pub outputs: Vec<Register>,
47}
48
49#[derive(Debug, Clone)]
50pub struct InlineTrace {
51    pub function_idx: usize,
52    pub register_count: u8,
53    pub first_arg: Register,
54    pub arg_count: u8,
55    pub arg_registers: Vec<Register>,
56    pub body: Vec<TraceOp>,
57    pub return_register: Option<Register>,
58    pub is_closure: bool,
59    pub upvalues_ptr: Option<*const ()>,
60}
61
62#[derive(Debug, Clone)]
63pub enum TraceOp {
64    LoadConst {
65        dest: Register,
66        value: Value,
67    },
68    Move {
69        dest: Register,
70        src: Register,
71    },
72    Add {
73        dest: Register,
74        lhs: Register,
75        rhs: Register,
76        lhs_type: ValueType,
77        rhs_type: ValueType,
78    },
79    Sub {
80        dest: Register,
81        lhs: Register,
82        rhs: Register,
83        lhs_type: ValueType,
84        rhs_type: ValueType,
85    },
86    Mul {
87        dest: Register,
88        lhs: Register,
89        rhs: Register,
90        lhs_type: ValueType,
91        rhs_type: ValueType,
92    },
93    Div {
94        dest: Register,
95        lhs: Register,
96        rhs: Register,
97        lhs_type: ValueType,
98        rhs_type: ValueType,
99    },
100    Mod {
101        dest: Register,
102        lhs: Register,
103        rhs: Register,
104        lhs_type: ValueType,
105        rhs_type: ValueType,
106    },
107    Neg {
108        dest: Register,
109        src: Register,
110    },
111    Eq {
112        dest: Register,
113        lhs: Register,
114        rhs: Register,
115    },
116    Ne {
117        dest: Register,
118        lhs: Register,
119        rhs: Register,
120    },
121    Lt {
122        dest: Register,
123        lhs: Register,
124        rhs: Register,
125    },
126    Le {
127        dest: Register,
128        lhs: Register,
129        rhs: Register,
130    },
131    Gt {
132        dest: Register,
133        lhs: Register,
134        rhs: Register,
135    },
136    Ge {
137        dest: Register,
138        lhs: Register,
139        rhs: Register,
140    },
141    And {
142        dest: Register,
143        lhs: Register,
144        rhs: Register,
145    },
146    Or {
147        dest: Register,
148        lhs: Register,
149        rhs: Register,
150    },
151    Not {
152        dest: Register,
153        src: Register,
154    },
155    Concat {
156        dest: Register,
157        lhs: Register,
158        rhs: Register,
159    },
160    GetIndex {
161        dest: Register,
162        array: Register,
163        index: Register,
164    },
165    ArrayLen {
166        dest: Register,
167        array: Register,
168    },
169    GuardNativeFunction {
170        register: Register,
171        function: TracedNativeFn,
172    },
173    GuardFunction {
174        register: Register,
175        function_idx: usize,
176    },
177    GuardClosure {
178        register: Register,
179        function_idx: usize,
180        upvalues_ptr: *const (),
181    },
182    CallNative {
183        dest: Register,
184        callee: Register,
185        function: TracedNativeFn,
186        first_arg: Register,
187        arg_count: u8,
188    },
189    CallFunction {
190        dest: Register,
191        callee: Register,
192        function_idx: usize,
193        first_arg: Register,
194        arg_count: u8,
195        is_closure: bool,
196        upvalues_ptr: Option<*const ()>,
197    },
198    InlineCall {
199        dest: Register,
200        callee: Register,
201        trace: InlineTrace,
202    },
203    CallMethod {
204        dest: Register,
205        object: Register,
206        method_name: String,
207        first_arg: Register,
208        arg_count: u8,
209    },
210    GetField {
211        dest: Register,
212        object: Register,
213        field_name: String,
214        field_index: Option<usize>,
215        value_type: Option<ValueType>,
216        is_weak: bool,
217    },
218    SetField {
219        object: Register,
220        field_name: String,
221        value: Register,
222        field_index: Option<usize>,
223        value_type: Option<ValueType>,
224        is_weak: bool,
225    },
226    NewArray {
227        dest: Register,
228        first_element: Register,
229        count: u8,
230    },
231    NewStruct {
232        dest: Register,
233        struct_name: String,
234        field_names: Vec<String>,
235        field_registers: Vec<Register>,
236    },
237    NewEnumUnit {
238        dest: Register,
239        enum_name: String,
240        variant_name: String,
241    },
242    NewEnumVariant {
243        dest: Register,
244        enum_name: String,
245        variant_name: String,
246        value_registers: Vec<Register>,
247    },
248    IsEnumVariant {
249        dest: Register,
250        value: Register,
251        enum_name: String,
252        variant_name: String,
253    },
254    GetEnumValue {
255        dest: Register,
256        enum_reg: Register,
257        index: u8,
258    },
259    Guard {
260        register: Register,
261        expected_type: ValueType,
262    },
263    GuardLoopContinue {
264        condition_register: Register,
265        expect_truthy: bool,
266        bailout_ip: usize,
267    },
268    NestedLoopCall {
269        function_idx: usize,
270        loop_start_ip: usize,
271        bailout_ip: usize,
272    },
273    Return {
274        value: Option<Register>,
275    },
276    /// Unbox a Value into specialized representation
277    Unbox {
278        specialized_id: usize,
279        source_reg: Register,
280        layout: crate::jit::specialization::SpecializedLayout,
281    },
282    /// Rebox a specialized value back to Value
283    Rebox {
284        dest_reg: Register,
285        specialized_id: usize,
286        layout: crate::jit::specialization::SpecializedLayout,
287    },
288    /// Drop a specialized value without reboxing (cleanup for leaked specializations)
289    DropSpecialized {
290        specialized_id: usize,
291        layout: crate::jit::specialization::SpecializedLayout,
292    },
293    /// Operation on specialized values
294    SpecializedOp {
295        op: SpecializedOpKind,
296        operands: Vec<Operand>,
297    },
298}
299
300/// Operand for specialized operations
301#[derive(Debug, Clone, PartialEq, Eq)]
302pub enum Operand {
303    Register(u8),
304    Specialized(usize),
305    Immediate(i64),
306}
307
308/// Types of operations on specialized values
309#[derive(Debug, Clone, PartialEq, Eq)]
310pub enum SpecializedOpKind {
311    // Vector operations
312    VecPush,
313    VecPop,
314    VecGet,
315    VecSet,
316    VecLen,
317
318    // Map operations
319    MapInsert,
320    MapGet,
321    MapRemove,
322
323    // Struct operations
324    StructGetField { field_index: usize },
325    StructSetField { field_index: usize },
326
327    // Arithmetic on unboxed values
328    Add,
329    Sub,
330    Mul,
331    Div,
332    Mod,
333    Neg,
334
335    // Comparison
336    Eq,
337    Ne,
338    Lt,
339    Le,
340    Gt,
341    Ge,
342}
343
344#[derive(Debug, Clone, Copy, PartialEq, Eq)]
345pub enum ValueType {
346    Int,
347    Float,
348    Bool,
349    String,
350    Array,
351    Tuple,
352    Struct,
353}
354
355pub struct TraceRecorder {
356    pub trace: Trace,
357    max_length: usize,
358    recording: bool,
359    guarded_registers: HashSet<Register>,
360    inline_stack: Vec<InlineContext>,
361    op_count: usize,
362    /// Track which registers contain specialized values (register -> (specialized_id, layout))
363    specialized_registers:
364        HashMap<Register, (usize, crate::jit::specialization::SpecializedLayout)>,
365    /// Counter for generating specialized IDs
366    next_specialized_id: usize,
367    /// Registry for type specializations
368    specialization_registry: crate::jit::specialization::SpecializationRegistry,
369    /// Track how many times we've seen each loop backedge to enable unrolling
370    loop_iterations: HashMap<(usize, usize), usize>,
371    /// Track specialized values that were unboxed but later invalidated (need cleanup/drop)
372    leaked_specialized_values: Vec<(usize, crate::jit::specialization::SpecializedLayout)>,
373}
374
375#[derive(Debug, Clone)]
376struct InlineContext {
377    function_idx: usize,
378    register_count: u8,
379    dest: Register,
380    callee_reg: Register,
381    first_arg: Register,
382    arg_count: u8,
383    arg_registers: Vec<Register>,
384    ops: Vec<TraceOp>,
385    guarded_registers: HashSet<Register>,
386    return_register: Option<Register>,
387    is_closure: bool,
388    upvalues_ptr: Option<*const ()>,
389}
390
391impl TraceRecorder {
392    pub fn new(function_idx: usize, start_ip: usize, max_length: usize) -> Self {
393        Self {
394            trace: Trace {
395                function_idx,
396                start_ip,
397                preamble: Vec::new(),
398                ops: Vec::new(),
399                postamble: Vec::new(),
400                inputs: Vec::new(),
401                outputs: Vec::new(),
402            },
403            max_length,
404            recording: true,
405            guarded_registers: HashSet::new(),
406            inline_stack: Vec::new(),
407            op_count: 0,
408            specialized_registers: HashMap::new(),
409            next_specialized_id: 0,
410            specialization_registry: crate::jit::specialization::SpecializationRegistry::new(),
411            loop_iterations: HashMap::new(),
412            leaked_specialized_values: Vec::new(),
413        }
414    }
415
416    /// Scan live registers at trace entry and specialize any loop-invariant arrays
417    /// This should be called right after trace recording starts
418    pub fn specialize_trace_inputs(
419        &mut self,
420        registers: &[Value; 256],
421        function: &crate::bytecode::Function,
422    ) {
423        crate::jit::log(|| format!("🔍 JIT: Scanning trace inputs for specialization..."));
424
425        // Scan all registers for arrays (not just ones with type info)
426        for reg in 0u8..=255 {
427            // Check if this register contains an Array at runtime
428            if let Value::Array(ref arr_rc) = registers[reg as usize] {
429                crate::jit::log(|| format!("🔍 JIT: Found array in reg {}", reg));
430                // Try to determine the element type by inspecting the array contents
431                // If we have type info, use it; otherwise infer from first element
432                let element_type = function.register_types.get(&reg).cloned().or_else(|| {
433                    // Try to infer from array contents
434                    let arr = arr_rc.borrow();
435                    if arr.is_empty() {
436                        None
437                    } else {
438                        match &arr[0] {
439                            Value::Int(_) => Some(crate::ast::TypeKind::Int),
440                            Value::Float(_) => Some(crate::ast::TypeKind::Float),
441                            Value::Bool(_) => Some(crate::ast::TypeKind::Bool),
442                            _ => None,
443                        }
444                    }
445                });
446
447                if let Some(elem_type) = element_type {
448                    // Build the full Array<T> type
449                    // Check if elem_type is already a full Array type (stored in register_types)
450                    // vs just the element type
451                    use crate::ast::{Span, Type};
452                    let (array_type, _actual_elem_type) =
453                        if matches!(elem_type, crate::ast::TypeKind::Array(_)) {
454                            // Already a full array type - use it directly
455                            (elem_type.clone(), elem_type.clone())
456                        } else {
457                            // Element type only, wrap in Array
458                            let arr_type = crate::ast::TypeKind::Array(Box::new(Type::new(
459                                elem_type.clone(),
460                                Span::dummy(),
461                            )));
462                            (arr_type, elem_type.clone())
463                        };
464
465                    // Check if this array type is specializable
466                    if let Some(layout) =
467                        self.specialization_registry.get_specialization(&array_type)
468                    {
469                        crate::jit::log(|| {
470                            format!(
471                                "🔬 JIT: Specializing trace input reg {} ({:?})",
472                                reg, array_type
473                            )
474                        });
475
476                        // Emit Unbox in PREAMBLE (executes once at trace entry, not in loop)
477                        let specialized_id = self.next_specialized_id;
478                        self.next_specialized_id += 1;
479
480                        self.trace.preamble.push(TraceOp::Unbox {
481                            specialized_id,
482                            source_reg: reg,
483                            layout: layout.clone(),
484                        });
485
486                        // Track this specialized value
487                        self.specialized_registers
488                            .insert(reg, (specialized_id, layout));
489                    }
490                }
491            }
492        }
493    }
494
495    fn current_function_idx(&self) -> usize {
496        self.inline_stack
497            .last()
498            .map(|ctx| ctx.function_idx)
499            .unwrap_or(self.trace.function_idx)
500    }
501
502    fn current_guard_set(&self) -> &HashSet<Register> {
503        self.inline_stack
504            .last()
505            .map(|ctx| &ctx.guarded_registers)
506            .unwrap_or(&self.guarded_registers)
507    }
508
509    fn current_guard_set_mut(&mut self) -> &mut HashSet<Register> {
510        self.inline_stack
511            .last_mut()
512            .map(|ctx| &mut ctx.guarded_registers)
513            .unwrap_or(&mut self.guarded_registers)
514    }
515
516    fn is_guarded(&self, register: Register) -> bool {
517        self.current_guard_set().contains(&register)
518    }
519
520    fn mark_guarded(&mut self, register: Register) {
521        let set = self.current_guard_set_mut();
522        set.insert(register);
523    }
524
525    fn push_op(&mut self, op: TraceOp) {
526        self.op_count += 1;
527        if let Some(ctx) = self.inline_stack.last_mut() {
528            ctx.ops.push(op);
529        } else {
530            self.trace.ops.push(op);
531        }
532    }
533
534    /// Finalize the trace by adding postamble operations (rebox all specialized values)
535    fn finalize_trace(&mut self) {
536        crate::jit::log(|| {
537            format!(
538                "🏁 JIT: Finalizing trace - reboxing {} specialized values, dropping {} leaked values",
539                self.specialized_registers.len(),
540                self.leaked_specialized_values.len()
541            )
542        });
543
544        // NOTE: We do NOT emit drops for leaked_specialized_values!
545        // Those values were invalidated during trace RECORDING, so they never
546        // actually exist on the JIT stack during trace EXECUTION.
547        // The arrays are still managed by their Rc<RefCell<>> wrappers.
548
549        // Rebox all remaining specialized values in the postamble
550        for (&register, &(specialized_id, ref layout)) in self.specialized_registers.iter() {
551            crate::jit::log(|| {
552                format!(
553                    "📦 JIT: Adding rebox to postamble for specialized #{} in reg {}",
554                    specialized_id, register
555                )
556            });
557
558            self.trace.postamble.push(TraceOp::Rebox {
559                dest_reg: register,
560                specialized_id,
561                layout: layout.clone(),
562            });
563        }
564    }
565
566    /// Stop recording and finalize the trace
567    fn stop_recording(&mut self) {
568        crate::jit::log(|| {
569            format!(
570                "🛑 JIT: stop_recording called, recording={}, specialized_regs={}",
571                self.recording,
572                self.specialized_registers.len()
573            )
574        });
575        if self.recording {
576            self.finalize_trace();
577            self.recording = false;
578        }
579    }
580
581    /// Rebox all currently active specialized values
582    /// This must be called before any side exit to restore interpreter-compatible state
583    fn rebox_all_specialized_values(&mut self) {
584        // Collect all specialized values that need reboxing
585        let to_rebox: Vec<(
586            Register,
587            usize,
588            crate::jit::specialization::SpecializedLayout,
589        )> = self
590            .specialized_registers
591            .iter()
592            .map(|(&reg, &(id, ref layout))| (reg, id, layout.clone()))
593            .collect();
594
595        // Emit Rebox operations
596        for (register, specialized_id, layout) in to_rebox {
597            crate::jit::log(|| {
598                format!(
599                    "📦 JIT: Reboxing specialized #{} back to reg {} before side exit",
600                    specialized_id, register
601                )
602            });
603
604            self.push_op(TraceOp::Rebox {
605                dest_reg: register,
606                specialized_id,
607                layout,
608            });
609
610            // Remove from tracking
611            self.specialized_registers.remove(&register);
612        }
613    }
614
615    /// Invalidate specialization for a register that's about to be overwritten
616    /// The specialized Vec data needs to be dropped since it won't be reboxed
617    fn invalidate_specialization(&mut self, register: Register) {
618        if let Some((specialized_id, layout)) = self.specialized_registers.remove(&register) {
619            crate::jit::log(|| {
620                format!(
621                    "🚫 JIT: Invalidating specialization for reg {} (being overwritten) - will drop specialized #{}",
622                    register, specialized_id
623                )
624            });
625            // Track this for cleanup in postamble - the Vec data needs to be dropped
626            self.leaked_specialized_values
627                .push((specialized_id, layout));
628        }
629    }
630
631    /// Remove specialization tracking if register is about to be overwritten
632    /// The Vec data becomes "leaked" on the JIT stack but that's fine - it's cleaned
633    /// up when the stack frame is destroyed. The array is still managed by Rc<RefCell<>>.
634    fn remove_specialization_tracking(&mut self, register: Register) {
635        if let Some((_specialized_id, _layout)) = self.specialized_registers.remove(&register) {
636            crate::jit::log(|| {
637                format!(
638                    "🗑️  JIT: Removing specialization tracking for reg {} (being overwritten)",
639                    register
640                )
641            });
642            // Don't emit rebox - the Vec data stays on JIT stack but that's OK
643            // It will be cleaned up when the JIT stack frame is destroyed
644        }
645    }
646
647    fn rebox_specialized_register(&mut self, register: Register, context: &str) {
648        if let Some((specialized_id, layout)) = self.specialized_registers.remove(&register) {
649            crate::jit::log(|| {
650                format!(
651                    "📦 JIT: Reboxing specialized #{} from reg {} before {}",
652                    specialized_id, register, context
653                )
654            });
655            self.push_op(TraceOp::Rebox {
656                dest_reg: register,
657                specialized_id,
658                layout,
659            });
660        }
661    }
662
663    fn should_inline(&self, function_idx: usize, callee_fn: &crate::bytecode::Function) -> bool {
664        if function_idx == self.trace.function_idx {
665            return false;
666        }
667
668        if self
669            .inline_stack
670            .iter()
671            .any(|ctx| ctx.function_idx == function_idx)
672        {
673            return false;
674        }
675
676        // Disable inlining when specialized values are active to avoid
677        // stack layout conflicts between inline frames and specialized storage
678        if !self.specialized_registers.is_empty() {
679            return false;
680        }
681
682        if callee_fn.chunk.instructions.iter().any(|inst| {
683            matches!(
684                inst,
685                Instruction::Jump(_) | Instruction::JumpIf(..) | Instruction::JumpIfNot(..)
686            )
687        }) {
688            return false;
689        }
690
691        true
692    }
693
694    fn push_inline_context(
695        &mut self,
696        function_idx: usize,
697        register_count: u8,
698        dest: Register,
699        callee_reg: Register,
700        first_arg: Register,
701        arg_count: u8,
702        arg_registers: Vec<Register>,
703        is_closure: bool,
704        upvalues_ptr: Option<*const ()>,
705    ) {
706        self.inline_stack.push(InlineContext {
707            function_idx,
708            register_count,
709            dest,
710            callee_reg,
711            first_arg,
712            arg_count,
713            arg_registers,
714            ops: Vec::new(),
715            guarded_registers: HashSet::new(),
716            return_register: None,
717            is_closure,
718            upvalues_ptr,
719        });
720    }
721
722    fn finalize_inline_context(&mut self) -> Option<TraceOp> {
723        let context = self.inline_stack.pop()?;
724        let trace = InlineTrace {
725            function_idx: context.function_idx,
726            register_count: context.register_count,
727            first_arg: context.first_arg,
728            arg_count: context.arg_count,
729            arg_registers: context.arg_registers,
730            body: context.ops,
731            return_register: context.return_register,
732            is_closure: context.is_closure,
733            upvalues_ptr: context.upvalues_ptr,
734        };
735        Some(TraceOp::InlineCall {
736            dest: context.dest,
737            callee: context.callee_reg,
738            trace,
739        })
740    }
741
742    pub fn record_instruction(
743        &mut self,
744        instruction: Instruction,
745        current_ip: usize,
746        registers: &[Value; 256],
747        function: &crate::bytecode::Function,
748        function_idx: usize,
749        functions: &[crate::bytecode::Function],
750    ) -> Result<(), LustError> {
751        if !self.recording {
752            return Ok(());
753        }
754
755        if function_idx != self.current_function_idx() {
756            return Ok(());
757        }
758
759        let outcome: Result<(), LustError> = match instruction {
760            Instruction::LoadConst(dest, _) => {
761                // Rebox specialized value if dest contains one
762                self.remove_specialization_tracking(dest);
763
764                if let Some(_ty) = Self::get_value_type(&registers[dest as usize]) {
765                    self.mark_guarded(dest);
766                }
767
768                self.push_op(TraceOp::LoadConst {
769                    dest,
770                    value: registers[dest as usize].clone(),
771                });
772                Ok(())
773            }
774
775            Instruction::LoadGlobal(dest, _) => {
776                // Remove specialization tracking if dest contains a specialized value
777                self.remove_specialization_tracking(dest);
778
779                if let Some(_ty) = Self::get_value_type(&registers[dest as usize]) {
780                    self.mark_guarded(dest);
781                }
782
783                self.push_op(TraceOp::LoadConst {
784                    dest,
785                    value: registers[dest as usize].clone(),
786                });
787                Ok(())
788            }
789
790            Instruction::StoreGlobal(_, _) => Ok(()),
791
792            Instruction::Move(dest, src) => {
793                // If dest contains a specialized value, rebox it first before overwriting
794                self.remove_specialization_tracking(dest);
795
796                // Check if we're moving a specialized value
797                if let Some(&(specialized_id, ref layout)) = self.specialized_registers.get(&src) {
798                    crate::jit::log(|| {
799                        format!(
800                            "📦 JIT: Moving specialized #{} from reg {} to reg {}",
801                            specialized_id, src, dest
802                        )
803                    });
804                    // Track that dest now contains the specialized value
805                    self.specialized_registers
806                        .insert(dest, (specialized_id, layout.clone()));
807                    // Remove from source register
808                    self.specialized_registers.remove(&src);
809                }
810
811                self.push_op(TraceOp::Move { dest, src });
812                Ok(())
813            }
814
815            Instruction::Add(dest, lhs, rhs) => {
816                // Rebox specialized value if dest contains one
817                self.remove_specialization_tracking(dest);
818
819                self.add_type_guards(lhs, rhs, registers, function)?;
820                let lhs_type =
821                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
822                let rhs_type =
823                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
824                self.push_op(TraceOp::Add {
825                    dest,
826                    lhs,
827                    rhs,
828                    lhs_type,
829                    rhs_type,
830                });
831                Ok(())
832            }
833
834            Instruction::Sub(dest, lhs, rhs) => {
835                self.remove_specialization_tracking(dest);
836                self.add_type_guards(lhs, rhs, registers, function)?;
837                let lhs_type =
838                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
839                let rhs_type =
840                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
841                self.push_op(TraceOp::Sub {
842                    dest,
843                    lhs,
844                    rhs,
845                    lhs_type,
846                    rhs_type,
847                });
848                Ok(())
849            }
850
851            Instruction::Mul(dest, lhs, rhs) => {
852                self.add_type_guards(lhs, rhs, registers, function)?;
853                let lhs_type =
854                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
855                let rhs_type =
856                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
857                self.push_op(TraceOp::Mul {
858                    dest,
859                    lhs,
860                    rhs,
861                    lhs_type,
862                    rhs_type,
863                });
864                Ok(())
865            }
866
867            Instruction::Div(dest, lhs, rhs) => {
868                self.add_type_guards(lhs, rhs, registers, function)?;
869                let lhs_type =
870                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
871                let rhs_type =
872                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
873                self.push_op(TraceOp::Div {
874                    dest,
875                    lhs,
876                    rhs,
877                    lhs_type,
878                    rhs_type,
879                });
880                Ok(())
881            }
882
883            Instruction::Mod(dest, lhs, rhs) => {
884                self.add_type_guards(lhs, rhs, registers, function)?;
885                let lhs_type =
886                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
887                let rhs_type =
888                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
889                self.push_op(TraceOp::Mod {
890                    dest,
891                    lhs,
892                    rhs,
893                    lhs_type,
894                    rhs_type,
895                });
896                Ok(())
897            }
898
899            Instruction::Neg(dest, src) => {
900                self.push_op(TraceOp::Neg { dest, src });
901                Ok(())
902            }
903
904            Instruction::Eq(dest, lhs, rhs) => {
905                self.push_op(TraceOp::Eq { dest, lhs, rhs });
906                Ok(())
907            }
908
909            Instruction::Ne(dest, lhs, rhs) => {
910                self.push_op(TraceOp::Ne { dest, lhs, rhs });
911                Ok(())
912            }
913
914            Instruction::Lt(dest, lhs, rhs) => {
915                self.push_op(TraceOp::Lt { dest, lhs, rhs });
916                Ok(())
917            }
918
919            Instruction::Le(dest, lhs, rhs) => {
920                self.push_op(TraceOp::Le { dest, lhs, rhs });
921                Ok(())
922            }
923
924            Instruction::Gt(dest, lhs, rhs) => {
925                self.push_op(TraceOp::Gt { dest, lhs, rhs });
926                Ok(())
927            }
928
929            Instruction::Ge(dest, lhs, rhs) => {
930                self.push_op(TraceOp::Ge { dest, lhs, rhs });
931                Ok(())
932            }
933
934            Instruction::And(dest, lhs, rhs) => {
935                self.push_op(TraceOp::And { dest, lhs, rhs });
936                Ok(())
937            }
938
939            Instruction::Or(dest, lhs, rhs) => {
940                self.push_op(TraceOp::Or { dest, lhs, rhs });
941                Ok(())
942            }
943
944            Instruction::Not(dest, src) => {
945                self.push_op(TraceOp::Not { dest, src });
946                Ok(())
947            }
948
949            Instruction::Concat(dest, lhs, rhs) => {
950                if let Some(ty) = Self::get_value_type(&registers[lhs as usize]) {
951                    if !self.is_guarded(lhs) {
952                        self.push_op(TraceOp::Guard {
953                            register: lhs,
954                            expected_type: ty,
955                        });
956                        self.mark_guarded(lhs);
957                    }
958                }
959
960                if let Some(ty) = Self::get_value_type(&registers[rhs as usize]) {
961                    if !self.is_guarded(rhs) {
962                        self.push_op(TraceOp::Guard {
963                            register: rhs,
964                            expected_type: ty,
965                        });
966                        self.mark_guarded(rhs);
967                    }
968                }
969
970                self.push_op(TraceOp::Concat { dest, lhs, rhs });
971                Ok(())
972            }
973
974            Instruction::GetIndex(dest, array, index) => {
975                if let Some(ty) = Self::get_value_type(&registers[array as usize]) {
976                    if !self.is_guarded(array) {
977                        self.push_op(TraceOp::Guard {
978                            register: array,
979                            expected_type: ty,
980                        });
981                        self.mark_guarded(array);
982                    }
983                }
984
985                if let Some(ty) = Self::get_value_type(&registers[index as usize]) {
986                    if !self.is_guarded(index) {
987                        self.push_op(TraceOp::Guard {
988                            register: index,
989                            expected_type: ty,
990                        });
991                        self.mark_guarded(index);
992                    }
993                }
994
995                self.push_op(TraceOp::GetIndex { dest, array, index });
996                Ok(())
997            }
998
999            Instruction::ArrayLen(dest, array) => {
1000                if let Some(ty) = Self::get_value_type(&registers[array as usize]) {
1001                    if !self.is_guarded(array) {
1002                        self.push_op(TraceOp::Guard {
1003                            register: array,
1004                            expected_type: ty,
1005                        });
1006                        self.mark_guarded(array);
1007                    }
1008                }
1009
1010                self.push_op(TraceOp::ArrayLen { dest, array });
1011                Ok(())
1012            }
1013
1014            Instruction::CallMethod(obj_reg, method_name_idx, first_arg, arg_count, dest_reg) => {
1015                // Rebox specialized value if dest_reg contains one
1016                self.remove_specialization_tracking(dest_reg);
1017
1018                let method_name = function.chunk.constants[method_name_idx as usize]
1019                    .as_string()
1020                    .unwrap_or("unknown")
1021                    .to_string();
1022
1023                // Check if this is a method on a specialized value
1024                if let Some(&(specialized_id, _)) = self.specialized_registers.get(&obj_reg) {
1025                    // This is a method call on a specialized value
1026                    match method_name.as_str() {
1027                        "push" if arg_count == 1 => {
1028                            // Specialized array push
1029                            crate::jit::log(|| {
1030                                format!(
1031                                    "⚡ JIT: Specializing push on reg {} (specialized #{})",
1032                                    obj_reg, specialized_id
1033                                )
1034                            });
1035
1036                            // Guard the argument
1037                            let value_reg = first_arg;
1038                            if let Some(ty) = Self::get_value_type(&registers[value_reg as usize]) {
1039                                if !self.is_guarded(value_reg) {
1040                                    self.push_op(TraceOp::Guard {
1041                                        register: value_reg,
1042                                        expected_type: ty,
1043                                    });
1044                                    self.mark_guarded(value_reg);
1045                                }
1046                            }
1047
1048                            // Emit specialized push operation
1049                            self.push_op(TraceOp::SpecializedOp {
1050                                op: SpecializedOpKind::VecPush,
1051                                operands: vec![
1052                                    Operand::Specialized(specialized_id),
1053                                    Operand::Register(value_reg),
1054                                ],
1055                            });
1056
1057                            return Ok(());
1058                        }
1059                        "len" if arg_count == 0 => {
1060                            // Specialized array len
1061                            crate::jit::log(|| {
1062                                format!(
1063                                    "⚡ JIT: Specializing len on reg {} (specialized #{})",
1064                                    obj_reg, specialized_id
1065                                )
1066                            });
1067
1068                            // Emit specialized len operation
1069                            self.push_op(TraceOp::SpecializedOp {
1070                                op: SpecializedOpKind::VecLen,
1071                                operands: vec![
1072                                    Operand::Specialized(specialized_id),
1073                                    Operand::Register(dest_reg),
1074                                ],
1075                            });
1076
1077                            return Ok(());
1078                        }
1079                        _ => {
1080                            // Other methods on specialized values - need to rebox first
1081                            // For now, fall through to normal handling (will be wrong!)
1082                            crate::jit::log(|| {
1083                                format!(
1084                                    "⚠️  JIT: Method '{}' on specialized value not supported, will be incorrect!",
1085                                    method_name
1086                                )
1087                            });
1088                        }
1089                    }
1090                }
1091
1092                // Normal (non-specialized) method call
1093                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
1094                    if !self.is_guarded(obj_reg) {
1095                        self.push_op(TraceOp::Guard {
1096                            register: obj_reg,
1097                            expected_type: ty,
1098                        });
1099                        self.mark_guarded(obj_reg);
1100                    }
1101                }
1102
1103                for i in 0..arg_count {
1104                    let arg_reg = first_arg + i;
1105                    if let Some(ty) = Self::get_value_type(&registers[arg_reg as usize]) {
1106                        if !self.is_guarded(arg_reg) {
1107                            self.push_op(TraceOp::Guard {
1108                                register: arg_reg,
1109                                expected_type: ty,
1110                            });
1111                            self.mark_guarded(arg_reg);
1112                        }
1113                    }
1114                }
1115
1116                self.push_op(TraceOp::CallMethod {
1117                    dest: dest_reg,
1118                    object: obj_reg,
1119                    method_name,
1120                    first_arg,
1121                    arg_count,
1122                });
1123                Ok(())
1124            }
1125
1126            Instruction::GetField(dest, obj_reg, field_name_idx) => {
1127                let field_name = function.chunk.constants[field_name_idx as usize]
1128                    .as_string()
1129                    .unwrap_or("unknown")
1130                    .to_string();
1131                let (field_index, is_weak_field) = match &registers[obj_reg as usize] {
1132                    Value::Struct { layout, .. } => {
1133                        let idx = layout.index_of_str(&field_name);
1134                        let is_weak = idx.map(|i| layout.is_weak(i)).unwrap_or(false);
1135                        (idx, is_weak)
1136                    }
1137
1138                    _ => (None, false),
1139                };
1140                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
1141                    if !self.is_guarded(obj_reg) {
1142                        self.push_op(TraceOp::Guard {
1143                            register: obj_reg,
1144                            expected_type: ty,
1145                        });
1146                        self.mark_guarded(obj_reg);
1147                    }
1148                }
1149
1150                let value_type = Self::get_value_type(&registers[dest as usize]);
1151                self.push_op(TraceOp::GetField {
1152                    dest,
1153                    object: obj_reg,
1154                    field_name,
1155                    field_index,
1156                    value_type,
1157                    is_weak: is_weak_field,
1158                });
1159                Ok(())
1160            }
1161
1162            Instruction::SetField(obj_reg, field_name_idx, value_reg) => {
1163                let field_name = function.chunk.constants[field_name_idx as usize]
1164                    .as_string()
1165                    .unwrap_or("unknown")
1166                    .to_string();
1167                let (field_index, is_weak_field) = match &registers[obj_reg as usize] {
1168                    Value::Struct { layout, .. } => {
1169                        let idx = layout.index_of_str(&field_name);
1170                        let is_weak = idx.map(|i| layout.is_weak(i)).unwrap_or(false);
1171                        (idx, is_weak)
1172                    }
1173
1174                    _ => (None, false),
1175                };
1176                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
1177                    if !self.is_guarded(obj_reg) {
1178                        self.push_op(TraceOp::Guard {
1179                            register: obj_reg,
1180                            expected_type: ty,
1181                        });
1182                        self.mark_guarded(obj_reg);
1183                    }
1184                }
1185
1186                let value_type = Self::get_value_type(&registers[value_reg as usize]);
1187                if let Some(ty) = value_type {
1188                    if !self.is_guarded(value_reg) {
1189                        self.push_op(TraceOp::Guard {
1190                            register: value_reg,
1191                            expected_type: ty,
1192                        });
1193                        self.mark_guarded(value_reg);
1194                    }
1195                }
1196
1197                self.rebox_specialized_register(value_reg, "SetField");
1198
1199                self.push_op(TraceOp::SetField {
1200                    object: obj_reg,
1201                    field_name,
1202                    value: value_reg,
1203                    field_index,
1204                    value_type,
1205                    is_weak: is_weak_field,
1206                });
1207                Ok(())
1208            }
1209
1210            Instruction::NewStruct(
1211                dest,
1212                struct_name_idx,
1213                first_field_name_idx,
1214                first_field_reg,
1215                field_count,
1216            ) => {
1217                let struct_name = function.chunk.constants[struct_name_idx as usize]
1218                    .as_string()
1219                    .unwrap_or("unknown")
1220                    .to_string();
1221                let mut field_names = Vec::new();
1222                for i in 0..field_count {
1223                    let field_name_idx = first_field_name_idx + (i as u16);
1224                    let field_name = function.chunk.constants[field_name_idx as usize]
1225                        .as_string()
1226                        .unwrap_or("unknown")
1227                        .to_string();
1228                    field_names.push(field_name);
1229                }
1230
1231                let mut field_registers = Vec::new();
1232                for i in 0..field_count {
1233                    let field_reg = first_field_reg + i;
1234                    field_registers.push(field_reg);
1235                    if let Some(ty) = Self::get_value_type(&registers[field_reg as usize]) {
1236                        if !self.is_guarded(field_reg) {
1237                            self.push_op(TraceOp::Guard {
1238                                register: field_reg,
1239                                expected_type: ty,
1240                            });
1241                            self.mark_guarded(field_reg);
1242                        }
1243                    }
1244                }
1245
1246                for &field_reg in &field_registers {
1247                    self.rebox_specialized_register(field_reg, "struct literal field");
1248                }
1249
1250                self.push_op(TraceOp::NewStruct {
1251                    dest,
1252                    struct_name,
1253                    field_names,
1254                    field_registers,
1255                });
1256                Ok(())
1257            }
1258
1259            Instruction::NewEnumUnit(dest, enum_name_idx, variant_idx) => {
1260                let enum_name = function.chunk.constants[enum_name_idx as usize]
1261                    .as_string()
1262                    .unwrap_or("unknown")
1263                    .to_string();
1264                let variant_name = function.chunk.constants[variant_idx as usize]
1265                    .as_string()
1266                    .unwrap_or("unknown")
1267                    .to_string();
1268                self.push_op(TraceOp::NewEnumUnit {
1269                    dest,
1270                    enum_name,
1271                    variant_name,
1272                });
1273                Ok(())
1274            }
1275
1276            Instruction::NewEnumVariant(
1277                dest,
1278                enum_name_idx,
1279                variant_idx,
1280                first_value,
1281                value_count,
1282            ) => {
1283                let enum_name = function.chunk.constants[enum_name_idx as usize]
1284                    .as_string()
1285                    .unwrap_or("unknown")
1286                    .to_string();
1287                let variant_name = function.chunk.constants[variant_idx as usize]
1288                    .as_string()
1289                    .unwrap_or("unknown")
1290                    .to_string();
1291                let mut value_registers = Vec::new();
1292                for i in 0..value_count {
1293                    value_registers.push(first_value + i);
1294                }
1295
1296                for &value_reg in &value_registers {
1297                    self.rebox_specialized_register(value_reg, "enum variant value");
1298                }
1299
1300                self.push_op(TraceOp::NewEnumVariant {
1301                    dest,
1302                    enum_name,
1303                    variant_name,
1304                    value_registers,
1305                });
1306                Ok(())
1307            }
1308
1309            Instruction::IsEnumVariant(dest, value_reg, enum_name_idx, variant_idx) => {
1310                let enum_name = function.chunk.constants[enum_name_idx as usize]
1311                    .as_string()
1312                    .unwrap_or("unknown")
1313                    .to_string();
1314                let variant_name = function.chunk.constants[variant_idx as usize]
1315                    .as_string()
1316                    .unwrap_or("unknown")
1317                    .to_string();
1318                self.push_op(TraceOp::IsEnumVariant {
1319                    dest,
1320                    value: value_reg,
1321                    enum_name,
1322                    variant_name,
1323                });
1324                Ok(())
1325            }
1326
1327            Instruction::GetEnumValue(dest, enum_reg, index) => {
1328                self.push_op(TraceOp::GetEnumValue {
1329                    dest,
1330                    enum_reg,
1331                    index,
1332                });
1333                Ok(())
1334            }
1335
1336            Instruction::Call(func_reg, first_arg, arg_count, dest_reg) => {
1337                // Rebox specialized value if dest_reg contains one
1338                self.remove_specialization_tracking(dest_reg);
1339
1340                match &registers[func_reg as usize] {
1341                    Value::NativeFunction(native_fn) => {
1342                        let traced = TracedNativeFn::new(native_fn.clone());
1343                        if !self.is_guarded(func_reg) {
1344                            self.push_op(TraceOp::GuardNativeFunction {
1345                                register: func_reg,
1346                                function: traced.clone(),
1347                            });
1348                            self.mark_guarded(func_reg);
1349                        }
1350
1351                        self.push_op(TraceOp::CallNative {
1352                            dest: dest_reg,
1353                            callee: func_reg,
1354                            function: traced,
1355                            first_arg,
1356                            arg_count,
1357                        });
1358                        Ok(())
1359                    }
1360
1361                    Value::Function(function_idx) => {
1362                        if !self.is_guarded(func_reg) {
1363                            self.push_op(TraceOp::GuardFunction {
1364                                register: func_reg,
1365                                function_idx: *function_idx,
1366                            });
1367                            self.mark_guarded(func_reg);
1368                        }
1369
1370                        let mut did_inline = false;
1371                        if let Some(callee_fn) = functions.get(*function_idx) {
1372                            if self.should_inline(*function_idx, callee_fn)
1373                                && (arg_count as usize) <= callee_fn.register_count as usize
1374                            {
1375                                let mut arg_registers = Vec::with_capacity(arg_count as usize);
1376                                for i in 0..arg_count {
1377                                    arg_registers.push(first_arg + i);
1378                                }
1379                                self.push_inline_context(
1380                                    *function_idx,
1381                                    callee_fn.register_count,
1382                                    dest_reg,
1383                                    func_reg,
1384                                    first_arg,
1385                                    arg_count,
1386                                    arg_registers,
1387                                    false,
1388                                    None,
1389                                );
1390                                did_inline = true;
1391                            }
1392                        }
1393
1394                        if !did_inline {
1395                            self.push_op(TraceOp::CallFunction {
1396                                dest: dest_reg,
1397                                callee: func_reg,
1398                                function_idx: *function_idx,
1399                                first_arg,
1400                                arg_count,
1401                                is_closure: false,
1402                                upvalues_ptr: None,
1403                            });
1404                        }
1405
1406                        Ok(())
1407                    }
1408
1409                    Value::Closure {
1410                        function_idx,
1411                        upvalues,
1412                    } => {
1413                        let upvalues_ptr = Rc::as_ptr(upvalues) as *const ();
1414                        if !self.is_guarded(func_reg) {
1415                            self.push_op(TraceOp::GuardClosure {
1416                                register: func_reg,
1417                                function_idx: *function_idx,
1418                                upvalues_ptr,
1419                            });
1420                            self.mark_guarded(func_reg);
1421                        }
1422
1423                        let mut did_inline = false;
1424                        if let Some(callee_fn) = functions.get(*function_idx) {
1425                            if self.should_inline(*function_idx, callee_fn)
1426                                && (arg_count as usize) <= callee_fn.register_count as usize
1427                            {
1428                                let mut arg_registers = Vec::with_capacity(arg_count as usize);
1429                                for i in 0..arg_count {
1430                                    arg_registers.push(first_arg + i);
1431                                }
1432                                self.push_inline_context(
1433                                    *function_idx,
1434                                    callee_fn.register_count,
1435                                    dest_reg,
1436                                    func_reg,
1437                                    first_arg,
1438                                    arg_count,
1439                                    arg_registers,
1440                                    true,
1441                                    Some(upvalues_ptr),
1442                                );
1443                                did_inline = true;
1444                            }
1445                        }
1446
1447                        if !did_inline {
1448                            self.push_op(TraceOp::CallFunction {
1449                                dest: dest_reg,
1450                                callee: func_reg,
1451                                function_idx: *function_idx,
1452                                first_arg,
1453                                arg_count,
1454                                is_closure: true,
1455                                upvalues_ptr: Some(upvalues_ptr),
1456                            });
1457                        }
1458
1459                        Ok(())
1460                    }
1461
1462                    _ => {
1463                        self.stop_recording();
1464                        crate::jit::log(|| {
1465                            format!(
1466                                "Trace aborted: unsupported call operation on register {} (value {:?})",
1467                                func_reg,
1468                                registers[func_reg as usize].tag()
1469                            )
1470                        });
1471                        Err(LustError::RuntimeError {
1472                            message: "Trace aborted: unsupported call operation".to_string(),
1473                        })
1474                    }
1475                }
1476            }
1477
1478            Instruction::NewArray(dest, first_elem, count) => {
1479                // Rebox specialized value if dest contains one
1480                self.remove_specialization_tracking(dest);
1481
1482                // Disable specialization inside inlined functions for now to avoid
1483                // register aliasing issues between inline frames and the parent trace.
1484                if !self.inline_stack.is_empty() {
1485                    self.push_op(TraceOp::NewArray {
1486                        dest,
1487                        first_element: first_elem,
1488                        count,
1489                    });
1490                    return Ok(());
1491                }
1492
1493                // Check if we can specialize this array based on element type
1494                // register_types contains the element type, not the array type
1495                let can_specialize = if let Some(element_type) = function.register_types.get(&dest)
1496                {
1497                    // Build the full Array type from the element type
1498                    use crate::ast::{Span, Type};
1499                    let array_type = crate::ast::TypeKind::Array(Box::new(Type::new(
1500                        element_type.clone(),
1501                        Span::dummy(),
1502                    )));
1503
1504                    // Check if this array type is specializable
1505                    self.specialization_registry
1506                        .get_specialization(&array_type)
1507                        .is_some()
1508                } else {
1509                    false
1510                };
1511
1512                if can_specialize {
1513                    let element_type = function.register_types.get(&dest).unwrap().clone();
1514
1515                    // Build the full Array type
1516                    use crate::ast::{Span, Type};
1517                    let array_type = crate::ast::TypeKind::Array(Box::new(Type::new(
1518                        element_type.clone(),
1519                        Span::dummy(),
1520                    )));
1521
1522                    let layout = self
1523                        .specialization_registry
1524                        .get_specialization(&array_type)
1525                        .unwrap();
1526
1527                    crate::jit::log(|| {
1528                        format!(
1529                            "🔬 JIT: Specializing NewArray for reg {} with element type {:?}",
1530                            dest, element_type
1531                        )
1532                    });
1533
1534                    // First create the array normally (for initial values if any)
1535                    if count > 0 {
1536                        self.push_op(TraceOp::NewArray {
1537                            dest,
1538                            first_element: first_elem,
1539                            count,
1540                        });
1541                    } else {
1542                        // Empty array
1543                        self.push_op(TraceOp::NewArray {
1544                            dest,
1545                            first_element: 0,
1546                            count: 0,
1547                        });
1548                    }
1549
1550                    // Then unbox it for specialized operations
1551                    let specialized_id = self.next_specialized_id;
1552                    self.next_specialized_id += 1;
1553
1554                    self.push_op(TraceOp::Unbox {
1555                        specialized_id,
1556                        source_reg: dest,
1557                        layout: layout.clone(),
1558                    });
1559
1560                    // Track that this register now contains a specialized value
1561                    self.specialized_registers
1562                        .insert(dest, (specialized_id, layout));
1563                } else {
1564                    // Normal non-specialized array
1565                    self.push_op(TraceOp::NewArray {
1566                        dest,
1567                        first_element: first_elem,
1568                        count,
1569                    });
1570                }
1571                Ok(())
1572            }
1573
1574            Instruction::NewMap(_) | Instruction::SetIndex(_, _, _) => {
1575                self.stop_recording();
1576                Err(LustError::RuntimeError {
1577                    message: "Trace aborted: unsupported index operation".to_string(),
1578                })
1579            }
1580
1581            Instruction::Return(value_reg) => {
1582                let return_reg = if value_reg == 255 {
1583                    None
1584                } else {
1585                    Some(value_reg)
1586                };
1587
1588                // Rebox any specialized values before return
1589                if let Some(reg) = return_reg {
1590                    if let Some(&(specialized_id, ref layout)) =
1591                        self.specialized_registers.get(&reg)
1592                    {
1593                        crate::jit::log(|| {
1594                            format!(
1595                                "📦 JIT: Reboxing specialized #{} in reg {} before return",
1596                                specialized_id, reg
1597                            )
1598                        });
1599
1600                        self.push_op(TraceOp::Rebox {
1601                            dest_reg: reg,
1602                            specialized_id,
1603                            layout: layout.clone(),
1604                        });
1605
1606                        self.specialized_registers.remove(&reg);
1607                    }
1608                }
1609
1610                // Ensure no specialized values leak past the return
1611                self.rebox_all_specialized_values();
1612
1613                if let Some(ctx) = self.inline_stack.last_mut() {
1614                    ctx.return_register = return_reg;
1615                    crate::jit::log(|| {
1616                        format!(
1617                            "🔧 JIT: Inline return detected, return_reg={:?}",
1618                            return_reg
1619                        )
1620                    });
1621                    if let Some(inline_op) = self.finalize_inline_context() {
1622                        self.push_op(inline_op);
1623                    }
1624                    Ok(())
1625                } else if function_idx == self.trace.function_idx {
1626                    self.stop_recording();
1627                    Ok(())
1628                } else {
1629                    self.push_op(TraceOp::Return { value: return_reg });
1630                    Ok(())
1631                }
1632            }
1633
1634            Instruction::Jump(offset) => {
1635                if offset < 0 {
1636                    let target_calc = (current_ip as isize) + (offset as isize);
1637                    if target_calc < 0 {
1638                        self.stop_recording();
1639                        Err(LustError::RuntimeError {
1640                            message: format!(
1641                                "Invalid jump target: offset={}, current_ip={}, target={}",
1642                                offset, current_ip, target_calc
1643                            ),
1644                        })
1645                    } else {
1646                        let jump_target = target_calc as usize;
1647                        let loop_key = (function_idx, jump_target);
1648
1649                        // Track how many times we've seen this loop backedge
1650                        let iteration_count = self.loop_iterations.entry(loop_key).or_insert(0);
1651                        *iteration_count += 1;
1652
1653                        if function_idx == self.trace.function_idx
1654                            && jump_target == self.trace.start_ip
1655                        {
1656                            // This is our main trace loop closing - check if we should unroll more
1657                            if *iteration_count < crate::jit::LOOP_UNROLL_COUNT {
1658                                crate::jit::log(|| {
1659                                    format!(
1660                                        "🔄 JIT: Unrolling main loop (iteration {}/{})",
1661                                        iteration_count,
1662                                        crate::jit::LOOP_UNROLL_COUNT
1663                                    )
1664                                });
1665                                // Continue recording to unroll the loop
1666                                Ok(())
1667                            } else {
1668                                crate::jit::log(|| {
1669                                    format!(
1670                                        "✅ JIT: Loop unrolled {} times, stopping trace",
1671                                        iteration_count
1672                                    )
1673                                });
1674                                self.stop_recording();
1675                                Ok(())
1676                            }
1677                        } else {
1678                            // This is a nested loop that should be compiled as a separate trace
1679                            // Following LuaJIT's approach: don't inline loops, compile them separately
1680                            let bailout_ip = current_ip.saturating_sub(1);
1681
1682                            crate::jit::log(|| {
1683                                format!(
1684                                    "🔄 JIT: Nested loop detected at func {} ip {} - will call as separate trace",
1685                                    function_idx, jump_target
1686                                )
1687                            });
1688
1689                            // Rebox all specialized values before calling nested trace
1690                            self.rebox_all_specialized_values();
1691
1692                            // Emit NestedLoopCall which will eventually call the compiled inner trace
1693                            self.push_op(TraceOp::NestedLoopCall {
1694                                function_idx,
1695                                loop_start_ip: jump_target,
1696                                bailout_ip,
1697                            });
1698                            Ok(())
1699                        }
1700                    }
1701                } else {
1702                    Ok(())
1703                }
1704            }
1705
1706            Instruction::JumpIf(cond, offset) => {
1707                let condition = &registers[cond as usize];
1708                let is_truthy = condition.is_truthy();
1709                let target_offset = (current_ip as isize) + (offset as isize);
1710                let target = if target_offset < 0 {
1711                    0
1712                } else {
1713                    target_offset as usize
1714                };
1715                let bailout_ip = if is_truthy { current_ip } else { target };
1716                self.push_op(TraceOp::GuardLoopContinue {
1717                    condition_register: cond,
1718                    expect_truthy: is_truthy,
1719                    bailout_ip,
1720                });
1721                Ok(())
1722            }
1723
1724            Instruction::JumpIfNot(cond, offset) => {
1725                let condition = &registers[cond as usize];
1726                let is_truthy = condition.is_truthy();
1727                let target_offset = (current_ip as isize) + (offset as isize);
1728                let target = if target_offset < 0 {
1729                    0
1730                } else {
1731                    target_offset as usize
1732                };
1733                let bailout_ip = if !is_truthy { current_ip } else { target };
1734                self.push_op(TraceOp::GuardLoopContinue {
1735                    condition_register: cond,
1736                    expect_truthy: is_truthy,
1737                    bailout_ip,
1738                });
1739                Ok(())
1740            }
1741
1742            _ => {
1743                self.stop_recording();
1744                crate::jit::log(|| {
1745                    format!(
1746                        "Trace aborted: unsupported instruction {:?}",
1747                        instruction.opcode()
1748                    )
1749                });
1750                Err(LustError::RuntimeError {
1751                    message: "Trace aborted: unsupported instruction".to_string(),
1752                })
1753            }
1754        };
1755
1756        outcome?;
1757
1758        if self.op_count >= self.max_length {
1759            self.stop_recording();
1760            return Err(LustError::RuntimeError {
1761                message: "Trace too long".to_string(),
1762            });
1763        }
1764
1765        Ok(())
1766    }
1767    fn add_type_guards(
1768        &mut self,
1769        lhs: Register,
1770        rhs: Register,
1771        registers: &[Value; 256],
1772        function: &crate::bytecode::Function,
1773    ) -> Result<(), LustError> {
1774        if let Some(ty) = Self::get_value_type(&registers[lhs as usize]) {
1775            let needs_guard = if self.is_guarded(lhs) {
1776                false
1777            } else if let Some(static_type) = function.register_types.get(&lhs) {
1778                !Self::type_kind_matches_value_type(static_type, ty)
1779            } else {
1780                true
1781            };
1782            if needs_guard {
1783                self.push_op(TraceOp::Guard {
1784                    register: lhs,
1785                    expected_type: ty,
1786                });
1787                self.mark_guarded(lhs);
1788            } else {
1789                self.mark_guarded(lhs);
1790            }
1791        }
1792
1793        if let Some(ty) = Self::get_value_type(&registers[rhs as usize]) {
1794            let needs_guard = if self.is_guarded(rhs) {
1795                false
1796            } else if let Some(static_type) = function.register_types.get(&rhs) {
1797                !Self::type_kind_matches_value_type(static_type, ty)
1798            } else {
1799                true
1800            };
1801            if needs_guard {
1802                self.push_op(TraceOp::Guard {
1803                    register: rhs,
1804                    expected_type: ty,
1805                });
1806                self.mark_guarded(rhs);
1807            } else {
1808                self.mark_guarded(rhs);
1809            }
1810        }
1811
1812        Ok(())
1813    }
1814
1815    fn type_kind_matches_value_type(
1816        type_kind: &crate::ast::TypeKind,
1817        value_type: ValueType,
1818    ) -> bool {
1819        use crate::ast::TypeKind;
1820        match (type_kind, value_type) {
1821            (TypeKind::Int, ValueType::Int) => true,
1822            (TypeKind::Float, ValueType::Float) => true,
1823            (TypeKind::Bool, ValueType::Bool) => true,
1824            (TypeKind::String, ValueType::String) => true,
1825            (TypeKind::Array(_), ValueType::Array) => true,
1826            (TypeKind::Tuple(_), ValueType::Tuple) => true,
1827            _ => false,
1828        }
1829    }
1830
1831    fn get_value_type(value: &Value) -> Option<ValueType> {
1832        match value {
1833            Value::Int(_) => Some(ValueType::Int),
1834            Value::Float(_) => Some(ValueType::Float),
1835            Value::Bool(_) => Some(ValueType::Bool),
1836            Value::String(_) => Some(ValueType::String),
1837            Value::Array(_) => Some(ValueType::Array),
1838            Value::Tuple(_) => Some(ValueType::Tuple),
1839            Value::Struct { .. } => Some(ValueType::Struct),
1840            _ => None,
1841        }
1842    }
1843
1844    pub fn finish(mut self) -> Trace {
1845        #[cfg(feature = "std")]
1846        if std::env::var("LUST_TRACE_DEBUG").is_ok() {
1847            eprintln!(
1848                "🧵 Trace dump (func {}, start_ip {}):",
1849                self.trace.function_idx, self.trace.start_ip
1850            );
1851            if !self.trace.preamble.is_empty() {
1852                eprintln!("  Preamble:");
1853                for (idx, op) in self.trace.preamble.iter().enumerate() {
1854                    eprintln!("    {:03}: {:?}", idx, op);
1855                }
1856            }
1857            eprintln!("  Body:");
1858            for (idx, op) in self.trace.ops.iter().enumerate() {
1859                eprintln!("    {:03}: {:?}", idx, op);
1860            }
1861            if !self.trace.postamble.is_empty() {
1862                eprintln!("  Postamble:");
1863                for (idx, op) in self.trace.postamble.iter().enumerate() {
1864                    eprintln!("    {:03}: {:?}", idx, op);
1865                }
1866            }
1867        }
1868
1869        // Finalize before returning (add rebox ops to postamble)
1870        self.finalize_trace();
1871        self.trace
1872    }
1873
1874    pub fn is_recording(&self) -> bool {
1875        self.recording
1876    }
1877
1878    pub fn abort(&mut self) {
1879        self.stop_recording();
1880    }
1881}
lust/jit/trace.rs

lust/jit/
trace.rs
