From 945721f4220006efeb2e028d439f757920616d95 Mon Sep 17 00:00:00 2001
From: Patric Bucher <patricbucher@gmail.com>
Date: Sun, 5 Nov 2023 21:18:55 +0100
Subject: [PATCH] Added support for up to 2^32 constants

Also simplified writing constants to a block
---
 src/main.rs         | 11 ++-----
 src/vm/block.rs     | 76 ++++++++++++++++++++++++++++++++++++++-------
 src/vm/constants.rs |  8 +++--
 3 files changed, 74 insertions(+), 21 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 00e32cf..968e109 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,15 +15,10 @@ fn main() {
 
     let mut block = Block::new("ZeBlock");
 
-    let constant_index = block.push_constant(1234.56, 2);
-    block.push_op_code(OpCode::Constant, 2);
-    block.write_byte(constant_index);
+    block.write_constant(1234.56, 2);
+    block.write_constant(789.10, 4);
 
-    let constant_index = block.push_constant(789.10, 4);
-    block.push_op_code(OpCode::Constant, 4);
-    block.write_byte(constant_index);
-
-    block.push_op_code(OpCode::Return, 4);
+    block.write_op_code(OpCode::Return, 4);
 
     #[cfg(feature = "disassemble")]
     block.disassemble_block();
diff --git a/src/vm/block.rs b/src/vm/block.rs
index 2fcdcc7..e88e389 100644
--- a/src/vm/block.rs
+++ b/src/vm/block.rs
@@ -4,10 +4,12 @@ use enum_primitive_derive::Primitive;
 use num_traits::FromPrimitive;
 
 #[repr(u8)]
-#[derive(Debug, Primitive)]
+#[derive(Debug, PartialEq, Primitive)]
 pub enum OpCode {
-    Constant = 0x00,
-    Return = 0x01,
+    Return = 0x00,
+    Constant = 0x01,
+    Constant2 = 0x02,
+    Constant4 = 0x03,
 }
 
 #[allow(dead_code)]
@@ -28,18 +30,38 @@ impl Block {
         }
     }
 
-    pub fn push_op_code(&mut self, op_code: OpCode, line: usize) {
+    pub fn write_op_code(&mut self, op_code: OpCode, line: usize) {
         self.lines.push(line);
         self.instructions.push(op_code as u8)
     }
 
-    pub fn push_constant(&mut self, value: f64, line: usize) -> i8 {
-        self.lines.push(line);
-        self.constants.push_value(value)
+    pub fn write_constant(&mut self, value: f64, line: usize) {
+        let constant_index = self.constants.push_value(value);
+
+        if constant_index <= 0xFF {
+            self.write_op_code(OpCode::Constant, line);
+            self.write_u8(constant_index as u8)
+        } else if constant_index <= 0xFFFF {
+            self.write_op_code(OpCode::Constant2, line);
+            self.write_u16(constant_index as u16)
+        } else {
+            self.write_op_code(OpCode::Constant4, line);
+            self.write_u32(constant_index)
+        }
     }
 
-    pub(crate) fn write_byte(&mut self, byte: i8) {
-        self.instructions.push(byte as u8)
+    fn write_u8(&mut self, value: u8) {
+        self.instructions.push(value)
+    }
+    fn write_u16(&mut self, value: u16) {
+        self.instructions.push((value) as u8);
+        self.instructions.push((value >> 8) as u8);
+    }
+    fn write_u32(&mut self, value: u32) {
+        self.instructions.push((value) as u8);
+        self.instructions.push((value >> 8) as u8);
+        self.instructions.push((value >> 16) as u8);
+        self.instructions.push((value >> 24) as u8);
     }
 }
 
@@ -76,8 +98,10 @@ impl BlockDbg for Block {
 
         let instruction = OpCode::from_u8(self.instructions[offset]).unwrap();
         return match instruction {
-            OpCode::Constant => self.constant_instruction(OpCode::Constant, offset),
             OpCode::Return => self.simple_instruction(OpCode::Return, offset),
+            OpCode::Constant => self.constant_instruction(instruction, offset),
+            OpCode::Constant2 => self.constant_instruction(instruction, offset),
+            OpCode::Constant4 => self.constant_instruction(instruction, offset),
         };
     }
 
@@ -110,9 +134,39 @@ mod tests {
     #[test]
     fn op_code_can_be_pushed_to_an_block() {
         let mut block = Block::new("jenny");
-        block.push_op_code(OpCode::Return, 123);
+        block.write_op_code(OpCode::Return, 123);
 
         assert_eq!(1, block.instructions.len());
         assert_eq!(OpCode::Return as u8, block.instructions[0]);
     }
+
+    #[test]
+    fn can_write_more_then_256_constants() {
+        let mut block = Block::new("maggie");
+        for i in 0..258 {
+            block.write_constant(i as f64, i);
+        }
+
+        assert_eq!(2 * 256 + 6, block.instructions.len());
+        assert_eq!(
+            OpCode::Constant2,
+            OpCode::from_u8(block.instructions[2 * 256]).unwrap()
+        );
+
+        let byte1 = block.instructions[2 * 256 + 1] as u16;
+        let byte2 = block.instructions[2 * 256 + 2] as u16;
+        let constant_index: u16 = (byte2 << 8) | byte1;
+        assert_eq!(256, constant_index);
+
+        assert_eq!(
+            OpCode::Constant2,
+            OpCode::from_u8(block.instructions[2 * 256 + 3]).unwrap()
+        );
+        let byte1 = block.instructions[2 * 256 + 4] as u16;
+        let byte2 = block.instructions[2 * 256 + 5] as u16;
+        let constant_index: u16 = (byte2 << 8) | byte1;
+        assert_eq!(257, constant_index);
+
+        assert_eq!(257f64, block.constants.get_value(constant_index as u32));
+    }
 }
diff --git a/src/vm/constants.rs b/src/vm/constants.rs
index 403ef4a..852fa45 100644
--- a/src/vm/constants.rs
+++ b/src/vm/constants.rs
@@ -9,9 +9,13 @@ impl Constants {
         Constants { values: Vec::new() }
     }
 
-    pub fn push_value(&mut self, value: Value) -> i8 {
+    pub fn push_value(&mut self, value: Value) -> u32 {
         self.values.push(value);
-        self.values.len() as i8 - 1
+        (self.values.len() - 1) as u32
+    }
+
+    pub fn get_value(&self, index: u32) -> Value {
+        self.values[index as usize]
     }
 }