Unify builtin and user-defined function values

src/environment.rs

@@ -3,41 +3,50 @@ use super::*;
 #[derive(Clone, Debug, Default)]
 pub struct Environment<'src> {
   pub config: Config,
-  pub functions: HashMap<&'src str, Function<'src>>,
   pub parent: Option<Box<Environment<'src>>>,
-  pub variables: HashMap<&'src str, Value<'src>>,
+  symbols: HashMap<&'src str, Symbol<'src>>,
+}
+
+#[derive(Clone, Debug, Default)]
+struct Symbol<'src> {
+  function: Option<Function<'src>>,
+  value: Option<Value<'src>>,
 }
 
 impl<'src> Environment<'src> {
   pub fn new(config: Config) -> Self {
     let mut environment = Self {
       config: config.clone(),
-      functions: HashMap::new(),
       parent: None,
-      variables: HashMap::new(),
+      symbols: HashMap::new(),
     };
 
     for builtin in BUILTINS {
       match builtin {
         Builtin::Constant { value, .. } => {
-          environment.add_variable(builtin.name(), value(config.clone()));
+          environment.add_symbol(builtin.name(), value(config.clone()));
         }
         Builtin::Function { function, .. } => {
-          environment
-            .add_function(builtin.name(), Function::Builtin(*function));
+          environment.add_function(
+            builtin.name(),
+            Function::Builtin {
+              function: *function,
+              name: builtin.name(),
+            },
+          );
         }
       }
     }
 
     environment
   }
 
-  pub fn add_function(&mut self, name: &'src str, function: Function<'src>) {
-    self.functions.insert(name, function);
+  pub fn add_symbol(&mut self, name: &'src str, value: Value<'src>) {
+    self.symbols.entry(name).or_default().value = Some(value);
   }
 
-  pub fn add_variable(&mut self, name: &'src str, value: Value<'src>) {
-    self.variables.insert(name, value);
+  pub fn add_function(&mut self, name: &'src str, function: Function<'src>) {
+    self.symbols.entry(name).or_default().function = Some(function);
   }
 
   pub fn call_function(
@@ -46,75 +55,10 @@ impl<'src> Environment<'src> {
     arguments: Vec<Value<'src>>,
     span: Span,
   ) -> Result<Value<'src>, Error> {
-    if let Some(function) = self.functions.get(name) {
-      match function {
-        Function::Builtin(function) => function(BuiltinFunctionPayload {
-          arguments,
-          config: self.config.clone(),
-          span,
-        }),
-        Function::UserDefined(Value::Function(
-          name,
-          parameters,
-          body,
-          environment,
-        )) => {
-          if parameters.len() != arguments.len() {
-            return Err(Error::new(
-              span,
-              format!(
-                "Function `{}` expects {} arguments, got {}",
-                name,
-                parameters.len(),
-                arguments.len()
-              ),
-            ));
-          }
-
-          let mut call_environment =
-            Environment::with_parent(environment.clone());
-
-          call_environment.add_function(name, function.clone());
-
-          for (parameter, argument) in parameters.iter().zip(arguments.iter()) {
-            let parameter_name = *parameter;
-
-            match argument {
-              Value::Function(_, _, _, _) => {
-                call_environment.add_function(
-                  parameter_name,
-                  Function::UserDefined(argument.clone()),
-                );
-              }
-              Value::BuiltinFunction(_, builtin) => {
-                call_environment
-                  .add_function(parameter_name, Function::Builtin(*builtin));
-              }
-              _ => {
-                call_environment.add_variable(parameter_name, argument.clone());
-              }
-            }
-          }
-
-          let mut evaluator = Evaluator::from(call_environment);
-          evaluator.inside_function = true;
-
-          if body.is_empty() {
-            return Ok(Value::Null);
-          }
-
-          for statement in body.iter() {
-            let result = evaluator.eval_statement(statement)?;
-
-            if result.is_return() {
-              return Ok(result.unwrap());
-            }
-          }
-
-          Ok(Value::Null)
-        }
-        _ => Err(Error::new(span, format!("`{}` is not a function", name))),
-      }
+    if let Some(function) = self.resolve_function(name) {
+      function.call(arguments, self.config.clone(), span)
+    } else if self.resolve_symbol(name).is_some() {
+      Err(Error::new(span, format!("`{}` is not a function", name)))
     } else {
       Err(Error::new(
         span,
@@ -123,16 +67,30 @@ impl<'src> Environment<'src> {
     }
   }
 
+  fn resolve_function(&self, name: &str) -> Option<Function<'src>> {
+    if let Some(symbol) = self.symbols.get(name) {
+      if let Some(function) = &symbol.function {
+        Some(function.clone())
+      } else if let Some(Value::Function(function)) = &symbol.value {
+        Some(function.clone())
+      } else if let Some(parent) = &self.parent {
+        parent.resolve_function(name)
+      } else {
+        None
+      }
+    } else if let Some(parent) = &self.parent {
+      parent.resolve_function(name)
+    } else {
+      None
+    }
+  }
+
   pub fn resolve_symbol(&self, symbol: &str) -> Option<Value<'src>> {
-    if let Some(value) = self.variables.get(symbol) {
-      Some(value.clone())
-    } else if let Some((name, function)) = self.functions.get_key_value(symbol)
-    {
-      match function {
-        Function::UserDefined(value) => Some(value.clone()),
-        Function::Builtin(builtin) => {
-          Some(Value::BuiltinFunction(name, *builtin))
-        }
+    if let Some(symbol) = self.symbols.get(symbol) {
+      if let Some(value) = &symbol.value {
+        Some(value.clone())
+      } else {
+        symbol.function.clone().map(Value::Function)
       }
     } else if let Some(parent) = &self.parent {
       parent.resolve_symbol(symbol)
@@ -144,9 +102,8 @@ impl<'src> Environment<'src> {
   pub fn with_parent(parent: Environment<'src>) -> Self {
     Self {
       config: parent.config.clone(),
-      functions: parent.functions.clone(),
       parent: Some(Box::new(parent)),
-      variables: HashMap::new(),
+      symbols: HashMap::new(),
     }
   }
 }

src/evaluator.rs

@@ -57,7 +57,7 @@ impl<'a> Evaluator<'a> {
 
         match &lhs.0 {
           Expression::Identifier(name) => {
-            self.environment.add_variable(name, value.clone());
+            self.environment.add_symbol(name, value.clone());
           }
           Expression::ListAccess(base_box, index_box) => {
             let (list_name, list_span) = match &base_box.0 {
@@ -117,7 +117,7 @@ impl<'a> Evaluator<'a> {
 
             list[index] = value.clone();
 
-            self.environment.add_variable(list_name, Value::List(list));
+            self.environment.add_symbol(list_name, Value::List(list));
           }
 
           _ => {
@@ -179,7 +179,7 @@ impl<'a> Evaluator<'a> {
         self.inside_loop = true;
 
         for item in list {
-          self.environment.add_variable(name, item);
+          self.environment.add_symbol(name, item);
 
           for statement in body {
             let eval_result = self.eval_statement(statement)?;
@@ -203,18 +203,16 @@ impl<'a> Evaluator<'a> {
         Ok(EvalResult::Value(result))
       }
       Statement::Function(name, params, body) => {
-        let function = Value::Function(
+        let function = Function::UserDefined {
+          body: body.clone(),
+          environment: self.environment.clone(),
           name,
-          params.clone(),
-          body.clone(),
-          self.environment.clone(),
-        );
+          parameters: params.clone(),
+        };
 
-        self
-          .environment
-          .add_function(name, Function::UserDefined(function.clone()));
+        self.environment.add_function(name, function.clone());
 
-        Ok(EvalResult::Value(function))
+        Ok(EvalResult::Value(Value::Function(function)))
       }
       Statement::If(condition, then_branch, else_branch) => {
         if self.eval_expression(condition)?.boolean(condition.1)? {

src/function.rs

@@ -1,7 +1,82 @@
 use super::*;
 
 #[derive(Clone, Debug)]
-pub enum Function<'a> {
-  Builtin(BuiltinFunction),
-  UserDefined(Value<'a>),
+pub enum Function<'src> {
+  Builtin {
+    function: BuiltinFunction,
+    name: &'src str,
+  },
+  UserDefined {
+    body: Vec<Spanned<Statement<'src>>>,
+    environment: Environment<'src>,
+    name: &'src str,
+    parameters: Vec<&'src str>,
+  },
+}
+
+impl<'src> Function<'src> {
+  pub fn call(
+    &self,
+    arguments: Vec<Value<'src>>,
+    config: Config,
+    span: Span,
+  ) -> Result<Value<'src>, Error> {
+    match self {
+      Self::Builtin { function, .. } => function(BuiltinFunctionPayload {
+        arguments,
+        config,
+        span,
+      }),
+      Self::UserDefined {
+        body,
+        environment,
+        name,
+        parameters,
+      } => {
+        if parameters.len() != arguments.len() {
+          return Err(Error::new(
+            span,
+            format!(
+              "Function `{}` expects {} arguments, got {}",
+              name,
+              parameters.len(),
+              arguments.len()
+            ),
+          ));
+        }
+
+        let mut call_environment =
+          Environment::with_parent(environment.clone());
+
+        call_environment.add_function(name, self.clone());
+
+        for (parameter, argument) in parameters.iter().zip(arguments.iter()) {
+          call_environment.add_symbol(parameter, argument.clone());
+        }
+
+        let mut evaluator = Evaluator::from(call_environment);
+        evaluator.inside_function = true;
+
+        if body.is_empty() {
+          return Ok(Value::Null);
+        }
+
+        for statement in body.iter() {
+          let result = evaluator.eval_statement(statement)?;
+
+          if result.is_return() {
+            return Ok(result.unwrap());
+          }
+        }
+
+        Ok(Value::Null)
+      }
+    }
+  }
+
+  pub fn name(&self) -> &'src str {
+    match self {
+      Self::Builtin { name, .. } | Self::UserDefined { name, .. } => name,
+    }
+  }
 }

src/value.rs

@@ -3,13 +3,7 @@ use super::*;
 #[derive(Clone, Debug)]
 pub enum Value<'src> {
   Boolean(bool),
-  BuiltinFunction(&'src str, BuiltinFunction),
-  Function(
-    &'src str,
-    Vec<&'src str>,
-    Vec<Spanned<Statement<'src>>>,
-    Environment<'src>,
-  ),
+  Function(Function<'src>),
   List(Vec<Self>),
   Null,
   Number(Float),
@@ -20,8 +14,7 @@ impl Display for Value<'_> {
   fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
     match self {
       Value::Boolean(boolean) => write!(f, "{boolean}"),
-      Value::BuiltinFunction(name, _) => write!(f, "<function: {name}>"),
-      Value::Function(name, _, _, _) => write!(f, "<function: {name}>"),
+      Value::Function(function) => write!(f, "<function: {}>", function.name()),
       Value::List(list) => write!(
         f,
         "[{}]",
@@ -45,13 +38,7 @@ impl PartialEq for Value<'_> {
   fn eq(&self, other: &Self) -> bool {
     match (self, other) {
       (Value::Boolean(a), Value::Boolean(b)) => a == b,
-      (
-        Value::BuiltinFunction(a_name, _),
-        Value::BuiltinFunction(b_name, _),
-      ) => a_name == b_name,
-      (Value::Function(a_name, _, _, _), Value::Function(b_name, _, _, _)) => {
-        a_name == b_name
-      }
+      (Value::Function(a), Value::Function(b)) => a.name() == b.name(),
       (Value::List(a), Value::List(b)) => {
         a.len() == b.len() && a.iter().zip(b.iter()).all(|(a, b)| a == b)
       }
@@ -111,8 +98,7 @@ impl<'a> Value<'a> {
   pub fn type_name(&self) -> &'static str {
     match self {
       Value::Boolean(_) => "boolean",
-      Value::BuiltinFunction(_, _) => "function",
-      Value::Function(_, _, _, _) => "function",
+      Value::Function(_) => "function",
       Value::List(_) => "list",
       Value::Null => "null",
       Value::Number(_) => "number",