lust/jit/
trace.rs

1use crate::bytecode::value::NativeFn;
2use crate::bytecode::Instruction;
3use crate::bytecode::{Register, Value};
4use crate::LustError;
5use alloc::{
6    format,
7    rc::Rc,
8    string::{String, ToString},
9    vec::Vec,
10};
11use core::fmt;
12use hashbrown::HashSet;
13
14#[derive(Clone)]
15pub struct TracedNativeFn {
16    function: NativeFn,
17}
18
19impl TracedNativeFn {
20    pub fn new(function: NativeFn) -> Self {
21        Self { function }
22    }
23
24    pub fn pointer(&self) -> *const () {
25        Rc::as_ptr(&self.function) as *const ()
26    }
27}
28
29impl fmt::Debug for TracedNativeFn {
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        write!(f, "NativeFn({:p})", Rc::as_ptr(&self.function))
32    }
33}
34
35#[derive(Debug, Clone)]
36pub struct Trace {
37    pub function_idx: usize,
38    pub start_ip: usize,
39    pub ops: Vec<TraceOp>,
40    pub inputs: Vec<Register>,
41    pub outputs: Vec<Register>,
42}
43
44#[derive(Debug, Clone)]
45pub struct InlineTrace {
46    pub function_idx: usize,
47    pub register_count: u8,
48    pub first_arg: Register,
49    pub arg_count: u8,
50    pub arg_registers: Vec<Register>,
51    pub body: Vec<TraceOp>,
52    pub return_register: Option<Register>,
53    pub is_closure: bool,
54    pub upvalues_ptr: Option<*const ()>,
55}
56
57#[derive(Debug, Clone)]
58pub enum TraceOp {
59    LoadConst {
60        dest: Register,
61        value: Value,
62    },
63    Move {
64        dest: Register,
65        src: Register,
66    },
67    Add {
68        dest: Register,
69        lhs: Register,
70        rhs: Register,
71        lhs_type: ValueType,
72        rhs_type: ValueType,
73    },
74    Sub {
75        dest: Register,
76        lhs: Register,
77        rhs: Register,
78        lhs_type: ValueType,
79        rhs_type: ValueType,
80    },
81    Mul {
82        dest: Register,
83        lhs: Register,
84        rhs: Register,
85        lhs_type: ValueType,
86        rhs_type: ValueType,
87    },
88    Div {
89        dest: Register,
90        lhs: Register,
91        rhs: Register,
92        lhs_type: ValueType,
93        rhs_type: ValueType,
94    },
95    Mod {
96        dest: Register,
97        lhs: Register,
98        rhs: Register,
99        lhs_type: ValueType,
100        rhs_type: ValueType,
101    },
102    Neg {
103        dest: Register,
104        src: Register,
105    },
106    Eq {
107        dest: Register,
108        lhs: Register,
109        rhs: Register,
110    },
111    Ne {
112        dest: Register,
113        lhs: Register,
114        rhs: Register,
115    },
116    Lt {
117        dest: Register,
118        lhs: Register,
119        rhs: Register,
120    },
121    Le {
122        dest: Register,
123        lhs: Register,
124        rhs: Register,
125    },
126    Gt {
127        dest: Register,
128        lhs: Register,
129        rhs: Register,
130    },
131    Ge {
132        dest: Register,
133        lhs: Register,
134        rhs: Register,
135    },
136    And {
137        dest: Register,
138        lhs: Register,
139        rhs: Register,
140    },
141    Or {
142        dest: Register,
143        lhs: Register,
144        rhs: Register,
145    },
146    Not {
147        dest: Register,
148        src: Register,
149    },
150    Concat {
151        dest: Register,
152        lhs: Register,
153        rhs: Register,
154    },
155    GetIndex {
156        dest: Register,
157        array: Register,
158        index: Register,
159    },
160    ArrayLen {
161        dest: Register,
162        array: Register,
163    },
164    GuardNativeFunction {
165        register: Register,
166        function: TracedNativeFn,
167    },
168    GuardFunction {
169        register: Register,
170        function_idx: usize,
171    },
172    GuardClosure {
173        register: Register,
174        function_idx: usize,
175        upvalues_ptr: *const (),
176    },
177    CallNative {
178        dest: Register,
179        callee: Register,
180        function: TracedNativeFn,
181        first_arg: Register,
182        arg_count: u8,
183    },
184    CallFunction {
185        dest: Register,
186        callee: Register,
187        function_idx: usize,
188        first_arg: Register,
189        arg_count: u8,
190        is_closure: bool,
191        upvalues_ptr: Option<*const ()>,
192    },
193    InlineCall {
194        dest: Register,
195        callee: Register,
196        trace: InlineTrace,
197    },
198    CallMethod {
199        dest: Register,
200        object: Register,
201        method_name: String,
202        first_arg: Register,
203        arg_count: u8,
204    },
205    GetField {
206        dest: Register,
207        object: Register,
208        field_name: String,
209        field_index: Option<usize>,
210        value_type: Option<ValueType>,
211        is_weak: bool,
212    },
213    SetField {
214        object: Register,
215        field_name: String,
216        value: Register,
217        field_index: Option<usize>,
218        value_type: Option<ValueType>,
219        is_weak: bool,
220    },
221    NewStruct {
222        dest: Register,
223        struct_name: String,
224        field_names: Vec<String>,
225        field_registers: Vec<Register>,
226    },
227    NewEnumUnit {
228        dest: Register,
229        enum_name: String,
230        variant_name: String,
231    },
232    NewEnumVariant {
233        dest: Register,
234        enum_name: String,
235        variant_name: String,
236        value_registers: Vec<Register>,
237    },
238    IsEnumVariant {
239        dest: Register,
240        value: Register,
241        enum_name: String,
242        variant_name: String,
243    },
244    GetEnumValue {
245        dest: Register,
246        enum_reg: Register,
247        index: u8,
248    },
249    Guard {
250        register: Register,
251        expected_type: ValueType,
252    },
253    GuardLoopContinue {
254        condition_register: Register,
255        expect_truthy: bool,
256        bailout_ip: usize,
257    },
258    NestedLoopCall {
259        function_idx: usize,
260        loop_start_ip: usize,
261        bailout_ip: usize,
262    },
263    Return {
264        value: Option<Register>,
265    },
266}
267
268#[derive(Debug, Clone, Copy, PartialEq, Eq)]
269pub enum ValueType {
270    Int,
271    Float,
272    Bool,
273    String,
274    Array,
275    Tuple,
276    Struct,
277}
278
279pub struct TraceRecorder {
280    pub trace: Trace,
281    max_length: usize,
282    recording: bool,
283    guarded_registers: HashSet<Register>,
284    inline_stack: Vec<InlineContext>,
285    op_count: usize,
286}
287
288#[derive(Debug, Clone)]
289struct InlineContext {
290    function_idx: usize,
291    register_count: u8,
292    dest: Register,
293    callee_reg: Register,
294    first_arg: Register,
295    arg_count: u8,
296    arg_registers: Vec<Register>,
297    ops: Vec<TraceOp>,
298    guarded_registers: HashSet<Register>,
299    return_register: Option<Register>,
300    is_closure: bool,
301    upvalues_ptr: Option<*const ()>,
302}
303
304impl TraceRecorder {
305    pub fn new(function_idx: usize, start_ip: usize, max_length: usize) -> Self {
306        Self {
307            trace: Trace {
308                function_idx,
309                start_ip,
310                ops: Vec::new(),
311                inputs: Vec::new(),
312                outputs: Vec::new(),
313            },
314            max_length,
315            recording: true,
316            guarded_registers: HashSet::new(),
317            inline_stack: Vec::new(),
318            op_count: 0,
319        }
320    }
321
322    fn current_function_idx(&self) -> usize {
323        self.inline_stack
324            .last()
325            .map(|ctx| ctx.function_idx)
326            .unwrap_or(self.trace.function_idx)
327    }
328
329    fn current_guard_set(&self) -> &HashSet<Register> {
330        self.inline_stack
331            .last()
332            .map(|ctx| &ctx.guarded_registers)
333            .unwrap_or(&self.guarded_registers)
334    }
335
336    fn current_guard_set_mut(&mut self) -> &mut HashSet<Register> {
337        self.inline_stack
338            .last_mut()
339            .map(|ctx| &mut ctx.guarded_registers)
340            .unwrap_or(&mut self.guarded_registers)
341    }
342
343    fn is_guarded(&self, register: Register) -> bool {
344        self.current_guard_set().contains(&register)
345    }
346
347    fn mark_guarded(&mut self, register: Register) {
348        let set = self.current_guard_set_mut();
349        set.insert(register);
350    }
351
352    fn push_op(&mut self, op: TraceOp) {
353        self.op_count += 1;
354        if let Some(ctx) = self.inline_stack.last_mut() {
355            ctx.ops.push(op);
356        } else {
357            self.trace.ops.push(op);
358        }
359    }
360
361    fn should_inline(&self, function_idx: usize) -> bool {
362        // return false;
363        if function_idx == self.trace.function_idx {
364            return false;
365        }
366
367        if self
368            .inline_stack
369            .iter()
370            .any(|ctx| ctx.function_idx == function_idx)
371        {
372            return false;
373        }
374
375        true
376    }
377
378    fn push_inline_context(
379        &mut self,
380        function_idx: usize,
381        register_count: u8,
382        dest: Register,
383        callee_reg: Register,
384        first_arg: Register,
385        arg_count: u8,
386        arg_registers: Vec<Register>,
387        is_closure: bool,
388        upvalues_ptr: Option<*const ()>,
389    ) {
390        self.inline_stack.push(InlineContext {
391            function_idx,
392            register_count,
393            dest,
394            callee_reg,
395            first_arg,
396            arg_count,
397            arg_registers,
398            ops: Vec::new(),
399            guarded_registers: HashSet::new(),
400            return_register: None,
401            is_closure,
402            upvalues_ptr,
403        });
404    }
405
406    fn finalize_inline_context(&mut self) -> Option<TraceOp> {
407        let context = self.inline_stack.pop()?;
408        let trace = InlineTrace {
409            function_idx: context.function_idx,
410            register_count: context.register_count,
411            first_arg: context.first_arg,
412            arg_count: context.arg_count,
413            arg_registers: context.arg_registers,
414            body: context.ops,
415            return_register: context.return_register,
416            is_closure: context.is_closure,
417            upvalues_ptr: context.upvalues_ptr,
418        };
419        Some(TraceOp::InlineCall {
420            dest: context.dest,
421            callee: context.callee_reg,
422            trace,
423        })
424    }
425
426    pub fn record_instruction(
427        &mut self,
428        instruction: Instruction,
429        current_ip: usize,
430        registers: &[Value; 256],
431        function: &crate::bytecode::Function,
432        function_idx: usize,
433        functions: &[crate::bytecode::Function],
434    ) -> Result<(), LustError> {
435        if !self.recording {
436            return Ok(());
437        }
438
439        if function_idx != self.current_function_idx() {
440            return Ok(());
441        }
442
443        let outcome: Result<(), LustError> = match instruction {
444            Instruction::LoadConst(dest, _) => {
445                if let Some(_ty) = Self::get_value_type(&registers[dest as usize]) {
446                    self.mark_guarded(dest);
447                }
448
449                self.push_op(TraceOp::LoadConst {
450                    dest,
451                    value: registers[dest as usize].clone(),
452                });
453                Ok(())
454            }
455
456            Instruction::LoadGlobal(dest, _) => {
457                if let Some(_ty) = Self::get_value_type(&registers[dest as usize]) {
458                    self.mark_guarded(dest);
459                }
460
461                self.push_op(TraceOp::LoadConst {
462                    dest,
463                    value: registers[dest as usize].clone(),
464                });
465                Ok(())
466            }
467
468            Instruction::StoreGlobal(_, _) => Ok(()),
469
470            Instruction::Move(dest, src) => {
471                self.push_op(TraceOp::Move { dest, src });
472                Ok(())
473            }
474
475            Instruction::Add(dest, lhs, rhs) => {
476                self.add_type_guards(lhs, rhs, registers, function)?;
477                let lhs_type =
478                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
479                let rhs_type =
480                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
481                self.push_op(TraceOp::Add {
482                    dest,
483                    lhs,
484                    rhs,
485                    lhs_type,
486                    rhs_type,
487                });
488                Ok(())
489            }
490
491            Instruction::Sub(dest, lhs, rhs) => {
492                self.add_type_guards(lhs, rhs, registers, function)?;
493                let lhs_type =
494                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
495                let rhs_type =
496                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
497                self.push_op(TraceOp::Sub {
498                    dest,
499                    lhs,
500                    rhs,
501                    lhs_type,
502                    rhs_type,
503                });
504                Ok(())
505            }
506
507            Instruction::Mul(dest, lhs, rhs) => {
508                self.add_type_guards(lhs, rhs, registers, function)?;
509                let lhs_type =
510                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
511                let rhs_type =
512                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
513                self.push_op(TraceOp::Mul {
514                    dest,
515                    lhs,
516                    rhs,
517                    lhs_type,
518                    rhs_type,
519                });
520                Ok(())
521            }
522
523            Instruction::Div(dest, lhs, rhs) => {
524                self.add_type_guards(lhs, rhs, registers, function)?;
525                let lhs_type =
526                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
527                let rhs_type =
528                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
529                self.push_op(TraceOp::Div {
530                    dest,
531                    lhs,
532                    rhs,
533                    lhs_type,
534                    rhs_type,
535                });
536                Ok(())
537            }
538
539            Instruction::Mod(dest, lhs, rhs) => {
540                self.add_type_guards(lhs, rhs, registers, function)?;
541                let lhs_type =
542                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
543                let rhs_type =
544                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
545                self.push_op(TraceOp::Mod {
546                    dest,
547                    lhs,
548                    rhs,
549                    lhs_type,
550                    rhs_type,
551                });
552                Ok(())
553            }
554
555            Instruction::Neg(dest, src) => {
556                self.push_op(TraceOp::Neg { dest, src });
557                Ok(())
558            }
559
560            Instruction::Eq(dest, lhs, rhs) => {
561                self.push_op(TraceOp::Eq { dest, lhs, rhs });
562                Ok(())
563            }
564
565            Instruction::Ne(dest, lhs, rhs) => {
566                self.push_op(TraceOp::Ne { dest, lhs, rhs });
567                Ok(())
568            }
569
570            Instruction::Lt(dest, lhs, rhs) => {
571                self.push_op(TraceOp::Lt { dest, lhs, rhs });
572                Ok(())
573            }
574
575            Instruction::Le(dest, lhs, rhs) => {
576                self.push_op(TraceOp::Le { dest, lhs, rhs });
577                Ok(())
578            }
579
580            Instruction::Gt(dest, lhs, rhs) => {
581                self.push_op(TraceOp::Gt { dest, lhs, rhs });
582                Ok(())
583            }
584
585            Instruction::Ge(dest, lhs, rhs) => {
586                self.push_op(TraceOp::Ge { dest, lhs, rhs });
587                Ok(())
588            }
589
590            Instruction::And(dest, lhs, rhs) => {
591                self.push_op(TraceOp::And { dest, lhs, rhs });
592                Ok(())
593            }
594
595            Instruction::Or(dest, lhs, rhs) => {
596                self.push_op(TraceOp::Or { dest, lhs, rhs });
597                Ok(())
598            }
599
600            Instruction::Not(dest, src) => {
601                self.push_op(TraceOp::Not { dest, src });
602                Ok(())
603            }
604
605            Instruction::Concat(dest, lhs, rhs) => {
606                if let Some(ty) = Self::get_value_type(&registers[lhs as usize]) {
607                    if !self.is_guarded(lhs) {
608                        self.push_op(TraceOp::Guard {
609                            register: lhs,
610                            expected_type: ty,
611                        });
612                        self.mark_guarded(lhs);
613                    }
614                }
615
616                if let Some(ty) = Self::get_value_type(&registers[rhs as usize]) {
617                    if !self.is_guarded(rhs) {
618                        self.push_op(TraceOp::Guard {
619                            register: rhs,
620                            expected_type: ty,
621                        });
622                        self.mark_guarded(rhs);
623                    }
624                }
625
626                self.push_op(TraceOp::Concat { dest, lhs, rhs });
627                Ok(())
628            }
629
630            Instruction::GetIndex(dest, array, index) => {
631                if let Some(ty) = Self::get_value_type(&registers[array as usize]) {
632                    if !self.is_guarded(array) {
633                        self.push_op(TraceOp::Guard {
634                            register: array,
635                            expected_type: ty,
636                        });
637                        self.mark_guarded(array);
638                    }
639                }
640
641                if let Some(ty) = Self::get_value_type(&registers[index as usize]) {
642                    if !self.is_guarded(index) {
643                        self.push_op(TraceOp::Guard {
644                            register: index,
645                            expected_type: ty,
646                        });
647                        self.mark_guarded(index);
648                    }
649                }
650
651                self.push_op(TraceOp::GetIndex { dest, array, index });
652                Ok(())
653            }
654
655            Instruction::ArrayLen(dest, array) => {
656                if let Some(ty) = Self::get_value_type(&registers[array as usize]) {
657                    if !self.is_guarded(array) {
658                        self.push_op(TraceOp::Guard {
659                            register: array,
660                            expected_type: ty,
661                        });
662                        self.mark_guarded(array);
663                    }
664                }
665
666                self.push_op(TraceOp::ArrayLen { dest, array });
667                Ok(())
668            }
669
670            Instruction::CallMethod(obj_reg, method_name_idx, first_arg, arg_count, dest_reg) => {
671                let method_name = function.chunk.constants[method_name_idx as usize]
672                    .as_string()
673                    .unwrap_or("unknown")
674                    .to_string();
675                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
676                    if !self.is_guarded(obj_reg) {
677                        self.push_op(TraceOp::Guard {
678                            register: obj_reg,
679                            expected_type: ty,
680                        });
681                        self.mark_guarded(obj_reg);
682                    }
683                }
684
685                for i in 0..arg_count {
686                    let arg_reg = first_arg + i;
687                    if let Some(ty) = Self::get_value_type(&registers[arg_reg as usize]) {
688                        if !self.is_guarded(arg_reg) {
689                            self.push_op(TraceOp::Guard {
690                                register: arg_reg,
691                                expected_type: ty,
692                            });
693                            self.mark_guarded(arg_reg);
694                        }
695                    }
696                }
697
698                self.push_op(TraceOp::CallMethod {
699                    dest: dest_reg,
700                    object: obj_reg,
701                    method_name,
702                    first_arg,
703                    arg_count,
704                });
705                Ok(())
706            }
707
708            Instruction::GetField(dest, obj_reg, field_name_idx) => {
709                let field_name = function.chunk.constants[field_name_idx as usize]
710                    .as_string()
711                    .unwrap_or("unknown")
712                    .to_string();
713                let (field_index, is_weak_field) = match &registers[obj_reg as usize] {
714                    Value::Struct { layout, .. } => {
715                        let idx = layout.index_of_str(&field_name);
716                        let is_weak = idx.map(|i| layout.is_weak(i)).unwrap_or(false);
717                        (idx, is_weak)
718                    }
719
720                    _ => (None, false),
721                };
722                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
723                    if !self.is_guarded(obj_reg) {
724                        self.push_op(TraceOp::Guard {
725                            register: obj_reg,
726                            expected_type: ty,
727                        });
728                        self.mark_guarded(obj_reg);
729                    }
730                }
731
732                let value_type = Self::get_value_type(&registers[dest as usize]);
733                self.push_op(TraceOp::GetField {
734                    dest,
735                    object: obj_reg,
736                    field_name,
737                    field_index,
738                    value_type,
739                    is_weak: is_weak_field,
740                });
741                Ok(())
742            }
743
744            Instruction::SetField(obj_reg, field_name_idx, value_reg) => {
745                let field_name = function.chunk.constants[field_name_idx as usize]
746                    .as_string()
747                    .unwrap_or("unknown")
748                    .to_string();
749                let (field_index, is_weak_field) = match &registers[obj_reg as usize] {
750                    Value::Struct { layout, .. } => {
751                        let idx = layout.index_of_str(&field_name);
752                        let is_weak = idx.map(|i| layout.is_weak(i)).unwrap_or(false);
753                        (idx, is_weak)
754                    }
755
756                    _ => (None, false),
757                };
758                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
759                    if !self.is_guarded(obj_reg) {
760                        self.push_op(TraceOp::Guard {
761                            register: obj_reg,
762                            expected_type: ty,
763                        });
764                        self.mark_guarded(obj_reg);
765                    }
766                }
767
768                let value_type = Self::get_value_type(&registers[value_reg as usize]);
769                if let Some(ty) = value_type {
770                    if !self.is_guarded(value_reg) {
771                        self.push_op(TraceOp::Guard {
772                            register: value_reg,
773                            expected_type: ty,
774                        });
775                        self.mark_guarded(value_reg);
776                    }
777                }
778
779                self.push_op(TraceOp::SetField {
780                    object: obj_reg,
781                    field_name,
782                    value: value_reg,
783                    field_index,
784                    value_type,
785                    is_weak: is_weak_field,
786                });
787                Ok(())
788            }
789
790            Instruction::NewStruct(
791                dest,
792                struct_name_idx,
793                first_field_name_idx,
794                first_field_reg,
795                field_count,
796            ) => {
797                let struct_name = function.chunk.constants[struct_name_idx as usize]
798                    .as_string()
799                    .unwrap_or("unknown")
800                    .to_string();
801                let mut field_names = Vec::new();
802                for i in 0..field_count {
803                    let field_name_idx = first_field_name_idx + (i as u16);
804                    let field_name = function.chunk.constants[field_name_idx as usize]
805                        .as_string()
806                        .unwrap_or("unknown")
807                        .to_string();
808                    field_names.push(field_name);
809                }
810
811                let mut field_registers = Vec::new();
812                for i in 0..field_count {
813                    let field_reg = first_field_reg + i;
814                    field_registers.push(field_reg);
815                    if let Some(ty) = Self::get_value_type(&registers[field_reg as usize]) {
816                        if !self.is_guarded(field_reg) {
817                            self.push_op(TraceOp::Guard {
818                                register: field_reg,
819                                expected_type: ty,
820                            });
821                            self.mark_guarded(field_reg);
822                        }
823                    }
824                }
825
826                self.push_op(TraceOp::NewStruct {
827                    dest,
828                    struct_name,
829                    field_names,
830                    field_registers,
831                });
832                Ok(())
833            }
834
835            Instruction::NewEnumUnit(dest, enum_name_idx, variant_idx) => {
836                let enum_name = function.chunk.constants[enum_name_idx as usize]
837                    .as_string()
838                    .unwrap_or("unknown")
839                    .to_string();
840                let variant_name = function.chunk.constants[variant_idx as usize]
841                    .as_string()
842                    .unwrap_or("unknown")
843                    .to_string();
844                self.push_op(TraceOp::NewEnumUnit {
845                    dest,
846                    enum_name,
847                    variant_name,
848                });
849                Ok(())
850            }
851
852            Instruction::NewEnumVariant(
853                dest,
854                enum_name_idx,
855                variant_idx,
856                first_value,
857                value_count,
858            ) => {
859                let enum_name = function.chunk.constants[enum_name_idx as usize]
860                    .as_string()
861                    .unwrap_or("unknown")
862                    .to_string();
863                let variant_name = function.chunk.constants[variant_idx as usize]
864                    .as_string()
865                    .unwrap_or("unknown")
866                    .to_string();
867                let mut value_registers = Vec::new();
868                for i in 0..value_count {
869                    value_registers.push(first_value + i);
870                }
871
872                self.push_op(TraceOp::NewEnumVariant {
873                    dest,
874                    enum_name,
875                    variant_name,
876                    value_registers,
877                });
878                Ok(())
879            }
880
881            Instruction::IsEnumVariant(dest, value_reg, enum_name_idx, variant_idx) => {
882                let enum_name = function.chunk.constants[enum_name_idx as usize]
883                    .as_string()
884                    .unwrap_or("unknown")
885                    .to_string();
886                let variant_name = function.chunk.constants[variant_idx as usize]
887                    .as_string()
888                    .unwrap_or("unknown")
889                    .to_string();
890                self.push_op(TraceOp::IsEnumVariant {
891                    dest,
892                    value: value_reg,
893                    enum_name,
894                    variant_name,
895                });
896                Ok(())
897            }
898
899            Instruction::GetEnumValue(dest, enum_reg, index) => {
900                self.push_op(TraceOp::GetEnumValue {
901                    dest,
902                    enum_reg,
903                    index,
904                });
905                Ok(())
906            }
907
908            Instruction::Call(func_reg, first_arg, arg_count, dest_reg) => {
909                match &registers[func_reg as usize] {
910                    Value::NativeFunction(native_fn) => {
911                        let traced = TracedNativeFn::new(native_fn.clone());
912                        if !self.is_guarded(func_reg) {
913                            self.push_op(TraceOp::GuardNativeFunction {
914                                register: func_reg,
915                                function: traced.clone(),
916                            });
917                            self.mark_guarded(func_reg);
918                        }
919
920                        self.push_op(TraceOp::CallNative {
921                            dest: dest_reg,
922                            callee: func_reg,
923                            function: traced,
924                            first_arg,
925                            arg_count,
926                        });
927                        Ok(())
928                    }
929
930                    Value::Function(function_idx) => {
931                        if !self.is_guarded(func_reg) {
932                            self.push_op(TraceOp::GuardFunction {
933                                register: func_reg,
934                                function_idx: *function_idx,
935                            });
936                            self.mark_guarded(func_reg);
937                        }
938
939                        let mut did_inline = false;
940                        if let Some(callee_fn) = functions.get(*function_idx) {
941                            if self.should_inline(*function_idx)
942                                && (arg_count as usize) <= callee_fn.register_count as usize
943                            {
944                                let mut arg_registers = Vec::with_capacity(arg_count as usize);
945                                for i in 0..arg_count {
946                                    arg_registers.push(first_arg + i);
947                                }
948                                self.push_inline_context(
949                                    *function_idx,
950                                    callee_fn.register_count,
951                                    dest_reg,
952                                    func_reg,
953                                    first_arg,
954                                    arg_count,
955                                    arg_registers,
956                                    false,
957                                    None,
958                                );
959                                did_inline = true;
960                            }
961                        }
962
963                        if !did_inline {
964                            self.push_op(TraceOp::CallFunction {
965                                dest: dest_reg,
966                                callee: func_reg,
967                                function_idx: *function_idx,
968                                first_arg,
969                                arg_count,
970                                is_closure: false,
971                                upvalues_ptr: None,
972                            });
973                        }
974
975                        Ok(())
976                    }
977
978                    Value::Closure {
979                        function_idx,
980                        upvalues,
981                    } => {
982                        let upvalues_ptr = Rc::as_ptr(upvalues) as *const ();
983                        if !self.is_guarded(func_reg) {
984                            self.push_op(TraceOp::GuardClosure {
985                                register: func_reg,
986                                function_idx: *function_idx,
987                                upvalues_ptr,
988                            });
989                            self.mark_guarded(func_reg);
990                        }
991
992                        let mut did_inline = false;
993                        if let Some(callee_fn) = functions.get(*function_idx) {
994                            if self.should_inline(*function_idx)
995                                && (arg_count as usize) <= callee_fn.register_count as usize
996                            {
997                                let mut arg_registers = Vec::with_capacity(arg_count as usize);
998                                for i in 0..arg_count {
999                                    arg_registers.push(first_arg + i);
1000                                }
1001                                self.push_inline_context(
1002                                    *function_idx,
1003                                    callee_fn.register_count,
1004                                    dest_reg,
1005                                    func_reg,
1006                                    first_arg,
1007                                    arg_count,
1008                                    arg_registers,
1009                                    true,
1010                                    Some(upvalues_ptr),
1011                                );
1012                                did_inline = true;
1013                            }
1014                        }
1015
1016                        if !did_inline {
1017                            self.push_op(TraceOp::CallFunction {
1018                                dest: dest_reg,
1019                                callee: func_reg,
1020                                function_idx: *function_idx,
1021                                first_arg,
1022                                arg_count,
1023                                is_closure: true,
1024                                upvalues_ptr: Some(upvalues_ptr),
1025                            });
1026                        }
1027
1028                        Ok(())
1029                    }
1030
1031                    _ => {
1032                        self.recording = false;
1033                        crate::jit::log(|| {
1034                            format!(
1035                                "Trace aborted: unsupported call operation on register {} (value {:?})",
1036                                func_reg,
1037                                registers[func_reg as usize].tag()
1038                            )
1039                        });
1040                        Err(LustError::RuntimeError {
1041                            message: "Trace aborted: unsupported call operation".to_string(),
1042                        })
1043                    }
1044                }
1045            }
1046
1047            Instruction::NewArray(_, _, _)
1048            | Instruction::NewMap(_)
1049            | Instruction::SetIndex(_, _, _) => {
1050                self.recording = false;
1051                Err(LustError::RuntimeError {
1052                    message: "Trace aborted: unsupported index operation".to_string(),
1053                })
1054            }
1055
1056            Instruction::Return(value_reg) => {
1057                let return_reg = if value_reg == 255 {
1058                    None
1059                } else {
1060                    Some(value_reg)
1061                };
1062
1063                if let Some(ctx) = self.inline_stack.last_mut() {
1064                    ctx.return_register = return_reg;
1065                    if let Some(inline_op) = self.finalize_inline_context() {
1066                        self.push_op(inline_op);
1067                    }
1068                    Ok(())
1069                } else if function_idx == self.trace.function_idx {
1070                    self.recording = false;
1071                    Ok(())
1072                } else {
1073                    self.push_op(TraceOp::Return { value: return_reg });
1074                    Ok(())
1075                }
1076            }
1077
1078            Instruction::Jump(offset) => {
1079                if offset < 0 {
1080                    let target_calc = (current_ip as isize) + (offset as isize);
1081                    if target_calc < 0 {
1082                        self.recording = false;
1083                        Err(LustError::RuntimeError {
1084                            message: format!(
1085                                "Invalid jump target: offset={}, current_ip={}, target={}",
1086                                offset, current_ip, target_calc
1087                            ),
1088                        })
1089                    } else {
1090                        let jump_target = target_calc as usize;
1091                        if function_idx == self.trace.function_idx
1092                            && jump_target == self.trace.start_ip
1093                        {
1094                            self.recording = false;
1095                            Ok(())
1096                        } else {
1097                            let bailout_ip = current_ip.saturating_sub(1);
1098                            self.push_op(TraceOp::NestedLoopCall {
1099                                function_idx,
1100                                loop_start_ip: jump_target,
1101                                bailout_ip,
1102                            });
1103                            Ok(())
1104                        }
1105                    }
1106                } else {
1107                    Ok(())
1108                }
1109            }
1110
1111            Instruction::JumpIf(cond, offset) => {
1112                let condition = &registers[cond as usize];
1113                let is_truthy = condition.is_truthy();
1114                let target_offset = (current_ip as isize) + (offset as isize);
1115                let target = if target_offset < 0 {
1116                    0
1117                } else {
1118                    target_offset as usize
1119                };
1120                let bailout_ip = if is_truthy { current_ip } else { target };
1121                self.push_op(TraceOp::GuardLoopContinue {
1122                    condition_register: cond,
1123                    expect_truthy: is_truthy,
1124                    bailout_ip,
1125                });
1126                Ok(())
1127            }
1128
1129            Instruction::JumpIfNot(cond, offset) => {
1130                let condition = &registers[cond as usize];
1131                let is_truthy = condition.is_truthy();
1132                let target_offset = (current_ip as isize) + (offset as isize);
1133                let target = if target_offset < 0 {
1134                    0
1135                } else {
1136                    target_offset as usize
1137                };
1138                let bailout_ip = if !is_truthy { current_ip } else { target };
1139                self.push_op(TraceOp::GuardLoopContinue {
1140                    condition_register: cond,
1141                    expect_truthy: is_truthy,
1142                    bailout_ip,
1143                });
1144                Ok(())
1145            }
1146
1147            _ => {
1148                self.recording = false;
1149                crate::jit::log(|| {
1150                    format!(
1151                        "Trace aborted: unsupported instruction {:?}",
1152                        instruction.opcode()
1153                    )
1154                });
1155                Err(LustError::RuntimeError {
1156                    message: "Trace aborted: unsupported instruction".to_string(),
1157                })
1158            }
1159        };
1160
1161        outcome?;
1162
1163        if self.op_count >= self.max_length {
1164            self.recording = false;
1165            return Err(LustError::RuntimeError {
1166                message: "Trace too long".to_string(),
1167            });
1168        }
1169
1170        Ok(())
1171    }
1172    fn add_type_guards(
1173        &mut self,
1174        lhs: Register,
1175        rhs: Register,
1176        registers: &[Value; 256],
1177        function: &crate::bytecode::Function,
1178    ) -> Result<(), LustError> {
1179        if let Some(ty) = Self::get_value_type(&registers[lhs as usize]) {
1180            let needs_guard = if self.is_guarded(lhs) {
1181                false
1182            } else if let Some(static_type) = function.register_types.get(&lhs) {
1183                !Self::type_kind_matches_value_type(static_type, ty)
1184            } else {
1185                true
1186            };
1187            if needs_guard {
1188                self.push_op(TraceOp::Guard {
1189                    register: lhs,
1190                    expected_type: ty,
1191                });
1192                self.mark_guarded(lhs);
1193            } else {
1194                self.mark_guarded(lhs);
1195            }
1196        }
1197
1198        if let Some(ty) = Self::get_value_type(&registers[rhs as usize]) {
1199            let needs_guard = if self.is_guarded(rhs) {
1200                false
1201            } else if let Some(static_type) = function.register_types.get(&rhs) {
1202                !Self::type_kind_matches_value_type(static_type, ty)
1203            } else {
1204                true
1205            };
1206            if needs_guard {
1207                self.push_op(TraceOp::Guard {
1208                    register: rhs,
1209                    expected_type: ty,
1210                });
1211                self.mark_guarded(rhs);
1212            } else {
1213                self.mark_guarded(rhs);
1214            }
1215        }
1216
1217        Ok(())
1218    }
1219
1220    fn type_kind_matches_value_type(
1221        type_kind: &crate::ast::TypeKind,
1222        value_type: ValueType,
1223    ) -> bool {
1224        use crate::ast::TypeKind;
1225        match (type_kind, value_type) {
1226            (TypeKind::Int, ValueType::Int) => true,
1227            (TypeKind::Float, ValueType::Float) => true,
1228            (TypeKind::Bool, ValueType::Bool) => true,
1229            (TypeKind::String, ValueType::String) => true,
1230            (TypeKind::Array(_), ValueType::Array) => true,
1231            (TypeKind::Tuple(_), ValueType::Tuple) => true,
1232            _ => false,
1233        }
1234    }
1235
1236    fn get_value_type(value: &Value) -> Option<ValueType> {
1237        match value {
1238            Value::Int(_) => Some(ValueType::Int),
1239            Value::Float(_) => Some(ValueType::Float),
1240            Value::Bool(_) => Some(ValueType::Bool),
1241            Value::String(_) => Some(ValueType::String),
1242            Value::Array(_) => Some(ValueType::Array),
1243            Value::Tuple(_) => Some(ValueType::Tuple),
1244            Value::Struct { .. } => Some(ValueType::Struct),
1245            _ => None,
1246        }
1247    }
1248
1249    pub fn finish(self) -> Trace {
1250        self.trace
1251    }
1252
1253    pub fn is_recording(&self) -> bool {
1254        self.recording
1255    }
1256
1257    pub fn abort(&mut self) {
1258        self.recording = false;
1259    }
1260}
lust/jit/trace.rs

lust/jit/
trace.rs
