lust/bytecode/
chunk.rs

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

lust/bytecode/
chunk.rs
