snix_eval/
chunk.rs

1use crate::opcode::{CodeIdx, ConstantIdx, Op, OpArg};
2use crate::value::Value;
3use crate::{CoercionKind, SourceCode};
4use std::io::Write;
5
6/// Maximum size of a u64 encoded in the vu128 varint encoding.
7const U64_VARINT_SIZE: usize = 9;
8
9/// Represents a source location from which one or more operations
10/// were compiled.
11///
12/// The span itself is an index into a [codemap::CodeMap], and the
13/// structure tracks the number of operations that were yielded from
14/// the same span.
15///
16/// At error reporting time, it becomes possible to either just fetch
17/// the textual representation of that span from the codemap, or to
18/// even re-parse the AST using rnix to create more semantically
19/// interesting errors.
20#[derive(Clone, Debug, PartialEq)]
21struct SourceSpan {
22    /// Span into the [codemap::CodeMap].
23    span: codemap::Span,
24
25    /// Index of the first operation covered by this span.
26    start: usize,
27}
28
29/// A chunk is a representation of a sequence of bytecode
30/// instructions, associated constants and additional metadata as
31/// emitted by the compiler.
32#[derive(Debug, Default)]
33pub struct Chunk {
34    pub code: Vec<u8>,
35    pub constants: Vec<Value>,
36
37    /// Spans for AttrSet's keys. Needed for `unsafeGetAttrPos` builtin.
38    ///
39    /// This code assumes that each Chunk can contain at most one AttrSet.
40    ///
41    /// This assumption is based on how AttrSets compiled now. Each
42    /// AttrSet emits its own thunk and each thunk creates its own Chunk.
43    pub attrsets_pos_spans: Vec<codemap::Span>,
44
45    spans: Vec<SourceSpan>,
46
47    /// Index of the last operation (i.e. not data) written to the code vector.
48    /// Some operations (e.g. jump patching) need to know this.
49    last_op: usize,
50}
51
52impl Chunk {
53    pub fn push_op(&mut self, op: Op, span: codemap::Span) -> usize {
54        self.last_op = self.code.len();
55        self.code.push(op as u8);
56        self.push_span(span, self.last_op);
57        self.last_op
58    }
59
60    pub fn push_uvarint(&mut self, data: u64) {
61        let mut encoded = [0u8; U64_VARINT_SIZE];
62        let bytes_written = vu128::encode_u64(&mut encoded, data);
63        self.code.extend_from_slice(&encoded[..bytes_written]);
64    }
65
66    pub fn read_uvarint(&self, idx: usize) -> (u64, usize) {
67        debug_assert!(
68            idx < self.code.len(),
69            "invalid bytecode (missing varint operand)",
70        );
71
72        if self.code.len() - idx >= U64_VARINT_SIZE {
73            vu128::decode_u64(
74                &self.code[idx..idx + U64_VARINT_SIZE]
75                    .try_into()
76                    .expect("size statically checked"),
77            )
78        } else {
79            let mut tmp = [0u8; U64_VARINT_SIZE];
80            tmp[..self.code.len() - idx].copy_from_slice(&self.code[idx..]);
81            vu128::decode_u64(&tmp)
82        }
83    }
84
85    pub fn push_u16(&mut self, data: u16) {
86        self.code.extend_from_slice(&data.to_le_bytes())
87    }
88
89    /// Patches the argument to the jump operand of the jump at the given index
90    /// to point to the *next* instruction that will be emitted.
91    pub fn patch_jump(&mut self, idx: usize) {
92        let offset = (self.code.len() - idx - /* arg idx = */ 1 - /* jump arg size = */ 2) as u16;
93        self.code[idx + 1..idx + 3].copy_from_slice(&offset.to_le_bytes())
94    }
95
96    pub fn read_u16(&self, idx: usize) -> u16 {
97        if idx + 2 > self.code.len() {
98            panic!("Snix bug: invalid bytecode (expected u16 operand not found)")
99        }
100
101        let byte_array: &[u8; 2] = &self.code[idx..idx + 2]
102            .try_into()
103            .expect("fixed-size slice can not fail to convert to array");
104
105        u16::from_le_bytes(*byte_array)
106    }
107
108    /// Get the first span of a chunk, no questions asked.
109    pub fn first_span(&self) -> codemap::Span {
110        self.spans[0].span
111    }
112
113    /// Return the last op in the chunk together with its index, if any.
114    pub fn last_op(&self) -> Option<(Op, usize)> {
115        if self.code.is_empty() {
116            return None;
117        }
118
119        Some((self.code[self.last_op].into(), self.last_op))
120    }
121
122    pub fn push_constant(&mut self, data: Value) -> ConstantIdx {
123        let idx = self.constants.len();
124        self.constants.push(data);
125        ConstantIdx(idx)
126    }
127
128    /// Return a reference to the constant at the given [`ConstantIdx`]
129    pub fn get_constant(&self, constant: ConstantIdx) -> Option<&Value> {
130        self.constants.get(constant.0)
131    }
132
133    fn push_span(&mut self, span: codemap::Span, start: usize) {
134        match self.spans.last_mut() {
135            // We do not need to insert the same span again, as this
136            // instruction was compiled from the same span as the last
137            // one.
138            Some(last) if last.span == span => {}
139
140            // In all other cases, this is a new source span.
141            _ => self.spans.push(SourceSpan { span, start }),
142        }
143    }
144
145    /// Retrieve the [codemap::Span] from which the instruction at
146    /// `offset` was compiled.
147    pub fn get_span(&self, offset: CodeIdx) -> codemap::Span {
148        let position = self
149            .spans
150            .binary_search_by(|span| span.start.cmp(&offset.0));
151
152        let span = match position {
153            Ok(index) => &self.spans[index],
154            Err(index) => {
155                if index == 0 {
156                    &self.spans[0]
157                } else {
158                    &self.spans[index - 1]
159                }
160            }
161        };
162
163        span.span
164    }
165
166    /// Write the disassembler representation of the operation at
167    /// `idx` to the specified writer, and return how many bytes in the code to
168    /// skip for the next op.
169    ///
170    /// # Format
171    ///
172    /// The format of the disassembly is shown in the diagram below.
173    ///
174    /// ```text
175    /// 0x0  1  OpClosure(BP @ 0, 1 upvalues)
176    /// ^    ^  ^        ^
177    /// |    |  |        |
178    /// |    |  |        \ One of the following, depending on the operation:
179    /// |    |  |          * Nothing
180    /// |    |  |          * Single argument
181    /// |    |  |          * Blueprint's index, whether `with`
182    /// |    |  |            was captured and number of upvalues
183    /// |    |  |
184    /// |    |  \ Operation identifier
185    /// |    \ Line number in source code, or continuation character (`|`)
186    /// \ Index of the instruction in the code chunk
187    /// ```
188    pub fn disassemble_op<W: Write>(
189        &self,
190        writer: &mut W,
191        source: &SourceCode,
192        width: usize,
193        idx: CodeIdx,
194    ) -> Result<usize, std::io::Error> {
195        write!(writer, "{:#width$x}\t ", idx.0, width = width)?;
196
197        // Print continuation character if the previous operation was at
198        // the same line, otherwise print the line.
199        let line = source.get_line(self.get_span(idx));
200        if idx.0 > 0 && source.get_line(self.get_span(idx - 1)) == line {
201            write!(writer, "   |\t")?;
202        } else {
203            write!(writer, "{line:4}\t")?;
204        }
205
206        let _fmt_constant = |idx: ConstantIdx| match &self.constants[idx.0] {
207            Value::Thunk(t) => t.debug_repr(),
208            Value::Closure(c) => format!("closure({:p})", c.lambda),
209            Value::Blueprint(b) => format!("blueprint({b:p})"),
210            val => format!("{val}"),
211        };
212
213        let op: Op = self.code[idx.0].into();
214
215        match op.arg_type() {
216            OpArg::None => {
217                writeln!(writer, "Op{op:?}")?;
218                Ok(1)
219            }
220
221            OpArg::Fixed => {
222                let arg = self.read_u16(idx.0 + 1);
223                writeln!(writer, "Op{op:?}({arg})")?;
224                Ok(3)
225            }
226
227            OpArg::Uvarint => {
228                let (arg, size) = self.read_uvarint(idx.0 + 1);
229                write!(writer, "Op{op:?}({arg})")?;
230                if let Op::Constant = &op {
231                    write!(writer, " (={})", self.constants[arg as usize])?;
232                }
233                writeln!(writer)?;
234                Ok(1 + size)
235            }
236
237            OpArg::Custom => match op {
238                Op::CoerceToString => {
239                    let kind: CoercionKind = self.code[idx.0 + 1].into();
240                    writeln!(writer, "Op{op:?}({kind:?})")?;
241                    Ok(2)
242                }
243
244                Op::Closure | Op::ThunkClosure | Op::ThunkSuspended => {
245                    let mut cidx = idx.0 + 1;
246
247                    let (bp_idx, size) = self.read_uvarint(cidx);
248                    cidx += size;
249
250                    let (packed_count, size) = self.read_uvarint(cidx);
251                    cidx += size;
252
253                    let captures_with = packed_count & 0b1 == 1;
254                    let count = packed_count >> 1;
255
256                    write!(writer, "Op{op:?}(BP @ {bp_idx}, ")?;
257                    if captures_with {
258                        write!(writer, "captures with, ")?;
259                    }
260                    writeln!(writer, "{count} upvalues)")?;
261
262                    for _ in 0..count {
263                        let (_, size) = self.read_uvarint(cidx);
264                        cidx += size;
265                    }
266
267                    Ok(cidx - idx.0)
268                }
269                _ => panic!("Snix bug: don't know how to format argument for Op{op:?}"),
270            },
271        }
272    }
273
274    /// Count the number of Opcodes. This function has O(n) time complexity.
275    pub fn op_count(&self) -> usize {
276        if self.code.is_empty() {
277            return 0;
278        }
279
280        let mut idx = 0;
281        let mut count = 0;
282        while idx <= self.last_op {
283            let op: Op = self.code[idx].into();
284            let op_len = match op.arg_type() {
285                OpArg::None => 1,
286                OpArg::Uvarint => {
287                    let (_, len) = self.read_uvarint(idx + 1);
288                    len + 1
289                }
290                OpArg::Fixed => 3,
291                OpArg::Custom => match op {
292                    Op::CoerceToString => 2,
293
294                    Op::Closure | Op::ThunkClosure | Op::ThunkSuspended => {
295                        let mut len = 1;
296
297                        let (_, size) = self.read_uvarint(idx + len);
298                        len += size;
299
300                        let (packed_count, size) = self.read_uvarint(idx + len);
301                        len += size;
302
303                        let count = packed_count >> 1;
304
305                        for _ in 0..count {
306                            let (_, size) = self.read_uvarint(idx + len);
307                            len += size;
308                        }
309
310                        len
311                    }
312                    _ => panic!("Snix bug: arg_type returned Custom for wrong Op"),
313                },
314            };
315            idx += op_len;
316            count += 1;
317        }
318        count
319    }
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325    use crate::test_utils::dummy_span;
326
327    // Note: These tests are about the functionality of the `Chunk` type, the
328    // opcodes used below do *not* represent valid, executable Snix code (and
329    // don't need to).
330
331    #[test]
332    fn push_op() {
333        let mut chunk = Chunk::default();
334        let idx = chunk.push_op(Op::Add, dummy_span());
335        assert_eq!(*chunk.code.last().unwrap(), Op::Add as u8);
336        assert_eq!(chunk.code[idx], Op::Add as u8);
337    }
338
339    #[test]
340    fn push_op_with_arg() {
341        let mut chunk = Chunk::default();
342        let mut idx = chunk.push_op(Op::Constant, dummy_span());
343        chunk.push_uvarint(42);
344
345        assert_eq!(chunk.code[idx], Op::Constant as u8);
346
347        idx += 1;
348        let (arg, size) = chunk.read_uvarint(idx);
349        assert_eq!(idx + size, chunk.code.len());
350        assert_eq!(arg, 42);
351    }
352
353    #[test]
354    fn push_jump() {
355        let mut chunk = Chunk::default();
356
357        chunk.push_op(Op::Constant, dummy_span());
358        chunk.push_uvarint(0);
359
360        let idx = chunk.push_op(Op::Jump, dummy_span());
361        chunk.push_u16(0);
362
363        chunk.push_op(Op::Constant, dummy_span());
364        chunk.push_uvarint(1);
365
366        chunk.patch_jump(idx);
367        chunk.push_op(Op::Return, dummy_span());
368
369        #[rustfmt::skip]
370        let expected: Vec<u8> = vec![
371            Op::Constant as u8, 0,
372            Op::Jump as u8, 2, 0,
373            Op::Constant as u8, 1,
374            Op::Return as u8,
375        ];
376
377        assert_eq!(chunk.code, expected);
378    }
379
380    #[test]
381    fn op_count() {
382        let mut chunk = Chunk::default();
383        assert_eq!(chunk.op_count(), 0);
384
385        chunk.push_op(Op::Constant, dummy_span());
386        chunk.push_uvarint(0);
387
388        let idx = chunk.push_op(Op::Jump, dummy_span());
389        chunk.push_u16(0);
390
391        chunk.push_op(Op::Constant, dummy_span());
392        chunk.push_uvarint(1);
393
394        chunk.patch_jump(idx);
395        chunk.push_op(Op::Return, dummy_span());
396
397        assert_eq!(chunk.op_count(), 4);
398        assert_eq!(chunk.code.len(), 8);
399    }
400}
snix_eval/chunk.rs

snix_eval/
chunk.rs
