lunar_wave/src/state.rs

use std::{
	collections::HashMap,
	rc::Rc,
};

use crate::{
	instruction::Instruction,
	value::{
		BogusClosure,
		Table,
		Value,
	},
};

pub struct Block {
	pub instructions: Vec <Instruction>,
	pub constants: Vec <Value>,
	pub upvalue_count: usize,
}

pub struct Chunk {
	pub blocks: Vec <Block>,
}

#[derive (Clone, Debug)]
struct StackFrame {
	// i32 makes it a little easier to implement jumps
	// Starts at 0 right after OP_CALL
	
	program_counter: i32,
	
	// Starts from 0 for main and 1 for the first closure
	block_idx: usize,
	
	register_offset: usize,
}

#[derive (Debug)]
pub struct Breakpoint {
	pub block_idx: usize,
	pub program_counter: i32,
}

#[derive (Debug)]
pub struct State {
	registers: Vec <Value>,
	stack: Vec <StackFrame>,
	
	pub debug_print: bool,
	pub breakpoints: Vec <Breakpoint>,
	step_count: u32,
}

impl Default for State {
	fn default () -> Self {
		Self {
			registers: vec! [Value::Nil; 16],
			stack: vec! [
				StackFrame {
					program_counter: 0,
					block_idx: 0,
					register_offset: 0,
				},
			],
			debug_print: false,
			breakpoints: Default::default(),
			step_count: 0,
		}
	}
}

impl State {
	pub fn upvalues_from_args <I: Iterator <Item = String>> (args: I) -> Vec <Value>
	{
		let arg: Vec <_> = args.map (|s| s.to_string ()).collect ();
		
		let env = HashMap::from_iter ([
			("arg", Value::BogusArg (arg.into ())),
			("print", Value::BogusPrint),
		].map (|(k, v)| (k.to_string (), v)));
		
		vec! [
			Value::BogusEnv (env.into ()),
		]
	}
	
	fn register_window (&self) -> &[Value] {
		let frame = self.stack.last ().unwrap ();
		&self.registers [frame.register_offset..]
	}
	
	/// Short form to get access to a register within our window
	
	fn reg (&self, i: u8) -> &Value {
		let frame = self.stack.last ().unwrap ();
		&self.registers [frame.register_offset + i as usize]
	}
	
	fn reg_mut (&mut self, i: u8) -> &mut Value {
		let frame = self.stack.last ().unwrap ();
		&mut self.registers [frame.register_offset + i as usize]
	}
	
	pub fn execute_chunk (&mut self, chunk: &Chunk, upvalues: &[Value]) 
	-> Vec <Value> {
		let max_iters = 2000;
		
		for _ in 0..max_iters {
			self.step_count += 1;
			
			let frame = self.stack.last_mut ().unwrap ().clone ();
			let block = chunk.blocks.get (frame.block_idx).unwrap ();
			
			for bp in &self.breakpoints {
				if frame.block_idx == bp.block_idx && frame.program_counter == bp.program_counter
				{
					dbg! (&self);
				}
			}
			
			let mut next_pc = frame.program_counter;
			
			let pc = usize::try_from (frame.program_counter).expect ("program_counter is not a valid usize");
			let instruction = match block.instructions.get (pc) {
				Some (x) => x,
				None => {
					dbg! (&self.stack);
					panic! ("program_counter went out of bounds");
				}
			};
			
			// let r = &mut self.registers [frame.register_offset..];
			let k = &block.constants;
			
			match instruction {
				Instruction::Add (a, b, c) => {
					let v_b = self.reg (*b).as_float ().unwrap ();
					let v_c = self.reg (*c).as_float ().unwrap ();
					
					*self.reg_mut (*a) = Value::from (v_b + v_c);
				},
				Instruction::Call (a, _b, c) => {
					// Take arguments from registers [a + 1, a + b)
					// Call the function in register [a]
					// Return values in registers [a, a + c - 1)
					// 
					// That is, call a with b - 1 arguments and expect c returns
					// 
					// e.g. CALL 0 2 1 mean "Call 0 with 1 argument, return 1 value", like for printing a constant
					
					// TODO: Only implement printing values for now
					
					let v_a = self.reg (*a);
					
					match v_a {
						Value::BogusClosure (rc) => {
							let idx = rc.idx;
							
							let block_idx = frame.block_idx;
							let target_block = idx;
							
							let current_frame = self.stack.last ().unwrap ();
							
							self.stack.push (StackFrame {
								program_counter: 0,
								block_idx: target_block,
								register_offset: current_frame.register_offset + *a as usize + 1,
							});
							
							if self.debug_print {
								println! ("Inst {block_idx}:{pc} calls {target_block}:0");
								let stack_depth = self.stack.len ();
								println! ("stack_depth: {stack_depth}");
							}
							
							// Skip the PC increment at the bottom of the loop
							continue;
						},
						Value::BogusPrint => {
							// In real Lua, print is a function inside
							// the runtime. Here it's bogus.
							
							// assert_eq! (*b, 2);
							assert_eq! (*c, 1);
							
							let value = self.reg (a + 1);
							println! ("{}", value);
							
							*self.reg_mut (*a) = self.reg_mut (*a + 1).take ();
						},
						_ => {
							let stack = &self.stack;
							panic! ("Cannot call value {a:?}. backtrace: {stack:?}");
						},
					}
				},
				Instruction::Closure (a, b) => {
					let b = usize::try_from (*b).unwrap ();
					
					*self.reg_mut (*a) = Value::BogusClosure (BogusClosure {
						idx: b + frame.block_idx + 1,
						upvalues: vec! [],
					}.into ());
				},
				Instruction::EqK (a, b, k_flag) => {
					let b = usize::from (*b);
					
					if (*self.reg (*a) == k [b]) != *k_flag {
						next_pc += 1;
					}
				},
				Instruction::ExtraArg (ax) => {
					// This is used for NewTable. Maybe it's for reserving
					// capacity in the array or something?
					
					assert_eq! (*ax, 0, "implemented only for ax == 0");
				},
				Instruction::GetField (a, b, c) => {
					let t = match self.reg (*b) {
						Value::Table (t) => t,
						_ => panic! ("R[B] must be a table"),
					};
					
					let key = match &k [usize::from (*c)] {
						Value::String (s) => s,
						_ => panic! ("K[C] must be a string"),
					};
					
					*self.reg_mut (*a) = t.get (Value::String (Rc::clone (key)));
				},
				Instruction::GetTable (a, b, c) => {
					let t = match self.reg (*b) {
						Value::Table (t) => t,
						_ => panic! ("R[B] must be a table"),
					};
					
					let key = self.reg (*c);
					
					let val = t.get (key.clone ());
					
					*self.reg_mut (*a) = val;
				},
				Instruction::GetTabUp (a, b, c) => {
					let b = usize::try_from (*b).unwrap ();
					let c = usize::try_from (*c).unwrap ();
					
					let env = match upvalues.get (b).unwrap () {
						Value::BogusEnv (x) => x,
						_ => panic! ("Only allowed upvalue is BogusEnv"),
					};
					
					let key = match k.get (c).unwrap () {
						Value::String (s) => s.as_ref (),
						_ => panic! ("GetTabUp only supports string keys"),
					};
					
					let value = env.get (key).unwrap ();
					*self.reg_mut (*a) = value.clone ();
				},
				Instruction::GetI (a, b, c) => {
					let key = usize::try_from (*c).unwrap ();
					
					let table = self.reg (*b);
					let value = match table {
						Value::BogusArg (arg) => arg.get (key).map (|x| x.as_str().into ()).unwrap_or_default(),
						Value::Table (t) => {
							let key = Value::from (i64::try_from (key).unwrap ());
							t.get (key)
						},
						_ => unimplemented! (),
					};
					
					*self.reg_mut (*a) = value;
				},
				Instruction::GetUpVal (a, b) => {
					let this_func = self.stack.last ().unwrap ().register_offset - 1;
					let upvalues = match &self.registers [this_func] {
						Value::BogusClosure (rc) => &rc.upvalues,
						_ => panic! ("Can't do GetUpVal outside a closure"),
					};
					
					let b  = usize::try_from  (*b).unwrap ();
					
					*self.reg_mut (*a) = upvalues [b].clone ();
				},
				Instruction::Jmp (s_j) => next_pc += s_j,
				Instruction::LoadF (a, sbx) => {
					*self.reg_mut (*a) = Value::Float (*sbx as f64);
				}
				Instruction::LoadFalse (a) => {
					*self.reg_mut (*a) = false.into ();
				},
				Instruction::LoadI (a, sbx) => {
					*self.reg_mut (*a) = Value::Integer (*sbx as i64);
				},
				Instruction::LoadK (a, bx) => {
					let bx = usize::try_from (*bx).unwrap ();
					
					*self.reg_mut (*a) = k [bx].clone ();
				},
				Instruction::LoadNil (a) => {
					*self.reg_mut (*a) = Value::Nil;
				},
				Instruction::LoadTrue (a) => {
					*self.reg_mut (*a) = true.into ();
				},
				Instruction::MmBin (a, b, _c) => {
					let a = self.reg (*a);
					let b = self.reg (*b);
					
					if a.as_float().is_some() && b.as_float().is_some () {
						// No need for metamethods
					}
					else {
						panic! ("Not sure how to implememtn OP_MMBIN for these 2 values {a:?}, {b:?}");
					}
				},
				Instruction::Move (a, b) => {
					*self.reg_mut (*a) = self.reg (*b).clone ();
				},
				Instruction::Mul (_a, _b, _c) => unimplemented!(),
				Instruction::NewTable (a) => {
					*self.reg_mut (*a) = Value::Table (Default::default ());
				},
				Instruction::Not (a, b) => {
					*self.reg_mut (*a) = Value::Boolean (! self.reg (*b).is_truthy());
				}
				Instruction::Return (a, b, _c, k) => {
					let a = usize::try_from (*a).unwrap ();
					let b = usize::try_from (*b).unwrap ();
					
					let popped_frame = self.stack.pop ().unwrap ();
					
					// Build closure if needed
					if *k {
						
						let closure_idx = match &self.registers [popped_frame.register_offset + a] {
							Value::BogusClosure (rc) => rc.idx,
							_ => panic! ("Impossible"),
						};
						
						let upvalue_count = chunk.blocks [closure_idx].upvalue_count;
						
						let start_reg = a + popped_frame.register_offset - upvalue_count;
						
						let upvalues = self.registers [start_reg..start_reg+upvalue_count].iter ().cloned ().collect ();
						
						self.registers [a + popped_frame.register_offset] = Value::BogusClosure (BogusClosure {
							idx: closure_idx,
							upvalues,
						}.into ());
					}
					
					if self.debug_print {
						let old_block = popped_frame.block_idx;
						let old_pc = popped_frame.program_counter;
						println! ("Inst {old_block}:{old_pc} returns");
						let stack_depth = self.stack.len ();
						println! ("stack_depth: {stack_depth}");
					}
					
					if let Some (new_frame) = self.stack.last() {
						next_pc = new_frame.program_counter;
						
						// Shift our output registers down so the caller
						// can grab them
						// idk exactly why Lua does this
						
						// Register that our function was in before we
						// called it.
						
						let offset = popped_frame.register_offset - 1;
						for i in (offset)..(offset - 1 + b) {
							self.registers [i] = self.registers [i + 1 + a].take ();
						}
					}
					else {
						// Return from the entire program
						return self.registers [a..(a + b - 1)].to_vec();
					}
				},
				Instruction::Return0 => unimplemented! (),
				Instruction::Return1 (a) => {
					let a = usize::try_from (*a).unwrap ();
					let popped_frame = self.stack.pop ().unwrap ();
					
					
					self.registers [popped_frame.register_offset - 1] = self.register_window ()[a].clone ();
					
					let frame = self.stack.last ().unwrap ();
					let new_block = frame.block_idx;
					next_pc = frame.program_counter;
					
					if self.debug_print {
						let old_block = popped_frame.block_idx;
						let old_pc = popped_frame.program_counter;
						println! ("Inst {old_block}:{old_pc} returns to inst {new_block}:{next_pc}");
						let stack_depth = self.stack.len ();
						println! ("stack_depth: {stack_depth}");
					}
					
					// Shift output register down
					let offset = popped_frame.register_offset;
					self.registers [offset - 1] = self.registers [offset + a].take ();
				},
				Instruction::SetField (a, b, c, k_flag) => {
					let value = if *k_flag {
						&k [usize::from (*c)]
					}
					else {
						self.reg (*c)
					}
					.clone ();
					
					let b = usize::try_from (*b).unwrap ();
					
					let key = match k.get (b).unwrap () {
						Value::String (s) => s.as_ref (),
						_ => panic! ("GetTabUp only supports string keys"),
					};
					
					let dst = self.reg_mut (*a).as_table_mut ()
					.expect ("SetField only works on tables");
					
					Rc::get_mut (dst).expect ("shrug").insert (Value::from (key.as_str ()), value);
				},
				Instruction::SetI (a, b, c, k_flag) => {
					let value = if *k_flag {
						&k [usize::from (*c)]
					}
					else {
						self.reg (*c)
					}
					.clone ();
					
					let dst = self.reg_mut (*a).as_table_mut ()
					.expect ("SetI only works on tables");
					
					Rc::get_mut (dst).expect ("shrug").insert (i64::from (*b), value);
				},
				Instruction::SetList (a, b, c, k) => {
					if *b == 0 {
						panic! ("SetList with b == 0 not implemented");
					}
					if *k {
						panic! ("SetList with k = true not implemented");
					}
					
					// Temporarily move the table out
					
					let mut dst = Rc::into_inner (std::mem::replace (self.reg_mut (*a).as_table_mut ().expect ("SetList only works on tables"), Default::default ())).unwrap ();
					
					for i in 1..=*b {
						let src = self.reg (*a + i);
						dst.insert (Value::from (i64::from (*c + i)), src.clone ());
					}
					
					// Put the table back and pretend it's all okay
					
					*self.reg_mut (*a) = Value::Table (dst.into ());
				},
				Instruction::SetTabUp (_a, _b, _c) => unimplemented! (),
				Instruction::TailCall (_a, _b, _c, _k) => unimplemented! (),
				Instruction::Test (a, _k) => {
					if self.reg (*a).is_truthy() {
						next_pc += 1;
					}
				},
				Instruction::VarArgPrep (_) => (),
			}
			
			next_pc += 1;
			{
				let frame = self.stack.last_mut ().unwrap ();
				frame.program_counter = next_pc;
			}
		}
		
		panic! ("Hit max iterations before block returned");
	}
}