From c2137d85b82580839679d17528763d9b6f831c1b Mon Sep 17 00:00:00 2001
From: Nika Beriachvili <nikaberiachvilipro@gmail.com>
Date: Fri, 3 Jan 2025 13:28:56 +0100
Subject: [PATCH] fix(expl): fix bug in an optional optimization in MARCO
 (greedy correction sets)

---
 .../musmcs_enumeration/marco/simple_marco.rs  | 46 ++++++++-----------
 1 file changed, 20 insertions(+), 26 deletions(-)

diff --git a/explainability/src/musmcs_enumeration/marco/simple_marco.rs b/explainability/src/musmcs_enumeration/marco/simple_marco.rs
index 0db63bde..4c1feb99 100644
--- a/explainability/src/musmcs_enumeration/marco/simple_marco.rs
+++ b/explainability/src/musmcs_enumeration/marco/simple_marco.rs
@@ -262,40 +262,34 @@ impl<Lbl: Label> Marco<Lbl> for SimpleMarco<Lbl> {
             self.map_solver.block_down(&mss);
             mcses.push(mcs.unwrap());
         } else {
-            // from Ignace Bleukx's implementation:
-            // find more MCSes, **disjoint** from this one, similar to "optimal_mus" in mus.py
-            // can only be done when MCSes do not have to be returned as there is no guarantee
-            // the MCSes encountered during enumeration are "new" MCSes
-            let mut sat_subset_lits: BTreeSet<Lit> = self.soft_constrs_reif_lits
-                .iter()
-                .filter(|&&l| self.subset_solver.s.model.entails(!l))
-                .cloned()
-                .collect();
-            self.map_solver
+            // Small optimization inspired by the implementation of Ignace Bleukx (in python).
+            // If we are not going to return MCSes, we can try to greedily search for
+            // more correction subsets, disjoint from this one (the seed).
+
+            let mut sat_subset = self.seed.clone();
+            self.map_solver.block_down(&sat_subset);
+
+            while self
+                .subset_solver
                 .s
-                .enforce(or(
-                        sat_subset_lits
-                        .iter()
-                        .map(|&l| self.map_solver.lits_translate_in[&l])
-                        .collect_vec()
-                    ), []);
-            while self.subset_solver.check_seed_sat(&sat_subset_lits) {
-                let s2: BTreeSet<Lit> = self
+                .solve_with_assumptions(sat_subset.iter().cloned())
+                .expect("Solver interrupted...")
+                .is_ok()
+            {
+                let new_sat_subset: BTreeSet<Lit> = self
                     .soft_constrs_reif_lits
                     .iter()
                     .filter(|&&l| self.subset_solver.s.model.entails(l))
                     .cloned()
                     .collect();
-                let new_mcs = self.soft_constrs_reif_lits.difference(&s2).cloned();
-                sat_subset_lits.extend(new_mcs.clone());
-                self.map_solver.s
-                .enforce(or(
-                    new_mcs
-                    .map(|l| self.map_solver.lits_translate_in[&l])
-                    .collect_vec()
-                ), []);
+                self.map_solver.block_down(&new_sat_subset);
+
+                let new_corr_subset = self.soft_constrs_reif_lits.difference(&new_sat_subset);
+                sat_subset.extend(new_corr_subset);
 
+                self.subset_solver.s.reset();
             }
+            self.subset_solver.s.reset();
         }
     }