-- Copyright (c) Facebook, Inc. and its affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.
--
{-# LANGUAGE CPP #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeriveFunctor #-}
module Retrie.PatternMap.Bag where

import qualified Data.Map as M
import qualified Data.IntMap as I

import Retrie.AlphaEnv
import qualified Retrie.GHC as GHC
import Retrie.PatternMap.Class
import Retrie.Quantifiers
import Retrie.Substitution

data BoolMap a
  = EmptyBoolMap
  | BoolMap
      { forall a. BoolMap a -> MaybeMap a
bmTrue :: MaybeMap a
      , forall a. BoolMap a -> MaybeMap a
bmFalse :: MaybeMap a
      }
  deriving ((forall a b. (a -> b) -> BoolMap a -> BoolMap b)
-> (forall a b. a -> BoolMap b -> BoolMap a) -> Functor BoolMap
forall a b. a -> BoolMap b -> BoolMap a
forall a b. (a -> b) -> BoolMap a -> BoolMap b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> BoolMap a -> BoolMap b
fmap :: forall a b. (a -> b) -> BoolMap a -> BoolMap b
$c<$ :: forall a b. a -> BoolMap b -> BoolMap a
<$ :: forall a b. a -> BoolMap b -> BoolMap a
Functor)

instance PatternMap BoolMap where
  type Key BoolMap = Bool

  mEmpty :: BoolMap a
  mEmpty :: forall a. BoolMap a
mEmpty = BoolMap a
forall a. BoolMap a
EmptyBoolMap

  mUnion :: BoolMap a -> BoolMap a -> BoolMap a
  mUnion :: forall a. BoolMap a -> BoolMap a -> BoolMap a
mUnion BoolMap a
EmptyBoolMap BoolMap a
m = BoolMap a
m
  mUnion BoolMap a
m BoolMap a
EmptyBoolMap = BoolMap a
m
  mUnion BoolMap a
m1 BoolMap a
m2 = BoolMap
    { bmTrue :: MaybeMap a
bmTrue = (BoolMap a -> MaybeMap a) -> BoolMap a -> BoolMap a -> MaybeMap a
forall (m :: * -> *) a b.
PatternMap m =>
(a -> m b) -> a -> a -> m b
unionOn BoolMap a -> MaybeMap a
forall a. BoolMap a -> MaybeMap a
bmTrue BoolMap a
m1 BoolMap a
m2
    , bmFalse :: MaybeMap a
bmFalse = (BoolMap a -> MaybeMap a) -> BoolMap a -> BoolMap a -> MaybeMap a
forall (m :: * -> *) a b.
PatternMap m =>
(a -> m b) -> a -> a -> m b
unionOn BoolMap a -> MaybeMap a
forall a. BoolMap a -> MaybeMap a
bmFalse BoolMap a
m1 BoolMap a
m2
    }

  mAlter
    :: AlphaEnv -> Quantifiers -> Key BoolMap -> A a -> BoolMap a -> BoolMap a
  mAlter :: forall a.
AlphaEnv
-> Quantifiers -> Key BoolMap -> A a -> BoolMap a -> BoolMap a
mAlter AlphaEnv
env Quantifiers
qs Key BoolMap
b A a
f BoolMap a
EmptyBoolMap = AlphaEnv
-> Quantifiers -> Key BoolMap -> A a -> BoolMap a -> BoolMap a
forall a.
AlphaEnv
-> Quantifiers -> Key BoolMap -> A a -> BoolMap a -> BoolMap a
forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
mAlter AlphaEnv
env Quantifiers
qs Key BoolMap
b A a
f (MaybeMap a -> MaybeMap a -> BoolMap a
forall a. MaybeMap a -> MaybeMap a -> BoolMap a
BoolMap MaybeMap a
forall a. MaybeMap a
forall (m :: * -> *) a. PatternMap m => m a
mEmpty MaybeMap a
forall a. MaybeMap a
forall (m :: * -> *) a. PatternMap m => m a
mEmpty)
  mAlter AlphaEnv
env Quantifiers
qs Key BoolMap
b A a
f m :: BoolMap a
m@BoolMap{}
    | Bool
Key BoolMap
b = BoolMap a
m { bmTrue = mAlter env qs () f (bmTrue m) }
    | Bool
otherwise = BoolMap a
m { bmFalse = mAlter env qs () f (bmFalse m) }

  mMatch
    :: MatchEnv
    -> Key BoolMap
    -> (Substitution, BoolMap a)
    -> [(Substitution, a)]
  mMatch :: forall a.
MatchEnv
-> Key BoolMap -> (Substitution, BoolMap a) -> [(Substitution, a)]
mMatch MatchEnv
_ Key BoolMap
_ (Substitution
_, BoolMap a
EmptyBoolMap) = []
  mMatch MatchEnv
env Key BoolMap
b hs :: (Substitution, BoolMap a)
hs@(Substitution
_, BoolMap{})
    | Bool
Key BoolMap
b = (BoolMap a -> MaybeMap a)
-> (Substitution, BoolMap a) -> [(Substitution, MaybeMap a)]
forall b c a. (b -> c) -> (a, b) -> [(a, c)]
mapFor BoolMap a -> MaybeMap a
forall a. BoolMap a -> MaybeMap a
bmTrue (Substitution, BoolMap a)
hs [(Substitution, MaybeMap a)]
-> ((Substitution, MaybeMap a) -> [(Substitution, a)])
-> [(Substitution, a)]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MatchEnv
-> Key MaybeMap
-> (Substitution, MaybeMap a)
-> [(Substitution, a)]
forall a.
MatchEnv
-> Key MaybeMap
-> (Substitution, MaybeMap a)
-> [(Substitution, a)]
forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env ()
    | Bool
otherwise = (BoolMap a -> MaybeMap a)
-> (Substitution, BoolMap a) -> [(Substitution, MaybeMap a)]
forall b c a. (b -> c) -> (a, b) -> [(a, c)]
mapFor BoolMap a -> MaybeMap a
forall a. BoolMap a -> MaybeMap a
bmFalse (Substitution, BoolMap a)
hs [(Substitution, MaybeMap a)]
-> ((Substitution, MaybeMap a) -> [(Substitution, a)])
-> [(Substitution, a)]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MatchEnv
-> Key MaybeMap
-> (Substitution, MaybeMap a)
-> [(Substitution, a)]
forall a.
MatchEnv
-> Key MaybeMap
-> (Substitution, MaybeMap a)
-> [(Substitution, a)]
forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env ()

------------------------------------------------------------------------

newtype IntMap a = IntMap { forall a. IntMap a -> IntMap [a]
unIntMap :: I.IntMap [a] }
  deriving ((forall a b. (a -> b) -> IntMap a -> IntMap b)
-> (forall a b. a -> IntMap b -> IntMap a) -> Functor IntMap
forall a b. a -> IntMap b -> IntMap a
forall a b. (a -> b) -> IntMap a -> IntMap b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> IntMap a -> IntMap b
fmap :: forall a b. (a -> b) -> IntMap a -> IntMap b
$c<$ :: forall a b. a -> IntMap b -> IntMap a
<$ :: forall a b. a -> IntMap b -> IntMap a
Functor)

instance PatternMap IntMap where
  type Key IntMap = I.Key

  mEmpty :: IntMap a
  mEmpty :: forall a. IntMap a
mEmpty = IntMap [a] -> IntMap a
forall a. IntMap [a] -> IntMap a
IntMap IntMap [a]
forall a. IntMap a
I.empty

  mUnion :: IntMap a -> IntMap a -> IntMap a
  mUnion :: forall a. IntMap a -> IntMap a -> IntMap a
mUnion (IntMap IntMap [a]
m1) (IntMap IntMap [a]
m2) = IntMap [a] -> IntMap a
forall a. IntMap [a] -> IntMap a
IntMap (IntMap [a] -> IntMap a) -> IntMap [a] -> IntMap a
forall a b. (a -> b) -> a -> b
$ ([a] -> [a] -> [a]) -> IntMap [a] -> IntMap [a] -> IntMap [a]
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
I.unionWith [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
(++) IntMap [a]
m1 IntMap [a]
m2

  mAlter :: AlphaEnv -> Quantifiers -> Key IntMap -> A a -> IntMap a -> IntMap a
  mAlter :: forall a.
AlphaEnv
-> Quantifiers -> Key IntMap -> A a -> IntMap a -> IntMap a
mAlter AlphaEnv
_ Quantifiers
_ Key IntMap
i A a
f (IntMap IntMap [a]
m) = IntMap [a] -> IntMap a
forall a. IntMap [a] -> IntMap a
IntMap (IntMap [a] -> IntMap a) -> IntMap [a] -> IntMap a
forall a b. (a -> b) -> a -> b
$ (Maybe [a] -> Maybe [a]) -> Key -> IntMap [a] -> IntMap [a]
forall a. (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
I.alter (A a -> Maybe [a] -> Maybe [a]
forall a. A a -> A [a]
toAList A a
f) Key
Key IntMap
i IntMap [a]
m

  mMatch
    :: MatchEnv
    -> Key IntMap
    -> (Substitution, IntMap a)
    -> [(Substitution, a)]
  mMatch :: forall a.
MatchEnv
-> Key IntMap -> (Substitution, IntMap a) -> [(Substitution, a)]
mMatch MatchEnv
_ Key IntMap
i = (IntMap a -> Maybe [a])
-> (Substitution, IntMap a) -> [(Substitution, a)]
forall b c a. (b -> Maybe [c]) -> (a, b) -> [(a, c)]
maybeListMap (Key -> IntMap [a] -> Maybe [a]
forall a. Key -> IntMap a -> Maybe a
I.lookup Key
Key IntMap
i (IntMap [a] -> Maybe [a])
-> (IntMap a -> IntMap [a]) -> IntMap a -> Maybe [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap a -> IntMap [a]
forall a. IntMap a -> IntMap [a]
unIntMap)

------------------------------------------------------------------------

newtype Map k a = Map { forall k a. Map k a -> Map k [a]
unMap :: M.Map k [a] }
  deriving ((forall a b. (a -> b) -> Map k a -> Map k b)
-> (forall a b. a -> Map k b -> Map k a) -> Functor (Map k)
forall a b. a -> Map k b -> Map k a
forall a b. (a -> b) -> Map k a -> Map k b
forall k a b. a -> Map k b -> Map k a
forall k a b. (a -> b) -> Map k a -> Map k b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall k a b. (a -> b) -> Map k a -> Map k b
fmap :: forall a b. (a -> b) -> Map k a -> Map k b
$c<$ :: forall k a b. a -> Map k b -> Map k a
<$ :: forall a b. a -> Map k b -> Map k a
Functor)

mapAssocs :: Map k v -> [(k,v)]
mapAssocs :: forall k v. Map k v -> [(k, v)]
mapAssocs (Map Map k [v]
m) = [ (k
k,v
v) | (k
k,[v]
vs) <- Map k [v] -> [(k, [v])]
forall k a. Map k a -> [(k, a)]
M.assocs Map k [v]
m, v
v <- [v]
vs ]

instance Ord k => PatternMap (Map k) where
  type Key (Map k) = k

  mEmpty :: Map k a
  mEmpty :: forall a. Map k a
mEmpty = Map k [a] -> Map k a
forall k a. Map k [a] -> Map k a
Map Map k [a]
forall k a. Map k a
M.empty

  mUnion :: Map k a -> Map k a -> Map k a
  mUnion :: forall a. Map k a -> Map k a -> Map k a
mUnion (Map Map k [a]
m1) (Map Map k [a]
m2) = Map k [a] -> Map k a
forall k a. Map k [a] -> Map k a
Map (Map k [a] -> Map k a) -> Map k [a] -> Map k a
forall a b. (a -> b) -> a -> b
$ ([a] -> [a] -> [a]) -> Map k [a] -> Map k [a] -> Map k [a]
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
(++) Map k [a]
m1 Map k [a]
m2

  mAlter :: AlphaEnv -> Quantifiers -> Key (Map k) -> A a -> Map k a -> Map k a
  mAlter :: forall a.
AlphaEnv -> Quantifiers -> Key (Map k) -> A a -> Map k a -> Map k a
mAlter AlphaEnv
_ Quantifiers
_ Key (Map k)
k A a
f (Map Map k [a]
m) = Map k [a] -> Map k a
forall k a. Map k [a] -> Map k a
Map (Map k [a] -> Map k a) -> Map k [a] -> Map k a
forall a b. (a -> b) -> a -> b
$ (Maybe [a] -> Maybe [a]) -> k -> Map k [a] -> Map k [a]
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter (A a -> Maybe [a] -> Maybe [a]
forall a. A a -> A [a]
toAList A a
f) k
Key (Map k)
k Map k [a]
m

  mMatch
    :: MatchEnv
    -> Key (Map k)
    -> (Substitution, Map k a)
    -> [(Substitution, a)]
  mMatch :: forall a.
MatchEnv
-> Key (Map k) -> (Substitution, Map k a) -> [(Substitution, a)]
mMatch MatchEnv
_ Key (Map k)
k = (Map k a -> Maybe [a])
-> (Substitution, Map k a) -> [(Substitution, a)]
forall b c a. (b -> Maybe [c]) -> (a, b) -> [(a, c)]
maybeListMap (k -> Map k [a] -> Maybe [a]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
Key (Map k)
k (Map k [a] -> Maybe [a])
-> (Map k a -> Map k [a]) -> Map k a -> Maybe [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k a -> Map k [a]
forall k a. Map k a -> Map k [a]
unMap)

------------------------------------------------------------------------

-- Note [OccEnv]
--
-- We avoid using OccEnv because the Uniquable instance for OccName
-- takes the NameSpace of the OccName into account, which we rarely actually
-- want. (Doing so requires creating new RdrNames with the proper namespace,
-- which is a bunch of fiddling for no obvious gain for our uses.) Instead
-- we just use a map based on the FastString name.

newtype FSEnv a =
  FSEnv { forall a. FSEnv a -> UniqFM a
_unFSEnv :: UniqFM a } -- this is the UniqFM below, NOT GHC's UniqFM
  deriving ((forall a b. (a -> b) -> FSEnv a -> FSEnv b)
-> (forall a b. a -> FSEnv b -> FSEnv a) -> Functor FSEnv
forall a b. a -> FSEnv b -> FSEnv a
forall a b. (a -> b) -> FSEnv a -> FSEnv b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> FSEnv a -> FSEnv b
fmap :: forall a b. (a -> b) -> FSEnv a -> FSEnv b
$c<$ :: forall a b. a -> FSEnv b -> FSEnv a
<$ :: forall a b. a -> FSEnv b -> FSEnv a
Functor)

instance PatternMap FSEnv where
  type Key FSEnv = GHC.FastString

  mEmpty :: FSEnv a
  mEmpty :: forall a. FSEnv a
mEmpty = UniqFM a -> FSEnv a
forall a. UniqFM a -> FSEnv a
FSEnv UniqFM a
forall a. UniqFM a
forall (m :: * -> *) a. PatternMap m => m a
mEmpty

  mUnion :: FSEnv a -> FSEnv a -> FSEnv a
  mUnion :: forall a. FSEnv a -> FSEnv a -> FSEnv a
mUnion (FSEnv UniqFM a
m1) (FSEnv UniqFM a
m2) = UniqFM a -> FSEnv a
forall a. UniqFM a -> FSEnv a
FSEnv (UniqFM a -> UniqFM a -> UniqFM a
forall a. UniqFM a -> UniqFM a -> UniqFM a
forall (m :: * -> *) a. PatternMap m => m a -> m a -> m a
mUnion UniqFM a
m1 UniqFM a
m2)

  mAlter :: AlphaEnv -> Quantifiers -> Key FSEnv -> A a -> FSEnv a -> FSEnv a
  mAlter :: forall a.
AlphaEnv -> Quantifiers -> Key FSEnv -> A a -> FSEnv a -> FSEnv a
mAlter AlphaEnv
env Quantifiers
qs Key FSEnv
fs A a
f (FSEnv UniqFM a
m) = UniqFM a -> FSEnv a
forall a. UniqFM a -> FSEnv a
FSEnv (AlphaEnv
-> Quantifiers -> Key UniqFM -> A a -> UniqFM a -> UniqFM a
forall a.
AlphaEnv
-> Quantifiers -> Key UniqFM -> A a -> UniqFM a -> UniqFM a
forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
mAlter AlphaEnv
env Quantifiers
qs (FastString -> Unique
forall a. Uniquable a => a -> Unique
GHC.getUnique FastString
Key FSEnv
fs) A a
f UniqFM a
m)

  mMatch :: MatchEnv -> Key FSEnv -> (Substitution, FSEnv a) -> [(Substitution, a)]
  mMatch :: forall a.
MatchEnv
-> Key FSEnv -> (Substitution, FSEnv a) -> [(Substitution, a)]
mMatch MatchEnv
env Key FSEnv
fs (Substitution
hs, FSEnv UniqFM a
m) = MatchEnv
-> Key UniqFM -> (Substitution, UniqFM a) -> [(Substitution, a)]
forall a.
MatchEnv
-> Key UniqFM -> (Substitution, UniqFM a) -> [(Substitution, a)]
forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env (FastString -> Unique
forall a. Uniquable a => a -> Unique
GHC.getUnique FastString
Key FSEnv
fs) (Substitution
hs, UniqFM a
m)

------------------------------------------------------------------------

newtype UniqFM a = UniqFM { forall a. UniqFM a -> UniqFM Unique [a]
unUniqFM :: GHC.UniqFM GHC.Unique [a] }
  deriving ((forall a b. (a -> b) -> UniqFM a -> UniqFM b)
-> (forall a b. a -> UniqFM b -> UniqFM a) -> Functor UniqFM
forall a b. a -> UniqFM b -> UniqFM a
forall a b. (a -> b) -> UniqFM a -> UniqFM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> UniqFM a -> UniqFM b
fmap :: forall a b. (a -> b) -> UniqFM a -> UniqFM b
$c<$ :: forall a b. a -> UniqFM b -> UniqFM a
<$ :: forall a b. a -> UniqFM b -> UniqFM a
Functor)

instance PatternMap UniqFM where
  type Key UniqFM = GHC.Unique

  mEmpty :: UniqFM a
  mEmpty :: forall a. UniqFM a
mEmpty = UniqFM Unique [a] -> UniqFM a
forall a. UniqFM Unique [a] -> UniqFM a
UniqFM UniqFM Unique [a]
forall key elt. UniqFM key elt
GHC.emptyUFM

  mUnion :: UniqFM a -> UniqFM a -> UniqFM a
  mUnion :: forall a. UniqFM a -> UniqFM a -> UniqFM a
mUnion (UniqFM UniqFM Unique [a]
m1) (UniqFM UniqFM Unique [a]
m2) = UniqFM Unique [a] -> UniqFM a
forall a. UniqFM Unique [a] -> UniqFM a
UniqFM (UniqFM Unique [a] -> UniqFM a) -> UniqFM Unique [a] -> UniqFM a
forall a b. (a -> b) -> a -> b
$ ([a] -> [a] -> [a])
-> UniqFM Unique [a] -> UniqFM Unique [a] -> UniqFM Unique [a]
forall elt key.
(elt -> elt -> elt)
-> UniqFM key elt -> UniqFM key elt -> UniqFM key elt
GHC.plusUFM_C [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
(++) UniqFM Unique [a]
m1 UniqFM Unique [a]
m2

  mAlter :: AlphaEnv -> Quantifiers -> Key UniqFM -> A a -> UniqFM a -> UniqFM a
  mAlter :: forall a.
AlphaEnv
-> Quantifiers -> Key UniqFM -> A a -> UniqFM a -> UniqFM a
mAlter AlphaEnv
_ Quantifiers
_ Key UniqFM
k A a
f (UniqFM UniqFM Unique [a]
m) = UniqFM Unique [a] -> UniqFM a
forall a. UniqFM Unique [a] -> UniqFM a
UniqFM (UniqFM Unique [a] -> UniqFM a) -> UniqFM Unique [a] -> UniqFM a
forall a b. (a -> b) -> a -> b
$ (Maybe [a] -> Maybe [a])
-> UniqFM Unique [a] -> Unique -> UniqFM Unique [a]
forall key elt.
Uniquable key =>
(Maybe elt -> Maybe elt) -> UniqFM key elt -> key -> UniqFM key elt
GHC.alterUFM (A a -> Maybe [a] -> Maybe [a]
forall a. A a -> A [a]
toAList A a
f) UniqFM Unique [a]
m Unique
Key UniqFM
k

  mMatch
    :: MatchEnv
    -> Key UniqFM
    -> (Substitution, UniqFM a)
    -> [(Substitution, a)]
  mMatch :: forall a.
MatchEnv
-> Key UniqFM -> (Substitution, UniqFM a) -> [(Substitution, a)]
mMatch MatchEnv
_ Key UniqFM
k = (UniqFM a -> Maybe [a])
-> (Substitution, UniqFM a) -> [(Substitution, a)]
forall b c a. (b -> Maybe [c]) -> (a, b) -> [(a, c)]
maybeListMap ((UniqFM Unique [a] -> Unique -> Maybe [a])
-> Unique -> UniqFM Unique [a] -> Maybe [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip UniqFM Unique [a] -> Unique -> Maybe [a]
forall key elt. UniqFM key elt -> Unique -> Maybe elt
GHC.lookupUFM_Directly Unique
Key UniqFM
k (UniqFM Unique [a] -> Maybe [a])
-> (UniqFM a -> UniqFM Unique [a]) -> UniqFM a -> Maybe [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UniqFM a -> UniqFM Unique [a]
forall a. UniqFM a -> UniqFM Unique [a]
unUniqFM)

------------------------------------------------------------------------