lust/bytecode/
chunk.rs

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}