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],
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    #[allow(dead_code)]
620    fn invalidate_specialization(&mut self, register: Register) {
621        if let Some((specialized_id, layout)) = self.specialized_registers.remove(&register) {
622            crate::jit::log(|| {
623                format!(
624                    "🚫 JIT: Invalidating specialization for reg {} (being overwritten) - will drop specialized #{}",
625                    register, specialized_id
626                )
627            });
628            // Track this for cleanup in postamble - the Vec data needs to be dropped
629            self.leaked_specialized_values
630                .push((specialized_id, layout));
631        }
632    }
633
634    /// Remove specialization tracking if register is about to be overwritten
635    /// The Vec data becomes "leaked" on the JIT stack but that's fine - it's cleaned
636    /// up when the stack frame is destroyed. The array is still managed by Rc<RefCell<>>.
637    fn remove_specialization_tracking(&mut self, register: Register) {
638        if let Some((_specialized_id, _layout)) = self.specialized_registers.remove(&register) {
639            crate::jit::log(|| {
640                format!(
641                    "🗑️  JIT: Removing specialization tracking for reg {} (being overwritten)",
642                    register
643                )
644            });
645            // Don't emit rebox - the Vec data stays on JIT stack but that's OK
646            // It will be cleaned up when the JIT stack frame is destroyed
647        }
648    }
649
650    fn rebox_specialized_register(&mut self, register: Register, context: &str) {
651        if let Some((specialized_id, layout)) = self.specialized_registers.remove(&register) {
652            crate::jit::log(|| {
653                format!(
654                    "📦 JIT: Reboxing specialized #{} from reg {} before {}",
655                    specialized_id, register, context
656                )
657            });
658            self.push_op(TraceOp::Rebox {
659                dest_reg: register,
660                specialized_id,
661                layout,
662            });
663        }
664    }
665
666    fn should_inline(&self, function_idx: usize, callee_fn: &crate::bytecode::Function) -> bool {
667        if function_idx == self.trace.function_idx {
668            return false;
669        }
670
671        if self
672            .inline_stack
673            .iter()
674            .any(|ctx| ctx.function_idx == function_idx)
675        {
676            return false;
677        }
678
679        // Disable inlining when specialized values are active to avoid
680        // stack layout conflicts between inline frames and specialized storage
681        if !self.specialized_registers.is_empty() {
682            return false;
683        }
684
685        if callee_fn.chunk.instructions.iter().any(|inst| {
686            matches!(
687                inst,
688                Instruction::Jump(_) | Instruction::JumpIf(..) | Instruction::JumpIfNot(..)
689            )
690        }) {
691            return false;
692        }
693
694        true
695    }
696
697    fn push_inline_context(
698        &mut self,
699        function_idx: usize,
700        register_count: u8,
701        dest: Register,
702        callee_reg: Register,
703        first_arg: Register,
704        arg_count: u8,
705        arg_registers: Vec<Register>,
706        is_closure: bool,
707        upvalues_ptr: Option<*const ()>,
708    ) {
709        self.inline_stack.push(InlineContext {
710            function_idx,
711            register_count,
712            dest,
713            callee_reg,
714            first_arg,
715            arg_count,
716            arg_registers,
717            ops: Vec::new(),
718            guarded_registers: HashSet::new(),
719            return_register: None,
720            is_closure,
721            upvalues_ptr,
722        });
723    }
724
725    fn finalize_inline_context(&mut self) -> Option<TraceOp> {
726        let context = self.inline_stack.pop()?;
727        let trace = InlineTrace {
728            function_idx: context.function_idx,
729            register_count: context.register_count,
730            first_arg: context.first_arg,
731            arg_count: context.arg_count,
732            arg_registers: context.arg_registers,
733            body: context.ops,
734            return_register: context.return_register,
735            is_closure: context.is_closure,
736            upvalues_ptr: context.upvalues_ptr,
737        };
738        Some(TraceOp::InlineCall {
739            dest: context.dest,
740            callee: context.callee_reg,
741            trace,
742        })
743    }
744
745    pub fn record_instruction(
746        &mut self,
747        instruction: Instruction,
748        current_ip: usize,
749        registers: &[Value],
750        function: &crate::bytecode::Function,
751        function_idx: usize,
752        functions: &[crate::bytecode::Function],
753    ) -> Result<(), LustError> {
754        if !self.recording {
755            return Ok(());
756        }
757
758        if function_idx != self.current_function_idx() {
759            return Ok(());
760        }
761
762        let outcome: Result<(), LustError> = match instruction {
763            Instruction::LoadConst(dest, _) => {
764                // Rebox specialized value if dest contains one
765                self.remove_specialization_tracking(dest);
766
767                if let Some(_ty) = Self::get_value_type(&registers[dest as usize]) {
768                    self.mark_guarded(dest);
769                }
770
771                self.push_op(TraceOp::LoadConst {
772                    dest,
773                    value: registers[dest as usize].clone(),
774                });
775                Ok(())
776            }
777
778            Instruction::LoadGlobal(dest, _) => {
779                // Remove specialization tracking if dest contains a specialized value
780                self.remove_specialization_tracking(dest);
781
782                if let Some(_ty) = Self::get_value_type(&registers[dest as usize]) {
783                    self.mark_guarded(dest);
784                }
785
786                self.push_op(TraceOp::LoadConst {
787                    dest,
788                    value: registers[dest as usize].clone(),
789                });
790                Ok(())
791            }
792
793            Instruction::StoreGlobal(_, _) => Ok(()),
794
795            Instruction::Move(dest, src) => {
796                // If dest contains a specialized value, rebox it first before overwriting
797                self.remove_specialization_tracking(dest);
798
799                // Check if we're moving a specialized value
800                if let Some(&(specialized_id, ref layout)) = self.specialized_registers.get(&src) {
801                    crate::jit::log(|| {
802                        format!(
803                            "📦 JIT: Moving specialized #{} from reg {} to reg {}",
804                            specialized_id, src, dest
805                        )
806                    });
807                    // Track that dest now contains the specialized value
808                    self.specialized_registers
809                        .insert(dest, (specialized_id, layout.clone()));
810                    // Remove from source register
811                    self.specialized_registers.remove(&src);
812                }
813
814                self.push_op(TraceOp::Move { dest, src });
815                Ok(())
816            }
817
818            Instruction::Add(dest, lhs, rhs) => {
819                // Rebox specialized value if dest contains one
820                self.remove_specialization_tracking(dest);
821
822                self.add_type_guards(lhs, rhs, registers, function)?;
823                let lhs_type =
824                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
825                let rhs_type =
826                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
827                self.push_op(TraceOp::Add {
828                    dest,
829                    lhs,
830                    rhs,
831                    lhs_type,
832                    rhs_type,
833                });
834                Ok(())
835            }
836
837            Instruction::Sub(dest, lhs, rhs) => {
838                self.remove_specialization_tracking(dest);
839                self.add_type_guards(lhs, rhs, registers, function)?;
840                let lhs_type =
841                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
842                let rhs_type =
843                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
844                self.push_op(TraceOp::Sub {
845                    dest,
846                    lhs,
847                    rhs,
848                    lhs_type,
849                    rhs_type,
850                });
851                Ok(())
852            }
853
854            Instruction::Mul(dest, lhs, rhs) => {
855                self.add_type_guards(lhs, rhs, registers, function)?;
856                let lhs_type =
857                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
858                let rhs_type =
859                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
860                self.push_op(TraceOp::Mul {
861                    dest,
862                    lhs,
863                    rhs,
864                    lhs_type,
865                    rhs_type,
866                });
867                Ok(())
868            }
869
870            Instruction::Div(dest, lhs, rhs) => {
871                self.add_type_guards(lhs, rhs, registers, function)?;
872                let lhs_type =
873                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
874                let rhs_type =
875                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
876                self.push_op(TraceOp::Div {
877                    dest,
878                    lhs,
879                    rhs,
880                    lhs_type,
881                    rhs_type,
882                });
883                Ok(())
884            }
885
886            Instruction::Mod(dest, lhs, rhs) => {
887                self.add_type_guards(lhs, rhs, registers, function)?;
888                let lhs_type =
889                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
890                let rhs_type =
891                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
892                self.push_op(TraceOp::Mod {
893                    dest,
894                    lhs,
895                    rhs,
896                    lhs_type,
897                    rhs_type,
898                });
899                Ok(())
900            }
901
902            Instruction::Neg(dest, src) => {
903                self.push_op(TraceOp::Neg { dest, src });
904                Ok(())
905            }
906
907            Instruction::Eq(dest, lhs, rhs) => {
908                self.push_op(TraceOp::Eq { dest, lhs, rhs });
909                Ok(())
910            }
911
912            Instruction::Ne(dest, lhs, rhs) => {
913                self.push_op(TraceOp::Ne { dest, lhs, rhs });
914                Ok(())
915            }
916
917            Instruction::Lt(dest, lhs, rhs) => {
918                self.push_op(TraceOp::Lt { dest, lhs, rhs });
919                Ok(())
920            }
921
922            Instruction::Le(dest, lhs, rhs) => {
923                self.push_op(TraceOp::Le { dest, lhs, rhs });
924                Ok(())
925            }
926
927            Instruction::Gt(dest, lhs, rhs) => {
928                self.push_op(TraceOp::Gt { dest, lhs, rhs });
929                Ok(())
930            }
931
932            Instruction::Ge(dest, lhs, rhs) => {
933                self.push_op(TraceOp::Ge { dest, lhs, rhs });
934                Ok(())
935            }
936
937            Instruction::And(dest, lhs, rhs) => {
938                self.push_op(TraceOp::And { dest, lhs, rhs });
939                Ok(())
940            }
941
942            Instruction::Or(dest, lhs, rhs) => {
943                self.push_op(TraceOp::Or { dest, lhs, rhs });
944                Ok(())
945            }
946
947            Instruction::Not(dest, src) => {
948                self.push_op(TraceOp::Not { dest, src });
949                Ok(())
950            }
951
952            Instruction::Concat(dest, lhs, rhs) => {
953                if let Some(ty) = Self::get_value_type(&registers[lhs as usize]) {
954                    if !self.is_guarded(lhs) {
955                        self.push_op(TraceOp::Guard {
956                            register: lhs,
957                            expected_type: ty,
958                        });
959                        self.mark_guarded(lhs);
960                    }
961                }
962
963                if let Some(ty) = Self::get_value_type(&registers[rhs as usize]) {
964                    if !self.is_guarded(rhs) {
965                        self.push_op(TraceOp::Guard {
966                            register: rhs,
967                            expected_type: ty,
968                        });
969                        self.mark_guarded(rhs);
970                    }
971                }
972
973                self.push_op(TraceOp::Concat { dest, lhs, rhs });
974                Ok(())
975            }
976
977            Instruction::GetIndex(dest, array, index) => {
978                if let Some(ty) = Self::get_value_type(&registers[array as usize]) {
979                    if !self.is_guarded(array) {
980                        self.push_op(TraceOp::Guard {
981                            register: array,
982                            expected_type: ty,
983                        });
984                        self.mark_guarded(array);
985                    }
986                }
987
988                if let Some(ty) = Self::get_value_type(&registers[index as usize]) {
989                    if !self.is_guarded(index) {
990                        self.push_op(TraceOp::Guard {
991                            register: index,
992                            expected_type: ty,
993                        });
994                        self.mark_guarded(index);
995                    }
996                }
997
998                self.push_op(TraceOp::GetIndex { dest, array, index });
999                Ok(())
1000            }
1001
1002            Instruction::ArrayLen(dest, array) => {
1003                if let Some(ty) = Self::get_value_type(&registers[array as usize]) {
1004                    if !self.is_guarded(array) {
1005                        self.push_op(TraceOp::Guard {
1006                            register: array,
1007                            expected_type: ty,
1008                        });
1009                        self.mark_guarded(array);
1010                    }
1011                }
1012
1013                self.push_op(TraceOp::ArrayLen { dest, array });
1014                Ok(())
1015            }
1016
1017            Instruction::CallMethod(obj_reg, method_name_idx, first_arg, arg_count, dest_reg) => {
1018                // Rebox specialized value if dest_reg contains one
1019                self.remove_specialization_tracking(dest_reg);
1020
1021                let method_name = function.chunk.constants[method_name_idx as usize]
1022                    .as_string()
1023                    .unwrap_or("unknown")
1024                    .to_string();
1025
1026                // Check if this is a method on a specialized value
1027                if let Some(&(specialized_id, _)) = self.specialized_registers.get(&obj_reg) {
1028                    // This is a method call on a specialized value
1029                    match method_name.as_str() {
1030                        "push" if arg_count == 1 => {
1031                            // Specialized array push
1032                            crate::jit::log(|| {
1033                                format!(
1034                                    "⚡ JIT: Specializing push on reg {} (specialized #{})",
1035                                    obj_reg, specialized_id
1036                                )
1037                            });
1038
1039                            // Guard the argument
1040                            let value_reg = first_arg;
1041                            if let Some(ty) = Self::get_value_type(&registers[value_reg as usize]) {
1042                                if !self.is_guarded(value_reg) {
1043                                    self.push_op(TraceOp::Guard {
1044                                        register: value_reg,
1045                                        expected_type: ty,
1046                                    });
1047                                    self.mark_guarded(value_reg);
1048                                }
1049                            }
1050
1051                            // Emit specialized push operation
1052                            self.push_op(TraceOp::SpecializedOp {
1053                                op: SpecializedOpKind::VecPush,
1054                                operands: vec![
1055                                    Operand::Specialized(specialized_id),
1056                                    Operand::Register(value_reg),
1057                                ],
1058                            });
1059
1060                            return Ok(());
1061                        }
1062                        "len" if arg_count == 0 => {
1063                            // Specialized array len
1064                            crate::jit::log(|| {
1065                                format!(
1066                                    "⚡ JIT: Specializing len on reg {} (specialized #{})",
1067                                    obj_reg, specialized_id
1068                                )
1069                            });
1070
1071                            // Emit specialized len operation
1072                            self.push_op(TraceOp::SpecializedOp {
1073                                op: SpecializedOpKind::VecLen,
1074                                operands: vec![
1075                                    Operand::Specialized(specialized_id),
1076                                    Operand::Register(dest_reg),
1077                                ],
1078                            });
1079
1080                            return Ok(());
1081                        }
1082                        _ => {
1083                            // Other methods on specialized values - need to rebox first
1084                            // For now, fall through to normal handling (will be wrong!)
1085                            crate::jit::log(|| {
1086                                format!(
1087                                    "⚠️  JIT: Method '{}' on specialized value not supported, will be incorrect!",
1088                                    method_name
1089                                )
1090                            });
1091                        }
1092                    }
1093                }
1094
1095                // Normal (non-specialized) method call
1096                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
1097                    if !self.is_guarded(obj_reg) {
1098                        self.push_op(TraceOp::Guard {
1099                            register: obj_reg,
1100                            expected_type: ty,
1101                        });
1102                        self.mark_guarded(obj_reg);
1103                    }
1104                }
1105
1106                for i in 0..arg_count {
1107                    let arg_reg = first_arg + i;
1108                    if let Some(ty) = Self::get_value_type(&registers[arg_reg as usize]) {
1109                        if !self.is_guarded(arg_reg) {
1110                            self.push_op(TraceOp::Guard {
1111                                register: arg_reg,
1112                                expected_type: ty,
1113                            });
1114                            self.mark_guarded(arg_reg);
1115                        }
1116                    }
1117                }
1118
1119                self.push_op(TraceOp::CallMethod {
1120                    dest: dest_reg,
1121                    object: obj_reg,
1122                    method_name,
1123                    first_arg,
1124                    arg_count,
1125                });
1126                Ok(())
1127            }
1128
1129            Instruction::GetField(dest, obj_reg, field_name_idx) => {
1130                let field_name = function.chunk.constants[field_name_idx as usize]
1131                    .as_string()
1132                    .unwrap_or("unknown")
1133                    .to_string();
1134                let (field_index, is_weak_field) = match &registers[obj_reg as usize] {
1135                    Value::Struct { layout, .. } => {
1136                        let idx = layout.index_of_str(&field_name);
1137                        let is_weak = idx.map(|i| layout.is_weak(i)).unwrap_or(false);
1138                        (idx, is_weak)
1139                    }
1140
1141                    _ => (None, false),
1142                };
1143                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
1144                    if !self.is_guarded(obj_reg) {
1145                        self.push_op(TraceOp::Guard {
1146                            register: obj_reg,
1147                            expected_type: ty,
1148                        });
1149                        self.mark_guarded(obj_reg);
1150                    }
1151                }
1152
1153                let value_type = Self::get_value_type(&registers[dest as usize]);
1154                self.push_op(TraceOp::GetField {
1155                    dest,
1156                    object: obj_reg,
1157                    field_name,
1158                    field_index,
1159                    value_type,
1160                    is_weak: is_weak_field,
1161                });
1162                Ok(())
1163            }
1164
1165            Instruction::SetField(obj_reg, field_name_idx, value_reg) => {
1166                let field_name = function.chunk.constants[field_name_idx as usize]
1167                    .as_string()
1168                    .unwrap_or("unknown")
1169                    .to_string();
1170                let (field_index, is_weak_field) = match &registers[obj_reg as usize] {
1171                    Value::Struct { layout, .. } => {
1172                        let idx = layout.index_of_str(&field_name);
1173                        let is_weak = idx.map(|i| layout.is_weak(i)).unwrap_or(false);
1174                        (idx, is_weak)
1175                    }
1176
1177                    _ => (None, false),
1178                };
1179                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
1180                    if !self.is_guarded(obj_reg) {
1181                        self.push_op(TraceOp::Guard {
1182                            register: obj_reg,
1183                            expected_type: ty,
1184                        });
1185                        self.mark_guarded(obj_reg);
1186                    }
1187                }
1188
1189                let value_type = Self::get_value_type(&registers[value_reg as usize]);
1190                if let Some(ty) = value_type {
1191                    if !self.is_guarded(value_reg) {
1192                        self.push_op(TraceOp::Guard {
1193                            register: value_reg,
1194                            expected_type: ty,
1195                        });
1196                        self.mark_guarded(value_reg);
1197                    }
1198                }
1199
1200                self.rebox_specialized_register(value_reg, "SetField");
1201
1202                self.push_op(TraceOp::SetField {
1203                    object: obj_reg,
1204                    field_name,
1205                    value: value_reg,
1206                    field_index,
1207                    value_type,
1208                    is_weak: is_weak_field,
1209                });
1210                Ok(())
1211            }
1212
1213            Instruction::NewStruct(
1214                dest,
1215                struct_name_idx,
1216                first_field_name_idx,
1217                first_field_reg,
1218                field_count,
1219            ) => {
1220                let struct_name = function.chunk.constants[struct_name_idx as usize]
1221                    .as_string()
1222                    .unwrap_or("unknown")
1223                    .to_string();
1224                let mut field_names = Vec::new();
1225                for i in 0..field_count {
1226                    let field_name_idx = first_field_name_idx + (i as u16);
1227                    let field_name = function.chunk.constants[field_name_idx as usize]
1228                        .as_string()
1229                        .unwrap_or("unknown")
1230                        .to_string();
1231                    field_names.push(field_name);
1232                }
1233
1234                let mut field_registers = Vec::new();
1235                for i in 0..field_count {
1236                    let field_reg = first_field_reg + i;
1237                    field_registers.push(field_reg);
1238                    if let Some(ty) = Self::get_value_type(&registers[field_reg as usize]) {
1239                        if !self.is_guarded(field_reg) {
1240                            self.push_op(TraceOp::Guard {
1241                                register: field_reg,
1242                                expected_type: ty,
1243                            });
1244                            self.mark_guarded(field_reg);
1245                        }
1246                    }
1247                }
1248
1249                for &field_reg in &field_registers {
1250                    self.rebox_specialized_register(field_reg, "struct literal field");
1251                }
1252
1253                self.push_op(TraceOp::NewStruct {
1254                    dest,
1255                    struct_name,
1256                    field_names,
1257                    field_registers,
1258                });
1259                Ok(())
1260            }
1261
1262            Instruction::NewEnumUnit(dest, enum_name_idx, variant_idx) => {
1263                let enum_name = function.chunk.constants[enum_name_idx as usize]
1264                    .as_string()
1265                    .unwrap_or("unknown")
1266                    .to_string();
1267                let variant_name = function.chunk.constants[variant_idx as usize]
1268                    .as_string()
1269                    .unwrap_or("unknown")
1270                    .to_string();
1271                self.push_op(TraceOp::NewEnumUnit {
1272                    dest,
1273                    enum_name,
1274                    variant_name,
1275                });
1276                Ok(())
1277            }
1278
1279            Instruction::NewEnumVariant(
1280                dest,
1281                enum_name_idx,
1282                variant_idx,
1283                first_value,
1284                value_count,
1285            ) => {
1286                let enum_name = function.chunk.constants[enum_name_idx as usize]
1287                    .as_string()
1288                    .unwrap_or("unknown")
1289                    .to_string();
1290                let variant_name = function.chunk.constants[variant_idx as usize]
1291                    .as_string()
1292                    .unwrap_or("unknown")
1293                    .to_string();
1294                let mut value_registers = Vec::new();
1295                for i in 0..value_count {
1296                    value_registers.push(first_value + i);
1297                }
1298
1299                for &value_reg in &value_registers {
1300                    self.rebox_specialized_register(value_reg, "enum variant value");
1301                }
1302
1303                self.push_op(TraceOp::NewEnumVariant {
1304                    dest,
1305                    enum_name,
1306                    variant_name,
1307                    value_registers,
1308                });
1309                Ok(())
1310            }
1311
1312            Instruction::IsEnumVariant(dest, value_reg, enum_name_idx, variant_idx) => {
1313                let enum_name = function.chunk.constants[enum_name_idx as usize]
1314                    .as_string()
1315                    .unwrap_or("unknown")
1316                    .to_string();
1317                let variant_name = function.chunk.constants[variant_idx as usize]
1318                    .as_string()
1319                    .unwrap_or("unknown")
1320                    .to_string();
1321                self.push_op(TraceOp::IsEnumVariant {
1322                    dest,
1323                    value: value_reg,
1324                    enum_name,
1325                    variant_name,
1326                });
1327                Ok(())
1328            }
1329
1330            Instruction::GetEnumValue(dest, enum_reg, index) => {
1331                self.push_op(TraceOp::GetEnumValue {
1332                    dest,
1333                    enum_reg,
1334                    index,
1335                });
1336                Ok(())
1337            }
1338
1339            Instruction::Call(func_reg, first_arg, arg_count, dest_reg) => {
1340                // Rebox specialized value if dest_reg contains one
1341                self.remove_specialization_tracking(dest_reg);
1342
1343                match &registers[func_reg as usize] {
1344                    Value::NativeFunction(native_fn) => {
1345                        let traced = TracedNativeFn::new(native_fn.clone());
1346                        if !self.is_guarded(func_reg) {
1347                            self.push_op(TraceOp::GuardNativeFunction {
1348                                register: func_reg,
1349                                function: traced.clone(),
1350                            });
1351                            self.mark_guarded(func_reg);
1352                        }
1353
1354                        self.push_op(TraceOp::CallNative {
1355                            dest: dest_reg,
1356                            callee: func_reg,
1357                            function: traced,
1358                            first_arg,
1359                            arg_count,
1360                        });
1361                        Ok(())
1362                    }
1363
1364                    Value::Function(function_idx) => {
1365                        if !self.is_guarded(func_reg) {
1366                            self.push_op(TraceOp::GuardFunction {
1367                                register: func_reg,
1368                                function_idx: *function_idx,
1369                            });
1370                            self.mark_guarded(func_reg);
1371                        }
1372
1373                        let mut did_inline = false;
1374                        if let Some(callee_fn) = functions.get(*function_idx) {
1375                            if self.should_inline(*function_idx, callee_fn)
1376                                && (arg_count as usize) <= callee_fn.register_count as usize
1377                            {
1378                                let mut arg_registers = Vec::with_capacity(arg_count as usize);
1379                                for i in 0..arg_count {
1380                                    arg_registers.push(first_arg + i);
1381                                }
1382                                self.push_inline_context(
1383                                    *function_idx,
1384                                    callee_fn.register_count,
1385                                    dest_reg,
1386                                    func_reg,
1387                                    first_arg,
1388                                    arg_count,
1389                                    arg_registers,
1390                                    false,
1391                                    None,
1392                                );
1393                                did_inline = true;
1394                            }
1395                        }
1396
1397                        if !did_inline {
1398                            self.push_op(TraceOp::CallFunction {
1399                                dest: dest_reg,
1400                                callee: func_reg,
1401                                function_idx: *function_idx,
1402                                first_arg,
1403                                arg_count,
1404                                is_closure: false,
1405                                upvalues_ptr: None,
1406                            });
1407                        }
1408
1409                        Ok(())
1410                    }
1411
1412                    Value::Closure {
1413                        function_idx,
1414                        upvalues,
1415                    } => {
1416                        let upvalues_ptr = Rc::as_ptr(upvalues) as *const ();
1417                        if !self.is_guarded(func_reg) {
1418                            self.push_op(TraceOp::GuardClosure {
1419                                register: func_reg,
1420                                function_idx: *function_idx,
1421                                upvalues_ptr,
1422                            });
1423                            self.mark_guarded(func_reg);
1424                        }
1425
1426                        let mut did_inline = false;
1427                        if let Some(callee_fn) = functions.get(*function_idx) {
1428                            if self.should_inline(*function_idx, callee_fn)
1429                                && (arg_count as usize) <= callee_fn.register_count as usize
1430                            {
1431                                let mut arg_registers = Vec::with_capacity(arg_count as usize);
1432                                for i in 0..arg_count {
1433                                    arg_registers.push(first_arg + i);
1434                                }
1435                                self.push_inline_context(
1436                                    *function_idx,
1437                                    callee_fn.register_count,
1438                                    dest_reg,
1439                                    func_reg,
1440                                    first_arg,
1441                                    arg_count,
1442                                    arg_registers,
1443                                    true,
1444                                    Some(upvalues_ptr),
1445                                );
1446                                did_inline = true;
1447                            }
1448                        }
1449
1450                        if !did_inline {
1451                            self.push_op(TraceOp::CallFunction {
1452                                dest: dest_reg,
1453                                callee: func_reg,
1454                                function_idx: *function_idx,
1455                                first_arg,
1456                                arg_count,
1457                                is_closure: true,
1458                                upvalues_ptr: Some(upvalues_ptr),
1459                            });
1460                        }
1461
1462                        Ok(())
1463                    }
1464
1465                    _ => {
1466                        self.stop_recording();
1467                        crate::jit::log(|| {
1468                            format!(
1469                                "Trace aborted: unsupported call operation on register {} (value {:?})",
1470                                func_reg,
1471                                registers[func_reg as usize].tag()
1472                            )
1473                        });
1474                        Err(LustError::RuntimeError {
1475                            message: "Trace aborted: unsupported call operation".to_string(),
1476                        })
1477                    }
1478                }
1479            }
1480
1481            Instruction::NewArray(dest, first_elem, count) => {
1482                // Rebox specialized value if dest contains one
1483                self.remove_specialization_tracking(dest);
1484
1485                // Disable specialization inside inlined functions for now to avoid
1486                // register aliasing issues between inline frames and the parent trace.
1487                if !self.inline_stack.is_empty() {
1488                    self.push_op(TraceOp::NewArray {
1489                        dest,
1490                        first_element: first_elem,
1491                        count,
1492                    });
1493                    return Ok(());
1494                }
1495
1496                // Check if we can specialize this array based on element type
1497                // register_types contains the element type, not the array type
1498                let can_specialize = if let Some(element_type) = function.register_types.get(&dest)
1499                {
1500                    // Build the full Array type from the element type
1501                    use crate::ast::{Span, Type};
1502                    let array_type = crate::ast::TypeKind::Array(Box::new(Type::new(
1503                        element_type.clone(),
1504                        Span::dummy(),
1505                    )));
1506
1507                    // Check if this array type is specializable
1508                    self.specialization_registry
1509                        .get_specialization(&array_type)
1510                        .is_some()
1511                } else {
1512                    false
1513                };
1514
1515                if can_specialize {
1516                    let element_type = function.register_types.get(&dest).unwrap().clone();
1517
1518                    // Build the full Array type
1519                    use crate::ast::{Span, Type};
1520                    let array_type = crate::ast::TypeKind::Array(Box::new(Type::new(
1521                        element_type.clone(),
1522                        Span::dummy(),
1523                    )));
1524
1525                    let layout = self
1526                        .specialization_registry
1527                        .get_specialization(&array_type)
1528                        .unwrap();
1529
1530                    crate::jit::log(|| {
1531                        format!(
1532                            "🔬 JIT: Specializing NewArray for reg {} with element type {:?}",
1533                            dest, element_type
1534                        )
1535                    });
1536
1537                    // First create the array normally (for initial values if any)
1538                    if count > 0 {
1539                        self.push_op(TraceOp::NewArray {
1540                            dest,
1541                            first_element: first_elem,
1542                            count,
1543                        });
1544                    } else {
1545                        // Empty array
1546                        self.push_op(TraceOp::NewArray {
1547                            dest,
1548                            first_element: 0,
1549                            count: 0,
1550                        });
1551                    }
1552
1553                    // Then unbox it for specialized operations
1554                    let specialized_id = self.next_specialized_id;
1555                    self.next_specialized_id += 1;
1556
1557                    self.push_op(TraceOp::Unbox {
1558                        specialized_id,
1559                        source_reg: dest,
1560                        layout: layout.clone(),
1561                    });
1562
1563                    // Track that this register now contains a specialized value
1564                    self.specialized_registers
1565                        .insert(dest, (specialized_id, layout));
1566                } else {
1567                    // Normal non-specialized array
1568                    self.push_op(TraceOp::NewArray {
1569                        dest,
1570                        first_element: first_elem,
1571                        count,
1572                    });
1573                }
1574                Ok(())
1575            }
1576
1577            Instruction::NewMap(_) | Instruction::SetIndex(_, _, _) => {
1578                self.stop_recording();
1579                Err(LustError::RuntimeError {
1580                    message: "Trace aborted: unsupported index operation".to_string(),
1581                })
1582            }
1583
1584            Instruction::Return(value_reg) => {
1585                let return_reg = if value_reg == 255 {
1586                    None
1587                } else {
1588                    Some(value_reg)
1589                };
1590
1591                // Rebox any specialized values before return
1592                if let Some(reg) = return_reg {
1593                    if let Some(&(specialized_id, ref layout)) =
1594                        self.specialized_registers.get(&reg)
1595                    {
1596                        crate::jit::log(|| {
1597                            format!(
1598                                "📦 JIT: Reboxing specialized #{} in reg {} before return",
1599                                specialized_id, reg
1600                            )
1601                        });
1602
1603                        self.push_op(TraceOp::Rebox {
1604                            dest_reg: reg,
1605                            specialized_id,
1606                            layout: layout.clone(),
1607                        });
1608
1609                        self.specialized_registers.remove(&reg);
1610                    }
1611                }
1612
1613                // Ensure no specialized values leak past the return
1614                self.rebox_all_specialized_values();
1615
1616                if let Some(ctx) = self.inline_stack.last_mut() {
1617                    ctx.return_register = return_reg;
1618                    crate::jit::log(|| {
1619                        format!(
1620                            "🔧 JIT: Inline return detected, return_reg={:?}",
1621                            return_reg
1622                        )
1623                    });
1624                    if let Some(inline_op) = self.finalize_inline_context() {
1625                        self.push_op(inline_op);
1626                    }
1627                    Ok(())
1628                } else if function_idx == self.trace.function_idx {
1629                    self.stop_recording();
1630                    Ok(())
1631                } else {
1632                    self.push_op(TraceOp::Return { value: return_reg });
1633                    Ok(())
1634                }
1635            }
1636
1637            Instruction::Jump(offset) => {
1638                if offset < 0 {
1639                    let target_calc = (current_ip as isize) + (offset as isize);
1640                    if target_calc < 0 {
1641                        self.stop_recording();
1642                        Err(LustError::RuntimeError {
1643                            message: format!(
1644                                "Invalid jump target: offset={}, current_ip={}, target={}",
1645                                offset, current_ip, target_calc
1646                            ),
1647                        })
1648                    } else {
1649                        let jump_target = target_calc as usize;
1650                        let loop_key = (function_idx, jump_target);
1651
1652                        // Track how many times we've seen this loop backedge
1653                        let iteration_count = self.loop_iterations.entry(loop_key).or_insert(0);
1654                        *iteration_count += 1;
1655
1656                        if function_idx == self.trace.function_idx
1657                            && jump_target == self.trace.start_ip
1658                        {
1659                            // This is our main trace loop closing - check if we should unroll more
1660                            if *iteration_count < crate::jit::LOOP_UNROLL_COUNT {
1661                                crate::jit::log(|| {
1662                                    format!(
1663                                        "🔄 JIT: Unrolling main loop (iteration {}/{})",
1664                                        iteration_count,
1665                                        crate::jit::LOOP_UNROLL_COUNT
1666                                    )
1667                                });
1668                                // Continue recording to unroll the loop
1669                                Ok(())
1670                            } else {
1671                                crate::jit::log(|| {
1672                                    format!(
1673                                        "✅ JIT: Loop unrolled {} times, stopping trace",
1674                                        iteration_count
1675                                    )
1676                                });
1677                                self.stop_recording();
1678                                Ok(())
1679                            }
1680                        } else {
1681                            // This is a nested loop that should be compiled as a separate trace
1682                            // Following LuaJIT's approach: don't inline loops, compile them separately
1683                            let bailout_ip = current_ip.saturating_sub(1);
1684
1685                            crate::jit::log(|| {
1686                                format!(
1687                                    "🔄 JIT: Nested loop detected at func {} ip {} - will call as separate trace",
1688                                    function_idx, jump_target
1689                                )
1690                            });
1691
1692                            // Rebox all specialized values before calling nested trace
1693                            self.rebox_all_specialized_values();
1694
1695                            // Emit NestedLoopCall which will eventually call the compiled inner trace
1696                            self.push_op(TraceOp::NestedLoopCall {
1697                                function_idx,
1698                                loop_start_ip: jump_target,
1699                                bailout_ip,
1700                            });
1701                            Ok(())
1702                        }
1703                    }
1704                } else {
1705                    Ok(())
1706                }
1707            }
1708
1709            Instruction::JumpIf(cond, offset) => {
1710                let condition = &registers[cond as usize];
1711                let is_truthy = condition.is_truthy();
1712                let target_offset = (current_ip as isize) + (offset as isize);
1713                let target = if target_offset < 0 {
1714                    0
1715                } else {
1716                    target_offset as usize
1717                };
1718                let bailout_ip = if is_truthy { current_ip } else { target };
1719                self.push_op(TraceOp::GuardLoopContinue {
1720                    condition_register: cond,
1721                    expect_truthy: is_truthy,
1722                    bailout_ip,
1723                });
1724                Ok(())
1725            }
1726
1727            Instruction::JumpIfNot(cond, offset) => {
1728                let condition = &registers[cond as usize];
1729                let is_truthy = condition.is_truthy();
1730                let target_offset = (current_ip as isize) + (offset as isize);
1731                let target = if target_offset < 0 {
1732                    0
1733                } else {
1734                    target_offset as usize
1735                };
1736                let bailout_ip = if !is_truthy { current_ip } else { target };
1737                self.push_op(TraceOp::GuardLoopContinue {
1738                    condition_register: cond,
1739                    expect_truthy: is_truthy,
1740                    bailout_ip,
1741                });
1742                Ok(())
1743            }
1744
1745            _ => {
1746                self.stop_recording();
1747                crate::jit::log(|| {
1748                    format!(
1749                        "Trace aborted: unsupported instruction {:?}",
1750                        instruction.opcode()
1751                    )
1752                });
1753                Err(LustError::RuntimeError {
1754                    message: "Trace aborted: unsupported instruction".to_string(),
1755                })
1756            }
1757        };
1758
1759        outcome?;
1760
1761        if self.op_count >= self.max_length {
1762            self.stop_recording();
1763            return Err(LustError::RuntimeError {
1764                message: "Trace too long".to_string(),
1765            });
1766        }
1767
1768        Ok(())
1769    }
1770    fn add_type_guards(
1771        &mut self,
1772        lhs: Register,
1773        rhs: Register,
1774        registers: &[Value],
1775        function: &crate::bytecode::Function,
1776    ) -> Result<(), LustError> {
1777        if let Some(ty) = Self::get_value_type(&registers[lhs as usize]) {
1778            let needs_guard = if self.is_guarded(lhs) {
1779                false
1780            } else if let Some(static_type) = function.register_types.get(&lhs) {
1781                !Self::type_kind_matches_value_type(static_type, ty)
1782            } else {
1783                true
1784            };
1785            if needs_guard {
1786                self.push_op(TraceOp::Guard {
1787                    register: lhs,
1788                    expected_type: ty,
1789                });
1790                self.mark_guarded(lhs);
1791            } else {
1792                self.mark_guarded(lhs);
1793            }
1794        }
1795
1796        if let Some(ty) = Self::get_value_type(&registers[rhs as usize]) {
1797            let needs_guard = if self.is_guarded(rhs) {
1798                false
1799            } else if let Some(static_type) = function.register_types.get(&rhs) {
1800                !Self::type_kind_matches_value_type(static_type, ty)
1801            } else {
1802                true
1803            };
1804            if needs_guard {
1805                self.push_op(TraceOp::Guard {
1806                    register: rhs,
1807                    expected_type: ty,
1808                });
1809                self.mark_guarded(rhs);
1810            } else {
1811                self.mark_guarded(rhs);
1812            }
1813        }
1814
1815        Ok(())
1816    }
1817
1818    fn type_kind_matches_value_type(
1819        type_kind: &crate::ast::TypeKind,
1820        value_type: ValueType,
1821    ) -> bool {
1822        use crate::ast::TypeKind;
1823        match (type_kind, value_type) {
1824            (TypeKind::Int, ValueType::Int) => true,
1825            (TypeKind::Float, ValueType::Float) => true,
1826            (TypeKind::Bool, ValueType::Bool) => true,
1827            (TypeKind::String, ValueType::String) => true,
1828            (TypeKind::Array(_), ValueType::Array) => true,
1829            (TypeKind::Tuple(_), ValueType::Tuple) => true,
1830            _ => false,
1831        }
1832    }
1833
1834    fn get_value_type(value: &Value) -> Option<ValueType> {
1835        match value {
1836            Value::Int(_) => Some(ValueType::Int),
1837            Value::Float(_) => Some(ValueType::Float),
1838            Value::Bool(_) => Some(ValueType::Bool),
1839            Value::String(_) => Some(ValueType::String),
1840            Value::Array(_) => Some(ValueType::Array),
1841            Value::Tuple(_) => Some(ValueType::Tuple),
1842            Value::Struct { .. } => Some(ValueType::Struct),
1843            _ => None,
1844        }
1845    }
1846
1847    pub fn finish(mut self) -> Trace {
1848        #[cfg(feature = "std")]
1849        if std::env::var("LUST_TRACE_DEBUG").is_ok() {
1850            eprintln!(
1851                "🧵 Trace dump (func {}, start_ip {}):",
1852                self.trace.function_idx, self.trace.start_ip
1853            );
1854            if !self.trace.preamble.is_empty() {
1855                eprintln!("  Preamble:");
1856                for (idx, op) in self.trace.preamble.iter().enumerate() {
1857                    eprintln!("    {:03}: {:?}", idx, op);
1858                }
1859            }
1860            eprintln!("  Body:");
1861            for (idx, op) in self.trace.ops.iter().enumerate() {
1862                eprintln!("    {:03}: {:?}", idx, op);
1863            }
1864            if !self.trace.postamble.is_empty() {
1865                eprintln!("  Postamble:");
1866                for (idx, op) in self.trace.postamble.iter().enumerate() {
1867                    eprintln!("    {:03}: {:?}", idx, op);
1868                }
1869            }
1870        }
1871
1872        // Finalize before returning (add rebox ops to postamble)
1873        self.finalize_trace();
1874        self.trace
1875    }
1876
1877    pub fn is_recording(&self) -> bool {
1878        self.recording
1879    }
1880
1881    pub fn abort(&mut self) {
1882        self.stop_recording();
1883    }
1884}
lust/jit/trace.rs

lust/jit/
trace.rs
