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

use std::collections::HashMap;

use rebel_parse::ast::{self, pat};

use crate::{
	func::Func,
	typing::{Coerce, Type, TypeFamily},
	value::Value,
	Error, Result,
};

#[derive(Debug, Clone)]
pub struct Var {
	pub explicit_type: Type,
	inferred_type: Type,
	value: Value,
	initialized: bool,
}

impl Var {
	pub fn new(explicit_type: Type, value: Value) -> Result<Self> {
		Ok(Var {
			explicit_type: explicit_type.clone(),
			inferred_type: explicit_type.unify(value.typ()?, Coerce::Assign)?,
			value,
			initialized: true,
		})
	}

	pub fn new_uninitialized(explicit_type: Type) -> Self {
		Var {
			inferred_type: explicit_type.clone(),
			explicit_type,
			value: Value::Uninitialized,
			initialized: false,
		}
	}

	pub fn inferred_type(&self) -> Result<&Type> {
		if !self.initialized {
			return Err(Error::typ("uninitialized variable"));
		}
		Ok(&self.inferred_type)
	}

	pub fn value(&self) -> Result<&Value> {
		if !self.initialized {
			return Err(Error::typ("uninitialized variable"));
		}
		Ok(&self.value)
	}
}

#[derive(Debug, Clone)]
pub struct Context {
	pub vars: Module<Var>,
	pub types: Module<Type>,
	pub methods: HashMap<TypeFamily, HashMap<&'static str, Func>>,
}

impl Default for Context {
	fn default() -> Self {
		let vars = Module::default();
		let mut types = Module::default();
		let methods = HashMap::default();

		// Type "prelude"
		types.insert("bool", Type::Bool);
		types.insert("int", Type::Int);
		types.insert("str", Type::Str);

		Self {
			vars,
			types,
			methods,
		}
	}
}

impl Context {
	fn pat_ident<'a>(pat: &pat::Pat<'a>) -> ast::Ident<'a> {
		match pat {
			pat::Pat::Paren(subpat) => Self::pat_ident(subpat),
			pat::Pat::Ident(ident) => *ident,
		}
	}

	pub fn lookup_var(&self, path: &ast::Path) -> Result<&Var> {
		if path.root != ast::PathRoot::Relative {
			return Err(Error::lookup("invalid path"));
		}

		if path.components == [ast::Ident { name: "_" }] {
			return Err(Error::lookup("_ in evaluated expression"));
		}

		self.vars
			.lookup(&path.components)
			.ok_or(Error::lookup("undefined variable"))
	}

	pub fn lookup_type(&self, path: &ast::Path) -> Result<&Type> {
		if path.root != ast::PathRoot::Relative {
			return Err(Error::lookup("invalid path"));
		}

		if path.components == [ast::Ident { name: "_" }] {
			return Ok(&Type::Free);
		}

		self.types
			.lookup(&path.components)
			.ok_or(Error::lookup("undefined type"))
	}

	fn lookup_path_var_mut(&mut self, path: &ast::Path) -> Result<&mut Var> {
		if path.root != ast::PathRoot::Relative {
			return Err(Error::lookup("invalid path"));
		}

		self.vars
			.lookup_mut(&path.components)
			.ok_or(Error::lookup("undefined variable"))
	}

	fn lookup_destr_pat_var_type_mut(
		&mut self,
		pat: &pat::DestrPat,
	) -> Result<(&mut Type, &mut bool)> {
		Ok(match pat {
			pat::DestrPat::Index { base, index } => {
				match Type::ast_expr_type(self, index)? {
					Type::Int => {}
					_ => {
						return Err(Error::typ("invalid array index type"));
					}
				}
				let (typ, initialized) = self.lookup_destr_pat_var_type_mut(base)?;
				if !*initialized {
					return Err(Error::typ(
						"tried to assign field of uninitialized variable",
					));
				}
				(
					match typ {
						Type::Array(inner) => inner.as_mut(),
						_ => return Err(Error::typ("invalid array index base")),
					},
					initialized,
				)
			}
			pat::DestrPat::Field { base, field } => {
				let (typ, initialized) = self.lookup_destr_pat_var_type_mut(base)?;
				if !*initialized {
					return Err(Error::typ(
						"tried to assign field of uninitialized variable",
					));
				}
				(
					match typ {
						Type::Tuple(elems) => {
							let index: usize =
								field.name.parse().or(Err(Error::typ("no such field")))?;
							elems.get_mut(index).ok_or(Error::typ("no such field"))?
						}
						Type::Struct(entries) => entries
							.get_mut(field.name)
							.ok_or(Error::typ("no such field"))?,
						_ => return Err(Error::typ("invalid field access base type")),
					},
					initialized,
				)
			}
			pat::DestrPat::Paren(subpat) => self.lookup_destr_pat_var_type_mut(subpat)?,
			pat::DestrPat::Path(path) => {
				let var = self.lookup_path_var_mut(path)?;
				(&mut var.inferred_type, &mut var.initialized)
			}
		})
	}

	fn lookup_destr_pat_var_type_value_mut(
		&mut self,
		pat: &pat::DestrPat,
	) -> Result<(&mut Type, &mut Value, &mut bool)> {
		Ok(match pat {
			pat::DestrPat::Index { base, index } => {
				let index = usize::try_from(match Value::eval(self, index)? {
					Value::Int(index) => index,
					_ => {
						return Err(Error::typ("invalid array index type"));
					}
				})
				.or(Err(Error::eval("array index out of bounds")))?;
				let (typ, value, initialized) = self.lookup_destr_pat_var_type_value_mut(base)?;
				if !*initialized {
					return Err(Error::typ(
						"tried to assign field of uninitialized variable",
					));
				}
				let (inner_type, inner_value) = match (typ, value) {
					(Type::Array(inner_type), Value::Array(inner_value)) => (
						inner_type.as_mut(),
						inner_value
							.get_mut(index)
							.ok_or(Error::eval("array index out of bounds"))?,
					),
					_ => return Err(Error::typ("invalid array index base")),
				};
				(inner_type, inner_value, initialized)
			}
			pat::DestrPat::Field { base, field } => {
				let (typ, value, initialized) = self.lookup_destr_pat_var_type_value_mut(base)?;
				if !*initialized {
					return Err(Error::typ(
						"tried to assign field of uninitialized variable",
					));
				}
				let (inner_type, inner_value) = match (typ, value) {
					(Type::Tuple(type_elems), Value::Tuple(value_elems)) => {
						let index: usize =
							field.name.parse().or(Err(Error::typ("no such field")))?;
						(
							type_elems
								.get_mut(index)
								.ok_or(Error::typ("no such field"))?,
							value_elems
								.get_mut(index)
								.ok_or(Error::typ("no such field"))?,
						)
					}
					(Type::Struct(type_entries), Value::Struct(value_entries)) => (
						type_entries
							.get_mut(field.name)
							.ok_or(Error::typ("no such field"))?,
						value_entries
							.get_mut(field.name)
							.ok_or(Error::typ("no such field"))?,
					),
					_ => return Err(Error::typ("invalid field access base type")),
				};
				(inner_type, inner_value, initialized)
			}
			pat::DestrPat::Paren(subpat) => self.lookup_destr_pat_var_type_value_mut(subpat)?,
			pat::DestrPat::Path(path) => {
				let var = self.lookup_path_var_mut(path)?;
				(&mut var.inferred_type, &mut var.value, &mut var.initialized)
			}
		})
	}

	fn is_wildcard_destr_pat(pat: &pat::DestrPat) -> bool {
		match pat {
			pat::DestrPat::Index { .. } => false,
			pat::DestrPat::Field { .. } => false,
			pat::DestrPat::Paren(pat) => Self::is_wildcard_destr_pat(pat),
			pat::DestrPat::Path(path) => {
				path.root == ast::PathRoot::Relative
					&& path.components == [ast::Ident { name: "_" }]
			}
		}
	}

	fn assign_destr_pat_type(&mut self, dest: &pat::DestrPat, typ: Type) -> Result<Type> {
		if Self::is_wildcard_destr_pat(dest) {
			return Ok(Type::Unit);
		}

		let (dest_type, initialized) = self.lookup_destr_pat_var_type_mut(dest)?;

		let inferred_type = dest_type.clone().unify(typ, Coerce::Common)?;
		*dest_type = inferred_type.clone();
		*initialized = true;

		// TODO: Check explicit type

		Ok(inferred_type)
	}

	fn assign_destr_pat_value(&mut self, dest: &pat::DestrPat, value: Value) -> Result<Value> {
		let typ = value.typ()?;

		if Self::is_wildcard_destr_pat(dest) {
			return Ok(Value::Unit);
		}

		let (dest_type, dest_value, initialized) =
			self.lookup_destr_pat_var_type_value_mut(dest)?;

		*dest_type = dest_type.clone().unify(typ, Coerce::Common)?;
		*dest_value = value.clone();
		*initialized = true;

		// TODO: Check explicit type

		Ok(value)
	}

	pub fn record_type(&mut self, stmt: &ast::BlockStmt) -> Result<Type> {
		Ok(match stmt {
			ast::BlockStmt::Let { dest, expr } => {
				let ast::TypedPat { pat, typ } = dest.as_ref();

				let dest_ident = Self::pat_ident(pat);

				let explicit_type = if let Some(typ) = typ {
					Type::ast_type(self, typ)?
				} else {
					Type::Free
				};

				let expr_type = expr
					.as_ref()
					.map(|expr| {
						let expr_type = Type::ast_expr_type(self, expr)?;
						explicit_type
							.clone()
							.unify(expr_type.clone(), Coerce::Assign)?;
						Ok(expr_type)
					})
					.transpose()?;

				if dest_ident.name == "_" {
					return Ok(Type::Unit);
				}

				self.vars
					.insert(dest_ident.name, Var::new_uninitialized(explicit_type));

				let Some(expr_type) = expr_type else {
					return Ok(Type::Unit);
				};

				self.assign_destr_pat_type(&pat::DestrPat::from(pat), expr_type)?
			}
			ast::BlockStmt::Assign { dest, expr } => {
				let expr_type = Type::ast_expr_type(self, expr)?;
				self.assign_destr_pat_type(dest, expr_type)?
			}
			ast::BlockStmt::Expr { expr } => Type::ast_expr_type(self, expr)?,
			ast::BlockStmt::Empty => Type::Unit,
		})
	}

	pub fn execute(&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 = Self::pat_ident(pat);

				let explicit_type = if let Some(typ) = typ {
					Type::ast_type(self, typ)?
				} else {
					Type::Free
				};

				let value = expr
					.as_ref()
					.map(|expr| Value::eval(self, expr))
					.transpose()?;

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

				self.vars
					.insert(dest_ident.name, Var::new_uninitialized(explicit_type));

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

				self.assign_destr_pat_value(&pat::DestrPat::from(pat), value)?
			}
			ast::BlockStmt::Assign { dest, expr } => {
				let value = Value::eval(self, expr)?;
				self.assign_destr_pat_value(dest, value)?
			}
			ast::BlockStmt::Expr { expr } => Value::eval(self, expr)?,
			ast::BlockStmt::Empty => Value::Unit,
		})
	}
}

#[derive(Debug, Clone)]
pub struct Module<T>(pub HashMap<String, ModuleEntry<T>>);

impl<T> Module<T> {
	pub fn insert(&mut self, ident: &str, value: T) {
		self.0.insert(ident.to_owned(), ModuleEntry::Def(value));
	}

	pub fn lookup(&self, path: &[ast::Ident<'_>]) -> Option<&T> {
		let (ident, rest) = path.split_first()?;

		match self.0.get(ident.name)? {
			ModuleEntry::Module(module) => module.lookup(rest),
			ModuleEntry::Def(def) => {
				if rest.is_empty() {
					Some(def)
				} else {
					None
				}
			}
		}
	}

	pub fn lookup_mut(&mut self, path: &[ast::Ident<'_>]) -> Option<&mut T> {
		let (ident, rest) = path.split_first()?;

		match self.0.get_mut(ident.name)? {
			ModuleEntry::Module(module) => module.lookup_mut(rest),
			ModuleEntry::Def(def) => {
				if rest.is_empty() {
					Some(def)
				} else {
					None
				}
			}
		}
	}
}

impl<T> Default for Module<T> {
	fn default() -> Self {
		Self(HashMap::new())
	}
}

#[derive(Debug, Clone)]
pub enum ModuleEntry<T> {
	Module(Module<T>),
	Def(T),
}