lust/bytecode/
chunk.rs

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}