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