1use crate::opcode::{CodeIdx, ConstantIdx, Op, OpArg};
2use crate::value::Value;
3use crate::{CoercionKind, SourceCode};
4use std::io::Write;
5
6const U64_VARINT_SIZE: usize = 9;
8
9#[derive(Clone, Debug, PartialEq)]
21struct SourceSpan {
22 span: codemap::Span,
24
25 start: usize,
27}
28
29#[derive(Debug, Default)]
33pub struct Chunk {
34 pub code: Vec<u8>,
35 pub constants: Vec<Value>,
36 spans: Vec<SourceSpan>,
37
38 last_op: usize,
41}
42
43impl Chunk {
44 pub fn push_op(&mut self, op: Op, span: codemap::Span) -> usize {
45 self.last_op = self.code.len();
46 self.code.push(op as u8);
47 self.push_span(span, self.last_op);
48 self.last_op
49 }
50
51 pub fn push_uvarint(&mut self, data: u64) {
52 let mut encoded = [0u8; U64_VARINT_SIZE];
53 let bytes_written = vu128::encode_u64(&mut encoded, data);
54 self.code.extend_from_slice(&encoded[..bytes_written]);
55 }
56
57 pub fn read_uvarint(&self, idx: usize) -> (u64, usize) {
58 debug_assert!(
59 idx < self.code.len(),
60 "invalid bytecode (missing varint operand)",
61 );
62
63 if self.code.len() - idx >= U64_VARINT_SIZE {
64 vu128::decode_u64(
65 &self.code[idx..idx + U64_VARINT_SIZE]
66 .try_into()
67 .expect("size statically checked"),
68 )
69 } else {
70 let mut tmp = [0u8; U64_VARINT_SIZE];
71 tmp[..self.code.len() - idx].copy_from_slice(&self.code[idx..]);
72 vu128::decode_u64(&tmp)
73 }
74 }
75
76 pub fn push_u16(&mut self, data: u16) {
77 self.code.extend_from_slice(&data.to_le_bytes())
78 }
79
80 pub fn patch_jump(&mut self, idx: usize) {
83 let offset = (self.code.len() - idx - 1 - 2) as u16;
84 self.code[idx + 1..idx + 3].copy_from_slice(&offset.to_le_bytes())
85 }
86
87 pub fn read_u16(&self, idx: usize) -> u16 {
88 if idx + 2 > self.code.len() {
89 panic!("Snix bug: invalid bytecode (expected u16 operand not found)")
90 }
91
92 let byte_array: &[u8; 2] = &self.code[idx..idx + 2]
93 .try_into()
94 .expect("fixed-size slice can not fail to convert to array");
95
96 u16::from_le_bytes(*byte_array)
97 }
98
99 pub fn first_span(&self) -> codemap::Span {
101 self.spans[0].span
102 }
103
104 pub fn last_op(&self) -> Option<(Op, usize)> {
106 if self.code.is_empty() {
107 return None;
108 }
109
110 Some((self.code[self.last_op].into(), self.last_op))
111 }
112
113 pub fn push_constant(&mut self, data: Value) -> ConstantIdx {
114 let idx = self.constants.len();
115 self.constants.push(data);
116 ConstantIdx(idx)
117 }
118
119 pub fn get_constant(&self, constant: ConstantIdx) -> Option<&Value> {
121 self.constants.get(constant.0)
122 }
123
124 fn push_span(&mut self, span: codemap::Span, start: usize) {
125 match self.spans.last_mut() {
126 Some(last) if last.span == span => {}
130
131 _ => self.spans.push(SourceSpan { span, start }),
133 }
134 }
135
136 pub fn get_span(&self, offset: CodeIdx) -> codemap::Span {
139 let position = self
140 .spans
141 .binary_search_by(|span| span.start.cmp(&offset.0));
142
143 let span = match position {
144 Ok(index) => &self.spans[index],
145 Err(index) => {
146 if index == 0 {
147 &self.spans[0]
148 } else {
149 &self.spans[index - 1]
150 }
151 }
152 };
153
154 span.span
155 }
156
157 pub fn disassemble_op<W: Write>(
161 &self,
162 writer: &mut W,
163 source: &SourceCode,
164 width: usize,
165 idx: CodeIdx,
166 ) -> Result<usize, std::io::Error> {
167 write!(writer, "{:#width$x}\t ", idx.0, width = width)?;
168
169 let line = source.get_line(self.get_span(idx));
172 if idx.0 > 0 && source.get_line(self.get_span(idx - 1)) == line {
173 write!(writer, " |\t")?;
174 } else {
175 write!(writer, "{:4}\t", line)?;
176 }
177
178 let _fmt_constant = |idx: ConstantIdx| match &self.constants[idx.0] {
179 Value::Thunk(t) => t.debug_repr(),
180 Value::Closure(c) => format!("closure({:p})", c.lambda),
181 Value::Blueprint(b) => format!("blueprint({:p})", b),
182 val => format!("{}", val),
183 };
184
185 let op: Op = self.code[idx.0].into();
186
187 match op.arg_type() {
188 OpArg::None => {
189 writeln!(writer, "Op{:?}", op)?;
190 Ok(1)
191 }
192
193 OpArg::Fixed => {
194 let arg = self.read_u16(idx.0 + 1);
195 writeln!(writer, "Op{:?}({})", op, arg)?;
196 Ok(3)
197 }
198
199 OpArg::Uvarint => {
200 let (arg, size) = self.read_uvarint(idx.0 + 1);
201 writeln!(writer, "Op{:?}({})", op, arg)?;
202 Ok(1 + size)
203 }
204
205 _ => match op {
206 Op::CoerceToString => {
207 let kind: CoercionKind = self.code[idx.0 + 1].into();
208 writeln!(writer, "Op{:?}({:?})", op, kind)?;
209 Ok(2)
210 }
211
212 Op::Closure | Op::ThunkClosure | Op::ThunkSuspended => {
213 let mut cidx = idx.0 + 1;
214
215 let (bp_idx, size) = self.read_uvarint(cidx);
216 cidx += size;
217
218 let (packed_count, size) = self.read_uvarint(cidx);
219 cidx += size;
220
221 let captures_with = packed_count & 0b1 == 1;
222 let count = packed_count >> 1;
223
224 write!(writer, "Op{:?}(BP @ {}, ", op, bp_idx)?;
225 if captures_with {
226 write!(writer, "captures with, ")?;
227 }
228 writeln!(writer, "{} upvalues)", count)?;
229
230 for _ in 0..count {
231 let (_, size) = self.read_uvarint(cidx);
232 cidx += size;
233 }
234
235 Ok(cidx - idx.0)
236 }
237 _ => panic!("Snix bug: don't know how to format argument for Op{:?}", op),
238 },
239 }
240 }
241}
242
243#[cfg(test)]
244mod tests {
245 use super::*;
246 use crate::test_utils::dummy_span;
247
248 #[test]
253 fn push_op() {
254 let mut chunk = Chunk::default();
255 let idx = chunk.push_op(Op::Add, dummy_span());
256 assert_eq!(*chunk.code.last().unwrap(), Op::Add as u8);
257 assert_eq!(chunk.code[idx], Op::Add as u8);
258 }
259
260 #[test]
261 fn push_op_with_arg() {
262 let mut chunk = Chunk::default();
263 let mut idx = chunk.push_op(Op::Constant, dummy_span());
264 chunk.push_uvarint(42);
265
266 assert_eq!(chunk.code[idx], Op::Constant as u8);
267
268 idx += 1;
269 let (arg, size) = chunk.read_uvarint(idx);
270 assert_eq!(idx + size, chunk.code.len());
271 assert_eq!(arg, 42);
272 }
273
274 #[test]
275 fn push_jump() {
276 let mut chunk = Chunk::default();
277
278 chunk.push_op(Op::Constant, dummy_span());
279 chunk.push_uvarint(0);
280
281 let idx = chunk.push_op(Op::Jump, dummy_span());
282 chunk.push_u16(0);
283
284 chunk.push_op(Op::Constant, dummy_span());
285 chunk.push_uvarint(1);
286
287 chunk.patch_jump(idx);
288 chunk.push_op(Op::Return, dummy_span());
289
290 #[rustfmt::skip]
291 let expected: Vec<u8> = vec![
292 Op::Constant as u8, 0,
293 Op::Jump as u8, 2, 0,
294 Op::Constant as u8, 1,
295 Op::Return as u8,
296 ];
297
298 assert_eq!(chunk.code, expected);
299 }
300}