ruchy_notebook/vm/
interpreter.rs

1use anyhow::{Result, bail};
2use super::bytecode::{OpCode, Value, Instruction, BytecodeModule};
3use std::collections::HashMap;
4
5#[derive(Debug, Clone)]
6pub struct ExecutionResult {
7    pub value: Option<Value>,
8    pub output: Vec<String>,
9}
10
11pub struct VirtualMachine {
12    stack: Vec<Value>,
13    call_stack: Vec<usize>,
14    locals: HashMap<usize, Value>,
15    globals: HashMap<String, Value>,
16    output: Vec<String>,
17    pc: usize,
18}
19
20impl VirtualMachine {
21    pub fn new() -> Self {
22        Self {
23            stack: Vec::with_capacity(256),
24            call_stack: Vec::with_capacity(32),
25            locals: HashMap::new(),
26            globals: HashMap::new(),
27            output: Vec::new(),
28            pc: 0,
29        }
30    }
31    
32    pub fn execute(&mut self, module: &BytecodeModule) -> Result<ExecutionResult> {
33        self.pc = module.entry_point;
34        
35        while self.pc < module.instructions.len() {
36            let instruction = &module.instructions[self.pc];
37            
38            match instruction.opcode {
39                OpCode::Push => self.execute_push(instruction)?,
40                OpCode::Pop => self.execute_pop()?,
41                OpCode::Dup => self.execute_dup()?,
42                OpCode::Swap => self.execute_swap()?,
43                
44                OpCode::Add => self.execute_binary_op(|a, b| a + b)?,
45                OpCode::Sub => self.execute_binary_op(|a, b| a - b)?,
46                OpCode::Mul => self.execute_binary_op(|a, b| a * b)?,
47                OpCode::Div => self.execute_div()?,
48                OpCode::Mod => self.execute_mod()?,
49                OpCode::Neg => self.execute_neg()?,
50                
51                OpCode::Eq => self.execute_comparison(|a, b| a == b)?,
52                OpCode::Ne => self.execute_comparison(|a, b| a != b)?,
53                OpCode::Lt => self.execute_comparison(|a, b| a < b)?,
54                OpCode::Le => self.execute_comparison(|a, b| a <= b)?,
55                OpCode::Gt => self.execute_comparison(|a, b| a > b)?,
56                OpCode::Ge => self.execute_comparison(|a, b| a >= b)?,
57                
58                OpCode::And => self.execute_and()?,
59                OpCode::Or => self.execute_or()?,
60                OpCode::Not => self.execute_not()?,
61                
62                OpCode::Jump => self.execute_jump(instruction)?,
63                OpCode::JumpIf => self.execute_jump_if(instruction)?,
64                OpCode::JumpIfNot => self.execute_jump_if_not(instruction)?,
65                
66                OpCode::LoadLocal => self.execute_load_local(instruction)?,
67                OpCode::StoreLocal => self.execute_store_local(instruction)?,
68                OpCode::LoadGlobal => self.execute_load_global(instruction)?,
69                OpCode::StoreGlobal => self.execute_store_global(instruction)?,
70                
71                OpCode::BuildList => self.execute_build_list(instruction)?,
72                OpCode::BuildMap => self.execute_build_map(instruction)?,
73                OpCode::Index => self.execute_index()?,
74                OpCode::SetIndex => self.execute_set_index()?,
75                
76                OpCode::Call => self.execute_call(instruction)?,
77                OpCode::Return => self.execute_return()?,
78                
79                OpCode::Print => self.execute_print()?,
80                OpCode::Halt => break,
81                
82                _ => bail!("Unimplemented opcode: {:?}", instruction.opcode),
83            }
84            
85            self.pc += 1;
86        }
87        
88        Ok(ExecutionResult {
89            value: self.stack.last().cloned(),
90            output: self.output.clone(),
91        })
92    }
93    
94    fn execute_push(&mut self, inst: &Instruction) -> Result<()> {
95        let value = inst.operand.as_ref()
96            .ok_or_else(|| anyhow::anyhow!("Push requires an operand"))?;
97        self.stack.push(value.clone());
98        Ok(())
99    }
100    
101    fn execute_pop(&mut self) -> Result<()> {
102        self.stack.pop()
103            .ok_or_else(|| anyhow::anyhow!("Stack underflow"))?;
104        Ok(())
105    }
106    
107    fn execute_dup(&mut self) -> Result<()> {
108        let value = self.stack.last()
109            .ok_or_else(|| anyhow::anyhow!("Stack underflow"))?
110            .clone();
111        self.stack.push(value);
112        Ok(())
113    }
114    
115    fn execute_swap(&mut self) -> Result<()> {
116        let len = self.stack.len();
117        if len < 2 {
118            bail!("Stack underflow: need 2 values to swap");
119        }
120        self.stack.swap(len - 1, len - 2);
121        Ok(())
122    }
123    
124    fn execute_binary_op(&mut self, op: fn(i64, i64) -> i64) -> Result<()> {
125        let b = self.pop_int()?;
126        let a = self.pop_int()?;
127        self.stack.push(Value::Int(op(a, b)));
128        Ok(())
129    }
130    
131    fn execute_div(&mut self) -> Result<()> {
132        let b = self.pop_int()?;
133        let a = self.pop_int()?;
134        if b == 0 {
135            bail!("Division by zero");
136        }
137        self.stack.push(Value::Int(a / b));
138        Ok(())
139    }
140    
141    fn execute_mod(&mut self) -> Result<()> {
142        let b = self.pop_int()?;
143        let a = self.pop_int()?;
144        if b == 0 {
145            bail!("Modulo by zero");
146        }
147        self.stack.push(Value::Int(a % b));
148        Ok(())
149    }
150    
151    fn execute_neg(&mut self) -> Result<()> {
152        let value = self.pop_int()?;
153        self.stack.push(Value::Int(-value));
154        Ok(())
155    }
156    
157    fn execute_comparison(&mut self, op: fn(i64, i64) -> bool) -> Result<()> {
158        let b = self.pop_int()?;
159        let a = self.pop_int()?;
160        self.stack.push(Value::Bool(op(a, b)));
161        Ok(())
162    }
163    
164    fn execute_and(&mut self) -> Result<()> {
165        let b = self.pop_bool()?;
166        let a = self.pop_bool()?;
167        self.stack.push(Value::Bool(a && b));
168        Ok(())
169    }
170    
171    fn execute_or(&mut self) -> Result<()> {
172        let b = self.pop_bool()?;
173        let a = self.pop_bool()?;
174        self.stack.push(Value::Bool(a || b));
175        Ok(())
176    }
177    
178    fn execute_not(&mut self) -> Result<()> {
179        let value = self.pop_bool()?;
180        self.stack.push(Value::Bool(!value));
181        Ok(())
182    }
183    
184    fn execute_jump(&mut self, inst: &Instruction) -> Result<()> {
185        let target = self.get_jump_target(inst)?;
186        self.pc = target.saturating_sub(1); // -1 because pc++ happens after
187        Ok(())
188    }
189    
190    fn execute_jump_if(&mut self, inst: &Instruction) -> Result<()> {
191        let condition = self.pop_bool()?;
192        if condition {
193            let target = self.get_jump_target(inst)?;
194            self.pc = target.saturating_sub(1);
195        }
196        Ok(())
197    }
198    
199    fn execute_jump_if_not(&mut self, inst: &Instruction) -> Result<()> {
200        let condition = self.pop_bool()?;
201        if !condition {
202            let target = self.get_jump_target(inst)?;
203            self.pc = target.saturating_sub(1);
204        }
205        Ok(())
206    }
207    
208    fn execute_load_local(&mut self, inst: &Instruction) -> Result<()> {
209        let index = self.get_index(inst)?;
210        let value = self.locals.get(&index)
211            .ok_or_else(|| anyhow::anyhow!("Undefined local variable at index {}", index))?
212            .clone();
213        self.stack.push(value);
214        Ok(())
215    }
216    
217    fn execute_store_local(&mut self, inst: &Instruction) -> Result<()> {
218        let index = self.get_index(inst)?;
219        let value = self.stack.pop()
220            .ok_or_else(|| anyhow::anyhow!("Stack underflow"))?;
221        self.locals.insert(index, value);
222        Ok(())
223    }
224    
225    fn execute_load_global(&mut self, inst: &Instruction) -> Result<()> {
226        let name = self.get_string(inst)?;
227        let value = self.globals.get(&name)
228            .ok_or_else(|| anyhow::anyhow!("Undefined global variable: {}", name))?
229            .clone();
230        self.stack.push(value);
231        Ok(())
232    }
233    
234    fn execute_store_global(&mut self, inst: &Instruction) -> Result<()> {
235        let name = self.get_string(inst)?;
236        let value = self.stack.pop()
237            .ok_or_else(|| anyhow::anyhow!("Stack underflow"))?;
238        self.globals.insert(name, value);
239        Ok(())
240    }
241    
242    fn execute_print(&mut self) -> Result<()> {
243        let value = self.stack.pop()
244            .ok_or_else(|| anyhow::anyhow!("Stack underflow"))?;
245        self.output.push(format!("{:?}", value));
246        Ok(())
247    }
248    
249    fn execute_build_list(&mut self, inst: &Instruction) -> Result<()> {
250        let count = self.get_index(inst)?;
251        if self.stack.len() < count {
252            bail!("Stack underflow: need {} values for list", count);
253        }
254        
255        let mut elements = Vec::with_capacity(count);
256        for _ in 0..count {
257            elements.push(self.stack.pop().unwrap());
258        }
259        elements.reverse(); // Restore original order
260        
261        self.stack.push(Value::List(elements));
262        Ok(())
263    }
264    
265    fn execute_build_map(&mut self, inst: &Instruction) -> Result<()> {
266        let count = self.get_index(inst)?;
267        if self.stack.len() < count * 2 {
268            bail!("Stack underflow: need {} key-value pairs for map", count);
269        }
270        
271        let mut pairs = Vec::with_capacity(count);
272        for _ in 0..count {
273            let value = self.stack.pop().unwrap();
274            let key = match self.stack.pop().unwrap() {
275                Value::String(s) => s,
276                _ => bail!("Map keys must be strings"),
277            };
278            pairs.push((key, value));
279        }
280        pairs.reverse(); // Restore original order
281        
282        self.stack.push(Value::Map(pairs));
283        Ok(())
284    }
285    
286    fn execute_index(&mut self) -> Result<()> {
287        let index = self.stack.pop().ok_or_else(|| anyhow::anyhow!("Stack underflow"))?;
288        let container = self.stack.pop().ok_or_else(|| anyhow::anyhow!("Stack underflow"))?;
289        
290        match (container, index) {
291            (Value::List(list), Value::Int(idx)) => {
292                if idx < 0 || idx as usize >= list.len() {
293                    bail!("List index out of bounds: {}", idx);
294                }
295                self.stack.push(list[idx as usize].clone());
296            }
297            (Value::Map(map), Value::String(key)) => {
298                let value = map.iter()
299                    .find(|(k, _)| k == &key)
300                    .map(|(_, v)| v.clone())
301                    .unwrap_or(Value::Null);
302                self.stack.push(value);
303            }
304            _ => bail!("Invalid types for indexing operation"),
305        }
306        Ok(())
307    }
308    
309    fn execute_set_index(&mut self) -> Result<()> {
310        let value = self.stack.pop().ok_or_else(|| anyhow::anyhow!("Stack underflow"))?;
311        let index = self.stack.pop().ok_or_else(|| anyhow::anyhow!("Stack underflow"))?;
312        let mut container = self.stack.pop().ok_or_else(|| anyhow::anyhow!("Stack underflow"))?;
313        
314        match (&mut container, index) {
315            (Value::List(list), Value::Int(idx)) => {
316                if idx < 0 || idx as usize >= list.len() {
317                    bail!("List index out of bounds: {}", idx);
318                }
319                list[idx as usize] = value;
320            }
321            (Value::Map(map), Value::String(key)) => {
322                if let Some(entry) = map.iter_mut().find(|(k, _)| k == &key) {
323                    entry.1 = value;
324                } else {
325                    map.push((key, value));
326                }
327            }
328            _ => bail!("Invalid types for set index operation"),
329        }
330        
331        self.stack.push(container);
332        Ok(())
333    }
334    
335    fn execute_call(&mut self, inst: &Instruction) -> Result<()> {
336        let target = self.get_jump_target(inst)?;
337        self.call_stack.push(self.pc + 1); // Save return address
338        self.pc = target.saturating_sub(1); // -1 because pc++ happens after
339        Ok(())
340    }
341    
342    fn execute_return(&mut self) -> Result<()> {
343        let return_addr = self.call_stack.pop()
344            .ok_or_else(|| anyhow::anyhow!("Call stack underflow"))?;
345        self.pc = return_addr.saturating_sub(1); // -1 because pc++ happens after
346        Ok(())
347    }
348    
349    // Helper methods
350    fn pop_int(&mut self) -> Result<i64> {
351        match self.stack.pop() {
352            Some(Value::Int(n)) => Ok(n),
353            Some(_) => bail!("Expected integer on stack"),
354            None => bail!("Stack underflow"),
355        }
356    }
357    
358    fn pop_bool(&mut self) -> Result<bool> {
359        match self.stack.pop() {
360            Some(Value::Bool(b)) => Ok(b),
361            Some(_) => bail!("Expected boolean on stack"),
362            None => bail!("Stack underflow"),
363        }
364    }
365    
366    fn get_jump_target(&self, inst: &Instruction) -> Result<usize> {
367        match &inst.operand {
368            Some(Value::Int(target)) => Ok(*target as usize),
369            _ => bail!("Jump instruction requires integer target"),
370        }
371    }
372    
373    fn get_index(&self, inst: &Instruction) -> Result<usize> {
374        match &inst.operand {
375            Some(Value::Int(idx)) => Ok(*idx as usize),
376            _ => bail!("Index operand must be integer"),
377        }
378    }
379    
380    fn get_string(&self, inst: &Instruction) -> Result<String> {
381        match &inst.operand {
382            Some(Value::String(s)) => Ok(s.clone()),
383            _ => bail!("String operand required"),
384        }
385    }
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391    
392    #[test]
393    fn test_vm_arithmetic() {
394        let mut vm = VirtualMachine::new();
395        let mut module = BytecodeModule::new();
396        
397        // 2 + 3 = 5
398        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(2)));
399        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(3)));
400        module.add_instruction(Instruction::new(OpCode::Add));
401        module.add_instruction(Instruction::new(OpCode::Halt));
402        
403        let result = vm.execute(&module).unwrap();
404        assert_eq!(result.value, Some(Value::Int(5)));
405    }
406    
407    #[test]
408    fn test_vm_comparison() {
409        let mut vm = VirtualMachine::new();
410        let mut module = BytecodeModule::new();
411        
412        // 5 > 3 = true
413        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(5)));
414        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(3)));
415        module.add_instruction(Instruction::new(OpCode::Gt));
416        module.add_instruction(Instruction::new(OpCode::Halt));
417        
418        let result = vm.execute(&module).unwrap();
419        assert_eq!(result.value, Some(Value::Bool(true)));
420    }
421    
422    #[test]
423    fn test_vm_print() {
424        let mut vm = VirtualMachine::new();
425        let mut module = BytecodeModule::new();
426        
427        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::String("Hello".to_string())));
428        module.add_instruction(Instruction::new(OpCode::Print));
429        module.add_instruction(Instruction::new(OpCode::Halt));
430        
431        let result = vm.execute(&module).unwrap();
432        assert_eq!(result.output, vec!["String(\"Hello\")"]);
433    }
434    
435    #[test]
436    fn test_vm_variables() {
437        let mut vm = VirtualMachine::new();
438        let mut module = BytecodeModule::new();
439        
440        // Store 42 in local[0], then load it
441        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(42)));
442        module.add_instruction(Instruction::with_operand(OpCode::StoreLocal, Value::Int(0)));
443        module.add_instruction(Instruction::with_operand(OpCode::LoadLocal, Value::Int(0)));
444        module.add_instruction(Instruction::new(OpCode::Halt));
445        
446        let result = vm.execute(&module).unwrap();
447        assert_eq!(result.value, Some(Value::Int(42)));
448    }
449    
450    #[test]
451    fn test_vm_jump() {
452        let mut vm = VirtualMachine::new();
453        let mut module = BytecodeModule::new();
454        
455        // Jump over the push 99
456        module.add_instruction(Instruction::with_operand(OpCode::Jump, Value::Int(2)));
457        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(99)));
458        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(42)));
459        module.add_instruction(Instruction::new(OpCode::Halt));
460        
461        let result = vm.execute(&module).unwrap();
462        assert_eq!(result.value, Some(Value::Int(42)));
463    }
464    
465    #[test]
466    fn test_vm_build_list() {
467        let mut vm = VirtualMachine::new();
468        let mut module = BytecodeModule::new();
469        
470        // Push elements and build list [1, 2, 3]
471        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(1)));
472        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(2)));
473        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(3)));
474        module.add_instruction(Instruction::with_operand(OpCode::BuildList, Value::Int(3)));
475        module.add_instruction(Instruction::new(OpCode::Halt));
476        
477        let result = vm.execute(&module).unwrap();
478        assert_eq!(result.value, Some(Value::List(vec![
479            Value::Int(1), Value::Int(2), Value::Int(3)
480        ])));
481    }
482    
483    #[test]
484    fn test_vm_build_map() {
485        let mut vm = VirtualMachine::new();
486        let mut module = BytecodeModule::new();
487        
488        // Build map {"key1": 42, "key2": 100}
489        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::String("key1".to_string())));
490        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(42)));
491        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::String("key2".to_string())));
492        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(100)));
493        module.add_instruction(Instruction::with_operand(OpCode::BuildMap, Value::Int(2)));
494        module.add_instruction(Instruction::new(OpCode::Halt));
495        
496        let result = vm.execute(&module).unwrap();
497        assert_eq!(result.value, Some(Value::Map(vec![
498            ("key1".to_string(), Value::Int(42)),
499            ("key2".to_string(), Value::Int(100))
500        ])));
501    }
502    
503    #[test]
504    fn test_vm_index_list() {
505        let mut vm = VirtualMachine::new();
506        let mut module = BytecodeModule::new();
507        
508        // Create list [10, 20, 30] and get index 1
509        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(10)));
510        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(20)));
511        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(30)));
512        module.add_instruction(Instruction::with_operand(OpCode::BuildList, Value::Int(3)));
513        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(1)));
514        module.add_instruction(Instruction::new(OpCode::Index));
515        module.add_instruction(Instruction::new(OpCode::Halt));
516        
517        let result = vm.execute(&module).unwrap();
518        assert_eq!(result.value, Some(Value::Int(20)));
519    }
520    
521    #[test]
522    fn test_vm_index_map() {
523        let mut vm = VirtualMachine::new();
524        let mut module = BytecodeModule::new();
525        
526        // Create map {"test": 123} and get "test"
527        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::String("test".to_string())));
528        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(123)));
529        module.add_instruction(Instruction::with_operand(OpCode::BuildMap, Value::Int(1)));
530        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::String("test".to_string())));
531        module.add_instruction(Instruction::new(OpCode::Index));
532        module.add_instruction(Instruction::new(OpCode::Halt));
533        
534        let result = vm.execute(&module).unwrap();
535        assert_eq!(result.value, Some(Value::Int(123)));
536    }
537    
538    #[test]
539    fn test_vm_call_return() {
540        let mut vm = VirtualMachine::new();
541        let mut module = BytecodeModule::new();
542        
543        // Main: call function at index 4, then halt
544        module.add_instruction(Instruction::with_operand(OpCode::Call, Value::Int(4)));  // 0
545        module.add_instruction(Instruction::new(OpCode::Halt));                         // 1
546        module.add_instruction(Instruction::new(OpCode::Halt));                         // 2 - padding
547        module.add_instruction(Instruction::new(OpCode::Halt));                         // 3 - padding
548        
549        // Function: push 42 and return
550        module.add_instruction(Instruction::with_operand(OpCode::Push, Value::Int(42))); // 4
551        module.add_instruction(Instruction::new(OpCode::Return));                       // 5
552        
553        let result = vm.execute(&module).unwrap();
554        assert_eq!(result.value, Some(Value::Int(42)));
555    }
556}