hara_native/jit/
recorder.rs

1use super::trace_ir::{Trace, TraceOp, TraceValue};
2use crate::core::{IntrinsicOp, Value};
3use crate::vm::{Instruction, Program};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub enum RecordError {
7    InvalidRange,
8    InvalidStack,
9    TooLong,
10    UnsupportedInstruction(u32),
11    UnsupportedConstant(u32),
12    UnsupportedLocal(u16),
13}
14
15pub struct TraceRecorder {
16    max_operations: usize,
17}
18
19impl TraceRecorder {
20    pub fn new(max_operations: usize) -> Self {
21        Self { max_operations }
22    }
23
24    pub fn record_loop(
25        &self,
26        program: &Program,
27        function: u16,
28        header: u32,
29        backedge: u32,
30        locals: &[TraceValue],
31    ) -> Result<Trace, RecordError> {
32        let path = (header..=backedge).collect::<Vec<_>>();
33        self.record_path(program, function, header, &path, locals)
34    }
35
36    /// Lowers the concrete instruction path observed by the VM. Forward
37    /// branches disappear into the linear trace; their observed direction is
38    /// retained as a guard.
39    pub fn record_path(
40        &self,
41        program: &Program,
42        function: u16,
43        header: u32,
44        path: &[u32],
45        locals: &[TraceValue],
46    ) -> Result<Trace, RecordError> {
47        let prototype = program
48            .functions
49            .get(function as usize)
50            .ok_or(RecordError::InvalidRange)?;
51        if path.first() != Some(&header) || path.is_empty() {
52            return Err(RecordError::InvalidRange);
53        }
54        let mut operations = Vec::new();
55        let mut vectors = Vec::new();
56        for (index, absolute) in path.iter().copied().enumerate() {
57            let instruction = prototype
58                .code
59                .get(absolute as usize)
60                .ok_or(RecordError::InvalidRange)?;
61            let next = path.get(index + 1).copied().unwrap_or(header);
62            match instruction {
63                Instruction::LoadLocal(local) => {
64                    operations.push(match locals.get(usize::from(*local)) {
65                        Some(TraceValue::I64(_)) => TraceOp::GuardLocalI64 { local: *local },
66                        Some(TraceValue::Bool(_)) => TraceOp::GuardLocalBool { local: *local },
67                        Some(TraceValue::Nil) => TraceOp::GuardLocalNil { local: *local },
68                        Some(TraceValue::Indexed(value)) if numeric_vector(value).is_some() => {
69                            TraceOp::GuardLocalVectorI64 { local: *local }
70                        }
71                        _ => return Err(RecordError::UnsupportedLocal(*local)),
72                    });
73                    operations.push(TraceOp::LoadLocal { local: *local });
74                }
75                Instruction::StoreLocal(local) => {
76                    operations.push(TraceOp::StoreLocal { local: *local })
77                }
78                Instruction::Constant(index) => match program.constants.get(*index as usize) {
79                    Some(Value::Number(value)) => operations.push(TraceOp::ConstantI64(*value)),
80                    Some(Value::Bool(value)) => operations.push(TraceOp::ConstantBool(*value)),
81                    Some(Value::Nil) => operations.push(TraceOp::ConstantNil),
82                    Some(value @ (Value::Tuple(_) | Value::Vector(_))) => {
83                        let values = numeric_vector(value)
84                            .ok_or(RecordError::UnsupportedConstant(*index))?;
85                        let vector =
86                            u16::try_from(vectors.len()).map_err(|_| RecordError::TooLong)?;
87                        vectors.push(values);
88                        operations.push(TraceOp::ConstantVectorI64 { vector });
89                    }
90                    _ => return Err(RecordError::UnsupportedConstant(*index)),
91                },
92                Instruction::Nil => operations.push(TraceOp::ConstantNil),
93                Instruction::True => operations.push(TraceOp::ConstantBool(true)),
94                Instruction::False => operations.push(TraceOp::ConstantBool(false)),
95                Instruction::IntrinsicCall { target, argc: 2 }
96                    if intrinsic_op(program, *target).is_some_and(binary_i64) =>
97                {
98                    operations.push(TraceOp::BinaryI64(
99                        intrinsic_op(program, *target).expect("guarded intrinsic operator"),
100                    ));
101                }
102                Instruction::IntrinsicCall { target, argc: 1 }
103                    if vector_operation(program, *target).is_some() =>
104                {
105                    operations.push(vector_operation(program, *target).expect("guarded vector op"));
106                }
107                Instruction::ProtocolCall { target, argc: 1 }
108                    if vector_operation(program, *target).is_some() =>
109                {
110                    operations.push(vector_operation(program, *target).expect("guarded vector op"));
111                }
112                Instruction::ProtocolCall { target, argc: 2 }
113                    if vector_operation(program, *target).is_some() =>
114                {
115                    operations.push(vector_operation(program, *target).expect("guarded vector op"));
116                }
117                Instruction::JumpIfFalse(target) => {
118                    let expected = next != *target;
119                    if next != absolute + 1 && next != *target {
120                        return Err(RecordError::InvalidRange);
121                    }
122                    operations.push(TraceOp::GuardTruthy { expected })
123                }
124                Instruction::Pop => operations.push(TraceOp::Pop),
125                Instruction::Jump(target) if *target == next => {
126                    if *target == header {
127                        operations.push(TraceOp::LoopBackedge)
128                    }
129                }
130                _ => return Err(RecordError::UnsupportedInstruction(absolute)),
131            }
132            if operations.len() > self.max_operations {
133                return Err(RecordError::TooLong);
134            }
135        }
136        if !matches!(operations.last(), Some(TraceOp::LoopBackedge)) {
137            return Err(RecordError::InvalidRange);
138        }
139        if !valid_types(&operations, locals) {
140            return Err(RecordError::InvalidStack);
141        }
142        Ok(Trace {
143            function,
144            header,
145            resume_ip: header,
146            operations,
147            vectors,
148        })
149    }
150}
151
152#[derive(Clone, Copy, PartialEq, Eq)]
153enum TraceType {
154    I64,
155    Bool,
156    Nil,
157    Vector,
158    Slice,
159}
160
161fn valid_types(operations: &[TraceOp], entry_locals: &[TraceValue]) -> bool {
162    use TraceType::*;
163    let mut locals = entry_locals
164        .iter()
165        .enumerate()
166        .filter_map(|(index, value)| {
167            let kind = match value {
168                TraceValue::I64(_) => I64,
169                TraceValue::Bool(_) => Bool,
170                TraceValue::Nil => Nil,
171                TraceValue::Indexed(value) if numeric_vector(value).is_some() => Vector,
172                _ => return None,
173            };
174            Some((index as u16, kind))
175        })
176        .collect::<std::collections::HashMap<_, _>>();
177    let entry_types = locals.clone();
178    let mut stack = Vec::new();
179    for operation in operations {
180        match *operation {
181            TraceOp::GuardLocalI64 { local } => {
182                locals.insert(local, I64);
183            }
184            TraceOp::GuardLocalBool { local } => {
185                locals.insert(local, Bool);
186            }
187            TraceOp::GuardLocalNil { local } => {
188                locals.insert(local, Nil);
189            }
190            TraceOp::GuardLocalVectorI64 { local } => {
191                locals.insert(local, Vector);
192            }
193            TraceOp::LoadLocal { local } => {
194                let Some(kind) = locals.get(&local).copied() else {
195                    return false;
196                };
197                stack.push(kind);
198            }
199            TraceOp::ConstantI64(_) => stack.push(I64),
200            TraceOp::ConstantBool(_) => stack.push(Bool),
201            TraceOp::ConstantNil => stack.push(Nil),
202            TraceOp::ConstantVectorI64 { .. } => stack.push(Vector),
203            TraceOp::BinaryI64(op) => {
204                if stack.pop() != Some(I64) || stack.pop() != Some(I64) {
205                    return false;
206                }
207                stack.push(
208                    if matches!(
209                        op,
210                        IntrinsicOp::Equal
211                            | IntrinsicOp::Less
212                            | IntrinsicOp::LessOrEqual
213                            | IntrinsicOp::Greater
214                            | IntrinsicOp::GreaterOrEqual
215                    ) {
216                        Bool
217                    } else {
218                        I64
219                    },
220                );
221            }
222            TraceOp::VectorCountI64 => {
223                if !matches!(stack.pop(), Some(Vector | Slice)) {
224                    return false;
225                }
226                stack.push(I64);
227            }
228            TraceOp::VectorFirstI64 | TraceOp::VectorSecondI64 => {
229                if !matches!(stack.pop(), Some(Vector | Slice)) {
230                    return false;
231                }
232                stack.push(I64);
233            }
234            TraceOp::VectorRestI64 => {
235                if !matches!(stack.pop(), Some(Vector | Slice)) {
236                    return false;
237                }
238                stack.push(Slice);
239            }
240            TraceOp::VectorNthI64 => {
241                if stack.pop() != Some(I64) || !matches!(stack.pop(), Some(Vector | Slice)) {
242                    return false;
243                }
244                stack.push(I64);
245            }
246            TraceOp::StoreLocal { local } => {
247                let Some(kind) = stack.pop() else {
248                    return false;
249                };
250                if matches!(kind, Vector | Slice)
251                    || entry_types.get(&local).is_some_and(|entry| *entry != kind)
252                    || !entry_types.contains_key(&local)
253                {
254                    return false;
255                }
256                locals.insert(local, kind);
257            }
258            TraceOp::GuardTruthy { .. } => {
259                if !matches!(stack.pop(), Some(Bool | Nil)) {
260                    return false;
261                }
262            }
263            TraceOp::Pop => {
264                if stack.pop().is_none() {
265                    return false;
266                }
267            }
268            TraceOp::LoopBackedge => {
269                if !stack.is_empty() {
270                    return false;
271                }
272            }
273        }
274    }
275    stack.is_empty()
276}
277
278fn binary_i64(op: IntrinsicOp) -> bool {
279    matches!(
280        op,
281        IntrinsicOp::Add
282            | IntrinsicOp::Subtract
283            | IntrinsicOp::Multiply
284            | IntrinsicOp::Divide
285            | IntrinsicOp::Remainder
286            | IntrinsicOp::Modulo
287            | IntrinsicOp::Less
288            | IntrinsicOp::LessOrEqual
289            | IntrinsicOp::Greater
290            | IntrinsicOp::GreaterOrEqual
291            | IntrinsicOp::Equal
292    )
293}
294
295fn target_name(program: &Program, target: u32) -> Option<&str> {
296    match program.constants.get(target as usize) {
297        Some(Value::String(name)) => Some(name),
298        _ => None,
299    }
300}
301
302fn intrinsic_op(program: &Program, target: u32) -> Option<IntrinsicOp> {
303    target_name(program, target).and_then(IntrinsicOp::from_symbol)
304}
305
306fn vector_operation(program: &Program, target: u32) -> Option<TraceOp> {
307    let name = target_name(program, target)?;
308    if name == "first" || name.ends_with("/first") {
309        Some(TraceOp::VectorFirstI64)
310    } else if name == "rest" || name.ends_with("/rest") {
311        Some(TraceOp::VectorRestI64)
312    } else if name == "second" || name.ends_with("/second") {
313        Some(TraceOp::VectorSecondI64)
314    } else if name == "count" || name.ends_with("/count") {
315        Some(TraceOp::VectorCountI64)
316    } else if name == "nth" || name.ends_with("/nth") {
317        Some(TraceOp::VectorNthI64)
318    } else {
319        None
320    }
321}
322
323fn numeric_vector(value: &Value) -> Option<Vec<i64>> {
324    let values: Box<dyn Iterator<Item = &Value> + '_> = match value {
325        Value::Tuple(values) => Box::new(values.iter()),
326        Value::Vector(values) => Box::new(values.iter()),
327        _ => return None,
328    };
329    values
330        .map(|value| match value {
331            Value::Number(value) => Some(*value),
332            _ => None,
333        })
334        .collect()
335}
hara_native/jit/recorder.rs

hara_native/jit/
recorder.rs
