{-# LANGUAGE DeriveTraversable #-}

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

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

-- SPDX-License-Identifier: BSD-3-Clause

-- |
-- Module      :  Disco.Context
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- A *context* is a mapping from names to other things (such as types
-- or values).  This module defines a generic type of contexts which
-- is used in many different places throughout the disco codebase.
module Disco.Context (
  -- * Context type
  Ctx,

  -- * Construction
  emptyCtx,
  singleCtx,
  fromList,
  ctxForModule,
  localCtx,

  -- * Insertion
  insert,
  extend,
  extends,

  -- * Query
  null,
  lookup,
  lookup',
  lookupNonLocal,
  lookupNonLocal',
  lookupAll,
  lookupAll',

  -- * Conversion
  names,
  elems,
  assocs,
  keysSet,

  -- * Traversal
  coerceKeys,
  restrictKeys,

  -- * Combination
  joinCtx,
  joinCtxs,

  -- * Filter
  filter,
) where

import Control.Monad ((<=<))
import Data.Bifunctor (first, second)
import Data.Coerce
import Data.Map (Map)
import qualified Data.Map as M
import Data.Map.Merge.Lazy as MM
import Data.Set (Set)
import qualified Data.Set as S
import Prelude hiding (filter, lookup, null)

import Unbound.Generics.LocallyNameless (Name)

import Polysemy
import Polysemy.Reader

import Disco.Names (
  ModuleName,
  NameProvenance (..),
  QName (..),
 )

-- | A context maps qualified names to things.  In particular a @Ctx a
--   b@ maps qualified names for @a@s to values of type @b@.
newtype Ctx a b = Ctx {forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx :: M.Map NameProvenance (M.Map (Name a) b)}
  deriving (Ctx a b -> Ctx a b -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall a b. Eq b => Ctx a b -> Ctx a b -> Bool
/= :: Ctx a b -> Ctx a b -> Bool
$c/= :: forall a b. Eq b => Ctx a b -> Ctx a b -> Bool
== :: Ctx a b -> Ctx a b -> Bool
$c== :: forall a b. Eq b => Ctx a b -> Ctx a b -> Bool
Eq, Int -> Ctx a b -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall a b. Show b => Int -> Ctx a b -> ShowS
forall a b. Show b => [Ctx a b] -> ShowS
forall a b. Show b => Ctx a b -> String
showList :: [Ctx a b] -> ShowS
$cshowList :: forall a b. Show b => [Ctx a b] -> ShowS
show :: Ctx a b -> String
$cshow :: forall a b. Show b => Ctx a b -> String
showsPrec :: Int -> Ctx a b -> ShowS
$cshowsPrec :: forall a b. Show b => Int -> Ctx a b -> ShowS
Show, forall a b. a -> Ctx a b -> Ctx a a
forall a b. (a -> b) -> Ctx a a -> Ctx a b
forall a a b. a -> Ctx a b -> Ctx a a
forall a a b. (a -> b) -> Ctx a a -> Ctx a 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 -> Ctx a b -> Ctx a a
$c<$ :: forall a a b. a -> Ctx a b -> Ctx a a
fmap :: forall a b. (a -> b) -> Ctx a a -> Ctx a b
$cfmap :: forall a a b. (a -> b) -> Ctx a a -> Ctx a b
Functor, forall a. Ctx a a -> Bool
forall a a. Eq a => a -> Ctx a a -> Bool
forall a a. Num a => Ctx a a -> a
forall a a. Ord a => Ctx a a -> a
forall m a. Monoid m => (a -> m) -> Ctx a a -> m
forall a m. Monoid m => Ctx a m -> m
forall a a. Ctx a a -> Bool
forall a a. Ctx a a -> Int
forall a a. Ctx a a -> [a]
forall a b. (a -> b -> b) -> b -> Ctx a a -> b
forall a a. (a -> a -> a) -> Ctx a a -> a
forall a m a. Monoid m => (a -> m) -> Ctx a a -> m
forall a b a. (b -> a -> b) -> b -> Ctx a a -> b
forall a a b. (a -> b -> b) -> b -> Ctx a a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => Ctx a a -> a
$cproduct :: forall a a. Num a => Ctx a a -> a
sum :: forall a. Num a => Ctx a a -> a
$csum :: forall a a. Num a => Ctx a a -> a
minimum :: forall a. Ord a => Ctx a a -> a
$cminimum :: forall a a. Ord a => Ctx a a -> a
maximum :: forall a. Ord a => Ctx a a -> a
$cmaximum :: forall a a. Ord a => Ctx a a -> a
elem :: forall a. Eq a => a -> Ctx a a -> Bool
$celem :: forall a a. Eq a => a -> Ctx a a -> Bool
length :: forall a. Ctx a a -> Int
$clength :: forall a a. Ctx a a -> Int
null :: forall a. Ctx a a -> Bool
$cnull :: forall a a. Ctx a a -> Bool
toList :: forall a. Ctx a a -> [a]
$ctoList :: forall a a. Ctx a a -> [a]
foldl1 :: forall a. (a -> a -> a) -> Ctx a a -> a
$cfoldl1 :: forall a a. (a -> a -> a) -> Ctx a a -> a
foldr1 :: forall a. (a -> a -> a) -> Ctx a a -> a
$cfoldr1 :: forall a a. (a -> a -> a) -> Ctx a a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> Ctx a a -> b
$cfoldl' :: forall a b a. (b -> a -> b) -> b -> Ctx a a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Ctx a a -> b
$cfoldl :: forall a b a. (b -> a -> b) -> b -> Ctx a a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Ctx a a -> b
$cfoldr' :: forall a a b. (a -> b -> b) -> b -> Ctx a a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Ctx a a -> b
$cfoldr :: forall a a b. (a -> b -> b) -> b -> Ctx a a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> Ctx a a -> m
$cfoldMap' :: forall a m a. Monoid m => (a -> m) -> Ctx a a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Ctx a a -> m
$cfoldMap :: forall a m a. Monoid m => (a -> m) -> Ctx a a -> m
fold :: forall m. Monoid m => Ctx a m -> m
$cfold :: forall a m. Monoid m => Ctx a m -> m
Foldable, forall a. Functor (Ctx a)
forall a. Foldable (Ctx a)
forall a (m :: * -> *) a. Monad m => Ctx a (m a) -> m (Ctx a a)
forall a (f :: * -> *) a.
Applicative f =>
Ctx a (f a) -> f (Ctx a a)
forall a (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Ctx a a -> m (Ctx a b)
forall a (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Ctx a a -> f (Ctx a b)
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Ctx a a -> f (Ctx a b)
sequence :: forall (m :: * -> *) a. Monad m => Ctx a (m a) -> m (Ctx a a)
$csequence :: forall a (m :: * -> *) a. Monad m => Ctx a (m a) -> m (Ctx a a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Ctx a a -> m (Ctx a b)
$cmapM :: forall a (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Ctx a a -> m (Ctx a b)
sequenceA :: forall (f :: * -> *) a. Applicative f => Ctx a (f a) -> f (Ctx a a)
$csequenceA :: forall a (f :: * -> *) a.
Applicative f =>
Ctx a (f a) -> f (Ctx a a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Ctx a a -> f (Ctx a b)
$ctraverse :: forall a (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Ctx a a -> f (Ctx a b)
Traversable)

-- Note that we implement a context as a nested map from
-- NameProvenance to Name to b, rather than as a Map QName b.  They
-- are isomorphic, but this way it is easier to do name resolution,
-- because given an (unqualified) Name, we can look it up in each
-- inner map corresponding to modules that are in scope.

instance Semigroup (Ctx a b) where
  <> :: Ctx a b -> Ctx a b -> Ctx a b
(<>) = forall a b. Ctx a b -> Ctx a b -> Ctx a b
joinCtx

instance Monoid (Ctx a b) where
  mempty :: Ctx a b
mempty = forall a b. Ctx a b
emptyCtx
  mappend :: Ctx a b -> Ctx a b -> Ctx a b
mappend = forall a. Semigroup a => a -> a -> a
(<>)

------------------------------------------------------------
-- Construction
------------------------------------------------------------

-- | The empty context.
emptyCtx :: Ctx a b
emptyCtx :: forall a b. Ctx a b
emptyCtx = forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx forall k a. Map k a
M.empty

-- | A singleton context, mapping a qualified name to a thing.
singleCtx :: QName a -> b -> Ctx a b
singleCtx :: forall a b. QName a -> b -> Ctx a b
singleCtx (QName NameProvenance
p Name a
n) = forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. k -> a -> Map k a
M.singleton NameProvenance
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. k -> a -> Map k a
M.singleton Name a
n

-- | Create a context from a list of (qualified name, value) pairs.
fromList :: [(QName a, b)] -> Ctx a b
fromList :: forall a b. [(QName a, b)] -> Ctx a b
fromList = forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\(QName NameProvenance
p Name a
n, b
b) -> (NameProvenance
p, forall k a. k -> a -> Map k a
M.singleton Name a
n b
b))

-- | Create a context for bindings from a single module.
ctxForModule :: ModuleName -> [(Name a, b)] -> Ctx a b
ctxForModule :: forall a b. ModuleName -> [(Name a, b)] -> Ctx a b
ctxForModule ModuleName
m = forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. k -> a -> Map k a
M.singleton (ModuleName -> NameProvenance
QualifiedName ModuleName
m) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList

-- | Create a context with local bindings.
localCtx :: [(Name a, b)] -> Ctx a b
localCtx :: forall a b. [(Name a, b)] -> Ctx a b
localCtx = forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. k -> a -> Map k a
M.singleton NameProvenance
LocalName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList

------------------------------------------------------------
-- Insertion
------------------------------------------------------------

-- | Insert a new binding into a context.  The new binding shadows any
--   old binding for the same qualified name.
insert :: QName a -> b -> Ctx a b -> Ctx a b
insert :: forall a b. QName a -> b -> Ctx a b -> Ctx a b
insert (QName NameProvenance
p Name a
n) b
b = forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union NameProvenance
p (forall k a. k -> a -> Map k a
M.singleton Name a
n b
b) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx

-- | Run a computation under a context extended with a new binding.
--   The new binding shadows any old binding for the same name.
extend :: Member (Reader (Ctx a b)) r => QName a -> b -> Sem r c -> Sem r c
extend :: forall a b (r :: EffectRow) c.
Member (Reader (Ctx a b)) r =>
QName a -> b -> Sem r c -> Sem r c
extend QName a
qn b
b = forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local (forall a b. QName a -> b -> Ctx a b -> Ctx a b
insert QName a
qn b
b)

-- | Run a computation in a context extended with an additional
--   context.  Bindings in the additional context shadow any bindings
--   with the same names in the existing context.
extends :: Member (Reader (Ctx a b)) r => Ctx a b -> Sem r c -> Sem r c
extends :: forall a b (r :: EffectRow) c.
Member (Reader (Ctx a b)) r =>
Ctx a b -> Sem r c -> Sem r c
extends = forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Ctx a b -> Ctx a b
joinCtx

------------------------------------------------------------
-- Query
------------------------------------------------------------

-- | Check if a context is empty.
null :: Ctx a b -> Bool
null :: forall a a. Ctx a a -> Bool
null = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall k a. Map k a -> Bool
M.null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx

-- | Look up a qualified name in an ambient context.
lookup :: Member (Reader (Ctx a b)) r => QName a -> Sem r (Maybe b)
lookup :: forall a b (r :: EffectRow).
Member (Reader (Ctx a b)) r =>
QName a -> Sem r (Maybe b)
lookup QName a
x = forall a b. QName a -> Ctx a b -> Maybe b
lookup' QName a
x forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall i (r :: EffectRow). Member (Reader i) r => Sem r i
ask

-- | Look up a qualified name in a context.
lookup' :: QName a -> Ctx a b -> Maybe b
lookup' :: forall a b. QName a -> Ctx a b -> Maybe b
lookup' (QName NameProvenance
p Name a
n) = (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name a
n forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup NameProvenance
p) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx

-- | Look up all the non-local bindings of a name in an ambient context.
lookupNonLocal :: Member (Reader (Ctx a b)) r => Name a -> Sem r [(ModuleName, b)]
lookupNonLocal :: forall a b (r :: EffectRow).
Member (Reader (Ctx a b)) r =>
Name a -> Sem r [(ModuleName, b)]
lookupNonLocal Name a
n = forall a b. Name a -> Ctx a b -> [(ModuleName, b)]
lookupNonLocal' Name a
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall i (r :: EffectRow). Member (Reader i) r => Sem r i
ask

-- | Look up all the non-local bindings of a name in a context.
lookupNonLocal' :: Name a -> Ctx a b -> [(ModuleName, b)]
lookupNonLocal' :: forall a b. Name a -> Ctx a b -> [(ModuleName, b)]
lookupNonLocal' Name a
n = forall {a} {b}. [(QName a, b)] -> [(ModuleName, b)]
nonLocal forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Name a -> Ctx a b -> [(QName a, b)]
lookupAll' Name a
n
 where
  nonLocal :: [(QName a, b)] -> [(ModuleName, b)]
nonLocal [(QName a, b)]
bs = [(ModuleName
m, b
b) | (QName (QualifiedName ModuleName
m) Name a
_, b
b) <- [(QName a, b)]
bs]

-- | Look up all the bindings of an (unqualified) name in an ambient context.
lookupAll :: Member (Reader (Ctx a b)) r => Name a -> Sem r [(QName a, b)]
lookupAll :: forall a b (r :: EffectRow).
Member (Reader (Ctx a b)) r =>
Name a -> Sem r [(QName a, b)]
lookupAll Name a
n = forall a b. Name a -> Ctx a b -> [(QName a, b)]
lookupAll' Name a
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall i (r :: EffectRow). Member (Reader i) r => Sem r i
ask

-- | Look up all the bindings of an (unqualified) name in a context.
lookupAll' :: Name a -> Ctx a b -> [(QName a, b)]
lookupAll' :: forall a b. Name a -> Ctx a b -> [(QName a, b)]
lookupAll' Name a
n = forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a. NameProvenance -> Name a -> QName a
`QName` Name a
n)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.assocs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name a
n) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx

------------------------------------------------------------
-- Conversion
------------------------------------------------------------

-- | Return a list of the names defined by the context.
names :: Ctx a b -> [Name a]
names :: forall a b. Ctx a b -> [Name a]
names = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall k a. Map k a -> [k]
M.keys forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [a]
M.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx

-- | Return a list of all the values bound by the context.
elems :: Ctx a b -> [b]
elems :: forall a a. Ctx a a -> [a]
elems = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall k a. Map k a -> [a]
M.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [a]
M.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx

-- | Return a list of the qualified name-value associations in the
--   context.
assocs :: Ctx a b -> [(QName a, b)]
assocs :: forall a b. Ctx a b -> [(QName a, b)]
assocs = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a b. NameProvenance -> Map (Name a) b -> [(QName a, b)]
modAssocs) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.assocs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx
 where
  modAssocs :: NameProvenance -> Map (Name a) b -> [(QName a, b)]
  modAssocs :: forall a b. NameProvenance -> Map (Name a) b -> [(QName a, b)]
modAssocs NameProvenance
p = forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a. NameProvenance -> Name a -> QName a
QName NameProvenance
p)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.assocs

-- | Return a set of all qualified names in the context.
keysSet :: Ctx a b -> Set (QName a)
keysSet :: forall a b. Ctx a b -> Set (QName a)
keysSet = forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NameProvenance -> Name a -> QName a
QName) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall k a. Map k a -> Set k
M.keysSet) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.assocs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx

------------------------------------------------------------
-- Traversal
------------------------------------------------------------

-- | Coerce the type of the qualified name keys in a context.
coerceKeys :: Ctx a1 b -> Ctx a2 b
coerceKeys :: forall a1 b a2. Ctx a1 b -> Ctx a2 b
coerceKeys = forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys coerce :: forall a b. Coercible a b => a -> b
coerce) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx

-- | Restrict a context to only the keys in the given set.
restrictKeys :: Ctx a b -> Set (QName a) -> Ctx a b
restrictKeys :: forall a b. Ctx a b -> Set (QName a) -> Ctx a b
restrictKeys Ctx a b
ctx Set (QName a)
xs = forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a}.
Map NameProvenance (Set (Name a))
-> Map NameProvenance (Map (Name a) a)
-> Map NameProvenance (Map (Name a) a)
restrict Map NameProvenance (Set (Name a))
m forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx forall a b. (a -> b) -> a -> b
$ Ctx a b
ctx
 where
  restrict :: Map NameProvenance (Set (Name a))
-> Map NameProvenance (Map (Name a) a)
-> Map NameProvenance (Map (Name a) a)
restrict = forall k a c b.
Ord k =>
SimpleWhenMissing k a c
-> SimpleWhenMissing k b c
-> SimpleWhenMatched k a b c
-> Map k a
-> Map k b
-> Map k c
MM.merge forall (f :: * -> *) k x y. Applicative f => WhenMissing f k x y
MM.dropMissing forall (f :: * -> *) k x y. Applicative f => WhenMissing f k x y
MM.dropMissing (forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> z) -> WhenMatched f k x y z
MM.zipWithMatched (\NameProvenance
_ Set (Name a)
ns Map (Name a) a
m' -> forall k a. Ord k => Map k a -> Set k -> Map k a
M.restrictKeys Map (Name a) a
m' Set (Name a)
ns))
  m :: Map NameProvenance (Set (Name a))
m = forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith forall a. Ord a => Set a -> Set a -> Set a
S.union forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\(QName NameProvenance
p Name a
n) -> (NameProvenance
p, forall a. a -> Set a
S.singleton Name a
n)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$ Set (QName a)
xs

------------------------------------------------------------
-- Combination
------------------------------------------------------------

-- | Join two contexts (left-biased, /i.e./ if the same qualified name
--   exists in both contexts, the result will use the value from the
--   first context, and throw away the value from the second.).
joinCtx :: Ctx a b -> Ctx a b -> Ctx a b
joinCtx :: forall a b. Ctx a b -> Ctx a b -> Ctx a b
joinCtx Ctx a b
a Ctx a b
b = forall a b. [Ctx a b] -> Ctx a b
joinCtxs [Ctx a b
a, Ctx a b
b]

-- | Join a list of contexts (left-biased).
joinCtxs :: [Ctx a b] -> Ctx a b
joinCtxs :: forall a b. [Ctx a b] -> Ctx a b
joinCtxs = forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
M.unionsWith forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx

------------------------------------------------------------
-- Filter
------------------------------------------------------------

-- | Filter a context using a predicate.
filter :: (b -> Bool) -> Ctx a b -> Ctx a b
filter :: forall b a. (b -> Bool) -> Ctx a b -> Ctx a b
filter b -> Bool
p = forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter b -> Bool
p) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx