tinyscript 0.6.1

// Copyright © 2025 Stephan Kunz
//! Virtual machine implementation.

#[doc(hidden)]
#[cfg(feature = "std")]
extern crate std;

use core::str::FromStr;

use alloc::string::ToString;

use super::{Chunk, op_code::OpCode};

use crate::{environment::Environment, execution::ExecutionError, scripting_value::ScriptingValue};

/// Stack size is fixed to avoid cache misses, which drastically reduce performance.
/// For the intended purpose (short inline scripting) this size should be enough.
const STACK_SIZE: usize = 8;

/// Subset of [`OpCode`] understood by [`VM::arithmetic_operator`]. Using a
/// dedicated enum keeps the inner `match` exhaustive and removes the need for
/// a defensive `_ => Unreachable` arm.
#[derive(Clone, Copy)]
enum ArithOp {
	Add,
	Subtract,
	Multiply,
	Divide,
}

/// Subset of [`OpCode`] understood by [`VM::bitwise_operator`].
#[derive(Clone, Copy)]
enum BitOp {
	And,
	Or,
	Xor,
}

/// Subset of [`OpCode`] understood by [`VM::comparison_operator`].
#[derive(Clone, Copy)]
enum CmpOp {
	Greater,
	Less,
}

/// A stack based Virtual Machine.
///
/// The stack size is limited to avoid cache misses, which drastically reduce performance.
/// For the intended purpose (short inline scripting) this size should be enough.
pub struct VM {
	/// The `InstructionPointer` (sometimes called `ProgramCounter`)
	ip: usize,
	/// Stack for values
	stack: [ScriptingValue; STACK_SIZE],
	/// Pointer to the next free stack place
	stack_top: usize,
}

impl core::fmt::Debug for VM {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		f.debug_struct("VM")
			.field("ip", &self.ip)
			.field("stack", &self.stack)
			.field("stack_top", &self.stack_top)
			.finish()
	}
}

impl Default for VM {
	fn default() -> Self {
		Self {
			ip: 0,
			stack: [const { ScriptingValue::nil() }; STACK_SIZE],
			stack_top: 0,
		}
	}
}

impl VM {
	fn reset(&mut self) {
		self.ip = 0;
		self.stack = [const { ScriptingValue::nil() }; STACK_SIZE];
		self.stack_top = 0;
	}

	const fn peek(&self, distance: usize) -> &ScriptingValue {
		&self.stack[self.stack_top - distance - 1]
	}

	fn push(&mut self, value: ScriptingValue) -> Result<(), ExecutionError> {
		if self.stack_top == STACK_SIZE {
			return Err(ExecutionError::StackOverflow);
		}
		self.stack[self.stack_top] = value;
		self.stack_top += 1;
		Ok(())
	}

	fn pop(&mut self) -> ScriptingValue {
		self.stack_top -= 1;
		core::mem::replace(&mut self.stack[self.stack_top], ScriptingValue::nil())
	}

	fn read_jmp_address(&mut self, chunk: &Chunk) -> usize {
		let byte1 = chunk.code()[self.ip];
		let byte2 = chunk.code()[self.ip + 1];
		self.ip += 2;
		usize::from(u16::from_be_bytes([byte1, byte2]))
	}

	#[allow(clippy::cast_precision_loss)]
	fn arithmetic_operator(&mut self, operator: ArithOp) -> Result<(), ExecutionError> {
		let b_val = self.pop();
		let a_val = self.pop();
		match (a_val, b_val) {
			(ScriptingValue::Float64(a), ScriptingValue::Float64(b)) => {
				let res = match operator {
					ArithOp::Add => a + b,
					ArithOp::Subtract => a - b,
					ArithOp::Multiply => a * b,
					ArithOp::Divide => a / b,
				};
				self.push(ScriptingValue::Float64(res))
			}
			(ScriptingValue::Float64(a), ScriptingValue::Int64(b)) => {
				let res = match operator {
					ArithOp::Add => a + (b as f64),
					ArithOp::Subtract => a - (b as f64),
					ArithOp::Multiply => a * (b as f64),
					ArithOp::Divide => a / (b as f64),
				};
				self.push(ScriptingValue::Float64(res))
			}
			(ScriptingValue::Int64(a), ScriptingValue::Float64(b)) => {
				let res = match operator {
					ArithOp::Add => (a as f64) + b,
					ArithOp::Subtract => (a as f64) - b,
					ArithOp::Multiply => (a as f64) * b,
					ArithOp::Divide => (a as f64) / b,
				};
				self.push(ScriptingValue::Float64(res))
			}
			(ScriptingValue::Int64(a), ScriptingValue::Int64(b)) => {
				let res = match operator {
					ArithOp::Add => a + b,
					ArithOp::Subtract => a - b,
					ArithOp::Multiply => a * b,
					ArithOp::Divide => a / b,
				};
				self.push(ScriptingValue::Int64(res))
			}
			(ScriptingValue::String(a), b_val) => {
				let res = match operator {
					ArithOp::Add => a + &b_val.to_string(),
					_ => return Err(ExecutionError::OnlyAdd),
				};
				self.push(ScriptingValue::String(res))
			}
			(a_val, ScriptingValue::String(b)) => {
				let res = match operator {
					ArithOp::Add => a_val.to_string() + &b,
					_ => return Err(ExecutionError::OnlyAdd),
				};
				self.push(ScriptingValue::String(res))
			}
			(ScriptingValue::Nil(), _) | (_, ScriptingValue::Nil()) => Err(ExecutionError::NilValue),
			(ScriptingValue::Boolean(_), _) | (_, ScriptingValue::Boolean(_)) => Err(ExecutionError::BoolNoArithmetic),
		}
	}

	fn bitwise_operator(&mut self, operator: BitOp) -> Result<(), ExecutionError> {
		let b_val = self.pop();
		let mut a_val = self.pop();
		match (a_val, b_val) {
			(ScriptingValue::Int64(a), ScriptingValue::Int64(b)) => {
				let res = match operator {
					BitOp::And => a & b,
					BitOp::Or => a | b,
					BitOp::Xor => a ^ b,
				};
				a_val = ScriptingValue::Int64(res);
				self.push(a_val)
			}
			(a_val, b_val) => Err(ExecutionError::NoInteger {
				value: (a_val.to_string() + "/" + &b_val.to_string()).into(),
			}),
		}
	}

	#[allow(clippy::cast_precision_loss)]
	fn comparison_operator(&mut self, operator: CmpOp) -> Result<(), ExecutionError> {
		let b_val = self.pop();
		let mut a_val = self.pop();
		let res = match (a_val, b_val) {
			(ScriptingValue::Int64(a), ScriptingValue::Int64(b)) => match operator {
				CmpOp::Greater => a > b,
				CmpOp::Less => a < b,
			},
			(ScriptingValue::Int64(a), ScriptingValue::Float64(b)) => match operator {
				CmpOp::Greater => (a as f64) > b,
				CmpOp::Less => (a as f64) < b,
			},
			(ScriptingValue::Float64(a), ScriptingValue::Int64(b)) => match operator {
				CmpOp::Greater => a > (b as f64),
				CmpOp::Less => a < (b as f64),
			},
			(ScriptingValue::Float64(a), ScriptingValue::Float64(b)) => match operator {
				CmpOp::Greater => a > b,
				CmpOp::Less => a < b,
			},
			(ScriptingValue::String(s), ScriptingValue::Float64(b)) => {
				if let Ok(a) = f64::from_str(&s) {
					match operator {
						CmpOp::Greater => a > b,
						CmpOp::Less => a < b,
					}
				} else {
					return Err(ExecutionError::NoComparison);
				}
			}
			(ScriptingValue::String(s), ScriptingValue::Int64(b)) => {
				if let Ok(a) = i64::from_str(&s) {
					match operator {
						CmpOp::Greater => a > b,
						CmpOp::Less => a < b,
					}
				} else {
					return Err(ExecutionError::NoComparison);
				}
			}
			(ScriptingValue::Float64(a), ScriptingValue::String(s)) => {
				if let Ok(b) = f64::from_str(&s) {
					match operator {
						CmpOp::Greater => a > b,
						CmpOp::Less => a < b,
					}
				} else {
					return Err(ExecutionError::NoComparison);
				}
			}
			(ScriptingValue::Int64(a), ScriptingValue::String(s)) => {
				if let Ok(b) = i64::from_str(&s) {
					match operator {
						CmpOp::Greater => a > b,
						CmpOp::Less => a < b,
					}
				} else {
					return Err(ExecutionError::NoComparison);
				}
			}
			_ => return Err(ExecutionError::NoComparison),
		};
		a_val = ScriptingValue::Boolean(res);
		self.push(a_val)
	}

	fn constant(&mut self, chunk: &Chunk) -> Result<(), ExecutionError> {
		let pos = chunk.code()[self.ip];
		let constant = chunk.read_constant(pos).clone();
		self.ip += 1;
		self.push(constant)
	}

	#[allow(clippy::cast_precision_loss)]
	fn equal(&mut self) -> Result<(), ExecutionError> {
		let b_val = self.pop();
		let mut a_val = self.pop();
		let res = match (a_val, b_val) {
			(ScriptingValue::Boolean(a), ScriptingValue::Boolean(b)) => a == b,
			(ScriptingValue::Float64(a), ScriptingValue::Float64(b)) => {
				let delta = f64::abs(a - b);
				delta <= 0.000_000_000_000_002
			}
			(ScriptingValue::Float64(a), ScriptingValue::Int64(b)) => {
				let delta = f64::abs(a - (b as f64));
				delta <= 0.000_000_000_000_002
			}
			(ScriptingValue::Int64(a), ScriptingValue::Float64(b)) => {
				let delta = f64::abs((a as f64) - b);
				delta <= 0.000_000_000_000_002
			}
			(ScriptingValue::Int64(a), ScriptingValue::Int64(b)) => a == b,
			(ScriptingValue::String(a), ScriptingValue::String(b)) => a == b,
			(ScriptingValue::Nil(), ScriptingValue::Nil()) => true,
			(ScriptingValue::String(s), ScriptingValue::Int64(b)) => i64::from_str(&s).is_ok_and(|a| a == b),
			(ScriptingValue::Int64(a), ScriptingValue::String(s)) => i64::from_str(&s).is_ok_and(|b| a == b),
			(ScriptingValue::String(s), ScriptingValue::Float64(b)) => f64::from_str(&s).is_ok_and(|a| {
				let delta = f64::abs(a - b);
				delta <= 0.000_000_000_000_002
			}),
			(ScriptingValue::Float64(a), ScriptingValue::String(s)) => f64::from_str(&s).is_ok_and(|b| {
				let delta = f64::abs(a - b);
				delta <= 0.000_000_000_000_002
			}),
			_ => false,
		};
		a_val = ScriptingValue::Boolean(res);
		self.push(a_val)
	}

	fn negate(&mut self) -> Result<(), ExecutionError> {
		let val = self.pop();
		let res = match val {
			ScriptingValue::Int64(v) => ScriptingValue::Int64(-v),
			ScriptingValue::Float64(v) => ScriptingValue::Float64(-v),
			ScriptingValue::String(s) => {
				if let Ok(i) = i64::from_str(&s) {
					ScriptingValue::Int64(-i)
				} else if let Ok(f) = f64::from_str(&s) {
					ScriptingValue::Float64(-f)
				} else {
					return Err(ExecutionError::NoNumber { value: s.into() });
				}
			}
			_ => {
				return Err(ExecutionError::NoNumber {
					value: val.to_string().into(),
				});
			}
		};
		self.push(res)
	}

	fn bitwise_not(&mut self) -> Result<(), ExecutionError> {
		let val = self.pop();
		let res = match val {
			ScriptingValue::Int64(v) => ScriptingValue::Int64(!v),
			_ => {
				return Err(ExecutionError::NoNumber {
					value: val.to_string().into(),
				});
			}
		};
		self.push(res)
	}

	fn not(&mut self) -> Result<(), ExecutionError> {
		let val = self.pop();
		let res = match val {
			ScriptingValue::Boolean(b) => ScriptingValue::Boolean(!b),
			ScriptingValue::Nil() => ScriptingValue::Boolean(true),
			_ => ScriptingValue::Boolean(false),
		};
		self.push(res)
	}

	#[cfg(feature = "std")]
	fn print(&mut self, stdout: &mut impl std::io::Write) {
		// The parser always emits an expression before `Print`, so the stack
		// is guaranteed to be non-empty here.
		let value = self.pop();
		let _ = std::writeln!(stdout, "{value}");
	}

	fn define_global(&mut self, chunk: &Chunk, globals: &mut impl Environment) -> Result<(), ExecutionError> {
		let pos = chunk.code()[self.ip];
		self.ip += 1;
		let value_val = self.pop();
		match chunk.read_constant(pos) {
			ScriptingValue::String(name) => globals.define_env(name, value_val)?,
			other => globals.define_env(&other.to_string(), value_val)?,
		}
		Ok(())
	}

	fn get_global(&mut self, chunk: &Chunk, globals: &impl Environment) -> Result<(), ExecutionError> {
		let pos = chunk.code()[self.ip];
		self.ip += 1;
		let val = match chunk.read_constant(pos) {
			ScriptingValue::String(name) => globals.get_env(name)?,
			other => globals.get_env(&other.to_string())?,
		};
		self.push(val)?;
		Ok(())
	}

	fn set_global(&mut self, chunk: &Chunk, globals: &mut impl Environment) -> Result<(), ExecutionError> {
		let pos = chunk.code()[self.ip];
		self.ip += 1;
		let value_val = self.pop();
		match chunk.read_constant(pos) {
			ScriptingValue::String(name) => globals.set_env(name, value_val)?,
			other => globals.set_env(&other.to_string(), value_val)?,
		}
		Ok(())
	}

	/// Execute a [`Chunk`] with the virtual machine,
	/// Returns the topmost stack [`ScriptingValue`] if there is one, otherwise [`ScriptingValue::nil()`].
	/// # Errors
	/// - unknown `OpCode`
	pub fn run(
		&mut self,
		chunk: &Chunk,
		globals: &mut impl Environment,
		#[cfg(feature = "std")] stdout: &mut impl std::io::Write,
	) -> Result<ScriptingValue, ExecutionError> {
		self.reset();
		// ignore empty chunks
		if chunk.code().is_empty() {
			return Ok(ScriptingValue::nil());
		}

		loop {
			let instruction: OpCode = chunk.code()[self.ip].into();
			self.ip += 1;
			match instruction {
				OpCode::Add => self.arithmetic_operator(ArithOp::Add)?,
				OpCode::Subtract => self.arithmetic_operator(ArithOp::Subtract)?,
				OpCode::Multiply => self.arithmetic_operator(ArithOp::Multiply)?,
				OpCode::Divide => self.arithmetic_operator(ArithOp::Divide)?,
				OpCode::BitwiseAnd => self.bitwise_operator(BitOp::And)?,
				OpCode::BitwiseOr => self.bitwise_operator(BitOp::Or)?,
				OpCode::BitwiseXor => self.bitwise_operator(BitOp::Xor)?,
				OpCode::BitwiseNot => self.bitwise_not()?,
				OpCode::Constant => self.constant(chunk)?,
				OpCode::DefineExternal => self.define_global(chunk, globals)?,
				OpCode::Equal => self.equal()?,
				OpCode::False => self.push(ScriptingValue::Boolean(false))?,
				OpCode::GetExternal => self.get_global(chunk, globals)?,
				OpCode::Greater => self.comparison_operator(CmpOp::Greater)?,
				OpCode::Jmp => {
					let target = self.read_jmp_address(chunk);
					self.ip = target;
				}
				OpCode::JmpIfFalse => {
					let target = self.read_jmp_address(chunk);
					if !self.peek(0).as_bool()? {
						self.ip = target;
					}
				}
				OpCode::JmpIfTrue => {
					let target = self.read_jmp_address(chunk);
					if self.peek(0).as_bool()? {
						self.ip = target;
					}
				}
				OpCode::Less => self.comparison_operator(CmpOp::Less)?,
				OpCode::Negate => self.negate()?,
				OpCode::Nil => self.push(ScriptingValue::nil())?,
				OpCode::Not => self.not()?,
				OpCode::Pop => {
					self.pop();
				}
				#[cfg(feature = "std")]
				OpCode::Print => self.print(stdout),
				OpCode::Return => {
					let val = if self.stack_top > 0 {
						self.pop()
					} else {
						ScriptingValue::nil()
					};
					return Ok(val);
				}
				OpCode::SetExternal => self.set_global(chunk, globals)?,
				OpCode::True => self.push(ScriptingValue::Boolean(true))?,
			}
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;

	// check, that the auto traits are available
	const fn is_normal<T: Sized + Send + Sync>() {}

	#[test]
	const fn normal_types() {
		is_normal::<&VM>();
		is_normal::<VM>();
	}
}