lust/jit/
trace.rs

1use crate::bytecode::value::NativeFn;
2use crate::bytecode::Instruction;
3use crate::bytecode::{Register, Value};
4use crate::LustError;
5use std::fmt;
6use std::rc::Rc;
7
8#[derive(Clone)]
9pub struct TracedNativeFn {
10    function: NativeFn,
11}
12
13impl TracedNativeFn {
14    pub fn new(function: NativeFn) -> Self {
15        Self { function }
16    }
17
18    pub fn pointer(&self) -> *const () {
19        Rc::as_ptr(&self.function) as *const ()
20    }
21}
22
23impl fmt::Debug for TracedNativeFn {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        write!(f, "NativeFn({:p})", Rc::as_ptr(&self.function))
26    }
27}
28
29#[derive(Debug, Clone)]
30pub struct Trace {
31    pub function_idx: usize,
32    pub start_ip: usize,
33    pub ops: Vec<TraceOp>,
34    pub inputs: Vec<Register>,
35    pub outputs: Vec<Register>,
36}
37
38#[derive(Debug, Clone)]
39pub enum TraceOp {
40    LoadConst {
41        dest: Register,
42        value: Value,
43    },
44    Move {
45        dest: Register,
46        src: Register,
47    },
48    Add {
49        dest: Register,
50        lhs: Register,
51        rhs: Register,
52        lhs_type: ValueType,
53        rhs_type: ValueType,
54    },
55    Sub {
56        dest: Register,
57        lhs: Register,
58        rhs: Register,
59        lhs_type: ValueType,
60        rhs_type: ValueType,
61    },
62    Mul {
63        dest: Register,
64        lhs: Register,
65        rhs: Register,
66        lhs_type: ValueType,
67        rhs_type: ValueType,
68    },
69    Div {
70        dest: Register,
71        lhs: Register,
72        rhs: Register,
73        lhs_type: ValueType,
74        rhs_type: ValueType,
75    },
76    Mod {
77        dest: Register,
78        lhs: Register,
79        rhs: Register,
80        lhs_type: ValueType,
81        rhs_type: ValueType,
82    },
83    Neg {
84        dest: Register,
85        src: Register,
86    },
87    Eq {
88        dest: Register,
89        lhs: Register,
90        rhs: Register,
91    },
92    Ne {
93        dest: Register,
94        lhs: Register,
95        rhs: Register,
96    },
97    Lt {
98        dest: Register,
99        lhs: Register,
100        rhs: Register,
101    },
102    Le {
103        dest: Register,
104        lhs: Register,
105        rhs: Register,
106    },
107    Gt {
108        dest: Register,
109        lhs: Register,
110        rhs: Register,
111    },
112    Ge {
113        dest: Register,
114        lhs: Register,
115        rhs: Register,
116    },
117    And {
118        dest: Register,
119        lhs: Register,
120        rhs: Register,
121    },
122    Or {
123        dest: Register,
124        lhs: Register,
125        rhs: Register,
126    },
127    Not {
128        dest: Register,
129        src: Register,
130    },
131    Concat {
132        dest: Register,
133        lhs: Register,
134        rhs: Register,
135    },
136    GetIndex {
137        dest: Register,
138        array: Register,
139        index: Register,
140    },
141    ArrayLen {
142        dest: Register,
143        array: Register,
144    },
145    GuardNativeFunction {
146        register: Register,
147        function: TracedNativeFn,
148    },
149    CallNative {
150        dest: Register,
151        callee: Register,
152        function: TracedNativeFn,
153        first_arg: Register,
154        arg_count: u8,
155    },
156    CallMethod {
157        dest: Register,
158        object: Register,
159        method_name: String,
160        first_arg: Register,
161        arg_count: u8,
162    },
163    GetField {
164        dest: Register,
165        object: Register,
166        field_name: String,
167        field_index: Option<usize>,
168        value_type: Option<ValueType>,
169        is_weak: bool,
170    },
171    SetField {
172        object: Register,
173        field_name: String,
174        value: Register,
175        field_index: Option<usize>,
176        value_type: Option<ValueType>,
177        is_weak: bool,
178    },
179    NewStruct {
180        dest: Register,
181        struct_name: String,
182        field_names: Vec<String>,
183        field_registers: Vec<Register>,
184    },
185    Guard {
186        register: Register,
187        expected_type: ValueType,
188    },
189    GuardLoopContinue {
190        condition_register: Register,
191        bailout_ip: usize,
192    },
193    NestedLoopCall {
194        function_idx: usize,
195        loop_start_ip: usize,
196        bailout_ip: usize,
197    },
198    Return {
199        value: Option<Register>,
200    },
201}
202
203#[derive(Debug, Clone, Copy, PartialEq, Eq)]
204pub enum ValueType {
205    Int,
206    Float,
207    Bool,
208    String,
209    Array,
210    Tuple,
211    Struct,
212}
213
214pub struct TraceRecorder {
215    pub trace: Trace,
216    max_length: usize,
217    recording: bool,
218    guarded_registers: std::collections::HashSet<Register>,
219}
220
221impl TraceRecorder {
222    pub fn new(function_idx: usize, start_ip: usize, max_length: usize) -> Self {
223        Self {
224            trace: Trace {
225                function_idx,
226                start_ip,
227                ops: Vec::new(),
228                inputs: Vec::new(),
229                outputs: Vec::new(),
230            },
231            max_length,
232            recording: true,
233            guarded_registers: std::collections::HashSet::new(),
234        }
235    }
236
237    pub fn record_instruction(
238        &mut self,
239        instruction: Instruction,
240        current_ip: usize,
241        registers: &[Value; 256],
242        function: &crate::bytecode::Function,
243        function_idx: usize,
244    ) -> Result<(), LustError> {
245        if !self.recording {
246            return Ok(());
247        }
248
249        if function_idx != self.trace.function_idx {
250            return Ok(());
251        }
252
253        let trace_op = match instruction {
254            Instruction::LoadConst(dest, _) => {
255                if let Some(_ty) = Self::get_value_type(&registers[dest as usize]) {
256                    self.guarded_registers.insert(dest);
257                }
258
259                TraceOp::LoadConst {
260                    dest,
261                    value: registers[dest as usize].clone(),
262                }
263            }
264
265            Instruction::LoadGlobal(dest, _) => {
266                if let Some(_ty) = Self::get_value_type(&registers[dest as usize]) {
267                    self.guarded_registers.insert(dest);
268                }
269
270                TraceOp::LoadConst {
271                    dest,
272                    value: registers[dest as usize].clone(),
273                }
274            }
275
276            Instruction::StoreGlobal(_, _) => {
277                return Ok(());
278            }
279
280            Instruction::Move(dest, src) => TraceOp::Move { dest, src },
281            Instruction::Add(dest, lhs, rhs) => {
282                self.add_type_guards(lhs, rhs, registers, function)?;
283                let lhs_type =
284                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
285                let rhs_type =
286                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
287                TraceOp::Add {
288                    dest,
289                    lhs,
290                    rhs,
291                    lhs_type,
292                    rhs_type,
293                }
294            }
295
296            Instruction::Sub(dest, lhs, rhs) => {
297                self.add_type_guards(lhs, rhs, registers, function)?;
298                let lhs_type =
299                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
300                let rhs_type =
301                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
302                TraceOp::Sub {
303                    dest,
304                    lhs,
305                    rhs,
306                    lhs_type,
307                    rhs_type,
308                }
309            }
310
311            Instruction::Mul(dest, lhs, rhs) => {
312                self.add_type_guards(lhs, rhs, registers, function)?;
313                let lhs_type =
314                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
315                let rhs_type =
316                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
317                TraceOp::Mul {
318                    dest,
319                    lhs,
320                    rhs,
321                    lhs_type,
322                    rhs_type,
323                }
324            }
325
326            Instruction::Div(dest, lhs, rhs) => {
327                self.add_type_guards(lhs, rhs, registers, function)?;
328                let lhs_type =
329                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
330                let rhs_type =
331                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
332                TraceOp::Div {
333                    dest,
334                    lhs,
335                    rhs,
336                    lhs_type,
337                    rhs_type,
338                }
339            }
340
341            Instruction::Mod(dest, lhs, rhs) => {
342                self.add_type_guards(lhs, rhs, registers, function)?;
343                let lhs_type =
344                    Self::get_value_type(&registers[lhs as usize]).unwrap_or(ValueType::Int);
345                let rhs_type =
346                    Self::get_value_type(&registers[rhs as usize]).unwrap_or(ValueType::Int);
347                TraceOp::Mod {
348                    dest,
349                    lhs,
350                    rhs,
351                    lhs_type,
352                    rhs_type,
353                }
354            }
355
356            Instruction::Neg(dest, src) => TraceOp::Neg { dest, src },
357            Instruction::Eq(dest, lhs, rhs) => TraceOp::Eq { dest, lhs, rhs },
358            Instruction::Ne(dest, lhs, rhs) => TraceOp::Ne { dest, lhs, rhs },
359            Instruction::Lt(dest, lhs, rhs) => TraceOp::Lt { dest, lhs, rhs },
360            Instruction::Le(dest, lhs, rhs) => TraceOp::Le { dest, lhs, rhs },
361            Instruction::Gt(dest, lhs, rhs) => TraceOp::Gt { dest, lhs, rhs },
362            Instruction::Ge(dest, lhs, rhs) => TraceOp::Ge { dest, lhs, rhs },
363            Instruction::And(dest, lhs, rhs) => TraceOp::And { dest, lhs, rhs },
364            Instruction::Or(dest, lhs, rhs) => TraceOp::Or { dest, lhs, rhs },
365            Instruction::Not(dest, src) => TraceOp::Not { dest, src },
366            Instruction::Concat(dest, lhs, rhs) => {
367                if let Some(ty) = Self::get_value_type(&registers[lhs as usize]) {
368                    if !self.guarded_registers.contains(&lhs) {
369                        self.trace.ops.push(TraceOp::Guard {
370                            register: lhs,
371                            expected_type: ty,
372                        });
373                        self.guarded_registers.insert(lhs);
374                    }
375                }
376
377                if let Some(ty) = Self::get_value_type(&registers[rhs as usize]) {
378                    if !self.guarded_registers.contains(&rhs) {
379                        self.trace.ops.push(TraceOp::Guard {
380                            register: rhs,
381                            expected_type: ty,
382                        });
383                        self.guarded_registers.insert(rhs);
384                    }
385                }
386
387                TraceOp::Concat { dest, lhs, rhs }
388            }
389
390            Instruction::GetIndex(dest, array, index) => {
391                if let Some(ty) = Self::get_value_type(&registers[array as usize]) {
392                    if !self.guarded_registers.contains(&array) {
393                        self.trace.ops.push(TraceOp::Guard {
394                            register: array,
395                            expected_type: ty,
396                        });
397                        self.guarded_registers.insert(array);
398                    }
399                }
400
401                if let Some(ty) = Self::get_value_type(&registers[index as usize]) {
402                    if !self.guarded_registers.contains(&index) {
403                        self.trace.ops.push(TraceOp::Guard {
404                            register: index,
405                            expected_type: ty,
406                        });
407                        self.guarded_registers.insert(index);
408                    }
409                }
410
411                TraceOp::GetIndex { dest, array, index }
412            }
413
414            Instruction::ArrayLen(dest, array) => {
415                if let Some(ty) = Self::get_value_type(&registers[array as usize]) {
416                    if !self.guarded_registers.contains(&array) {
417                        self.trace.ops.push(TraceOp::Guard {
418                            register: array,
419                            expected_type: ty,
420                        });
421                        self.guarded_registers.insert(array);
422                    }
423                }
424
425                TraceOp::ArrayLen { dest, array }
426            }
427
428            Instruction::CallMethod(obj_reg, method_name_idx, first_arg, arg_count, dest_reg) => {
429                let method_name = function.chunk.constants[method_name_idx as usize]
430                    .as_string()
431                    .unwrap_or("unknown")
432                    .to_string();
433                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
434                    if !self.guarded_registers.contains(&obj_reg) {
435                        self.trace.ops.push(TraceOp::Guard {
436                            register: obj_reg,
437                            expected_type: ty,
438                        });
439                        self.guarded_registers.insert(obj_reg);
440                    }
441                }
442
443                for i in 0..arg_count {
444                    let arg_reg = first_arg + i;
445                    if let Some(ty) = Self::get_value_type(&registers[arg_reg as usize]) {
446                        if !self.guarded_registers.contains(&arg_reg) {
447                            self.trace.ops.push(TraceOp::Guard {
448                                register: arg_reg,
449                                expected_type: ty,
450                            });
451                            self.guarded_registers.insert(arg_reg);
452                        }
453                    }
454                }
455
456                TraceOp::CallMethod {
457                    dest: dest_reg,
458                    object: obj_reg,
459                    method_name,
460                    first_arg,
461                    arg_count,
462                }
463            }
464
465            Instruction::GetField(dest, obj_reg, field_name_idx) => {
466                let field_name = function.chunk.constants[field_name_idx as usize]
467                    .as_string()
468                    .unwrap_or("unknown")
469                    .to_string();
470                let (field_index, is_weak_field) = match &registers[obj_reg as usize] {
471                    Value::Struct { layout, .. } => {
472                        let idx = layout.index_of_str(&field_name);
473                        let is_weak = idx.map(|i| layout.is_weak(i)).unwrap_or(false);
474                        (idx, is_weak)
475                    }
476
477                    _ => (None, false),
478                };
479                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
480                    if !self.guarded_registers.contains(&obj_reg) {
481                        self.trace.ops.push(TraceOp::Guard {
482                            register: obj_reg,
483                            expected_type: ty,
484                        });
485                        self.guarded_registers.insert(obj_reg);
486                    }
487                }
488
489                let value_type = Self::get_value_type(&registers[dest as usize]);
490                TraceOp::GetField {
491                    dest,
492                    object: obj_reg,
493                    field_name,
494                    field_index,
495                    value_type,
496                    is_weak: is_weak_field,
497                }
498            }
499
500            Instruction::SetField(obj_reg, field_name_idx, value_reg) => {
501                let field_name = function.chunk.constants[field_name_idx as usize]
502                    .as_string()
503                    .unwrap_or("unknown")
504                    .to_string();
505                let (field_index, is_weak_field) = match &registers[obj_reg as usize] {
506                    Value::Struct { layout, .. } => {
507                        let idx = layout.index_of_str(&field_name);
508                        let is_weak = idx.map(|i| layout.is_weak(i)).unwrap_or(false);
509                        (idx, is_weak)
510                    }
511
512                    _ => (None, false),
513                };
514                if let Some(ty) = Self::get_value_type(&registers[obj_reg as usize]) {
515                    if !self.guarded_registers.contains(&obj_reg) {
516                        self.trace.ops.push(TraceOp::Guard {
517                            register: obj_reg,
518                            expected_type: ty,
519                        });
520                        self.guarded_registers.insert(obj_reg);
521                    }
522                }
523
524                let value_type = Self::get_value_type(&registers[value_reg as usize]);
525                if let Some(ty) = value_type {
526                    if !self.guarded_registers.contains(&value_reg) {
527                        self.trace.ops.push(TraceOp::Guard {
528                            register: value_reg,
529                            expected_type: ty,
530                        });
531                        self.guarded_registers.insert(value_reg);
532                    }
533                }
534
535                TraceOp::SetField {
536                    object: obj_reg,
537                    field_name,
538                    value: value_reg,
539                    field_index,
540                    value_type,
541                    is_weak: is_weak_field,
542                }
543            }
544
545            Instruction::NewStruct(
546                dest,
547                struct_name_idx,
548                first_field_name_idx,
549                first_field_reg,
550                field_count,
551            ) => {
552                let struct_name = function.chunk.constants[struct_name_idx as usize]
553                    .as_string()
554                    .unwrap_or("unknown")
555                    .to_string();
556                let mut field_names = Vec::new();
557                for i in 0..field_count {
558                    let field_name_idx = first_field_name_idx + (i as u16);
559                    let field_name = function.chunk.constants[field_name_idx as usize]
560                        .as_string()
561                        .unwrap_or("unknown")
562                        .to_string();
563                    field_names.push(field_name);
564                }
565
566                let mut field_registers = Vec::new();
567                for i in 0..field_count {
568                    let field_reg = first_field_reg + i;
569                    field_registers.push(field_reg);
570                    if let Some(ty) = Self::get_value_type(&registers[field_reg as usize]) {
571                        if !self.guarded_registers.contains(&field_reg) {
572                            self.trace.ops.push(TraceOp::Guard {
573                                register: field_reg,
574                                expected_type: ty,
575                            });
576                            self.guarded_registers.insert(field_reg);
577                        }
578                    }
579                }
580
581                TraceOp::NewStruct {
582                    dest,
583                    struct_name,
584                    field_names,
585                    field_registers,
586                }
587            }
588
589            Instruction::Call(func_reg, first_arg, arg_count, dest_reg) => {
590                match &registers[func_reg as usize] {
591                    Value::NativeFunction(native_fn) => {
592                        let traced = TracedNativeFn::new(native_fn.clone());
593                        if !self.guarded_registers.contains(&func_reg) {
594                            self.trace.ops.push(TraceOp::GuardNativeFunction {
595                                register: func_reg,
596                                function: traced.clone(),
597                            });
598                            self.guarded_registers.insert(func_reg);
599                        }
600
601                        self.trace.ops.push(TraceOp::CallNative {
602                            dest: dest_reg,
603                            callee: func_reg,
604                            function: traced,
605                            first_arg,
606                            arg_count,
607                        });
608                        return Ok(());
609                    }
610
611                    _ => {
612                        self.recording = false;
613                        return Err(LustError::RuntimeError {
614                            message: "Trace aborted: unsupported operation".to_string(),
615                        });
616                    }
617                }
618            }
619
620            Instruction::NewArray(_, _, _)
621            | Instruction::NewMap(_)
622            | Instruction::SetIndex(_, _, _) => {
623                self.recording = false;
624                return Err(LustError::RuntimeError {
625                    message: "Trace aborted: unsupported operation".to_string(),
626                });
627            }
628
629            Instruction::Return(value_reg) => {
630                if function_idx == self.trace.function_idx {
631                    self.recording = false;
632                    return Ok(());
633                } else {
634                    TraceOp::Return {
635                        value: if value_reg == 255 {
636                            None
637                        } else {
638                            Some(value_reg)
639                        },
640                    }
641                }
642            }
643
644            Instruction::Jump(offset) => {
645                if offset < 0 {
646                    let target_calc = (current_ip as isize) + (offset as isize);
647                    if target_calc < 0 {
648                        self.recording = false;
649                        return Err(LustError::RuntimeError {
650                            message: format!(
651                                "Invalid jump target: offset={}, current_ip={}, target={}",
652                                offset, current_ip, target_calc
653                            ),
654                        });
655                    }
656
657                    let jump_target = target_calc as usize;
658                    if function_idx == self.trace.function_idx && jump_target == self.trace.start_ip
659                    {
660                        self.recording = false;
661                        return Ok(());
662                    } else {
663                        let bailout_ip = current_ip.saturating_sub(1);
664                        TraceOp::NestedLoopCall {
665                            function_idx,
666                            loop_start_ip: jump_target,
667                            bailout_ip,
668                        }
669                    }
670                } else {
671                    return Ok(());
672                }
673            }
674
675            Instruction::JumpIf(cond, _) | Instruction::JumpIfNot(cond, _) => {
676                return Ok(());
677            }
678
679            _ => {
680                self.recording = false;
681                return Err(LustError::RuntimeError {
682                    message: "Trace aborted: unsupported instruction".to_string(),
683                });
684            }
685        };
686        self.trace.ops.push(trace_op);
687        if self.trace.ops.len() >= self.max_length {
688            self.recording = false;
689            return Err(LustError::RuntimeError {
690                message: "Trace too long".to_string(),
691            });
692        }
693
694        Ok(())
695    }
696
697    fn add_type_guards(
698        &mut self,
699        lhs: Register,
700        rhs: Register,
701        registers: &[Value; 256],
702        function: &crate::bytecode::Function,
703    ) -> Result<(), LustError> {
704        if let Some(ty) = Self::get_value_type(&registers[lhs as usize]) {
705            let needs_guard = if self.guarded_registers.contains(&lhs) {
706                false
707            } else if let Some(static_type) = function.register_types.get(&lhs) {
708                !Self::type_kind_matches_value_type(static_type, ty)
709            } else {
710                true
711            };
712            if needs_guard {
713                self.trace.ops.push(TraceOp::Guard {
714                    register: lhs,
715                    expected_type: ty,
716                });
717                self.guarded_registers.insert(lhs);
718            } else {
719                self.guarded_registers.insert(lhs);
720            }
721        }
722
723        if let Some(ty) = Self::get_value_type(&registers[rhs as usize]) {
724            let needs_guard = if self.guarded_registers.contains(&rhs) {
725                false
726            } else if let Some(static_type) = function.register_types.get(&rhs) {
727                !Self::type_kind_matches_value_type(static_type, ty)
728            } else {
729                true
730            };
731            if needs_guard {
732                self.trace.ops.push(TraceOp::Guard {
733                    register: rhs,
734                    expected_type: ty,
735                });
736                self.guarded_registers.insert(rhs);
737            } else {
738                self.guarded_registers.insert(rhs);
739            }
740        }
741
742        Ok(())
743    }
744
745    fn type_kind_matches_value_type(
746        type_kind: &crate::ast::TypeKind,
747        value_type: ValueType,
748    ) -> bool {
749        use crate::ast::TypeKind;
750        match (type_kind, value_type) {
751            (TypeKind::Int, ValueType::Int) => true,
752            (TypeKind::Float, ValueType::Float) => true,
753            (TypeKind::Bool, ValueType::Bool) => true,
754            (TypeKind::String, ValueType::String) => true,
755            (TypeKind::Array(_), ValueType::Array) => true,
756            (TypeKind::Tuple(_), ValueType::Tuple) => true,
757            _ => false,
758        }
759    }
760
761    fn get_value_type(value: &Value) -> Option<ValueType> {
762        match value {
763            Value::Int(_) => Some(ValueType::Int),
764            Value::Float(_) => Some(ValueType::Float),
765            Value::Bool(_) => Some(ValueType::Bool),
766            Value::String(_) => Some(ValueType::String),
767            Value::Array(_) => Some(ValueType::Array),
768            Value::Tuple(_) => Some(ValueType::Tuple),
769            Value::Struct { .. } => Some(ValueType::Struct),
770            _ => None,
771        }
772    }
773
774    pub fn finish(self) -> Trace {
775        self.trace
776    }
777
778    pub fn is_recording(&self) -> bool {
779        self.recording
780    }
781
782    pub fn abort(&mut self) {
783        self.recording = false;
784    }
785}
lust/jit/trace.rs

lust/jit/
trace.rs
