Fix function comparisons

crates/runtime/src/vm.rs

@@ -1860,21 +1860,9 @@ impl KotoVm {
             }
             (Object(o), _) => o.try_borrow()?.equal(rhs_value)?,
             (Function(a), Function(b)) => {
-                if a.chunk == b.chunk && a.ip == b.ip {
-                    match (&a.captures, &b.captures) {
-                        (None, None) => true,
-                        (Some(captures_a), Some(captures_b)) => {
-                            let captures_a = captures_a.clone();
-                            let captures_b = captures_b.clone();
-                            let data_a = captures_a.data();
-                            let data_b = captures_b.data();
-                            self.compare_value_ranges(&data_a, &data_b)?
-                        }
-                        _ => false,
-                    }
-                } else {
-                    false
-                }
+                let a = a.clone();
+                let b = b.clone();
+                self.compare_functions(a, b)?
             }
             _ => false,
         };
@@ -1925,21 +1913,9 @@ impl KotoVm {
             }
             (Object(o), _) => o.try_borrow()?.not_equal(rhs_value)?,
             (Function(a), Function(b)) => {
-                if a.chunk == b.chunk && a.ip == b.ip {
-                    match (&a.captures, &b.captures) {
-                        (None, None) => true,
-                        (Some(captures_a), Some(captures_b)) => {
-                            let captures_a = captures_a.clone();
-                            let captures_b = captures_b.clone();
-                            let data_a = captures_a.data();
-                            let data_b = captures_b.data();
-                            !self.compare_value_ranges(&data_a, &data_b)?
-                        }
-                        _ => false,
-                    }
-                } else {
-                    true
-                }
+                let a = a.clone();
+                let b = b.clone();
+                !self.compare_functions(a, b)?
             }
             _ => true,
         };
@@ -1948,6 +1924,24 @@ impl KotoVm {
         Ok(())
     }
 
+    fn compare_functions(&mut self, a: KFunction, b: KFunction) -> Result<bool> {
+        if a.chunk == b.chunk && a.ip == b.ip {
+            match (&a.captures, &b.captures) {
+                (None, None) => Ok(true),
+                (Some(captures_a), Some(captures_b)) => {
+                    let captures_a = captures_a.clone();
+                    let captures_b = captures_b.clone();
+                    let data_a = captures_a.data();
+                    let data_b = captures_b.data();
+                    self.compare_value_ranges(&data_a, &data_b)
+                }
+                _ => Ok(false),
+            }
+        } else {
+            Ok(false)
+        }
+    }
+
     // Called from run_equal / run_not_equal to compare the contents of lists and tuples
     fn compare_value_ranges(&mut self, range_a: &[KValue], range_b: &[KValue]) -> Result<bool> {
         if range_a.len() != range_b.len() {

crates/runtime/tests/vm_tests.rs

@@ -1612,6 +1612,57 @@ f().x
 ";
                 check_script_output(script, 99);
             }
+
+            mod comparisions {
+                use super::*;
+
+                #[test]
+                fn matching_functions_no_captures() {
+                    let script = "
+f = |x| x + x
+f2 = f
+f == f, f == f2, f != f, f != f2
+";
+                    check_script_output(
+                        script,
+                        tuple(&[true.into(), true.into(), false.into(), false.into()]),
+                    );
+                }
+
+                #[test]
+                fn different_functions() {
+                    let script = "
+f = |x| x + x
+g = |y| y * y
+f == g, f != g
+";
+                    check_script_output(script, tuple(&[false.into(), true.into()]));
+                }
+
+                #[test]
+                fn matching_captures() {
+                    let script = "
+f, g =
+  (1, 1)
+    .each |x| return || x * x
+    .to_tuple()
+f == g, f != g
+";
+                    check_script_output(script, tuple(&[true.into(), false.into()]));
+                }
+
+                #[test]
+                fn different_captures() {
+                    let script = "
+f, g =
+  (1, 2)
+    .each |x| return || x * x
+    .to_tuple()
+f == g, f != g
+";
+                    check_script_output(script, tuple(&[false.into(), true.into()]));
+                }
+            }
         }
 
         mod piped_calls {
@@ -3416,18 +3467,19 @@ b >= a
         #[test]
         fn equality_of_functions_with_overridden_captures() {
             let script = "
-# Make two functions which capture a different foo
-foos = (0, 1)
+# Make two functions which capture different values of `foo`
+f, g = (0, 1)
   .each |n|
     foo =
       x: n
-      @==: |other| self.x != other.x # inverting the usual behaviour to show its effect
-    || foo # The function returns its captured foo
-  .to_tuple()
+      # Invert the equality comparison so that its effect is visible
+      @==: |other| self.x != other.x
+    || foo
 
-foos[0] == foos[1]
+# The overridden operator on the captured value of `foo` inverts the usual comparison logic
+f == g, f != g
 ";
-            check_script_output(script, true);
+            check_script_output(script, tuple(&[true.into(), false.into()]));
         }
 
         #[test]