path: root/crates/rebel-lang/src/value.rs

use std::{
	cell::RefCell,
	fmt::{Display, Write},
	iter,
};

use rebel_parse::ast::{self, expr, pat, typ};
use rustc_hash::{FxHashMap, FxHashSet};

use crate::{
	func::{Func, FuncDef},
	scope::{self, Scope},
	typing::{Coerce, OrdType, Type, TypeContext, TypeFamily},
	Error, Result,
};

#[derive(Debug)]
pub struct Context<'scope>(pub &'scope mut Box<Scope<Value>>);

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Value {
	Uninitialized,
	None,
	Bool(bool),
	Int(i64),
	Str(String),
	Tuple(Vec<Value>),
	Map(FxHashMap<OrdValue, Value>),
	Array(Vec<Value>),
	Struct(FxHashMap<String, Value>),
	Fn(Box<Func>),
}

impl Value {
	pub fn unit() -> Self {
		Value::Tuple(Vec::new())
	}

	pub fn typ(&self) -> Result<Type> {
		Ok(match self {
			Value::Uninitialized => return Err(Error::typ("uninitialized value")),
			Value::None => Type::Option(Box::new(Type::Free)),
			Value::Bool(_) => Type::Bool,
			Value::Int(_) => Type::Int,
			Value::Str(_) => Type::Str,
			Value::Tuple(elems) => {
				Type::Tuple(elems.iter().map(|elem| elem.typ()).collect::<Result<_>>()?)
			}
			Value::Array(elems) => Type::Array(Box::new(Self::common_type(elems)?)),
			Value::Map(entries) => Type::Map(
				Self::common_ord_type(entries.keys())?,
				Box::new(Self::common_type(entries.values())?),
			),
			Value::Struct(entries) => Type::Struct(
				entries
					.iter()
					.map(|(k, v)| Ok((k.clone(), v.typ()?)))
					.collect::<Result<_>>()?,
			),
			Value::Fn(func) => Type::Fn(Box::new(func.typ.clone())),
		})
	}

	fn common_type<'a>(values: impl IntoIterator<Item = &'a Value>) -> Result<Type> {
		values.into_iter().try_fold(Type::Free, |acc, value| {
			acc.unify(value.typ()?, Coerce::Common)
		})
	}

	fn common_ord_type<'a>(values: impl IntoIterator<Item = &'a OrdValue>) -> Result<OrdType> {
		values.into_iter().try_fold(OrdType::Free, |acc, value| {
			acc.unify(value.typ()?, Coerce::Common)
		})
	}
}

impl From<OrdValue> for Value {
	fn from(value: OrdValue) -> Self {
		match value {
			OrdValue::Int(v) => Value::Int(v),
			OrdValue::Str(v) => Value::Str(v),
		}
	}
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum OrdValue {
	Int(i64),
	Str(String),
}

impl OrdValue {
	pub fn typ(&self) -> Result<OrdType> {
		Ok(match self {
			OrdValue::Int(_) => OrdType::Int,
			OrdValue::Str(_) => OrdType::Str,
		})
	}
}

impl Display for OrdValue {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self {
			OrdValue::Int(value) => value.fmt(f),
			OrdValue::Str(value) => write!(f, "{value:?}"),
		}
	}
}

impl TryFrom<Value> for OrdValue {
	type Error = Error;

	fn try_from(value: Value) -> Result<Self> {
		Ok(match value {
			Value::Int(v) => OrdValue::Int(v),
			Value::Str(v) => OrdValue::Str(v),
			_ => return Err(Error::bug("OrdValue from unordered value")),
		})
	}
}

impl<'scope> Context<'scope> {
	pub fn eval_block_stmt(&mut self, stmt: &ast::BlockStmt) -> Result<Value> {
		Ok(match stmt {
			ast::BlockStmt::Let { dest, expr } => {
				let ast::TypedPat { pat, typ: _ } = dest.as_ref();

				let dest_ident = scope::pat_ident(pat);

				let value = expr.as_ref().map(|expr| self.eval_expr(expr)).transpose()?;

				if dest_ident.name == "_" {
					return Ok(Value::unit());
				}

				self.0
					.defs
					.insert_value(dest_ident.name.as_ref(), Value::Uninitialized);

				let Some(value) = value else {
					return Ok(Value::unit());
				};

				self.assign_destr_pat_value(&pat::DestrPat::from(pat.borrowed()), value)?
			}
			ast::BlockStmt::Assign { dest, expr } => {
				let value = self.eval_expr(expr)?;
				self.assign_destr_pat_value(dest, value)?
			}
			ast::BlockStmt::Fn {
				ident,
				params,
				ret,
				block,
			} => {
				let value = Value::Fn(Box::new(self.eval_fn_def(params, ret.as_deref(), block)?));

				// TODO: Reject in validation?
				if ident.name == "_" {
					return Ok(Value::unit());
				}
				self.0.defs.insert_value(ident.name.as_ref(), value.clone());

				value
			}
			ast::BlockStmt::Expr { expr } => self.eval_expr(expr)?,
			ast::BlockStmt::Empty => Value::unit(),
		})
	}

	fn eval_fn_def(
		&mut self,
		params: &[ast::FuncParam],
		ret: Option<&typ::Type>,
		block: &ast::Block,
	) -> Result<Func> {
		let type_ctx = TypeContext(self.0);
		let typ = type_ctx.type_fn_def(params, ret)?;

		let param_names: Vec<_> = params
			.iter()
			.map(|param| param.name.name.as_ref().to_owned())
			.collect();

		Ok(Func {
			typ,
			def: FuncDef::Body {
				param_names,
				block: block.borrowed().into_owned(),
			},
		})
	}

	pub fn eval_expr(&mut self, expr: &expr::Expr<'_>) -> Result<Value> {
		use expr::Expr::*;

		match expr {
			Binary { left, op, right } => self.eval_binary_op(left, *op, right),
			Unary { op, expr } => self.eval_unary_op(*op, expr),
			Apply { expr, params } => self.eval_apply(expr, params),
			Method {
				expr,
				method,
				params,
			} => self.eval_method(expr, method, params),
			Index { base, index } => self.eval_index(base, index),
			Field { base, field } => self.eval_field(base, field),
			Block(block) => self.eval_block(block),
			IfElse {
				if_blocks,
				else_block,
			} => self.eval_ifelse(if_blocks, else_block),
			Paren(subexpr) => self.eval_expr(subexpr),
			Path(path) => self.eval_path(path),
			Literal(lit) => self.eval_literal(lit),
		}
	}

	fn eval_binary_op(
		&mut self,
		left: &expr::Expr<'_>,
		op: expr::OpBinary,
		right: &expr::Expr<'_>,
	) -> Result<Value> {
		use expr::OpBinary::*;
		use Value::*;

		let tl = self.eval_expr(left)?;
		let tr = self.eval_expr(right)?;

		Ok(match (tl, op, tr) {
			(Str(s1), Add, Str(s2)) => Str(s1 + &s2),
			(Int(i1), Add, Int(i2)) => Int(i1
				.checked_add(i2)
				.ok_or(Error::eval("integer over- or underflow"))?),
			(Array(elems1), Add, Array(elems2)) => Array([elems1, elems2].concat()),
			(Int(i1), Sub, Int(i2)) => Int(i1
				.checked_sub(i2)
				.ok_or(Error::eval("integer over- or underflow"))?),
			(Array(elems1), Sub, Array(elems2)) => Array(
				elems1
					.into_iter()
					.filter(|elem| !elems2.contains(elem))
					.collect(),
			),
			(Int(i1), Mul, Int(i2)) => Int(i1
				.checked_mul(i2)
				.ok_or(Error::eval("integer over- or underflow"))?),
			(Int(i1), Div, Int(i2)) => {
				Int(i1.checked_div(i2).ok_or(Error::eval("division by zero"))?)
			}
			(Int(i1), Rem, Int(i2)) => {
				Int(i1.checked_rem(i2).ok_or(Error::eval("division by zero"))?)
			}
			(Bool(b1), And, Bool(b2)) => Bool(b1 && b2),
			(Bool(b1), Or, Bool(b2)) => Bool(b1 || b2),
			(l, Eq, r) => Bool(l == r),
			(l, Ne, r) => Bool(l != r),
			(Int(i1), Lt, Int(i2)) => Bool(i1 < i2),
			(Int(i1), Le, Int(i2)) => Bool(i1 <= i2),
			(Int(i1), Ge, Int(i2)) => Bool(i1 >= i2),
			(Int(i1), Gt, Int(i2)) => Bool(i1 > i2),
			(Str(i1), Lt, Str(i2)) => Bool(i1 < i2),
			(Str(i1), Le, Str(i2)) => Bool(i1 <= i2),
			(Str(i1), Ge, Str(i2)) => Bool(i1 >= i2),
			(Str(i1), Gt, Str(i2)) => Bool(i1 > i2),
			_ => return Err(Error::typ("invalid types for operation")),
		})
	}

	fn eval_unary_op(&mut self, op: expr::OpUnary, expr: &expr::Expr<'_>) -> Result<Value> {
		use expr::OpUnary::*;
		use Value::*;

		let typ = self.eval_expr(expr)?;

		Ok(match (op, typ) {
			(Not, Bool(val)) => Bool(!val),
			(Neg, Int(val)) => Int(val
				.checked_neg()
				.ok_or(Error::eval("integer over- or underflow"))?),
			_ => return Err(Error::typ("invalid type for operation")),
		})
	}

	fn eval_index(&mut self, base: &expr::Expr<'_>, index: &expr::Expr<'_>) -> Result<Value> {
		use Value::*;

		let base_value = self.eval_expr(base)?;
		let index_value = self.eval_expr(index)?;

		if let Array(elems) = base_value {
			if let Int(index) = index_value {
				return usize::try_from(index)
					.ok()
					.and_then(|index| elems.into_iter().nth(index))
					.ok_or(Error::eval("array index out of bounds"));
			}
		} else if let Map(entries) = base_value {
			return entries
				.get(&index_value.try_into()?)
				.cloned()
				.ok_or(Error::eval("map key not found"));
		}

		Err(Error::typ("invalid types index operation"))
	}

	fn eval_func(&mut self, func: Func, call_param_values: Vec<Value>) -> Result<Value> {
		match func.def {
			FuncDef::Intrinsic(f) => f(&call_param_values),
			FuncDef::Body { param_names, block } => {
				let mut scope = Box::new(self.0.func());

				for (param_name, param_value) in param_names.iter().zip(call_param_values) {
					scope.defs.insert_value(param_name, param_value);
				}

				Context(&mut scope).eval_block(&block)
			}
		}
	}

	fn eval_apply(&mut self, expr: &expr::Expr<'_>, params: &[expr::Expr]) -> Result<Value> {
		use Value::*;

		let value = self.eval_expr(expr)?;

		let Fn(func) = value else {
			return Err(Error::typ("invalid type for function call"));
		};

		let param_values: Vec<_> = params
			.iter()
			.map(|param| self.eval_expr(param))
			.collect::<Result<_>>()?;

		self.eval_func(*func, param_values)
	}

	fn eval_method(
		&mut self,
		expr: &expr::Expr<'_>,
		method: &ast::Ident<'_>,
		params: &[expr::Expr],
	) -> Result<Value> {
		let self_value = self.eval_expr(expr)?;
		// TODO: Use inferred type for method lookup
		let self_type = self_value.typ()?;
		let type_family = TypeFamily::from(&self_type);

		let method = self
			.0
			.methods
			.get(&type_family)
			.and_then(|methods| methods.get(method.name.as_ref()))
			.ok_or(Error::lookup("undefined method"))?
			.clone();

		let param_values: Vec<_> = iter::once(Ok(self_value))
			.chain(params.iter().map(|param| self.eval_expr(param)))
			.collect::<Result<_>>()?;

		self.eval_func(method, param_values)
	}

	fn eval_field(&mut self, base: &expr::Expr<'_>, field: &ast::Ident<'_>) -> Result<Value> {
		use Value::*;

		let base_value = self.eval_expr(base)?;
		let name = field.name.as_ref();

		Ok(match base_value {
			Tuple(elems) => {
				let index: usize = name.parse().or(Err(Error::typ("no such field")))?;
				elems
					.into_iter()
					.nth(index)
					.ok_or(Error::typ("no such field"))?
			}
			Struct(mut entries) => entries.remove(name).ok_or(Error::typ("no such field"))?,
			_ => return Err(Error::typ("invalid field access base type")),
		})
	}

	fn eval_scope(&mut self, stmts: &[ast::BlockStmt]) -> Result<(Value, FxHashSet<String>)> {
		let (ret, upvalues) = self.scoped(|ctx| {
			let mut ret = Value::unit();
			for stmt in stmts {
				ret = ctx.eval_block_stmt(stmt)?;
			}
			Ok(ret)
		});
		Ok((ret?, upvalues))
	}

	fn eval_block(&mut self, block: &ast::Block) -> Result<Value> {
		let (ret, upvalues) = self.eval_scope(&block.0)?;
		self.0.initialize_all(upvalues);
		Ok(ret)
	}

	fn eval_ifelse(
		&mut self,
		if_blocks: &[(expr::Expr, ast::Block)],
		else_block: &Option<Box<ast::Block>>,
	) -> Result<Value> {
		for (cond, block) in if_blocks {
			if self.eval_expr(cond)? == Value::Bool(true) {
				return self.eval_block(block);
			}
		}
		if let Some(block) = else_block {
			self.eval_block(block)
		} else {
			Ok(Value::unit())
		}
	}

	fn eval_path(&self, path: &ast::Path<'_>) -> Result<Value> {
		Ok(self.0.lookup_value(path)?.clone())
	}

	fn eval_literal(&mut self, lit: &expr::Literal<'_>) -> Result<Value> {
		use expr::Literal;
		use Value::*;

		Ok(match lit {
			Literal::Unit => Value::unit(),
			Literal::None => Value::None,
			Literal::Bool(val) => Bool(*val),
			Literal::Int(val) => Int(*val),
			Literal::Str { pieces, kind } => Str(StrDisplay {
				pieces,
				kind: *kind,
				ctx: RefCell::new(self),
			}
			.to_string()),
			Literal::Tuple(elems) => Tuple(
				elems
					.iter()
					.map(|elem| self.eval_expr(elem))
					.collect::<Result<_>>()?,
			),
			Literal::Array(elems) => Array(
				elems
					.iter()
					.map(|elem| self.eval_expr(elem))
					.collect::<Result<_>>()?,
			),
			Literal::Map(entries) => {
				let map: FxHashMap<_, _> = entries
					.iter()
					.map(|expr::MapEntry { key, value }| {
						Ok((self.eval_expr(key)?.try_into()?, self.eval_expr(value)?))
					})
					.collect::<Result<_>>()?;
				if map.len() != entries.len() {
					return Err(Error::eval("duplicate map key"));
				}
				Map(map)
			}
			Literal::Struct(entries) => {
				if entries.is_empty() {
					return Ok(Value::unit());
				}
				Struct(
					entries
						.iter()
						.map(|expr::StructField { name, value }| {
							Ok((name.as_ref().to_owned(), self.eval_expr(value)?))
						})
						.collect::<Result<_>>()?,
				)
			}
		})
	}

	#[allow(clippy::type_complexity)]
	fn assign_destr_pat_value_with<'r, R: 'r>(
		&mut self,
		pat: &pat::DestrPat,
		f: Box<dyn FnOnce(&mut Value) -> Result<R> + 'r>,
	) -> Result<R> {
		match pat {
			pat::DestrPat::Index { base, index } => {
				let index_value = self.eval_expr(index)?;
				self.assign_destr_pat_value_with(
					base,
					Box::new(|base_value| match (base_value, index_value) {
						(Value::Array(inner), Value::Int(index)) => f(usize::try_from(index)
							.ok()
							.and_then(|index| inner.get_mut(index))
							.ok_or(Error::eval("array index out of bounds"))?),
						(Value::Map(entries), key) => {
							use std::collections::hash_map::Entry;
							match entries.entry(OrdValue::try_from(key)?) {
								Entry::Occupied(mut entry) => f(entry.get_mut()),
								Entry::Vacant(entry) => {
									let mut tmp = Value::Uninitialized;
									let ret =
										f(&mut tmp).or(Err(Error::eval("map key not found")))?;
									entry.insert(tmp);
									Ok(ret)
								}
							}
						}
						_ => Err(Error::typ("invalid types for index operation")),
					}),
				)
			}
			pat::DestrPat::Field { base, field } => self.assign_destr_pat_value_with(
				base,
				Box::new(|value| match value {
					Value::Tuple(elems) => {
						let index: Option<usize> = field.name.parse().ok();
						f(index
							.and_then(|index| elems.get_mut(index))
							.ok_or(Error::typ("no such field"))?)
					}
					Value::Struct(value_entries) => f(value_entries
						.get_mut(field.name.as_ref())
						.ok_or(Error::typ("no such field"))?),
					_ => Err(Error::typ("invalid field access base type")),
				}),
			),
			pat::DestrPat::Paren(subpat) => self.assign_destr_pat_value_with(subpat, f),
			pat::DestrPat::Path(path) => f(self.0.lookup_var_mut(path)?.0),
		}
	}

	fn assign_destr_pat_value(&mut self, dest: &pat::DestrPat, value: Value) -> Result<Value> {
		if scope::is_wildcard_destr_pat(dest) {
			return Ok(Value::unit());
		}

		self.assign_destr_pat_value_with(
			dest,
			Box::new(|dest_value| {
				*dest_value = value.clone();
				Ok(())
			}),
		)?;

		Ok(value)
	}

	fn scoped<F, R>(&mut self, f: F) -> (R, FxHashSet<String>)
	where
		F: FnOnce(&mut Context) -> R,
	{
		self.0.scoped(|scope| f(&mut Context(scope)))
	}
}

impl Display for Value {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self {
			Value::Uninitialized => f.write_str("_"),
			Value::None => f.write_str("none"),
			Value::Bool(value) => value.fmt(f),
			Value::Int(value) => value.fmt(f),
			Value::Str(value) => write!(f, "{value:?}"),
			Value::Tuple(elems) => {
				let mut first = true;
				f.write_str("(")?;
				for elem in elems {
					if !first {
						f.write_str(", ")?;
					}
					first = false;
					elem.fmt(f)?;
				}
				if elems.len() == 1 {
					f.write_str(",")?;
				}
				f.write_str(")")
			}
			Value::Array(elems) => {
				let mut first = true;
				f.write_str("[")?;
				for elem in elems {
					if !first {
						f.write_str(", ")?;
					}
					first = false;
					elem.fmt(f)?;
				}
				f.write_str("]")
			}
			Value::Map(entries) => {
				let mut first = true;
				let mut entries: Vec<_> = entries.iter().collect();
				entries.sort_unstable_by_key(|(key, _)| *key);
				f.write_str("map{")?;
				for (key, value) in entries {
					if !first {
						f.write_str(", ")?;
					}
					first = false;
					write!(f, "{key} => {value}")?;
				}
				f.write_str("}")
			}
			Value::Struct(entries) => {
				let mut first = true;
				f.write_str("{")?;
				for (key, value) in entries {
					if !first {
						f.write_str(", ")?;
					}
					first = false;
					write!(f, "{key}: {value}")?;
				}
				f.write_str("}")
			}
			Value::Fn(func) => func.typ.fmt(f),
		}
	}
}

#[derive(Debug)]
struct Stringify<'a>(&'a Value);

impl<'a> Display for Stringify<'a> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self.0 {
			Value::Bool(value) => value.fmt(f),
			Value::Int(value) => value.fmt(f),
			Value::Str(value) => value.fmt(f),
			_ => Err(std::fmt::Error),
		}
	}
}

#[derive(Debug)]
struct ScriptStringify<'a>(&'a Value);

impl<'a> ScriptStringify<'a> {
	fn fmt_list(elems: &'a [Value], f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		let mut first = true;
		for elem in elems {
			if !first {
				f.write_char(' ')?;
			}
			ScriptStringify(elem).fmt(f)?;
			first = false;
		}
		Ok(())
	}
}

impl<'a> Display for ScriptStringify<'a> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self.0 {
			Value::Bool(value) => {
				value.fmt(f)?;
			}
			Value::Int(value) => {
				value.fmt(f)?;
			}
			Value::Str(value) => {
				f.write_char('\'')?;
				f.write_str(&value.replace('\'', "'\\''"))?;
				f.write_char('\'')?;
			}
			Value::None => {}
			Value::Array(elems) => {
				Self::fmt_list(elems, f)?;
			}
			Value::Tuple(elems) => {
				Self::fmt_list(elems, f)?;
			}
			_ => return Err(std::fmt::Error),
		};
		Ok(())
	}
}

#[derive(Debug)]
struct StrDisplay<'a, 'scope> {
	pieces: &'a [expr::StrPiece<'a>],
	kind: expr::StrKind,
	ctx: RefCell<&'a mut Context<'scope>>,
}

impl<'a, 'scope> Display for StrDisplay<'a, 'scope> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		for piece in self.pieces {
			match piece {
				expr::StrPiece::Chars(chars) => f.write_str(chars)?,
				expr::StrPiece::Escape(c) => f.write_char(*c)?,
				expr::StrPiece::Interp(expr) => {
					let val = self
						.ctx
						.borrow_mut()
						.eval_expr(expr)
						.or(Err(std::fmt::Error))?;
					match self.kind {
						expr::StrKind::Regular => Stringify(&val).fmt(f),
						expr::StrKind::Raw => unreachable!(),
						expr::StrKind::Script => ScriptStringify(&val).fmt(f),
					}?;
				}
			};
		}
		Ok(())
	}
}