lust/jit/
trace.rs

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

lust/jit/
trace.rs
