xbasic 0.3.2

use crate::basic_io::BasicIO;
use crate::chunk::Chunk;
use crate::compiler::Compiler;
use crate::error::*;
use crate::error_handler::ErrorHandler;
use crate::expr::ExprValue;
use crate::function::{Function, FunctionDefinition, FunctionType};
use crate::native_function::{NativeFunction, NativeFunctionDefinition};
use crate::opcodes::OpCode;
use crate::parser::Parser;
use crate::resolver::Resolver;
use crate::scanner::Scanner;
use crate::std::Std;
use crate::tokens::TokenType;
use crate::vm::VirtualMachine;
use std::collections::HashMap;

/// Used for building a custom `XBasic` interpreter
/// There are two main reasons to use this:
/// 1. To define custom native functions
/// 2. To set a compute limit
pub struct XBasicBuilder<T: 'static> {
	stdio: T,
	native_functions: HashMap<String, NativeFunction<T>>,
	compute_limit: usize,
}

impl<T> XBasicBuilder<T>
where
	T: BasicIO,
{
	/// Creates a new `XBasicBuilder` struct with the given `BasicIO` instance.
	///
	/// # Arguments
	///
	/// * `stdio` - An instance of a struct which implements the `BasicIO` trait.
	pub fn new(stdio: T) -> Self {
		let mut xbb = Self {
			stdio,
			native_functions: HashMap::new(),
			compute_limit: 0, // No limit
		};
		Std::attach(&mut xbb);
		xbb
	}

	/// Defines a new native function.
	/// Returns `Err` if the function name is already in use by a native function, or if the native function's arity is more than 255. Returns `Ok` otherwise.
	///
	/// # Arguments
	///
	/// * `name` - The name of the native function.
	/// * `arity` - The arity of the function, i.e. the number of arguments it expects.
	/// * `function` - The closure that implements the function. This is passed a vec of arguments, as well as a mutable reference to the BasicIO object. It is ensured that `arguments.len() == arity`.
	pub fn define_function<F: 'static>(
		&mut self,
		name: String,
		arity: u8,
		function: F,
	) -> Result<(), DefineFunctionError>
	where
		F: Fn(Vec<ExprValue>, &mut T) -> ExprValue,
	{
		if self.native_functions.len() < 256 {
			// Prevent defining the same native function twice
			if self.native_functions.contains_key(&name) {
				return Err(DefineFunctionError::NameInUse);
			}

			self.native_functions
				.insert(name.clone(), NativeFunction::new(name, arity, function));
			Ok(())
		} else {
			Err(DefineFunctionError::FunctionLimitReached)
		}
	}

	/// Set the compute limit, in terms of number of instructions.
	/// Set to 0 for no limit (the default)
	pub fn compute_limit(&mut self, limit: usize) {
		self.compute_limit = limit;
	}

	/// Consume the builder, producing an XBasic instance.
	pub fn build(self) -> XBasic<T> {
		let native_function_definitions = self
			.native_functions
			.iter()
			.map(|(_, v)| NativeFunctionDefinition::new(v.name.clone(), v.arity))
			.collect();

		XBasic {
			error_handler: ErrorHandler::new(),
			functions: Vec::new(),
			native_functions: native_function_definitions,
			compiler: Compiler::new(),
			vm: VirtualMachine::new(
				self.stdio,
				self.native_functions.into_iter().map(|(_, v)| v).collect(),
				self.compute_limit,
			),
		}
	}
}

/// Represents an xBASIC Interpreter.
pub struct XBasic<T: 'static> {
	/// Keeps track of errors that have been encountered while interpreting source code.
	/// Errors are recorded as user-friendly strings so that they can be passed directly to the user.
	pub error_handler: ErrorHandler,
	functions: Vec<Function>,
	native_functions: Vec<NativeFunctionDefinition>,
	compiler: Compiler,
	vm: VirtualMachine<T>,
}

impl<T> XBasic<T>
where
	T: BasicIO,
{
	/// Creates a new `XBasic` struct with the given `BasicIO` instance.
	pub fn new(stdio: T) -> Self {
		let xbb = XBasicBuilder::new(stdio);
		xbb.build()
	}

	/// Runs a snippet of XBasic source code using the interpreter. Returns `Ok` if there are no errors. Otherwise returns `Err` and keeps track of errors in the `error_handler` field.
	///
	/// # Arguments
	///
	/// * source - The XBasic code. This should be terminated by a newline character.
	pub fn run(&mut self, source: &str) -> Result<(), RunError> {
		let mut sc = Scanner::new(source, &mut self.error_handler);
		let tokens = sc.scan_tokens();

		let tokens = match tokens {
			Ok(x) => x,
			Err(_) => return Err(RunError::Scan),
		};

		// If a program consists of only newlines and an Eof, we shouldn't run it
		let mut only_whitespace = true;
		for token in &tokens {
			if token.token_type != TokenType::Newline && token.token_type != TokenType::Eof {
				only_whitespace = false;
				break;
			}
		}
		if only_whitespace {
			return Ok(());
		}

		let mut function_definitions = {
			let mut resolver = Resolver::new(tokens.clone(), &mut self.error_handler);
			match resolver.resolve(
				&self
					.functions
					.iter()
					.map(|a| a.definition())
					.collect::<Vec<_>>(),
			) {
				Some(x) => x,
				None => return Err(RunError::ResolveFunctions),
			}
		};

		// Tack on native functions
		// O(n^2) but should still be fast enough
		for function in &self.native_functions {
			let mut unique = true;
			for other in &function_definitions {
				if function.name == other.name {
					unique = false;
					break;
				}
			}
			if unique {
				function_definitions.push(FunctionDefinition::new(
					function.name.to_string(),
					function.arity,
					0,
					FunctionType::Native,
				));
			}
		}

		let mut parser = Parser::new(tokens, function_definitions, &mut self.error_handler);
		let stmts = match parser.parse() {
			Some((stmts, functions)) => {
				// Combine new functions with functions that we already know about
				for func in functions {
					self.functions.push(func);
				}
				stmts
			}
			None => return Err(RunError::Parse),
		};

		let native_function_ids: HashMap<String, usize> = self
			.native_functions
			.iter()
			.enumerate()
			.map(|(i, f)| (f.name.clone(), i))
			.collect();

		match self.compiler.compile(
			stmts,
			&mut self.functions,
			native_function_ids,
			&mut self.error_handler,
		) {
			Some(chunk) => {
				self.vm.run(
					chunk,
					self.functions
						.iter()
						.cloned()
						.map(|func| func.chunk.unwrap())
						.collect(),
					&mut self.error_handler,
				);
			}
			None => return Err(RunError::Compile),
		}

		if self.error_handler.had_runtime_error {
			return Err(RunError::Runtime);
		}

		Ok(())
	}

	/// Clears all existing errors in the various stages of the interpreter.
	pub fn clear_errors(&mut self) {
		self.error_handler.reset();
		self.compiler.clear_errors();
		self.vm.clear_errors();
	}

	/// Calls an xBASIC function from Rust code.
	/// # Arguments
	///
	/// * name - The name of the function to call
	/// * arguments - The arguments of the function to be called
	pub fn call_function(
		&mut self,
		name: &str,
		arguments: &[ExprValue],
	) -> Result<ExprValue, CallFunctionError> {
		for (i, function) in self.functions.iter().enumerate() {
			if function.name == name {
				if function.parameters.len() != arguments.len() {
					return Err(CallFunctionError::ArityMismatch);
				}
				// Generate instructions for function call
				let mut instructions = Vec::new();

				// Add arguments as literals
				for i in 0..arguments.len() {
					instructions.push(OpCode::Literal8 as u8);
					instructions.push(i as u8);
				}

				// Add call opcode
				instructions.push(OpCode::Call as u8);
				instructions.push(i as u8);

				// 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, 0);

				// Run the instructions
				if !self.vm.run(
					chunk,
					self.functions
						.iter()
						.cloned()
						.map(|func| func.chunk.unwrap())
						.collect(),
					&mut self.error_handler,
				) {
					return Err(CallFunctionError::Runtime);
				}

				// Pop the return value
				return Ok(self.vm.pop());
			}
		}

		Err(CallFunctionError::NotFound)
	}

	/// Gets a reference to the IO object.
	/// Useful if you are using the IO object for some kind of state.
	pub fn get_io(&self) -> &T {
		&self.vm.stdio
	}

	/// Gets a mutable reference to the IO object.
	/// Useful if you are using the IO object for some kind of state.
	pub fn get_io_mut(&mut self) -> &mut T {
		&mut self.vm.stdio
	}
}