From b35dc346e87463e9029dec8a4c8f3a10c0b375f8 Mon Sep 17 00:00:00 2001
From: _ <_@_>
Date: Mon, 2 Oct 2023 19:14:54 -0500
Subject: [PATCH] :chart_with_upwards_trend: performance: caching the current
 block helps, barely

---
 lunar_wave_vm/src/loader.rs |  9 ++--
 lunar_wave_vm/src/state.rs  | 93 +++++++++++++++++++++++--------------
 2 files changed, 63 insertions(+), 39 deletions(-)

diff --git a/lunar_wave_vm/src/loader.rs b/lunar_wave_vm/src/loader.rs
index 0188e95..3969f6b 100644
--- a/lunar_wave_vm/src/loader.rs
+++ b/lunar_wave_vm/src/loader.rs
@@ -1,4 +1,7 @@
-use std::io::Read;
+use std::{
+	io::Read,
+	rc::Rc,
+};
 
 use crate::{
 	instruction::Instruction as Inst,
@@ -243,7 +246,7 @@ fn parse_i64 <R: Read> (rdr: &mut R) -> Option <i64> {
 // code, but I don't like recursion in general, and I don't know
 // why PUC wrote it that way.
 
-pub fn parse_block <R: Read> (rdr: &mut R, si: &mut Interner, blocks: &mut Vec <Block>) 
+pub fn parse_block <R: Read> (rdr: &mut R, si: &mut Interner, blocks: &mut Vec <Rc <Block>>) 
 -> Option <()> 
 {
 	// Ignore things I haven't implemented yet
@@ -306,7 +309,7 @@ pub fn parse_block <R: Read> (rdr: &mut R, si: &mut Interner, blocks: &mut Vec <
 		constants,
 		instructions,
 		upvalues,
-	});
+	}.into ());
 	
 	// Recursion
 	
diff --git a/lunar_wave_vm/src/state.rs b/lunar_wave_vm/src/state.rs
index 6863a6f..f4265e2 100644
--- a/lunar_wave_vm/src/state.rs
+++ b/lunar_wave_vm/src/state.rs
@@ -1,3 +1,5 @@
+use std::rc::Rc;
+
 use crate::{
 	instruction::Instruction,
 	string_interner::Interner,
@@ -23,7 +25,7 @@ pub struct Block {
 
 #[derive (Clone, Debug, Default)]
 pub struct Chunk {
-	pub blocks: Vec <Block>,
+	pub blocks: Vec <Rc <Block>>,
 }
 
 #[derive (Clone, Copy, Debug, Default)]
@@ -54,6 +56,7 @@ pub struct State {
 	
 	pub debug_print: bool,
 	chunk: Chunk,
+	current_block: Rc <Block>,
 	pub upvalues: Vec <Value>,
 	pub si: Interner,
 }
@@ -180,22 +183,7 @@ pub enum StepError {
 
 impl State {
 	pub fn new (chunk: Chunk, upvalues: Vec <Value>) -> Self {
-		Self {
-			// TODO: Stack is actually supposed to grow to a limit of
-			// idk 10,000. I thought it was fixed at 256.
-			registers: vec! [Value::Nil; 256],
-			top: 0,
-			stack: vec! [],
-			stack_top: Default::default (),
-			debug_print: false,
-			chunk,
-			upvalues,
-			si: Default::default (),
-		}
-	}
-	
-	pub fn new_with_args <I: Iterator <Item = String>> (chunk: Chunk, mut si: Interner, args: I) -> Self {
-		let upvalues = Self::upvalues_from_args (&mut si, args);
+		let current_block = Rc::clone (&chunk.blocks [0]);
 		
 		Self {
 			// TODO: Stack is actually supposed to grow to a limit of
@@ -206,6 +194,26 @@ impl State {
 			stack_top: Default::default (),
 			debug_print: false,
 			chunk,
+			current_block,
+			upvalues,
+			si: Default::default (),
+		}
+	}
+	
+	pub fn new_with_args <I: Iterator <Item = String>> (chunk: Chunk, mut si: Interner, args: I) -> Self {
+		let upvalues = Self::upvalues_from_args (&mut si, args);
+		let current_block = Rc::clone (&chunk.blocks [0]);
+		
+		Self {
+			// TODO: Stack is actually supposed to grow to a limit of
+			// idk 10,000. I thought it was fixed at 256.
+			registers: vec! [Value::Nil; 256],
+			top: 0,
+			stack: vec! [],
+			stack_top: Default::default (),
+			debug_print: false,
+			chunk,
+			current_block,
 			upvalues,
 			si,
 		}
@@ -321,11 +329,11 @@ impl State {
 				let target_block = idx;
 				
 				self.stack.push (self.stack_top);
-				self.stack_top = StackFrame {
+				self.set_stack_top (StackFrame {
 					program_counter: 0,
 					block_idx: target_block,
 					register_offset: self.stack_top.register_offset + a as usize + 1,
-				};
+				});
 				
 				// Skip the PC increment at the bottom of the loop
 				return true;
@@ -354,7 +362,8 @@ impl State {
 				let num_results = x (self, num_args);
 				
 				let popped_frame = self.stack_top;
-				self.stack_top = self.stack.pop ().unwrap ();
+				let x = self.stack.pop ().unwrap ();
+				self.set_stack_top (x);
 				self.stack_top.register_offset = old_offset;
 				let offset = old_offset + usize::from (a);
 				
@@ -391,7 +400,7 @@ impl State {
 	}
 	
 	fn op_get_field (&mut self, a: u8, b: u8, c: u8) {
-		let block = self.chunk.blocks.get (self.stack_top.block_idx).unwrap ();
+		let block = &self.current_block;
 		let constants = &block.constants;
 		
 		let key = match &constants [usize::from (c)] {
@@ -449,8 +458,7 @@ impl State {
 	}
 	
 	fn op_set_field (&mut self, a: u8, b: u8, c: u8, k: bool) {
-		let frame = &self.stack_top;
-		let block = self.chunk.blocks.get (frame.block_idx).unwrap ();
+		let block = &self.current_block;
 		let constants = &block.constants;
 		
 		let b = usize::try_from (b).unwrap ();
@@ -489,12 +497,21 @@ impl State {
 	}
 	
 	fn constants (&self) -> &[Value] {
-		&self.chunk.blocks.get (self.stack_top.block_idx).unwrap ().constants
+		&self.current_block.constants
 	}
 	
-	fn decode (&self) -> Instruction {
-		let block = self.chunk.blocks.get (self.stack_top.block_idx).unwrap ();
-		match block.instructions.get (self.stack_top.program_counter) {
+	fn set_block_idx (&mut self, block_idx: usize) {
+		self.stack_top.block_idx = block_idx;
+		self.current_block = Rc::clone (&self.chunk.blocks [block_idx]);
+	}
+	
+	fn set_stack_top (&mut self, frame: StackFrame) {
+		self.stack_top = frame;
+		self.current_block = Rc::clone (&self.chunk.blocks [frame.block_idx]);
+	}
+	
+	fn fetch (&self) -> Instruction {
+		match self.current_block.instructions.get (self.stack_top.program_counter) {
 			Some (x) => *x,
 			None => {
 				dbg! (&self.stack, &self.stack_top);
@@ -509,7 +526,7 @@ impl State {
 	
 	pub fn step (&mut self) -> Result <Option <StepOutput>, StepError> 
 	{
-		let instruction = self.decode ();
+		let instruction = self.fetch ();
 		
 		let make_step_error = |msg| {
 			self.make_step_error (msg, &instruction)
@@ -798,7 +815,7 @@ impl State {
 				}
 				
 				if let Some (new_frame) = self.stack.pop () {
-					self.stack_top = new_frame;
+					self.set_stack_top (new_frame);
 					
 					// Shift our output registers down so the caller
 					// can grab them
@@ -821,7 +838,8 @@ impl State {
 			},
 			Instruction::Return0 => {
 				let popped_frame = self.stack_top;
-				self.stack_top = self.stack.pop ().unwrap ();
+				let x = self.stack.pop ().unwrap ();
+				self.set_stack_top (x);
 				self.top = popped_frame.register_offset - 1 + 0;
 			},
 			Instruction::Return1 (a) => {
@@ -830,7 +848,8 @@ impl State {
 				
 				self.registers [popped_frame.register_offset - 1] = self.register_window ()[a].clone ();
 				
-				self.stack_top = self.stack.pop ().unwrap ();
+				let x = self.stack.pop ().unwrap ();
+				self.set_stack_top (x);
 				
 				// Shift output register down
 				let offset = popped_frame.register_offset;
@@ -917,9 +936,11 @@ impl State {
 						
 						// Jump into the other function
 						
-						self.stack_top.program_counter = 0;
-						self.stack_top.block_idx = closure.idx;
-						self.stack_top.program_counter = 0;
+						self.set_stack_top (StackFrame {
+							block_idx: closure.idx,
+							program_counter: 0,
+							register_offset: self.stack_top.register_offset,
+						});
 						
 						// Skip the PC increment
 						return Ok (None);
@@ -954,7 +975,7 @@ impl State {
 						let popped_frame = self.stack_top;
 						
 						if let Some (new_frame) = self.stack.pop () {
-							self.stack_top = new_frame;
+							self.set_stack_top (new_frame);
 							
 							// Set up top for the next call
 							if c == 0 {
@@ -1025,7 +1046,7 @@ impl State {
 	
 	pub fn set_chunk (&mut self, chunk: Chunk) {
 		self.stack = vec! [];
-		self.stack_top = Default::default ();
+		self.set_stack_top (Default::default ());
 		self.chunk = chunk;
 	}