Fix edge case

src/chunk.rs

@@ -5,6 +5,7 @@ pub struct Chunk {
 	pub(crate) instructions: Vec<u8>,
 	pub(crate) lines: Vec<usize>,
 	pub(crate) literals: Vec<ExprValue>,
+	pub(crate) num_variables: usize,
 	pub(crate) arity: usize,
 }
 
@@ -13,6 +14,7 @@ impl Chunk {
 		instructions: Vec<u8>,
 		lines: Vec<usize>,
 		literals: Vec<ExprValue>,
+		num_variables: usize,
 		arity: usize,
 	) -> Self {
 		assert_eq!(lines.len(), instructions.len());
@@ -20,6 +22,7 @@ impl Chunk {
 			instructions,
 			lines,
 			literals,
+			num_variables,
 			arity,
 		}
 	}

src/compiler.rs

@@ -58,6 +58,7 @@ impl Compiler {
 			self.instructions.clone(),
 			self.lines.clone(),
 			self.literals.clone(),
+			self.cur_variable,
 			0,
 		);
 
@@ -82,6 +83,7 @@ impl Compiler {
 					self.instructions.clone(),
 					self.lines.clone(),
 					self.literals.clone(),
+					self.cur_variable,
 					func.parameters.len(),
 				));
 			}
@@ -127,7 +129,6 @@ impl Compiler {
 			}
 
 			self.variables.insert(name.lexeme, self.cur_variable);
-			self.emit_opcode(OpCode::NewVar);
 			let index = self.cur_variable;
 			self.cur_variable += 1;
 			return index as u8;
@@ -231,7 +232,10 @@ impl StmtVisitor<()> for Compiler {
 				.error_token(&stmt.variable, "Too many variables");
 		}
 
+		// Allocate space for max_value
 		self.cur_variable += 1;
+		self.emit_opcode(OpCode::SetVar);
+		self.emit_byte(self.cur_variable as u8 - 1);
 
 		// set x to min_value
 		stmt.min_value.accept(self);

src/opcodes.rs

@@ -10,7 +10,6 @@ pub(crate) enum OpCode {
 	Nil,
 
 	Var,
-	NewVar,
 	SetVar,
 
 	Add,

src/vm.rs

@@ -70,6 +70,11 @@ where
 		function_chunks: Vec<Chunk>,
 		error_handler: &mut ErrorHandler,
 	) -> bool {
+		// Allocate space for variables
+		for _ in 0..chunk.num_variables {
+			self.push(ExprValue::Integer(0))
+		}
+
 		self.ip = 0;
 		self.cur_chunk = 0;
 		self.chunks = vec![chunk];
@@ -119,7 +124,7 @@ where
 			}
 			OpCode::Input => {
 				let var_id = self.read_byte() as usize;
-				self.stack[self.offset + var_id] = ExprValue::String(self.stdio.read_line());
+				self.stack[var_id] = ExprValue::String(self.stdio.read_line());
 			}
 			OpCode::Literal8 => {
 				let byte = self.read_byte();
@@ -263,9 +268,6 @@ where
 				let var_id = self.read_byte() as usize;
 				self.push(self.stack[self.offset + var_id].clone());
 			}
-			OpCode::NewVar => {
-				self.push(ExprValue::Integer(0));
-			}
 			OpCode::SetVar => {
 				let var_id = self.read_byte() as usize;
 				self.stack[self.offset + var_id] = self.pop();
@@ -377,6 +379,11 @@ where
 					.push(CallFrame::new(old_chunk, self.ip, self.offset));
 				self.offset = self.stack.len() - self.arity();
 				self.ip = 0; // Reset IP
+
+				// Allocate space on stack for variables
+				for _ in 0..self.chunks[self.cur_chunk].num_variables {
+					self.push(ExprValue::Integer(0));
+				}
 			}
 			OpCode::Ret => {
 				if self.call_stack.is_empty() {

src/xbasic.rs

@@ -263,7 +263,7 @@ where
 
 				// Generate line numbers(all set to 0, since they should never be used)
 				let lines = instructions.iter().map(|_| 0).collect();
-				let chunk = Chunk::new(instructions, lines, arguments.to_vec(), 0);
+				let chunk = Chunk::new(instructions, lines, arguments.to_vec(), 0, 0);
 
 				// Run the instructions
 				if !self.vm.run(

tests/control.rs

@@ -158,6 +158,121 @@ end if
 }
 
 #[test]
+fn if_missing_endif() {
+	let tio = common::TestIO::new("");
+
+	let mut xb = XBasic::new(tio);
+	assert!(xb
+		.run(
+			"if 5 = 3 then
+    print \"hi\"
+	"
+		)
+		.is_err());
+	assert!(xb.error_handler.had_errors);
+	assert_eq!(
+		xb.error_handler.errors,
+		["[line 3] Error at end: Expected END IF after IF statement."]
+	);
+
+	xb.get_io().check();
+}
+
+#[test]
+fn nested_if_missing_endif() {
+	let tio = common::TestIO::new("");
+
+	let mut xb = XBasic::new(tio);
+	assert!(xb
+		.run(
+			"
+if 6 = 5 then
+    if 5 = 3 then
+        print \"hi\"
+    end if
+    a = 5
+	"
+		)
+		.is_err());
+	assert!(xb.error_handler.had_errors);
+	assert_eq!(
+		xb.error_handler.errors,
+		["[line 7] Error at end: Expected END IF after IF statement."]
+	);
+
+	xb.get_io().check();
+}
+
+#[test]
+fn nested_if_in_for_missing_endif() {
+	let tio = common::TestIO::new("");
+
+	let mut xb = XBasic::new(tio);
+	assert!(xb
+		.run(
+			"
+for x = 0 to 5
+    for y = 0 to 5
+        if 6 = 5 then
+            if 5 = 3 then
+                print \"hi\"
+        end if
+        a = 5
+    next y
+next x
+	"
+		)
+		.is_err());
+	assert!(xb.error_handler.had_errors);
+	assert_eq!(
+		xb.error_handler.errors,
+		["[line 9] Error at 'y': Expected expression."]
+	);
+
+	xb.get_io().check();
+}
+
+#[test]
+fn if_assignment_edgecase_1() {
+	common::test_program(
+		"
+if 4 = 0 then
+    b = 0
+end if
+a = 4
+    ",
+		"",
+	);
+}
+
+#[test]
+fn if_assignment_edgecase_2() {
+	common::test_program(
+		"
+a = 0
+if 4 = 0 then
+    a = 1
+end if
+a = 4
+    ",
+		"",
+	);
+}
+
+#[test]
+fn if_assignment_edgecase_3() {
+	common::test_program(
+		"
+if 4 = 0 then
+    a = 0
+end if
+a = 4
+    ",
+		"",
+	);
+}
+
+#[test]
 fn while_loop() {
 	common::test_program(
 		"while a < 10

tests/limits.rs

@@ -31,7 +31,7 @@ fn above_compute_limit() {
 	assert!(xb
 		.run(
 			"
-for a = 0 to 90
+for a = 0 to 98
 b = b + a
 next a
 print b
@@ -40,7 +40,7 @@ print b
 		.is_err());
 	assert!(xb.error_handler.had_runtime_error);
 	assert_eq!(
-		"[line 2] Error : Compute limit exceeded.",
+		"[line 5] Error : Compute limit exceeded.",
 		xb.error_handler.errors[0]
 	);
 	xb.get_io().check();