♻️ refactor: trying to fix OP_RETURN so it can close over values

src/state.rs

@@ -66,7 +66,10 @@ pub enum Value {
 	// tables and function pointers yet
 	
 	BogusArg (Vec <String>),
-	BogusClosure (usize),
+	BogusClosure {
+		idx: usize,
+		upvalues: Vec <Value>,
+	},
 	BogusEnv (BTreeMap <String, Value>),
 	BogusPrint,
 }
@@ -131,7 +134,7 @@ pub struct Chunk {
 	pub blocks: Vec <Block>,
 }
 
-#[derive (Debug)]
+#[derive (Clone, Debug)]
 struct StackFrame {
 	// i32 makes it a little easier to implement jumps
 	// Starts at 0 right after OP_CALL
@@ -146,6 +149,7 @@ struct StackFrame {
 
 pub struct State {
 	registers: Vec <Value>,
+	stack: Vec <StackFrame>,
 	pub debug_print: bool,
 }
 
@@ -153,6 +157,13 @@ impl Default for State {
 	fn default () -> Self {
 		Self {
 			registers: vec! [Value::Nil; 256],
+			stack: vec! [
+				StackFrame {
+					program_counter: 0,
+					block_idx: 0,
+					register_offset: 0,
+				},
+			],
 			debug_print: false,
 		}
 	}
@@ -173,21 +184,22 @@ impl State {
 		]
 	}
 	
+	fn register_window (&self) -> &[Value] {
+		let frame = self.stack.last ().unwrap ();
+		&self.registers [frame.register_offset..]
+	}
+	
+	fn register_window_mut (&mut self) -> &mut [Value] {
+		let frame = self.stack.last_mut ().unwrap ();
+		&mut self.registers [frame.register_offset..]
+	}
+	
 	pub fn execute_chunk (&mut self, chunk: &Chunk, upvalues: &[Value]) 
 	-> Vec <Value> {
 		let max_iters = 2000;
 		
-		let mut stack = vec! [
-			StackFrame {
-				program_counter: 0,
-				block_idx: 0,
-				register_offset: 0,
-			},
-		];
-		
 		for _ in 0..max_iters {
-			let stack_idx = stack.len () - 1;
+			let frame = self.stack.last_mut ().unwrap ().clone ();
-			let frame = stack.get_mut (stack_idx).unwrap ();
 			let block = chunk.blocks.get (frame.block_idx).unwrap ();
 			
 			let mut next_pc = frame.program_counter;
@@ -196,16 +208,18 @@ impl State {
 			let instruction = match block.instructions.get (pc) {
 				Some (x) => x,
 				None => {
-					dbg! (&stack);
+					dbg! (&self.stack);
 					panic! ("program_counter went out of bounds");
 				}
 			};
 			
-			let r = &mut self.registers [frame.register_offset..];
+			// let r = &mut self.registers [frame.register_offset..];
 			let k = &block.constants;
 			
 			match instruction {
 				Instruction::Add (a, b, c) => {
+					let r = self.register_window_mut ();
+					
 					let a = usize::try_from (*a).unwrap ();
 					let b = usize::try_from (*b).unwrap ();
 					let c = usize::try_from (*c).unwrap ();
@@ -227,16 +241,20 @@ impl State {
 					// TODO: Only implement printing values for now
 					
 					let a = usize::try_from (*a).unwrap ();
+					let r = self.register_window ();
 					let v_a = r.get (a).unwrap ();
 					
 					match v_a {
-						Value::BogusClosure (idx) => {
+						Value::BogusClosure {
+							idx,
+							upvalues,
+						}=> {
 							let block_idx = frame.block_idx;
 							let target_block = *idx;
 							
-							let current_frame = &stack [stack.len () - 1];
+							let current_frame = self.stack.last ().unwrap ();
 							
-							stack.push (StackFrame {
+							self.stack.push (StackFrame {
 								program_counter: 0,
 								block_idx: target_block,
 								register_offset: current_frame.register_offset + a + 1,
@@ -244,7 +262,7 @@ impl State {
 							
 							if self.debug_print {
 								println! ("Inst {block_idx}:{pc} calls {target_block}:0");
-								let stack_depth = stack.len ();
+								let stack_depth = self.stack.len ();
 								println! ("stack_depth: {stack_depth}");
 							}
 							
@@ -260,18 +278,26 @@ impl State {
 							
 							println! ("{:?}", r.get (a + 1).unwrap ());
 						},
-						_ => panic! ("Cannot call value {a:?}. backtrace: {stack:?}"),
+						_ => {
+							let stack = &self.stack;
+							panic! ("Cannot call value {a:?}. backtrace: {stack:?}");
+						},
 					}
 				},
 				Instruction::Closure (a, b) => {
 					let a = usize::try_from (*a).unwrap ();
 					let b = usize::try_from (*b).unwrap ();
+					let r = self.register_window_mut ();
 					
-					r [a] = Value::BogusClosure (b + frame.block_idx + 1);
+					r [a] = Value::BogusClosure {
+						idx: b + frame.block_idx + 1,
+						upvalues: vec! [],
+					};
 				},
 				Instruction::EqK (a, b, c_k) => {
 					let a = usize::try_from (*a).unwrap ();
 					let b = usize::try_from (*b).unwrap ();
+					let r = self.register_window ();
 					
 					let equal = r [a] == k [b];
 					
@@ -299,13 +325,16 @@ impl State {
 					
 					let value = env.get (key).unwrap ();
 					
-					r [a] = value.clone ();
+					
+					self.register_window_mut() [a] = value.clone ();
 				},
 				Instruction::GetI (a, b, c) => {
 					let a = usize::try_from (*a).unwrap ();
 					let b = usize::try_from (*b).unwrap ();
 					let c = usize::try_from (*c).unwrap ();
 					
+					let r = self.register_window_mut ();
+					
 					let table = r.get (b).unwrap ();
 					let value = match table {
 						Value::BogusArg (arg) => arg.get (c).map (|x| x.as_str().into ()).unwrap_or_default(),
@@ -317,27 +346,28 @@ impl State {
 				Instruction::Jmp (s_j) => next_pc += s_j,
 				Instruction::LoadFalse (a) => {
 					let a  = usize::try_from  (*a).unwrap ();
-					r [a] = Value::Boolean (false);
+					self.register_window_mut ()[a] = Value::Boolean (false);
 				},
 				Instruction::LoadI (a, sbx) => {
 					let a  = usize::try_from  (*a).unwrap ();
 					
-					r [a] = (*sbx).into ();
+					self.register_window_mut ()[a] = (*sbx).into ();
 				},
 				Instruction::LoadK (a, bx) => {
 					let a  = usize::try_from  (*a).unwrap ();
 					let bx = usize::try_from (*bx).unwrap ();
 					
-					r [a] = k [bx].clone ();
+					self.register_window_mut ()[a] = k [bx].clone ();
 				},
 				Instruction::LoadTrue (a) => {
 					let a  = usize::try_from  (*a).unwrap ();
-					r [a] = Value::Boolean (true);
+					self.register_window_mut ()[a] = Value::Boolean (true);
 				},
 				Instruction::MmBin (a, b, _c) => {
 					let a = usize::try_from (*a).unwrap ();
 					let b = usize::try_from (*b).unwrap ();
 					
+					let r = self.register_window();
 					let a = &r [a];
 					let b = &r [b];
 					
@@ -352,43 +382,64 @@ impl State {
 					let a = usize::try_from (*a).unwrap ();
 					let b = usize::try_from (*b).unwrap ();
 					
+					let r = self.register_window_mut();
 					r [a] = r [b].clone ();
 				},
 				Instruction::Not (a, b) => {
 					let a = usize::try_from (*a).unwrap ();
 					let b = usize::try_from (*b).unwrap ();
-					
+					let r = self.register_window_mut();
 					r [a] = Value::Boolean (! r [b].is_truthy());
 				}
-				Instruction::Return (a, b, _c, _k) => {
+				Instruction::Return (a, b, _c, k) => {
 					let a = usize::try_from (*a).unwrap ();
 					let b = usize::try_from (*b).unwrap ();
 					
-					let popped_frame = stack.pop ().unwrap ();
+					let popped_frame = self.stack.pop ().unwrap ();
+					
+					// Build closure if needed
+					if *k {
+						
+						let closure_idx = match &self.register_window ()[a] {
+							Value::BogusClosure { idx, upvalues } => 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 {
+							idx: *closure_idx,
+							upvalues,
+						};
+					}
 					
 					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 = stack.len ();
+						let stack_depth = self.stack.len ();
 						println! ("stack_depth: {stack_depth}");
 					}
 					
-					if let Some (new_frame) = stack.last() {
+					if let Some (new_frame) = self.stack.last() {
 						next_pc = new_frame.program_counter;
 					}
 					else {
-						return r [a..(a + b - 1)].to_vec();
+						return self.register_window ()[a..(a + b - 1)].to_vec();
 					}
 				},
 				Instruction::Return1 (a) => {
 					let a = usize::try_from (*a).unwrap ();
-					let popped_frame = stack.pop ().unwrap ();
+					let popped_frame = self.stack.pop ().unwrap ();
 					
-					self.registers [popped_frame.register_offset - 1] = r [a].clone ();
 					
-					let stack_idx = stack.len () - 1;
+					self.registers [popped_frame.register_offset - 1] = self.register_window ()[a].clone ();
-					let frame = stack.get (stack_idx).unwrap ();
+					
+					let frame = self.stack.last ().unwrap ();
 					let new_block = frame.block_idx;
 					next_pc = frame.program_counter;
 					
@@ -396,14 +447,14 @@ impl State {
 						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 = stack.len ();
+						let stack_depth = self.stack.len ();
 						println! ("stack_depth: {stack_depth}");
 					}
 				},
 				Instruction::Test (a, _k) => {
 					let a = usize::try_from (*a).unwrap ();
 					
-					let a = r.get (a).unwrap ();
+					let a = self.register_window ().get (a).unwrap ();
 					
 					if self.debug_print {
 						println! ("Test {a:?}");
@@ -419,8 +470,7 @@ impl State {
 			
 			next_pc += 1;
 			{
-				let stack_idx = stack.len () - 1;
+				let frame = self.stack.last_mut ().unwrap ();
-				let frame = stack.get_mut (stack_idx).unwrap ();
 				frame.program_counter = next_pc;
 			}
 		}