Simplify function context and arity validation

src/context.rs

@@ -1,28 +1,27 @@
 #[derive(Default)]
 pub(crate) struct Context {
-  function_depth: usize,
+  inside_function: bool,
   loop_depth: usize,
 }
 
 impl Context {
-  pub(crate) fn enter_function(&mut self) {
-    self.function_depth += 1;
-  }
-
   pub(crate) fn enter_loop(&mut self) {
     self.loop_depth += 1;
   }
 
-  pub(crate) fn exit_function(&mut self) {
-    self.function_depth -= 1;
-  }
-
   pub(crate) fn exit_loop(&mut self) {
     self.loop_depth -= 1;
   }
 
+  pub(crate) fn for_function() -> Self {
+    Self {
+      inside_function: true,
+      ..Self::default()
+    }
+  }
+
   pub(crate) fn inside_function(&self) -> bool {
-    self.function_depth > 0
+    self.inside_function
   }
 
   pub(crate) fn inside_loop(&self) -> bool {

src/evaluator.rs

@@ -80,16 +80,6 @@ impl Evaluator {
     Ok(Value::List(list))
   }
 
-  pub(crate) fn enter_function<T>(
-    &mut self,
-    f: impl FnOnce(&mut Self) -> Result<T, Error>,
-  ) -> Result<T, Error> {
-    self.context.enter_function();
-    let result = f(self);
-    self.context.exit_function();
-    result
-  }
-
   fn enter_loop<T>(
     &mut self,
     f: impl FnOnce(&mut Self) -> Result<T, Error>,
@@ -529,6 +519,13 @@ impl Evaluator {
 
     Ok(Completion::Value(result))
   }
+
+  pub(crate) fn for_function(environment: Environment) -> Self {
+    Self {
+      context: Context::for_function(),
+      environment,
+    }
+  }
 }
 
 impl From<Environment> for Evaluator {

src/function.rs

@@ -49,12 +49,12 @@ impl Function {
           call_environment.add_symbol(parameter, argument);
         }
 
-        Evaluator::from(call_environment).enter_function(|evaluator| {
-          match evaluator.evaluate_statements(body)? {
-            Completion::Return(value) | Completion::Value(value) => Ok(value),
-            Completion::Break | Completion::Continue => Ok(Value::Null),
-          }
-        })
+        match Evaluator::for_function(call_environment)
+          .evaluate_statements(body)?
+        {
+          Completion::Return(value) | Completion::Value(value) => Ok(value),
+          Completion::Break | Completion::Continue => Ok(Value::Null),
+        }
       }
     }
   }
@@ -67,19 +67,7 @@ impl Function {
     match self {
       Self::Builtin { arity, name, .. } => arity.check(name, len, span),
       Self::UserDefined { parameters, .. } => {
-        if parameters.len() == len {
-          return Ok(());
-        }
-
-        Err(Error::new(
-          span,
-          format!(
-            "Function `{}` expects {} arguments, got {}",
-            self.name(),
-            parameters.len(),
-            len
-          ),
-        ))
+        BuiltinArity::Exact(parameters.len()).check(self.name(), len, span)
       }
     }
   }

tests/integration.rs

@@ -1364,6 +1364,25 @@ fn function_calling_builtin() -> Result {
     .run()
 }
 
+#[test]
+fn function_does_not_inherit_loop_context() -> Result {
+  #[track_caller]
+  fn case(statement: &str) -> Result {
+    Test::new()?
+      .program(&format!(
+        "fn foo() {{ {statement} }}\nfor bar in [0] {{ foo() }}"
+      ))
+      .expected_status(1)
+      .expected_stderr(Contains(&format!(
+        "Cannot use '{statement}' outside of a loop"
+      )))
+      .run()
+  }
+
+  case("break")?;
+  case("continue")
+}
+
 #[test]
 fn function_equality_uses_identity() -> Result {
   Test::new()?
@@ -1567,19 +1586,29 @@ fn function_with_return_value() -> Result {
 
 #[test]
 fn function_wrong_argument_count() -> Result {
-  Test::new()?
-    .program(indoc! {
-      "
-      fn add(a, b) {
-        return a + b
-      }
+  #[track_caller]
+  fn case(program: &str, expected: &str) -> Result {
+    Test::new()?
+      .program(program)
+      .expected_status(1)
+      .expected_stderr(Contains(expected))
+      .run()
+  }
 
-      println(add(1))
-      "
-    })
-    .expected_status(1)
-    .expected_stderr(Contains("Function `add` expects 2 arguments, got 1"))
-    .run()
+  case(
+    "fn foo(bar, baz) {}\nfoo(1)",
+    "Function `foo` expects 2 arguments, got 1",
+  )?;
+
+  case(
+    "fn foo(bar) {}\nfoo()",
+    "Function `foo` expects 1 argument, got 0",
+  )?;
+
+  case(
+    "(fn(foo) {})(1, 2)",
+    "Function `<anonymous>` expects 1 argument, got 2",
+  )
 }
 
 #[test]