UnionView.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}

-- |
-- Module      :   Grisette.Internal.Core.Data.Class.UnionView
-- Copyright   :   (c) Sirui Lu 2024
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Internal.Core.Data.Class.UnionView
  ( UnionView (..),
    IfViewResult (..),
    pattern Single,
    pattern If,
    simpleMerge,
    symIteMerge,
    (.#),
    onUnion,
    onUnion2,
    onUnion3,
    onUnion4,
    unionToCon,
    liftUnion,
    liftToMonadUnion,
  )
where

import Control.Monad.Identity (Identity (runIdentity))
import Data.Bifunctor (Bifunctor (first))
import Data.Kind (Type)
import Grisette.Internal.Core.Data.Class.AsKey (AsKey1 (AsKey1))
import Grisette.Internal.Core.Data.Class.Function (Function ((#)))
import Grisette.Internal.Core.Data.Class.ITEOp (ITEOp (symIte))
import Grisette.Internal.Core.Data.Class.LogicalOp
  ( LogicalOp (symNot, (.&&)),
  )
import Grisette.Internal.Core.Data.Class.Solvable (Solvable (con))
import Grisette.Internal.Internal.Decl.Core.Data.Class.Mergeable
  ( Mergeable,
  )
import Grisette.Internal.Internal.Decl.Core.Data.Class.SimpleMergeable
  ( SimpleMergeable,
    SymBranching,
    mrgIf,
  )
import Grisette.Internal.Internal.Decl.Core.Data.Class.ToCon (ToCon (toCon))
import Grisette.Internal.Internal.Decl.Core.Data.Class.TryMerge
  ( TryMerge,
    mrgSingle,
    tryMerge,
  )
import Grisette.Internal.SymPrim.SymBool (SymBool)

-- $setup
-- >>> import Grisette.Core
-- >>> import Grisette.SymPrim

-- | The result of 'ifView'.
data IfViewResult u a where
  IfViewResult :: (SymBranching u) => SymBool -> u a -> u a -> IfViewResult u a

instance (Show (u a)) => Show (IfViewResult u a) where
  showsPrec d (IfViewResult c l r) =
    showParen (d > 10) $
      showString "IfViewResult "
        . showsPrec 11 c
        . showString " "
        . showsPrec 11 l
        . showString " "
        . showsPrec 11 r

-- | Containers that can be projected back into single value or if-guarded
-- values. 'Identity' is an instance of this class as we can always project to
-- single value.
class (Applicative u, TryMerge u) => UnionView (u :: Type -> Type) where
  -- | Pattern match to extract single values.
  --
  -- >>> singleView (return 1 :: Union Integer)
  -- Just 1
  -- >>> singleView (mrgIfPropagatedStrategy "a" (return 1) (return 2) :: Union Integer)
  -- Nothing
  singleView :: u a -> Maybe a

  -- | Pattern match to extract if values.
  --
  -- >>> ifView (return 1 :: Union Integer)
  -- Nothing
  -- >>> ifView (mrgIfPropagatedStrategy "a" (return 1) (return 2) :: Union Integer)
  -- Just (IfViewResult a <1> <2>)
  -- >>> ifView (mrgIf "a" (return 1) (return 2) :: Union Integer)
  -- Just (IfViewResult a {1} {2})
  ifView :: u a -> Maybe (IfViewResult u a)

  -- | Convert the union to a guarded list.
  --
  -- >>> toGuardedList (mrgIf "a" (return 1) (mrgIf "b" (return 2) (return 3)) :: Union 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 (IfViewResult c l r)) ->
        fmap (first (.&& c)) (toGuardedList l)
          ++ fmap (first (.&& symNot c)) (toGuardedList r)
      _ -> error "Should not happen"

  -- | Return all possible values in the union. Drop the path conditions.
  --
  -- >>> overestimateUnionValues (return 1 :: Union Integer)
  -- [1]
  --
  -- >>> overestimateUnionValues (mrgIf "a" (return 1) (return 2) :: Union Integer)
  -- [1,2]
  overestimateUnionValues :: (Mergeable a) => u a -> [a]
  overestimateUnionValues (Single v) = [v]
  overestimateUnionValues (If _ l r) =
    overestimateUnionValues l ++ overestimateUnionValues r

instance UnionView Identity where
  singleView = Just . runIdentity
  ifView _ = Nothing

-- | Pattern match to extract single values with 'singleView'.
--
-- >>> case (return 1 :: Union Integer) of Single v -> v
-- 1
pattern Single :: (UnionView u, Mergeable a) => a -> u a
pattern Single x <-
  (singleView -> Just x)
  where
    Single x = mrgSingle x

-- | Pattern match to extract guard values with 'ifView'
--
-- >>> case (mrgIfPropagatedStrategy "a" (return 1) (return 2) :: Union Integer) of If c t f -> (c,t,f)
-- (a,<1>,<2>)
pattern If :: (UnionView u, Mergeable a) => (SymBranching u) => SymBool -> u a -> u a -> u a
pattern If c t f <-
  (ifView -> Just (IfViewResult c t f))
  where
    If c t f = mrgIf c t f

#if MIN_VERSION_base(4, 16, 4)
{-# COMPLETE Single, If #-}
#endif

-- | Merge the simply mergeable values in a union, and extract the merged value.
--
-- In the following example,
-- 'Grisette.Internal.Core.Data.Class.SimpleMergeable.mrgIfPropagatedStrategy'
-- will not merge the results, and 'simpleMerge' will merge it and extract the
-- single merged value.
--
-- >>> mrgIfPropagatedStrategy (ssym "a") (return $ ssym "b") (return $ ssym "c") :: Union SymBool
-- <If a b c>
-- >>> simpleMerge $ (mrgIfPropagatedStrategy (ssym "a") (return $ ssym "b") (return $ ssym "c") :: Union SymBool)
-- (ite a b c)
simpleMerge :: forall u a. (SimpleMergeable a, UnionView u) => u a -> a
simpleMerge u = case tryMerge u of
  Single x -> x
  _ -> error "Should not happen"
{-# INLINE simpleMerge #-}

-- | Merge the mergeable values in a union, using `symIte`, and extract the
-- merged value.
--
-- The reason why we provide this class is that for some types, we only have
-- `ITEOp` (which may throw an error), and we don't have a `SimpleMergeable`
-- instance. In this case, we can use `symIteMerge` to merge the values.
symIteMerge :: (ITEOp a, Mergeable a, UnionView u) => u a -> a
symIteMerge (Single x) = x
symIteMerge (If cond l r) = symIte cond (symIteMerge l) (symIteMerge r)
{-# INLINE symIteMerge #-}

-- | Helper for applying functions on 'UnionView' and 'SimpleMergeable'.
--
-- >>> let f :: Integer -> Union Integer = \x -> mrgIf (ssym "a") (mrgSingle $ x + 1) (mrgSingle $ x + 2)
-- >>> f .# (mrgIf (ssym "b" :: SymBool) (mrgSingle 0) (mrgSingle 2) :: Union Integer)
-- {If (&& b a) 1 (If b 2 (If a 3 4))}
(.#) ::
  (Function f a r, SimpleMergeable r, UnionView u) =>
  f ->
  u a ->
  r
(.#) f u = simpleMerge $ fmap (f #) u
{-# INLINE (.#) #-}

infixl 9 .#

-- | Lift a function to work on union values.
--
-- >>> sumU = onUnion sum :: Union [SymInteger] -> SymInteger
-- >>> sumU (mrgIfPropagatedStrategy "cond" (return ["a"]) (return ["b","c"]) :: Union [SymInteger])
-- (ite cond a (+ b c))
onUnion ::
  forall u a r.
  (SimpleMergeable r, SymBranching u, UnionView u, Mergeable a) =>
  (a -> r) ->
  (u a -> r)
onUnion f = simpleMerge . fmap f . tryMerge

-- | Lift a function to work on union values.
onUnion2 ::
  forall u a b r.
  ( SimpleMergeable r,
    SymBranching u,
    UnionView u,
    Mergeable a,
    Mergeable b
  ) =>
  (a -> b -> r) ->
  (u a -> u b -> r)
onUnion2 f ua ub = simpleMerge $ f <$> tryMerge ua <*> tryMerge ub

-- | Lift a function to work on union values.
onUnion3 ::
  forall u a b c r.
  ( SimpleMergeable r,
    SymBranching u,
    UnionView u,
    Mergeable a,
    Mergeable b,
    Mergeable c
  ) =>
  (a -> b -> c -> r) ->
  (u a -> u b -> u c -> r)
onUnion3 f ua ub uc =
  simpleMerge $ f <$> tryMerge ua <*> tryMerge ub <*> tryMerge uc

-- | Lift a function to work on union values.
onUnion4 ::
  forall u a b c d r.
  ( SimpleMergeable r,
    SymBranching u,
    UnionView u,
    Mergeable a,
    Mergeable b,
    Mergeable c,
    Mergeable d
  ) =>
  (a -> b -> c -> d -> r) ->
  (u a -> u b -> u c -> u d -> r)
onUnion4 f ua ub uc ud =
  simpleMerge $
    f <$> tryMerge ua <*> tryMerge ub <*> tryMerge uc <*> tryMerge ud

-- | Convert a plain union to concrete values.
--
-- >>> unionToCon (return 1 :: Union SymInteger) :: Maybe Integer
-- Just 1
-- >>> unionToCon (mrgIf "a" (return 1) (return 2) :: Union SymInteger) :: Maybe Integer
-- Nothing
-- >>> unionToCon (return "a" :: Union SymInteger) :: Maybe Integer
-- Nothing
unionToCon :: (ToCon a b, UnionView u) => u a -> Maybe b
unionToCon u =
  case (singleView u, ifView u) of
    (Just x, _) -> toCon x
    (_, Just (IfViewResult c l r)) -> do
      cl <- toCon c
      if cl then unionToCon l else unionToCon r
    _ -> Nothing
{-# INLINE unionToCon #-}

-- | Lift the 'UnionView' value to any Applicative 'SymBranching'.
liftUnion ::
  (Mergeable a, UnionView u, Applicative m, SymBranching m) => u a -> m a
liftUnion u = case u of
  Single x -> mrgSingle x
  If c l r -> mrgIf c (liftUnion l) (liftUnion r)

-- | Alias for 'liftUnion', but for monads.
liftToMonadUnion ::
  (Mergeable a, UnionView u, Monad m, SymBranching m) => u a -> m a
liftToMonadUnion = liftUnion

#if MIN_VERSION_base(4,16,0)
instance (UnionView u) => UnionView (AsKey1 u) where
  singleView (AsKey1 u) = singleView u
  ifView (AsKey1 u) = case ifView u of
    Just (IfViewResult c l r) -> Just (IfViewResult c (AsKey1 l) (AsKey1 r))
    Nothing -> Nothing
  toGuardedList (AsKey1 u) = toGuardedList u
  overestimateUnionValues (AsKey1 u) = overestimateUnionValues u
#endif