1use super::{Instruction, Value};
2use crate::ast::TypeKind;
3use alloc::{format, string::String, vec::Vec};
4use hashbrown::HashMap;
5#[derive(Debug, Clone)]
6pub struct Chunk {
7 pub instructions: Vec<Instruction>,
8 pub constants: Vec<Value>,
9 pub lines: Vec<usize>,
10}
11
12impl Chunk {
13 pub fn new() -> Self {
14 Self {
15 instructions: Vec::new(),
16 constants: Vec::new(),
17 lines: Vec::new(),
18 }
19 }
20
21 pub fn emit(&mut self, instruction: Instruction, line: usize) -> usize {
22 let idx = self.instructions.len();
23 self.instructions.push(instruction);
24 self.lines.push(line);
25 idx
26 }
27
28 pub fn add_constant(&mut self, value: Value) -> u16 {
29 if let Some(idx) = self.constants.iter().position(|v| v == &value) {
30 return idx as u16;
31 }
32
33 let idx = self.constants.len();
34 if idx > u16::MAX as usize {
35 panic!("Too many constants in chunk (max: {})", u16::MAX);
36 }
37
38 self.constants.push(value);
39 idx as u16
40 }
41
42 pub fn patch_jump(&mut self, jump_idx: usize, target_idx: usize) {
43 let offset = (target_idx as isize - jump_idx as isize - 1) as i16;
44 match &mut self.instructions[jump_idx] {
45 Instruction::Jump(ref mut off) => *off = offset,
46 Instruction::JumpIf(_, ref mut off) => *off = offset,
47 Instruction::JumpIfNot(_, ref mut off) => *off = offset,
48 _ => panic!("Attempted to patch non-jump instruction"),
49 }
50 }
51
52 pub fn disassemble(&self, name: &str) -> String {
53 let mut output = String::new();
54 output.push_str(&format!("===== {} =====\n", name));
55 output.push_str(&format!("Constants: {}\n", self.constants.len()));
56 for (i, constant) in self.constants.iter().enumerate() {
57 output.push_str(&format!(" K{}: {}\n", i, constant));
58 }
59
60 output.push_str("\nInstructions:\n");
61 for (i, instruction) in self.instructions.iter().enumerate() {
62 let line = self.lines.get(i).copied().unwrap_or(0);
63 output.push_str(&format!("{:04} [L{:03}] {}\n", i, line, instruction));
64 }
65
66 output
67 }
68}
69
70impl Default for Chunk {
71 fn default() -> Self {
72 Self::new()
73 }
74}
75
76#[derive(Debug, Clone)]
77pub struct Function {
78 pub name: String,
79 pub param_count: u8,
80 pub register_count: u8,
81 pub is_method: bool,
82 pub chunk: Chunk,
83 pub upvalues: Vec<(bool, u8)>,
84 pub register_types: HashMap<u8, TypeKind>,
85}
86
87impl Function {
88 pub fn new(name: impl Into<String>, param_count: u8, is_method: bool) -> Self {
89 Self {
90 name: name.into(),
91 param_count,
92 register_count: 0,
93 is_method,
94 chunk: Chunk::new(),
95 upvalues: Vec::new(),
96 register_types: HashMap::new(),
97 }
98 }
99
100 pub fn set_register_count(&mut self, count: u8) {
101 self.register_count = count;
102 }
103
104 pub fn add_upvalue(&mut self, is_local: bool, index: u8) -> u8 {
105 let idx = self.upvalues.len();
106 if idx > u8::MAX as usize {
107 panic!("Too many upvalues in function (max: 255)");
108 }
109
110 self.upvalues.push((is_local, index));
111 idx as u8
112 }
113
114 pub fn disassemble(&self) -> String {
115 let mut output = String::new();
116 output.push_str(&format!(
117 "Function: {} (params: {}, registers: {}, method: {})\n",
118 self.name, self.param_count, self.register_count, self.is_method
119 ));
120 if !self.upvalues.is_empty() {
121 output.push_str("Upvalues:\n");
122 for (i, (is_local, idx)) in self.upvalues.iter().enumerate() {
123 let kind = if *is_local { "local" } else { "upvalue" };
124 output.push_str(&format!(" U{}: {} {}\n", i, kind, idx));
125 }
126 }
127
128 output.push_str(&self.chunk.disassemble(&self.name));
129 output
130 }
131}