-- 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 -> 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
<$ :: forall a b. a -> BoolMap b -> BoolMap a
$c<$ :: forall a b. a -> BoolMap b -> BoolMap a
fmap :: forall a b. (a -> b) -> BoolMap a -> BoolMap b
$cfmap :: forall a b. (a -> b) -> BoolMap a -> BoolMap b
Functor)

instance PatternMap BoolMap where
  type Key BoolMap = Bool

  mEmpty :: BoolMap a
  mEmpty :: forall a. BoolMap a
mEmpty = 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 = forall (m :: * -> *) a b.
PatternMap m =>
(a -> m b) -> a -> a -> m b
unionOn forall a. BoolMap a -> MaybeMap a
bmTrue BoolMap a
m1 BoolMap a
m2
    , bmFalse :: MaybeMap a
bmFalse = forall (m :: * -> *) a b.
PatternMap m =>
(a -> m b) -> a -> a -> m b
unionOn 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 = 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 (forall a. MaybeMap a -> MaybeMap a -> BoolMap a
BoolMap forall (m :: * -> *) a. PatternMap m => m a
mEmpty forall (m :: * -> *) a. PatternMap m => m a
mEmpty)
  mAlter AlphaEnv
env Quantifiers
qs Key BoolMap
b A a
f m :: BoolMap a
m@BoolMap{}
    | Key BoolMap
b = BoolMap a
m { bmTrue :: MaybeMap a
bmTrue = forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
mAlter AlphaEnv
env Quantifiers
qs () A a
f (forall a. BoolMap a -> MaybeMap a
bmTrue BoolMap a
m) }
    | Bool
otherwise = BoolMap a
m { bmFalse :: MaybeMap a
bmFalse = forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
mAlter AlphaEnv
env Quantifiers
qs () A a
f (forall a. BoolMap a -> MaybeMap a
bmFalse BoolMap a
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{})
    | Key BoolMap
b = forall b c a. (b -> c) -> (a, b) -> [(a, c)]
mapFor forall a. BoolMap a -> MaybeMap a
bmTrue (Substitution, BoolMap a)
hs forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env ()
    | Bool
otherwise = forall b c a. (b -> c) -> (a, b) -> [(a, c)]
mapFor forall a. BoolMap a -> MaybeMap a
bmFalse (Substitution, BoolMap a)
hs forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 -> 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
<$ :: forall a b. a -> IntMap b -> IntMap a
$c<$ :: forall a b. a -> IntMap b -> IntMap a
fmap :: forall a b. (a -> b) -> IntMap a -> IntMap b
$cfmap :: forall a b. (a -> b) -> IntMap a -> IntMap b
Functor)

instance PatternMap IntMap where
  type Key IntMap = I.Key

  mEmpty :: IntMap a
  mEmpty :: forall a. IntMap a
mEmpty = forall a. IntMap [a] -> IntMap a
IntMap 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) = forall a. IntMap [a] -> IntMap a
IntMap forall a b. (a -> b) -> a -> b
$ forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
I.unionWith 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) = forall a. IntMap [a] -> IntMap a
IntMap forall a b. (a -> b) -> a -> b
$ forall a. (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
I.alter (forall a. A a -> A [a]
toAList A a
f) 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 = forall b c a. (b -> Maybe [c]) -> (a, b) -> [(a, c)]
maybeListMap (forall a. Key -> IntMap a -> Maybe a
I.lookup Key IntMap
i forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 -> 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
<$ :: forall a b. a -> Map k b -> Map k a
$c<$ :: forall k a b. a -> Map k b -> Map k a
fmap :: forall a b. (a -> b) -> Map k a -> Map k b
$cfmap :: forall k a b. (a -> b) -> Map k a -> Map k b
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) <- 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 = forall k a. Map k [a] -> Map k a
Map 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) = forall k a. Map k [a] -> Map k a
Map forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith 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) = forall k a. Map k [a] -> Map k a
Map forall a b. (a -> b) -> a -> b
$ forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter (forall a. A a -> A [a]
toAList A a
f) 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 = forall b c a. (b -> Maybe [c]) -> (a, b) -> [(a, c)]
maybeListMap (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Key (Map k)
k forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 -> 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
<$ :: forall a b. a -> FSEnv b -> FSEnv a
$c<$ :: forall a b. a -> FSEnv b -> FSEnv a
fmap :: forall a b. (a -> b) -> FSEnv a -> FSEnv b
$cfmap :: forall a b. (a -> b) -> FSEnv a -> FSEnv b
Functor)

instance PatternMap FSEnv where
  type Key FSEnv = GHC.FastString

  mEmpty :: FSEnv a
  mEmpty :: forall a. FSEnv a
mEmpty = forall a. UniqFM a -> FSEnv a
FSEnv 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) = forall a. UniqFM a -> FSEnv a
FSEnv (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) = forall a. UniqFM a -> FSEnv a
FSEnv (forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
mAlter AlphaEnv
env Quantifiers
qs (forall a. Uniquable a => a -> Unique
GHC.getUnique 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) = forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env (forall a. Uniquable a => a -> Unique
GHC.getUnique 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 -> 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
<$ :: forall a b. a -> UniqFM b -> UniqFM a
$c<$ :: forall a b. a -> UniqFM b -> UniqFM a
fmap :: forall a b. (a -> b) -> UniqFM a -> UniqFM b
$cfmap :: forall a b. (a -> b) -> UniqFM a -> UniqFM b
Functor)

instance PatternMap UniqFM where
  type Key UniqFM = GHC.Unique

  mEmpty :: UniqFM a
  mEmpty :: forall a. UniqFM a
mEmpty = forall a. UniqFM Unique [a] -> UniqFM a
UniqFM 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) = forall a. UniqFM Unique [a] -> UniqFM a
UniqFM forall a b. (a -> b) -> a -> b
$ forall elt key.
(elt -> elt -> elt)
-> UniqFM key elt -> UniqFM key elt -> UniqFM key elt
GHC.plusUFM_C 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) = forall a. UniqFM Unique [a] -> UniqFM a
UniqFM forall a b. (a -> b) -> a -> b
$ forall key elt.
Uniquable key =>
(Maybe elt -> Maybe elt) -> UniqFM key elt -> key -> UniqFM key elt
GHC.alterUFM (forall a. A a -> A [a]
toAList A a
f) UniqFM Unique [a]
m 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 = forall b c a. (b -> Maybe [c]) -> (a, b) -> [(a, c)]
maybeListMap (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall key elt. UniqFM key elt -> Unique -> Maybe elt
GHC.lookupUFM_Directly Key UniqFM
k forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. UniqFM a -> UniqFM Unique [a]
unUniqFM)

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