{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE UndecidableInstances #-}

{-
(c) The University of Glasgow 2006
(c) The GRASP/AQUA Project, Glasgow University, 1992-1998
-}
module GHC.Data.TrieMap(
   -- * Maps over 'Maybe' values
   MaybeMap,
   -- * Maps over 'List' values
   ListMap,
   -- * Maps over 'Literal's
   LiteralMap,
   -- * 'TrieMap' class
   TrieMap(..), insertTM, deleteTM, foldMapTM, isEmptyTM,

   -- * Things helpful for adding additional Instances.
   (>.>), (|>), (|>>), XT,
   foldMaybe, filterMaybe,
   -- * Map for leaf compression
   GenMap,
   lkG, xtG, mapG, fdG,
   xtList, lkList

 ) where

import GHC.Prelude

import GHC.Types.Literal
import GHC.Types.Unique.DFM
import GHC.Types.Unique( Uniquable )

import qualified Data.Map    as Map
import qualified Data.IntMap as IntMap
import GHC.Utils.Outputable
import Control.Monad( (>=>) )
import Data.Kind( Type )

import qualified Data.Semigroup as S

{-
This module implements TrieMaps, which are finite mappings
whose key is a structured value like a CoreExpr or Type.

This file implements tries over general data structures.
Implementation for tries over Core Expressions/Types are
available in GHC.Core.Map.Expr.

The regular pattern for handling TrieMaps on data structures was first
described (to my knowledge) in Connelly and Morris's 1995 paper "A
generalization of the Trie Data Structure"; there is also an accessible
description of the idea in Okasaki's book "Purely Functional Data
Structures", Section 10.3.2

************************************************************************
*                                                                      *
                   The TrieMap class
*                                                                      *
************************************************************************
-}

type XT a = Maybe a -> Maybe a  -- How to alter a non-existent elt (Nothing)
                                --               or an existing elt (Just)

class Functor m => TrieMap m where
   type Key m :: Type
   emptyTM  :: m a
   lookupTM :: forall b. Key m -> m b -> Maybe b
   alterTM  :: forall b. Key m -> XT b -> m b -> m b
   filterTM :: (a -> Bool) -> m a -> m a

   foldTM   :: (a -> b -> b) -> m a -> b -> b
      -- The unusual argument order here makes
      -- it easy to compose calls to foldTM;
      -- see for example fdE below

insertTM :: TrieMap m => Key m -> a -> m a -> m a
insertTM :: forall (m :: * -> *) a. TrieMap m => Key m -> a -> m a -> m a
insertTM Key m
k a
v m a
m = forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
alterTM Key m
k (\Maybe a
_ -> forall a. a -> Maybe a
Just a
v) m a
m

deleteTM :: TrieMap m => Key m -> m a -> m a
deleteTM :: forall (m :: * -> *) a. TrieMap m => Key m -> m a -> m a
deleteTM Key m
k m a
m = forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
alterTM Key m
k (\Maybe a
_ -> forall a. Maybe a
Nothing) m a
m

foldMapTM :: (TrieMap m, Monoid r) => (a -> r) -> m a -> r
foldMapTM :: forall (m :: * -> *) r a.
(TrieMap m, Monoid r) =>
(a -> r) -> m a -> r
foldMapTM a -> r
f m a
m = forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM (\ a
x r
r -> a -> r
f a
x forall a. Semigroup a => a -> a -> a
S.<> r
r) m a
m forall a. Monoid a => a
mempty

-- This looks inefficient.
isEmptyTM :: TrieMap m => m a -> Bool
isEmptyTM :: forall (m :: * -> *) a. TrieMap m => m a -> Bool
isEmptyTM m a
m = forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM (\ a
_ Bool
_ -> Bool
False) m a
m Bool
True

----------------------
-- Recall that
--   Control.Monad.(>=>) :: (a -> Maybe b) -> (b -> Maybe c) -> a -> Maybe c

(>.>) :: (a -> b) -> (b -> c) -> a -> c
-- Reverse function composition (do f first, then g)
infixr 1 >.>
(a -> b
f >.> :: forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> b -> c
g) a
x = b -> c
g (a -> b
f a
x)
infixr 1 |>, |>>

(|>) :: a -> (a->b) -> b     -- Reverse application
a
x |> :: forall a b. a -> (a -> b) -> b
|> a -> b
f = a -> b
f a
x

----------------------
(|>>) :: TrieMap m2
      => (XT (m2 a) -> m1 (m2 a) -> m1 (m2 a))
      -> (m2 a -> m2 a)
      -> m1 (m2 a) -> m1 (m2 a)
|>> :: forall (m2 :: * -> *) a (m1 :: * -> *).
TrieMap m2 =>
(XT (m2 a) -> m1 (m2 a) -> m1 (m2 a))
-> (m2 a -> m2 a) -> m1 (m2 a) -> m1 (m2 a)
(|>>) XT (m2 a) -> m1 (m2 a) -> m1 (m2 a)
f m2 a -> m2 a
g = XT (m2 a) -> m1 (m2 a) -> m1 (m2 a)
f (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. m2 a -> m2 a
g forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. TrieMap m => Maybe (m a) -> m a
deMaybe)

deMaybe :: TrieMap m => Maybe (m a) -> m a
deMaybe :: forall (m :: * -> *) a. TrieMap m => Maybe (m a) -> m a
deMaybe Maybe (m a)
Nothing  = forall (m :: * -> *) a. TrieMap m => m a
emptyTM
deMaybe (Just m a
m) = m a
m

{-
Note [Every TrieMap is a Functor]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Every TrieMap T admits
   fmap :: (a->b) -> T a -> T b
where (fmap f t) applies `f` to every element of the range of `t`.
Ergo, we make `Functor` a superclass of `TrieMap`.

Moreover it is almost invariably possible to /derive/ Functor for each
particular instance. E.g. in the list instance we have
    data ListMap m a
      = LM { lm_nil  :: Maybe a
           , lm_cons :: m (ListMap m a) }
      deriving (Functor)
    instance TrieMap m => TrieMap (ListMap m) where { .. }

Alas, we not yet derive `Functor` for reasons of performance; see #22292.
-}

{-
************************************************************************
*                                                                      *
                   IntMaps
*                                                                      *
************************************************************************
-}

instance TrieMap IntMap.IntMap where
  type Key IntMap.IntMap = Int
  emptyTM :: forall a. IntMap a
emptyTM = forall a. IntMap a
IntMap.empty
  lookupTM :: forall b. Key IntMap -> IntMap b -> Maybe b
lookupTM Key IntMap
k IntMap b
m = forall a. Key -> IntMap a -> Maybe a
IntMap.lookup Key IntMap
k IntMap b
m
  alterTM :: forall b. Key IntMap -> XT b -> IntMap b -> IntMap b
alterTM = forall a. Key -> XT a -> IntMap a -> IntMap a
xtInt
  foldTM :: forall a b. (a -> b -> b) -> IntMap a -> b -> b
foldTM a -> b -> b
k IntMap a
m b
z = forall a b. (a -> b -> b) -> b -> IntMap a -> b
IntMap.foldr a -> b -> b
k b
z IntMap a
m
  filterTM :: forall a. (a -> Bool) -> IntMap a -> IntMap a
filterTM a -> Bool
f IntMap a
m = forall a. (a -> Bool) -> IntMap a -> IntMap a
IntMap.filter a -> Bool
f IntMap a
m

xtInt :: Int -> XT a -> IntMap.IntMap a -> IntMap.IntMap a
xtInt :: forall a. Key -> XT a -> IntMap a -> IntMap a
xtInt Key
k XT a
f IntMap a
m = forall a. (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
IntMap.alter XT a
f Key
k IntMap a
m

instance Ord k => TrieMap (Map.Map k) where
  type Key (Map.Map k) = k
  emptyTM :: forall a. Map k a
emptyTM = forall k a. Map k a
Map.empty
  lookupTM :: forall b. Key (Map k) -> Map k b -> Maybe b
lookupTM = forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup
  alterTM :: forall b. Key (Map k) -> XT b -> Map k b -> Map k b
alterTM Key (Map k)
k XT b
f Map k b
m = forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter XT b
f Key (Map k)
k Map k b
m
  foldTM :: forall a b. (a -> b -> b) -> Map k a -> b -> b
foldTM a -> b -> b
k Map k a
m b
z = forall a b k. (a -> b -> b) -> b -> Map k a -> b
Map.foldr a -> b -> b
k b
z Map k a
m
  filterTM :: forall a. (a -> Bool) -> Map k a -> Map k a
filterTM a -> Bool
f Map k a
m = forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter a -> Bool
f Map k a
m


{-
Note [foldTM determinism]
~~~~~~~~~~~~~~~~~~~~~~~~~
We want foldTM to be deterministic, which is why we have an instance of
TrieMap for UniqDFM, but not for UniqFM. Here's an example of some things that
go wrong if foldTM is nondeterministic. Consider:

  f a b = return (a <> b)

Depending on the order that the typechecker generates constraints you
get either:

  f :: (Monad m, Monoid a) => a -> a -> m a

or:

  f :: (Monoid a, Monad m) => a -> a -> m a

The generated code will be different after desugaring as the dictionaries
will be bound in different orders, leading to potential ABI incompatibility.

One way to solve this would be to notice that the typeclasses could be
sorted alphabetically.

Unfortunately that doesn't quite work with this example:

  f a b = let x = a <> a; y = b <> b in x

where you infer:

  f :: (Monoid m, Monoid m1) => m1 -> m -> m1

or:

  f :: (Monoid m1, Monoid m) => m1 -> m -> m1

Here you could decide to take the order of the type variables in the type
according to depth first traversal and use it to order the constraints.

The real trouble starts when the user enables incoherent instances and
the compiler has to make an arbitrary choice. Consider:

  class T a b where
    go :: a -> b -> String

  instance (Show b) => T Int b where
    go a b = show a ++ show b

  instance (Show a) => T a Bool where
    go a b = show a ++ show b

  f = go 10 True

GHC is free to choose either dictionary to implement f, but for the sake of
determinism we'd like it to be consistent when compiling the same sources
with the same flags.

inert_dicts :: DictMap is implemented with a TrieMap. In getUnsolvedInerts it
gets converted to a bag of (Wanted) Cts using a fold. Then in
solve_simple_wanteds it's merged with other WantedConstraints. We want the
conversion to a bag to be deterministic. For that purpose we use UniqDFM
instead of UniqFM to implement the TrieMap.

See Note [Deterministic UniqFM] in GHC.Types.Unique.DFM for more details on how it's made
deterministic.
-}

instance forall key. Uniquable key => TrieMap (UniqDFM key) where
  type Key (UniqDFM key) = key
  emptyTM :: forall a. UniqDFM key a
emptyTM = forall key elt. UniqDFM key elt
emptyUDFM
  lookupTM :: forall b. Key (UniqDFM key) -> UniqDFM key b -> Maybe b
lookupTM Key (UniqDFM key)
k UniqDFM key b
m = forall key elt.
Uniquable key =>
UniqDFM key elt -> key -> Maybe elt
lookupUDFM UniqDFM key b
m Key (UniqDFM key)
k
  alterTM :: forall b.
Key (UniqDFM key) -> XT b -> UniqDFM key b -> UniqDFM key b
alterTM Key (UniqDFM key)
k XT b
f UniqDFM key b
m = forall key elt.
Uniquable key =>
(Maybe elt -> Maybe elt)
-> UniqDFM key elt -> key -> UniqDFM key elt
alterUDFM XT b
f UniqDFM key b
m Key (UniqDFM key)
k
  foldTM :: forall a b. (a -> b -> b) -> UniqDFM key a -> b -> b
foldTM a -> b -> b
k UniqDFM key a
m b
z = forall elt a key. (elt -> a -> a) -> a -> UniqDFM key elt -> a
foldUDFM a -> b -> b
k b
z UniqDFM key a
m
  filterTM :: forall a. (a -> Bool) -> UniqDFM key a -> UniqDFM key a
filterTM a -> Bool
f UniqDFM key a
m = forall elt key. (elt -> Bool) -> UniqDFM key elt -> UniqDFM key elt
filterUDFM a -> Bool
f UniqDFM key a
m

{-
************************************************************************
*                                                                      *
                   Maybes
*                                                                      *
************************************************************************

If              m is a map from k -> val
then (MaybeMap m) is a map from (Maybe k) -> val
-}

data MaybeMap m a = MM { forall (m :: * -> *) a. MaybeMap m a -> Maybe a
mm_nothing  :: Maybe a, forall (m :: * -> *) a. MaybeMap m a -> m a
mm_just :: m a }

-- TODO(22292): derive
instance Functor m => Functor (MaybeMap m) where
    fmap :: forall a b. (a -> b) -> MaybeMap m a -> MaybeMap m b
fmap a -> b
f MM { mm_nothing :: forall (m :: * -> *) a. MaybeMap m a -> Maybe a
mm_nothing = Maybe a
mn, mm_just :: forall (m :: * -> *) a. MaybeMap m a -> m a
mm_just = m a
mj } = MM
      { mm_nothing :: Maybe b
mm_nothing = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Maybe a
mn, mm_just :: m b
mm_just = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f m a
mj }

instance TrieMap m => TrieMap (MaybeMap m) where
   type Key (MaybeMap m) = Maybe (Key m)
   emptyTM :: forall a. MaybeMap m a
emptyTM  = MM { mm_nothing :: Maybe a
mm_nothing = forall a. Maybe a
Nothing, mm_just :: m a
mm_just = forall (m :: * -> *) a. TrieMap m => m a
emptyTM }
   lookupTM :: forall b. Key (MaybeMap m) -> MaybeMap m b -> Maybe b
lookupTM = forall k (m :: * -> *) a.
(forall b. k -> m b -> Maybe b)
-> Maybe k -> MaybeMap m a -> Maybe a
lkMaybe forall (m :: * -> *) b. TrieMap m => Key m -> m b -> Maybe b
lookupTM
   alterTM :: forall b. Key (MaybeMap m) -> XT b -> MaybeMap m b -> MaybeMap m b
alterTM  = forall k (m :: * -> *) a.
(forall b. k -> XT b -> m b -> m b)
-> Maybe k -> XT a -> MaybeMap m a -> MaybeMap m a
xtMaybe forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
alterTM
   foldTM :: forall a b. (a -> b -> b) -> MaybeMap m a -> b -> b
foldTM   = forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> MaybeMap m a -> b -> b
fdMaybe
   filterTM :: forall a. (a -> Bool) -> MaybeMap m a -> MaybeMap m a
filterTM = forall (m :: * -> *) a.
TrieMap m =>
(a -> Bool) -> MaybeMap m a -> MaybeMap m a
ftMaybe

instance TrieMap m => Foldable (MaybeMap m) where
  foldMap :: forall m a. Monoid m => (a -> m) -> MaybeMap m a -> m
foldMap = forall (m :: * -> *) r a.
(TrieMap m, Monoid r) =>
(a -> r) -> m a -> r
foldMapTM

lkMaybe :: (forall b. k -> m b -> Maybe b)
        -> Maybe k -> MaybeMap m a -> Maybe a
lkMaybe :: forall k (m :: * -> *) a.
(forall b. k -> m b -> Maybe b)
-> Maybe k -> MaybeMap m a -> Maybe a
lkMaybe forall b. k -> m b -> Maybe b
_  Maybe k
Nothing  = forall (m :: * -> *) a. MaybeMap m a -> Maybe a
mm_nothing
lkMaybe forall b. k -> m b -> Maybe b
lk (Just k
x) = forall (m :: * -> *) a. MaybeMap m a -> m a
mm_just forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall b. k -> m b -> Maybe b
lk k
x

xtMaybe :: (forall b. k -> XT b -> m b -> m b)
        -> Maybe k -> XT a -> MaybeMap m a -> MaybeMap m a
xtMaybe :: forall k (m :: * -> *) a.
(forall b. k -> XT b -> m b -> m b)
-> Maybe k -> XT a -> MaybeMap m a -> MaybeMap m a
xtMaybe forall b. k -> XT b -> m b -> m b
_  Maybe k
Nothing  XT a
f MaybeMap m a
m = MaybeMap m a
m { mm_nothing :: Maybe a
mm_nothing  = XT a
f (forall (m :: * -> *) a. MaybeMap m a -> Maybe a
mm_nothing MaybeMap m a
m) }
xtMaybe forall b. k -> XT b -> m b -> m b
tr (Just k
x) XT a
f MaybeMap m a
m = MaybeMap m a
m { mm_just :: m a
mm_just = forall (m :: * -> *) a. MaybeMap m a -> m a
mm_just MaybeMap m a
m forall a b. a -> (a -> b) -> b
|> forall b. k -> XT b -> m b -> m b
tr k
x XT a
f }

fdMaybe :: TrieMap m => (a -> b -> b) -> MaybeMap m a -> b -> b
fdMaybe :: forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> MaybeMap m a -> b -> b
fdMaybe a -> b -> b
k MaybeMap m a
m = forall a b. (a -> b -> b) -> Maybe a -> b -> b
foldMaybe a -> b -> b
k (forall (m :: * -> *) a. MaybeMap m a -> Maybe a
mm_nothing MaybeMap m a
m)
            forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM a -> b -> b
k (forall (m :: * -> *) a. MaybeMap m a -> m a
mm_just MaybeMap m a
m)

ftMaybe :: TrieMap m => (a -> Bool) -> MaybeMap m a -> MaybeMap m a
ftMaybe :: forall (m :: * -> *) a.
TrieMap m =>
(a -> Bool) -> MaybeMap m a -> MaybeMap m a
ftMaybe a -> Bool
f (MM { mm_nothing :: forall (m :: * -> *) a. MaybeMap m a -> Maybe a
mm_nothing = Maybe a
mn, mm_just :: forall (m :: * -> *) a. MaybeMap m a -> m a
mm_just = m a
mj })
  = MM { mm_nothing :: Maybe a
mm_nothing = forall a. (a -> Bool) -> Maybe a -> Maybe a
filterMaybe a -> Bool
f Maybe a
mn, mm_just :: m a
mm_just = forall (m :: * -> *) a. TrieMap m => (a -> Bool) -> m a -> m a
filterTM a -> Bool
f m a
mj }

foldMaybe :: (a -> b -> b) -> Maybe a -> b -> b
foldMaybe :: forall a b. (a -> b -> b) -> Maybe a -> b -> b
foldMaybe a -> b -> b
_ Maybe a
Nothing  b
b = b
b
foldMaybe a -> b -> b
k (Just a
a) b
b = a -> b -> b
k a
a b
b

filterMaybe :: (a -> Bool) -> Maybe a -> Maybe a
filterMaybe :: forall a. (a -> Bool) -> Maybe a -> Maybe a
filterMaybe a -> Bool
_ Maybe a
Nothing = forall a. Maybe a
Nothing
filterMaybe a -> Bool
f input :: Maybe a
input@(Just a
x) | a -> Bool
f a
x       = Maybe a
input
                             | Bool
otherwise = forall a. Maybe a
Nothing

{-
************************************************************************
*                                                                      *
                   Lists
*                                                                      *
************************************************************************
-}

data ListMap m a
  = LM { forall (m :: * -> *) a. ListMap m a -> Maybe a
lm_nil  :: Maybe a
       , forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lm_cons :: m (ListMap m a) }

-- TODO(22292): derive
instance Functor m => Functor (ListMap m) where
    fmap :: forall a b. (a -> b) -> ListMap m a -> ListMap m b
fmap a -> b
f LM { lm_nil :: forall (m :: * -> *) a. ListMap m a -> Maybe a
lm_nil = Maybe a
mnil, lm_cons :: forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lm_cons = m (ListMap m a)
mcons } = LM
      { lm_nil :: Maybe b
lm_nil = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Maybe a
mnil, lm_cons :: m (ListMap m b)
lm_cons = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) m (ListMap m a)
mcons }

instance TrieMap m => TrieMap (ListMap m) where
   type Key (ListMap m) = [Key m]
   emptyTM :: forall a. ListMap m a
emptyTM  = LM { lm_nil :: Maybe a
lm_nil = forall a. Maybe a
Nothing, lm_cons :: m (ListMap m a)
lm_cons = forall (m :: * -> *) a. TrieMap m => m a
emptyTM }
   lookupTM :: forall b. Key (ListMap m) -> ListMap m b -> Maybe b
lookupTM = forall (m :: * -> *) k a.
TrieMap m =>
(forall b. k -> m b -> Maybe b) -> [k] -> ListMap m a -> Maybe a
lkList forall (m :: * -> *) b. TrieMap m => Key m -> m b -> Maybe b
lookupTM
   alterTM :: forall b. Key (ListMap m) -> XT b -> ListMap m b -> ListMap m b
alterTM  = forall (m :: * -> *) k a.
TrieMap m =>
(forall b. k -> XT b -> m b -> m b)
-> [k] -> XT a -> ListMap m a -> ListMap m a
xtList forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
alterTM
   foldTM :: forall a b. (a -> b -> b) -> ListMap m a -> b -> b
foldTM   = forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> ListMap m a -> b -> b
fdList
   filterTM :: forall a. (a -> Bool) -> ListMap m a -> ListMap m a
filterTM = forall (m :: * -> *) a.
TrieMap m =>
(a -> Bool) -> ListMap m a -> ListMap m a
ftList

instance TrieMap m => Foldable (ListMap m) where
  foldMap :: forall m a. Monoid m => (a -> m) -> ListMap m a -> m
foldMap = forall (m :: * -> *) r a.
(TrieMap m, Monoid r) =>
(a -> r) -> m a -> r
foldMapTM

instance (TrieMap m, Outputable a) => Outputable (ListMap m a) where
  ppr :: ListMap m a -> SDoc
ppr ListMap m a
m = String -> SDoc
text String
"List elts" SDoc -> SDoc -> SDoc
<+> forall a. Outputable a => a -> SDoc
ppr (forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM (:) ListMap m a
m [])

lkList :: TrieMap m => (forall b. k -> m b -> Maybe b)
        -> [k] -> ListMap m a -> Maybe a
lkList :: forall (m :: * -> *) k a.
TrieMap m =>
(forall b. k -> m b -> Maybe b) -> [k] -> ListMap m a -> Maybe a
lkList forall b. k -> m b -> Maybe b
_  []     = forall (m :: * -> *) a. ListMap m a -> Maybe a
lm_nil
lkList forall b. k -> m b -> Maybe b
lk (k
x:[k]
xs) = forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lm_cons forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall b. k -> m b -> Maybe b
lk k
x forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> forall (m :: * -> *) k a.
TrieMap m =>
(forall b. k -> m b -> Maybe b) -> [k] -> ListMap m a -> Maybe a
lkList forall b. k -> m b -> Maybe b
lk [k]
xs

xtList :: TrieMap m => (forall b. k -> XT b -> m b -> m b)
        -> [k] -> XT a -> ListMap m a -> ListMap m a
xtList :: forall (m :: * -> *) k a.
TrieMap m =>
(forall b. k -> XT b -> m b -> m b)
-> [k] -> XT a -> ListMap m a -> ListMap m a
xtList forall b. k -> XT b -> m b -> m b
_  []     XT a
f ListMap m a
m = ListMap m a
m { lm_nil :: Maybe a
lm_nil  = XT a
f (forall (m :: * -> *) a. ListMap m a -> Maybe a
lm_nil ListMap m a
m) }
xtList forall b. k -> XT b -> m b -> m b
tr (k
x:[k]
xs) XT a
f ListMap m a
m = ListMap m a
m { lm_cons :: m (ListMap m a)
lm_cons = forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lm_cons ListMap m a
m forall a b. a -> (a -> b) -> b
|> forall b. k -> XT b -> m b -> m b
tr k
x forall (m2 :: * -> *) a (m1 :: * -> *).
TrieMap m2 =>
(XT (m2 a) -> m1 (m2 a) -> m1 (m2 a))
-> (m2 a -> m2 a) -> m1 (m2 a) -> m1 (m2 a)
|>> forall (m :: * -> *) k a.
TrieMap m =>
(forall b. k -> XT b -> m b -> m b)
-> [k] -> XT a -> ListMap m a -> ListMap m a
xtList forall b. k -> XT b -> m b -> m b
tr [k]
xs XT a
f }

fdList :: forall m a b. TrieMap m
       => (a -> b -> b) -> ListMap m a -> b -> b
fdList :: forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> ListMap m a -> b -> b
fdList a -> b -> b
k ListMap m a
m = forall a b. (a -> b -> b) -> Maybe a -> b -> b
foldMaybe a -> b -> b
k          (forall (m :: * -> *) a. ListMap m a -> Maybe a
lm_nil ListMap m a
m)
           forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM    (forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> ListMap m a -> b -> b
fdList a -> b -> b
k) (forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lm_cons ListMap m a
m)

ftList :: TrieMap m => (a -> Bool) -> ListMap m a -> ListMap m a
ftList :: forall (m :: * -> *) a.
TrieMap m =>
(a -> Bool) -> ListMap m a -> ListMap m a
ftList a -> Bool
f (LM { lm_nil :: forall (m :: * -> *) a. ListMap m a -> Maybe a
lm_nil = Maybe a
mnil, lm_cons :: forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lm_cons = m (ListMap m a)
mcons })
  = LM { lm_nil :: Maybe a
lm_nil = forall a. (a -> Bool) -> Maybe a -> Maybe a
filterMaybe a -> Bool
f Maybe a
mnil, lm_cons :: m (ListMap m a)
lm_cons = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *) a. TrieMap m => (a -> Bool) -> m a -> m a
filterTM a -> Bool
f) m (ListMap m a)
mcons }

{-
************************************************************************
*                                                                      *
                   Basic maps
*                                                                      *
************************************************************************
-}

type LiteralMap  a = Map.Map Literal a

{-
************************************************************************
*                                                                      *
                   GenMap
*                                                                      *
************************************************************************

Note [Compressed TrieMap]
~~~~~~~~~~~~~~~~~~~~~~~~~

The GenMap constructor augments TrieMaps with leaf compression.  This helps
solve the performance problem detailed in #9960: suppose we have a handful
H of entries in a TrieMap, each with a very large key, size K. If you fold over
such a TrieMap you'd expect time O(H). That would certainly be true of an
association list! But with TrieMap we actually have to navigate down a long
singleton structure to get to the elements, so it takes time O(K*H).  This
can really hurt on many type-level computation benchmarks:
see for example T9872d.

The point of a TrieMap is that you need to navigate to the point where only one
key remains, and then things should be fast.  So the point of a SingletonMap
is that, once we are down to a single (key,value) pair, we stop and
just use SingletonMap.

'EmptyMap' provides an even more basic (but essential) optimization: if there is
nothing in the map, don't bother building out the (possibly infinite) recursive
TrieMap structure!

Compressed triemaps are heavily used by GHC.Core.Map.Expr. So we have to mark some things
as INLINEABLE to permit specialization.
-}

data GenMap m a
   = EmptyMap
   | SingletonMap (Key m) a
   | MultiMap (m a)

instance (Outputable a, Outputable (m a)) => Outputable (GenMap m a) where
  ppr :: GenMap m a -> SDoc
ppr GenMap m a
EmptyMap = String -> SDoc
text String
"Empty map"
  ppr (SingletonMap Key m
_ a
v) = String -> SDoc
text String
"Singleton map" SDoc -> SDoc -> SDoc
<+> forall a. Outputable a => a -> SDoc
ppr a
v
  ppr (MultiMap m a
m) = forall a. Outputable a => a -> SDoc
ppr m a
m

-- TODO(22292): derive
instance Functor m => Functor (GenMap m) where
    fmap :: forall a b. (a -> b) -> GenMap m a -> GenMap m b
fmap = forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> GenMap m a -> GenMap m b
mapG
    {-# INLINE fmap #-}

-- TODO undecidable instance
instance (Eq (Key m), TrieMap m) => TrieMap (GenMap m) where
   type Key (GenMap m) = Key m
   emptyTM :: forall a. GenMap m a
emptyTM  = forall (m :: * -> *) a. GenMap m a
EmptyMap
   lookupTM :: forall b. Key (GenMap m) -> GenMap m b -> Maybe b
lookupTM = forall (m :: * -> *) a.
(Eq (Key m), TrieMap m) =>
Key m -> GenMap m a -> Maybe a
lkG
   alterTM :: forall b. Key (GenMap m) -> XT b -> GenMap m b -> GenMap m b
alterTM  = forall (m :: * -> *) a.
(Eq (Key m), TrieMap m) =>
Key m -> XT a -> GenMap m a -> GenMap m a
xtG
   foldTM :: forall a b. (a -> b -> b) -> GenMap m a -> b -> b
foldTM   = forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> GenMap m a -> b -> b
fdG
   filterTM :: forall a. (a -> Bool) -> GenMap m a -> GenMap m a
filterTM = forall (m :: * -> *) a.
TrieMap m =>
(a -> Bool) -> GenMap m a -> GenMap m a
ftG

instance (Eq (Key m), TrieMap m) => Foldable (GenMap m) where
  foldMap :: forall m a. Monoid m => (a -> m) -> GenMap m a -> m
foldMap = forall (m :: * -> *) r a.
(TrieMap m, Monoid r) =>
(a -> r) -> m a -> r
foldMapTM

--We want to be able to specialize these functions when defining eg
--tries over (GenMap CoreExpr) which requires INLINEABLE

{-# INLINEABLE lkG #-}
lkG :: (Eq (Key m), TrieMap m) => Key m -> GenMap m a -> Maybe a
lkG :: forall (m :: * -> *) a.
(Eq (Key m), TrieMap m) =>
Key m -> GenMap m a -> Maybe a
lkG Key m
_ GenMap m a
EmptyMap                         = forall a. Maybe a
Nothing
lkG Key m
k (SingletonMap Key m
k' a
v') | Key m
k forall a. Eq a => a -> a -> Bool
== Key m
k'   = forall a. a -> Maybe a
Just a
v'
                           | Bool
otherwise = forall a. Maybe a
Nothing
lkG Key m
k (MultiMap m a
m)                     = forall (m :: * -> *) b. TrieMap m => Key m -> m b -> Maybe b
lookupTM Key m
k m a
m

{-# INLINEABLE xtG #-}
xtG :: (Eq (Key m), TrieMap m) => Key m -> XT a -> GenMap m a -> GenMap m a
xtG :: forall (m :: * -> *) a.
(Eq (Key m), TrieMap m) =>
Key m -> XT a -> GenMap m a -> GenMap m a
xtG Key m
k XT a
f GenMap m a
EmptyMap
    = case XT a
f forall a. Maybe a
Nothing of
        Just a
v  -> forall (m :: * -> *) a. Key m -> a -> GenMap m a
SingletonMap Key m
k a
v
        Maybe a
Nothing -> forall (m :: * -> *) a. GenMap m a
EmptyMap
xtG Key m
k XT a
f m :: GenMap m a
m@(SingletonMap Key m
k' a
v')
    | Key m
k' forall a. Eq a => a -> a -> Bool
== Key m
k
    -- The new key matches the (single) key already in the tree.  Hence,
    -- apply @f@ to @Just v'@ and build a singleton or empty map depending
    -- on the 'Just'/'Nothing' response respectively.
    = case XT a
f (forall a. a -> Maybe a
Just a
v') of
        Just a
v'' -> forall (m :: * -> *) a. Key m -> a -> GenMap m a
SingletonMap Key m
k' a
v''
        Maybe a
Nothing  -> forall (m :: * -> *) a. GenMap m a
EmptyMap
    | Bool
otherwise
    -- We've hit a singleton tree for a different key than the one we are
    -- searching for. Hence apply @f@ to @Nothing@. If result is @Nothing@ then
    -- we can just return the old map. If not, we need a map with *two*
    -- entries. The easiest way to do that is to insert two items into an empty
    -- map of type @m a@.
    = case XT a
f forall a. Maybe a
Nothing of
        Maybe a
Nothing  -> GenMap m a
m
        Just a
v   -> forall (m :: * -> *) a. TrieMap m => m a
emptyTM forall a b. a -> (a -> b) -> b
|> forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
alterTM Key m
k' (forall a b. a -> b -> a
const (forall a. a -> Maybe a
Just a
v'))
                           forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
alterTM Key m
k  (forall a b. a -> b -> a
const (forall a. a -> Maybe a
Just a
v))
                           forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall (m :: * -> *) a. m a -> GenMap m a
MultiMap
xtG Key m
k XT a
f (MultiMap m a
m) = forall (m :: * -> *) a. m a -> GenMap m a
MultiMap (forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
alterTM Key m
k XT a
f m a
m)

{-# INLINEABLE mapG #-}
mapG :: Functor m => (a -> b) -> GenMap m a -> GenMap m b
mapG :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> GenMap m a -> GenMap m b
mapG a -> b
_ GenMap m a
EmptyMap = forall (m :: * -> *) a. GenMap m a
EmptyMap
mapG a -> b
f (SingletonMap Key m
k a
v) = forall (m :: * -> *) a. Key m -> a -> GenMap m a
SingletonMap Key m
k (a -> b
f a
v)
mapG a -> b
f (MultiMap m a
m) = forall (m :: * -> *) a. m a -> GenMap m a
MultiMap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f m a
m)

{-# INLINEABLE fdG #-}
fdG :: TrieMap m => (a -> b -> b) -> GenMap m a -> b -> b
fdG :: forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> GenMap m a -> b -> b
fdG a -> b -> b
_ GenMap m a
EmptyMap = \b
z -> b
z
fdG a -> b -> b
k (SingletonMap Key m
_ a
v) = \b
z -> a -> b -> b
k a
v b
z
fdG a -> b -> b
k (MultiMap m a
m) = forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM a -> b -> b
k m a
m

{-# INLINEABLE ftG #-}
ftG :: TrieMap m => (a -> Bool) -> GenMap m a -> GenMap m a
ftG :: forall (m :: * -> *) a.
TrieMap m =>
(a -> Bool) -> GenMap m a -> GenMap m a
ftG a -> Bool
_ GenMap m a
EmptyMap = forall (m :: * -> *) a. GenMap m a
EmptyMap
ftG a -> Bool
f input :: GenMap m a
input@(SingletonMap Key m
_ a
v)
  | a -> Bool
f a
v       = GenMap m a
input
  | Bool
otherwise = forall (m :: * -> *) a. GenMap m a
EmptyMap
ftG a -> Bool
f (MultiMap m a
m) = forall (m :: * -> *) a. m a -> GenMap m a
MultiMap (forall (m :: * -> *) a. TrieMap m => (a -> Bool) -> m a -> m a
filterTM a -> Bool
f m a
m)
  -- we don't have enough information to reconstruct the key to make
  -- a SingletonMap