From 2b02984b60a70edbdfa1c212b2973fba97fa598e Mon Sep 17 00:00:00 2001
From: Vitaly Bogdanov <vsbogd@gmail.com>
Date: Wed, 31 Jan 2024 12:22:36 +0300
Subject: [PATCH] Fix crash on retain all variables from single binding

---
 lib/src/atom/matcher.rs | 38 ++++++++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/lib/src/atom/matcher.rs b/lib/src/atom/matcher.rs
index 0bde8deb6..596195d31 100644
--- a/lib/src/atom/matcher.rs
+++ b/lib/src/atom/matcher.rs
@@ -181,7 +181,7 @@ impl Bindings {
                     self.bindings.remove(binding_id).atom
                 } else {
                     if binding.var == *var {
-                        self.rename_binding(binding_id, var);
+                        let _ = self.rename_binding(binding_id);
                     }
                     None
                 }
@@ -190,14 +190,15 @@ impl Bindings {
         }
     }
 
-    fn rename_binding(&mut self, binding_id: usize, old_var: &VariableAtom) {
+    fn rename_binding(&mut self, binding_id: usize) -> Result<(), ()> {
         let binding = &mut self.bindings[binding_id];
         let var = self.binding_by_var.iter()
-            .filter(|(v, i)| **i == binding.id && *v != old_var)
+            .filter(|(v, i)| **i == binding.id && **v != binding.var)
             .map(|(v, _i)| v)
             .next()
-            .expect("Unexpected state");
+            .ok_or(())?;
         binding.var = var.clone();
+        Ok(())
     }
 
     fn move_binding_to_binding(&mut self, from_binding_id: usize, to_binding_id: usize) {
@@ -1712,6 +1713,35 @@ mod test {
         Ok(())
     }
 
+    #[test]
+    fn bindings_retain_all() -> Result<(), &'static str> {
+        let mut bindings = Bindings::new()
+            .add_var_equality(&VariableAtom::new("a"), &VariableAtom::new("b"))?;
+        bindings.retain(|_| false);
+        assert!(bindings.is_empty());
+        Ok(())
+    }
+
+    #[test]
+    fn bindings_rename_binding() -> Result<(), &'static str> {
+        let a = &VariableAtom::new("a");
+        let b = &VariableAtom::new("b");
+        let mut bindings = Bindings::new().add_var_equality(a, b)?;
+        let binding_id = bindings.get_binding(a).unwrap().id;
+        assert_eq!(bindings.resolve(a), Some(Atom::Variable(a.clone())));
+        assert_eq!(bindings.resolve(b), Some(Atom::Variable(a.clone())));
+
+        assert_eq!(bindings.rename_binding(binding_id), Ok(()));
+        assert_eq!(bindings.resolve(a), Some(Atom::Variable(b.clone())));
+        assert_eq!(bindings.resolve(b), Some(Atom::Variable(b.clone())));
+
+        let mut bindings = Bindings::new().with_var_no_value(a);
+        let binding_id = bindings.get_binding(a).unwrap().id;
+        assert_eq!(bindings.rename_binding(binding_id), Err(()));
+        assert_eq!(bindings.resolve(a), Some(Atom::Variable(a.clone())));
+        Ok(())
+    }
+
     #[ignore = "Unstable because HashMap doesn't guarantee the order of entries"]
     #[test]
     fn bindings_into_entries() -> Result<(), &'static str> {