From 8f71171784e87a67d1eaa0036b1f43e1f67f1470 Mon Sep 17 00:00:00 2001
From: Vitaly Bogdanov <vsbogd@gmail.com>
Date: Fri, 14 Jun 2024 12:43:35 +0300
Subject: [PATCH] Rename wiped variables in Bindings::retain

---
 lib/src/atom/matcher.rs | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/lib/src/atom/matcher.rs b/lib/src/atom/matcher.rs
index 97b4b822f..d85cdb529 100644
--- a/lib/src/atom/matcher.rs
+++ b/lib/src/atom/matcher.rs
@@ -736,8 +736,17 @@ impl Bindings {
             Atom::Variable(var) => {
                 if to_remove.contains(&var) {
                     match self.resolve(var) {
-                        Some(Atom::Variable(v)) if to_remove.contains(&v) => { self.rename_var(var, &to_remove).map(|v| *atom = v); },
-                        Some(value) => { *atom = value; },
+                        Some(mut value) => {
+                            value.iter_mut().for_each(|atom| match atom {
+                                Atom::Variable(var) if to_remove.contains(&var) => {
+                                    self.rename_var(var, &to_remove).map(|var| {
+                                        *atom = var;
+                                    });
+                                },
+                                _ => {},
+                            });
+                            *atom = value;
+                        },
                         None => {},
                     }
                 }
@@ -1801,7 +1810,7 @@ mod test {
             .add_var_equality(&VariableAtom::new("a"), &VariableAtom::new("b"))?;
         bindings.apply_and_retain(&mut atom, |_| false);
         assert!(bindings.is_empty());
-        assert_eq!(atom, expr!(a b));
+        assert_eq!(atom, expr!(a a));
         Ok(())
     }
 
@@ -1815,6 +1824,18 @@ mod test {
         Ok(())
     }
 
+    #[test]
+    fn bindings_retain_apply_wiped_variable() -> Result<(), &'static str> {
+        let mut atom = expr!(b);
+        let mut bindings = Bindings::new()
+            .add_var_equality(&VariableAtom::new("a"), &VariableAtom::new("retained"))?
+            .add_var_binding_v2(&VariableAtom::new("b"), expr!((a)))?;
+        bindings.apply_and_retain(&mut atom, |v| *v == VariableAtom::new("retained"));
+        assert_eq!(bindings, bind!{ retained: expr!(retained) });
+        assert_eq!(atom, expr!((retained)));
+        Ok(())
+    }
+
     #[test]
     fn bindings_rename_binding() -> Result<(), &'static str> {
         let a = &VariableAtom::new("a");