From 6c9706f949c800f2463949525395299fc6e46bdc Mon Sep 17 00:00:00 2001
From: Sirui Lu <siruilu@cs.washington.edu>
Date: Sun, 1 Oct 2023 07:09:29 -0700
Subject: [PATCH] :sparkles: Add toGuardedList to UnionPrjOp

This commit introduces a toGuardedList function that converts a union to a list of values guarded by their path conditions.
---
 CHANGELOG.md                                    |  1 +
 src/Grisette/Core/Data/Class/SimpleMergeable.hs | 17 ++++++++++++++++-
 test/Grisette/Core/Control/Monad/UnionMTests.hs | 17 +++++++++++++++--
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5f65979c..3e351bc6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Added wrappers for state transformers. ([#132](https://github.com/lsrcz/grisette/pull/132))
+- Added `toGuardList` function. ([#137](https://github.com/lsrcz/grisette/pull/137))
 
 ### Removed
 
diff --git a/src/Grisette/Core/Data/Class/SimpleMergeable.hs b/src/Grisette/Core/Data/Class/SimpleMergeable.hs
index 2c21b967..6fd23b68 100644
--- a/src/Grisette/Core/Data/Class/SimpleMergeable.hs
+++ b/src/Grisette/Core/Data/Class/SimpleMergeable.hs
@@ -60,6 +60,7 @@ import Control.Monad.Trans.Cont (ContT (ContT))
 import Control.Monad.Trans.Maybe (MaybeT (MaybeT))
 import qualified Control.Monad.Writer.Lazy as WriterLazy
 import qualified Control.Monad.Writer.Strict as WriterStrict
+import Data.Bifunctor (Bifunctor (first))
 import Data.Kind (Type)
 import GHC.Generics
   ( Generic (Rep, from, to),
@@ -71,7 +72,7 @@ import GHC.Generics
   )
 import GHC.TypeNats (KnownNat, type (<=))
 import Generics.Deriving (Default (Default))
-import Grisette.Core.Data.Class.Bool (ITEOp (ites))
+import Grisette.Core.Data.Class.Bool (ITEOp (ites), LogicalOp (nots, (&&~)))
 import Grisette.Core.Data.Class.Function (Function (Arg, Ret, (#)))
 import Grisette.Core.Data.Class.Mergeable
   ( Mergeable (rootStrategy),
@@ -81,6 +82,7 @@ import Grisette.Core.Data.Class.Mergeable
     Mergeable3 (liftRootStrategy3),
     MergingStrategy (SimpleStrategy),
   )
+import Grisette.Core.Data.Class.Solvable (Solvable (con))
 import Grisette.IR.SymPrim.Data.Prim.InternedTerm.Term
   ( LinkedRep,
     SupportedPrim,
@@ -717,6 +719,19 @@ class (UnionLike u) => UnionPrjOp (u :: Type -> Type) where
   -- 1
   leftMost :: u a -> a
 
+  -- | Convert the union to a guarded list.
+  --
+  -- >>> toGuardedList (mrgIf "a" (single 1) (mrgIf "b" (single 2) (single 3)) :: UnionM Integer)
+  -- [(a,1),((&& b (! a)),2),((! (|| b a)),3)]
+  toGuardedList :: u a -> [(SymBool, a)]
+  toGuardedList u =
+    case (singleView u, ifView u) of
+      (Just x, _) -> [(con True, x)]
+      (_, Just (c, l, r)) ->
+        fmap (first (&&~ c)) (toGuardedList l)
+          ++ fmap (first (&&~ nots c)) (toGuardedList r)
+      _ -> error "Should not happen"
+
 -- | Pattern match to extract single values with 'singleView'.
 --
 -- >>> case (single 1 :: UnionM Integer) of Single v -> v
diff --git a/test/Grisette/Core/Control/Monad/UnionMTests.hs b/test/Grisette/Core/Control/Monad/UnionMTests.hs
index e00b2a3e..6509d344 100644
--- a/test/Grisette/Core/Control/Monad/UnionMTests.hs
+++ b/test/Grisette/Core/Control/Monad/UnionMTests.hs
@@ -42,7 +42,7 @@ import Grisette.Core.Data.Class.SOrd
 import Grisette.Core.Data.Class.SimpleMergeable
   ( SimpleMergeable (mrgIte),
     UnionLike (single, unionIf),
-    UnionPrjOp (ifView, leftMost, singleView),
+    UnionPrjOp (ifView, leftMost, singleView, toGuardedList),
     merge,
     mrgIf,
     mrgIte1,
@@ -61,6 +61,7 @@ import Grisette.IR.SymPrim.Data.Prim.Model
   ( ModelValuePair ((::=)),
   )
 import Grisette.IR.SymPrim.Data.SymPrim (SymBool)
+import Grisette.TestUtil.SymbolicAssertion ((@?=~))
 import Test.Framework (Test, testGroup)
 import Test.Framework.Providers.HUnit (testCase)
 import Test.HUnit (assertFailure, (@?=))
@@ -241,7 +242,19 @@ unionMTests =
             let r2 :: UnionM (Either SymBool SymBool) =
                   mrgIf "a" (mrgSingle $ Left "b") (mrgSingle $ Right "c")
             leftMost r1 @?= Left "b"
-            leftMost r2 @?= Left "b"
+            leftMost r2 @?= Left "b",
+          testCase "toGuardedList should work" $ do
+            let actual =
+                  toGuardedList
+                    ( mrgIf "a" (single 1) (mrgIf "b" (single 2) (single 3)) ::
+                        UnionM Integer
+                    )
+            let expected =
+                  [ ("a", 1),
+                    (nots "a" &&~ "b", 2),
+                    (nots "a" &&~ nots "b", 3)
+                  ]
+            actual @?=~ expected
         ],
       testGroup
         "MonadUnion"