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

lust/bytecode/
chunk.rs
