module Intensional.Guard where

import qualified GhcPlugins as GHC
import Binary 
import Data.Map (Map)
import Data.Set (Set)
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.IntSet as IntSet
import qualified Data.IntMap as IntMap

import Intensional.Types
import Intensional.Constructors

-- data Named a = Named {toPair :: (GHC.Name, a)}
--   deriving (Eq, Functor)

-- instance Semigroup a => Semigroup (Named a) where
--   Named (n, ks1) <> Named (_, ks2) = Named (n, ks1 <> ks2)

-- instance GHC.Uniquable (Named a) where
--   getUnique (Named (n, _)) = GHC.getUnique n

-- instance Binary a => Binary (Named a) where
--   put_ bh = put_ bh . toPair
--   get bh = Named <$> Binary.get bh


-- A set of simple inclusion constraints, i.e. k in X(d), grouped by X(d)
newtype Guard
  = Guard
      { 
        Guard -> Map (RVar, Name) (UniqSet Name)
groups :: Map (RVar, GHC.Name) (GHC.UniqSet GHC.Name)
      }
  deriving (Guard -> Guard -> Bool
(Guard -> Guard -> Bool) -> (Guard -> Guard -> Bool) -> Eq Guard
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Guard -> Guard -> Bool
$c/= :: Guard -> Guard -> Bool
== :: Guard -> Guard -> Bool
$c== :: Guard -> Guard -> Bool
Eq)

instance Semigroup Guard where
  Guard g :: Map (RVar, Name) (UniqSet Name)
g <> :: Guard -> Guard -> Guard
<> Guard g' :: Map (RVar, Name) (UniqSet Name)
g' = Map (RVar, Name) (UniqSet Name) -> Guard
Guard ((UniqSet Name -> UniqSet Name -> UniqSet Name)
-> Map (RVar, Name) (UniqSet Name)
-> Map (RVar, Name) (UniqSet Name)
-> Map (RVar, Name) (UniqSet Name)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith UniqSet Name -> UniqSet Name -> UniqSet Name
forall a. UniqSet a -> UniqSet a -> UniqSet a
GHC.unionUniqSets Map (RVar, Name) (UniqSet Name)
g Map (RVar, Name) (UniqSet Name)
g')

instance Monoid Guard where
  mempty :: Guard
mempty = Map (RVar, Name) (UniqSet Name) -> Guard
Guard Map (RVar, Name) (UniqSet Name)
forall a. Monoid a => a
mempty

instance GHC.Outputable Guard where
  ppr :: Guard -> SDoc
ppr = (RVar -> SDoc) -> Guard -> SDoc
forall t. Refined t => (RVar -> SDoc) -> t -> SDoc
prpr RVar -> SDoc
forall a. Outputable a => a -> SDoc
GHC.ppr

isEmpty :: Guard -> Bool
isEmpty :: Guard -> Bool
isEmpty (Guard g :: Map (RVar, Name) (UniqSet Name)
g) = Map (RVar, Name) (UniqSet Name) -> Bool
forall k a. Map k a -> Bool
Map.null Map (RVar, Name) (UniqSet Name)
g

toList :: Guard -> [(Int, GHC.Name, GHC.Name)]
toList :: Guard -> [(RVar, Name, Name)]
toList (Guard g :: Map (RVar, Name) (UniqSet Name)
g) =
  [ (RVar
x, Name
d, Name
k) | ((x :: RVar
x,d :: Name
d), ks :: UniqSet Name
ks) <- Map (RVar, Name) (UniqSet Name) -> [((RVar, Name), UniqSet Name)]
forall k a. Map k a -> [(k, a)]
Map.toList Map (RVar, Name) (UniqSet Name)
g, Name
k <- UniqSet Name -> [Name]
forall elt. UniqSet elt -> [elt]
GHC.nonDetEltsUniqSet UniqSet Name
ks ]

fromList :: [(Int, GHC.Name, GHC.Name)] -> Guard
fromList :: [(RVar, Name, Name)] -> Guard
fromList = ((RVar, Name, Name) -> Guard) -> [(RVar, Name, Name)] -> Guard
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\(x :: RVar
x, d :: Name
d, k :: Name
k) -> [Name] -> RVar -> Name -> Guard
singleton [Name
k] RVar
x Name
d)

typedVars :: Guard -> Set (RVar, GHC.Name)
typedVars :: Guard -> Set (RVar, Name)
typedVars (Guard g :: Map (RVar, Name) (UniqSet Name)
g) = Map (RVar, Name) (UniqSet Name) -> Set (RVar, Name)
forall k a. Map k a -> Set k
Map.keysSet Map (RVar, Name) (UniqSet Name)
g

instance Binary Guard where
  put_ :: BinHandle -> Guard -> IO ()
put_ bh :: BinHandle
bh = BinHandle -> [(RVar, Name, Name)] -> IO ()
forall a. Binary a => BinHandle -> a -> IO ()
put_ BinHandle
bh ([(RVar, Name, Name)] -> IO ())
-> (Guard -> [(RVar, Name, Name)]) -> Guard -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Guard -> [(RVar, Name, Name)]
toList
  get :: BinHandle -> IO Guard
get bh :: BinHandle
bh = [(RVar, Name, Name)] -> Guard
fromList ([(RVar, Name, Name)] -> Guard)
-> IO [(RVar, Name, Name)] -> IO Guard
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BinHandle -> IO [(RVar, Name, Name)]
forall a. Binary a => BinHandle -> IO a
get BinHandle
bh

instance Refined Guard where
  domain :: Guard -> Domain
domain (Guard g :: Map (RVar, Name) (UniqSet Name)
g) = 
      (Domain -> (RVar, Name) -> Domain)
-> Domain -> Set (RVar, Name) -> Domain
forall a b. (a -> b -> a) -> a -> Set b -> a
Set.foldl' (\acc :: Domain
acc (x :: RVar
x,_) -> RVar -> Domain -> Domain
IntSet.insert RVar
x Domain
acc) Domain
forall a. Monoid a => a
mempty (Map (RVar, Name) (UniqSet Name) -> Set (RVar, Name)
forall k a. Map k a -> Set k
Map.keysSet Map (RVar, Name) (UniqSet Name)
g)

  rename :: RVar -> RVar -> Guard -> Guard
rename x :: RVar
x y :: RVar
y (Guard g :: Map (RVar, Name) (UniqSet Name)
g) =
    Map (RVar, Name) (UniqSet Name) -> Guard
Guard (((RVar, Name)
 -> UniqSet Name
 -> Map (RVar, Name) (UniqSet Name)
 -> Map (RVar, Name) (UniqSet Name))
-> Map (RVar, Name) (UniqSet Name)
-> Map (RVar, Name) (UniqSet Name)
-> Map (RVar, Name) (UniqSet Name)
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\(z :: RVar
z,d :: Name
d) ks :: UniqSet Name
ks m :: Map (RVar, Name) (UniqSet Name)
m -> (UniqSet Name -> UniqSet Name -> UniqSet Name)
-> (RVar, Name)
-> UniqSet Name
-> Map (RVar, Name) (UniqSet Name)
-> Map (RVar, Name) (UniqSet Name)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith UniqSet Name -> UniqSet Name -> UniqSet Name
forall a. UniqSet a -> UniqSet a -> UniqSet a
GHC.unionUniqSets (if RVar
z RVar -> RVar -> Bool
forall a. Eq a => a -> a -> Bool
== RVar
x then RVar
y else RVar
z, Name
d) UniqSet Name
ks Map (RVar, Name) (UniqSet Name)
m) Map (RVar, Name) (UniqSet Name)
forall a. Monoid a => a
mempty Map (RVar, Name) (UniqSet Name)
g)

  prpr :: (RVar -> SDoc) -> Guard -> SDoc
prpr m :: RVar -> SDoc
m (Guard g :: Map (RVar, Name) (UniqSet Name)
g) = (((RVar, Name), [Name]) -> SDoc)
-> [((RVar, Name), [Name])] -> SDoc
forall a. (a -> SDoc) -> [a] -> SDoc
GHC.pprWithCommas ((RVar, Name), [Name]) -> SDoc
forall a. Outputable a => ((RVar, Name), a) -> SDoc
pprGuardAtom [((RVar, Name), [Name])]
guardList
    where
    pprGuardAtom :: ((RVar, Name), a) -> SDoc
pprGuardAtom ((x :: RVar
x,d :: Name
d), ks :: a
ks) = [SDoc] -> SDoc
GHC.hsep [a -> SDoc
forall a. Outputable a => a -> SDoc
GHC.ppr a
ks, String -> SDoc
GHC.text "in", (RVar -> SDoc) -> K Any -> SDoc
forall t. Refined t => (RVar -> SDoc) -> t -> SDoc
prpr RVar -> SDoc
m (DataType Name -> K Any
forall (s :: Side). DataType Name -> K s
Dom (RVar -> Name -> DataType Name
forall d. RVar -> d -> DataType d
Inj RVar
x Name
d))]
    guardList :: [((RVar, Name), [Name])]
guardList = (((RVar, Name), UniqSet Name) -> ((RVar, Name), [Name]))
-> [((RVar, Name), UniqSet Name)] -> [((RVar, Name), [Name])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(x :: (RVar, Name)
x,y :: UniqSet Name
y) -> ((RVar, Name)
x, UniqSet Name -> [Name]
forall elt. UniqSet elt -> [elt]
GHC.nonDetEltsUniqSet UniqSet Name
y)) (Map (RVar, Name) (UniqSet Name) -> [((RVar, Name), UniqSet Name)]
forall k a. Map k a -> [(k, a)]
Map.toList Map (RVar, Name) (UniqSet Name)
g)

lookup :: RVar -> GHC.Name -> Guard -> Maybe (GHC.UniqSet GHC.Name)
lookup :: RVar -> Name -> Guard -> Maybe (UniqSet Name)
lookup x :: RVar
x d :: Name
d (Guard g :: Map (RVar, Name) (UniqSet Name)
g) = (RVar, Name)
-> Map (RVar, Name) (UniqSet Name) -> Maybe (UniqSet Name)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (RVar
x,Name
d) Map (RVar, Name) (UniqSet Name)
g

delete :: GHC.Name -> RVar -> GHC.Name -> Guard -> Guard
delete :: Name -> RVar -> Name -> Guard -> Guard
delete k :: Name
k x :: RVar
x d :: Name
d (Guard g :: Map (RVar, Name) (UniqSet Name)
g) = Map (RVar, Name) (UniqSet Name) -> Guard
Guard ((Maybe (UniqSet Name) -> Maybe (UniqSet Name))
-> (RVar, Name)
-> Map (RVar, Name) (UniqSet Name)
-> Map (RVar, Name) (UniqSet Name)
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter Maybe (UniqSet Name) -> Maybe (UniqSet Name)
del (RVar
x,Name
d) Map (RVar, Name) (UniqSet Name)
g)
  where
    del :: Maybe (UniqSet Name) -> Maybe (UniqSet Name)
del Nothing = Maybe (UniqSet Name)
forall a. Maybe a
Nothing
    del (Just ks :: UniqSet Name
ks) =
      let ks' :: UniqSet Name
ks' = UniqSet Name -> Name -> UniqSet Name
forall a. Uniquable a => UniqSet a -> a -> UniqSet a
GHC.delOneFromUniqSet UniqSet Name
ks Name
k
      in if UniqSet Name -> Bool
forall a. UniqSet a -> Bool
GHC.isEmptyUniqSet UniqSet Name
ks' then Maybe (UniqSet Name)
forall a. Maybe a
Nothing else UniqSet Name -> Maybe (UniqSet Name)
forall a. a -> Maybe a
Just UniqSet Name
ks'

deleteAll :: [GHC.Name] -> RVar -> GHC.Name -> Guard -> Guard
deleteAll :: [Name] -> RVar -> Name -> Guard -> Guard
deleteAll ms :: [Name]
ms x :: RVar
x d :: Name
d (Guard g :: Map (RVar, Name) (UniqSet Name)
g) = Map (RVar, Name) (UniqSet Name) -> Guard
Guard ((Maybe (UniqSet Name) -> Maybe (UniqSet Name))
-> (RVar, Name)
-> Map (RVar, Name) (UniqSet Name)
-> Map (RVar, Name) (UniqSet Name)
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter Maybe (UniqSet Name) -> Maybe (UniqSet Name)
del (RVar
x,Name
d) Map (RVar, Name) (UniqSet Name)
g)
  where
    del :: Maybe (UniqSet Name) -> Maybe (UniqSet Name)
del Nothing = Maybe (UniqSet Name)
forall a. Maybe a
Nothing
    del (Just ks :: UniqSet Name
ks) =
      let ks' :: UniqSet Name
ks' = UniqSet Name -> [Name] -> UniqSet Name
forall a. Uniquable a => UniqSet a -> [a] -> UniqSet a
GHC.delListFromUniqSet UniqSet Name
ks [Name]
ms
      in if UniqSet Name -> Bool
forall a. UniqSet a -> Bool
GHC.isEmptyUniqSet UniqSet Name
ks' then Maybe (UniqSet Name)
forall a. Maybe a
Nothing else UniqSet Name -> Maybe (UniqSet Name)
forall a. a -> Maybe a
Just UniqSet Name
ks'

-- A guard literal
-- Ignorning possibly trivial guards (e.g. 1-constructor types has already
-- happened in InferM.branch)
singleton :: [GHC.Name] -> RVar -> GHC.Name -> Guard
singleton :: [Name] -> RVar -> Name -> Guard
singleton ks :: [Name]
ks x :: RVar
x d :: Name
d =
  Map (RVar, Name) (UniqSet Name) -> Guard
Guard ((RVar, Name) -> UniqSet Name -> Map (RVar, Name) (UniqSet Name)
forall k a. k -> a -> Map k a
Map.singleton (RVar
x, Name
d) (UniqSet Name -> [Name] -> UniqSet Name
forall a. Uniquable a => UniqSet a -> [a] -> UniqSet a
GHC.addListToUniqSet UniqSet Name
forall a. Monoid a => a
mempty [Name]
ks))

-- guardsFromList :: [GHC.Name] -> DataType GHC.Name -> Guard
-- guardsFromList ks (Inj x d) = foldr (\k gs -> singleton k (Inj x d) <> gs) mempty ks

impliedBy :: Guard -> Guard -> Bool
impliedBy :: Guard -> Guard -> Bool
impliedBy (Guard g :: Map (RVar, Name) (UniqSet Name)
g) (Guard g' :: Map (RVar, Name) (UniqSet Name)
g') =
    (UniqSet Name -> UniqSet Name -> Bool)
-> Map (RVar, Name) (UniqSet Name)
-> Map (RVar, Name) (UniqSet Name)
-> Bool
forall k a b.
Ord k =>
(a -> b -> Bool) -> Map k a -> Map k b -> Bool
Map.isSubmapOfBy UniqSet Name -> UniqSet Name -> Bool
forall a b. UniqSet a -> UniqSet b -> Bool
keyInclusion Map (RVar, Name) (UniqSet Name)
g' Map (RVar, Name) (UniqSet Name)
g
  where
    keyInclusion :: UniqSet a -> UniqSet b -> Bool
keyInclusion u1 :: UniqSet a
u1 u2 :: UniqSet b
u2 =
      {-# SCC keyInclusion #-}
      (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
forall a b. (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
IntMap.isSubmapOfBy (\_ _ -> Bool
True) (UniqFM a -> IntMap a
forall elt. UniqFM elt -> IntMap elt
GHC.ufmToIntMap (UniqFM a -> IntMap a) -> UniqFM a -> IntMap a
forall a b. (a -> b) -> a -> b
$ UniqSet a -> UniqFM a
forall a. UniqSet a -> UniqFM a
GHC.getUniqSet UniqSet a
u1) (UniqFM b -> IntMap b
forall elt. UniqFM elt -> IntMap elt
GHC.ufmToIntMap (UniqFM b -> IntMap b) -> UniqFM b -> IntMap b
forall a b. (a -> b) -> a -> b
$ UniqSet b -> UniqFM b
forall a. UniqSet a -> UniqFM a
GHC.getUniqSet UniqSet b
u2)