egglang/scope/
functions.rs

1use alloc::{collections::BTreeMap, format, string::ToString, vec::Vec};
2use arcstr::ArcStr;
3
4use crate::{
5	error::{EggError, EggResult},
6	evaluator::evaluate,
7	expression::{Expression, Value},
8	operators::Operator,
9};
10
11use super::Scope;
12
13/// A function defined withing the Egg script.
14///
15/// Functions are defined using the `fn` operator.
16/// They keep track of their argument names. These are later used to construct a local [`Scope`] during function invocation.
17/// The body of the function is an [`Expression`], which is [`evaluated`](evaluate) using the newly created local [`Scope`].
18pub struct FunctionDefinition {
19	/// The names of the parameters of the function.
20	pub parameter_names: Vec<ArcStr>,
21	/// The body of the function.
22	pub body: Expression,
23}
24
25impl core::fmt::Debug for FunctionDefinition {
26	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
27		let parameters = self.parameter_names.first().map(|s| s.as_str()).unwrap_or("");
28		let parameters = self.parameter_names.iter().skip(1).fold(parameters.to_string(), |acc, s| format!("{}, {}", acc, s.as_str()));
29
30		write!(f, "Function ({})", parameters)
31	}
32}
33
34fn get_parameter_name(expr: &Expression) -> EggResult<ArcStr> {
35	match expr {
36		Expression::Word { name } => Ok(name.clone()),
37		_ => Err(EggError::InvalidFunctionDefinition("Parameter name must be a word".to_string())),
38	}
39}
40
41impl super::Scope {
42	pub(crate) fn get_function(&self, name: &str) -> Option<usize> {
43		let value = self.get(name);
44		match value {
45			Some(Value::Function(idx)) => Some(*idx),
46			_ => None,
47		}
48	}
49
50	pub fn get_function_definition(&self, idx: usize) -> EggResult<&FunctionDefinition> {
51		self.extras()
52			.functions
53			.get(&idx)
54			.ok_or_else(|| EggError::InvalidFunctionCall(format!("Function with index {} not found", idx)))
55	}
56
57	pub fn get_function_definition_mut(&mut self, idx: usize) -> EggResult<&mut FunctionDefinition> {
58		self.extras_mut()
59			.functions
60			.get_mut(&idx)
61			.ok_or_else(|| EggError::InvalidFunctionCall(format!("Function with index {} not found", idx)))
62	}
63
64	pub fn call_function(&mut self, idx: usize, parameters: &[Expression]) -> EggResult<Value> {
65		let function = unsafe {
66			// SAFETY: We are not modifying the scope, only reading from the function definition from it
67			(*(self as *const Scope)).get_function_definition(idx)?
68		};
69
70		if parameters.len() != function.parameter_names.len() {
71			return Err(EggError::InvalidFunctionCall(format!(
72				"Function expects {} parameters, but {} were given",
73				function.parameter_names.len(),
74				parameters.len()
75			)));
76		}
77
78		let mut new_scope = BTreeMap::new();
79		for (name, expression) in function.parameter_names.iter().zip(parameters.iter()) {
80			let value = evaluate(expression, self)?;
81			new_scope.insert(name.clone(), value);
82		}
83
84		let mut new_scope = self.overlay(new_scope);
85		evaluate(&function.body, &mut new_scope)
86	}
87
88	pub fn delete_function(&mut self, idx: usize) -> Option<FunctionDefinition> {
89		self.extras_mut().functions.remove(&idx)
90	}
91}
92
93/// Create a new Function
94pub struct CreateFunction;
95
96impl Operator for CreateFunction {
97	fn evaluate(&self, args: &[Expression], scope: &mut super::Scope) -> EggResult<crate::expression::Value> {
98		if args.is_empty() {
99			return Err(EggError::InvalidFunctionDefinition("Function Definition requires at least a body".to_string()));
100		}
101
102		// assemble function parts
103		let body = args[args.len() - 1].clone();
104		let parameter_names = args.iter().take(args.len() - 1).map(get_parameter_name).collect::<EggResult<Vec<ArcStr>>>()?;
105
106		scope.extras_mut().counter += 1;
107		let index = scope.extras().counter;
108		scope.extras_mut().functions.insert(index, FunctionDefinition { parameter_names, body });
109
110		Ok(crate::expression::Value::Function(index))
111	}
112}