{-# LANGUAGE DeriveTraversable #-}

-----------------------------------------------------------------------------
-- |
-- 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.
--
-----------------------------------------------------------------------------

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

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 { 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
(Ctx a b -> Ctx a b -> Bool)
-> (Ctx a b -> Ctx a b -> Bool) -> Eq (Ctx a b)
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
[Ctx a b] -> ShowS
Ctx a b -> String
(Int -> Ctx a b -> ShowS)
-> (Ctx a b -> String) -> ([Ctx a b] -> ShowS) -> Show (Ctx a b)
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, a -> Ctx a b -> Ctx a a
(a -> b) -> Ctx a a -> Ctx a b
(forall a b. (a -> b) -> Ctx a a -> Ctx a b)
-> (forall a b. a -> Ctx a b -> Ctx a a) -> Functor (Ctx a)
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
<$ :: a -> Ctx a b -> Ctx a a
$c<$ :: forall a a b. a -> Ctx a b -> Ctx a a
fmap :: (a -> b) -> Ctx a a -> Ctx a b
$cfmap :: forall a a b. (a -> b) -> Ctx a a -> Ctx a b
Functor, Ctx a a -> Bool
(a -> m) -> Ctx a a -> m
(a -> b -> b) -> b -> Ctx a a -> b
(forall m. Monoid m => Ctx a m -> m)
-> (forall m a. Monoid m => (a -> m) -> Ctx a a -> m)
-> (forall m a. Monoid m => (a -> m) -> Ctx a a -> m)
-> (forall a b. (a -> b -> b) -> b -> Ctx a a -> b)
-> (forall a b. (a -> b -> b) -> b -> Ctx a a -> b)
-> (forall b a. (b -> a -> b) -> b -> Ctx a a -> b)
-> (forall b a. (b -> a -> b) -> b -> Ctx a a -> b)
-> (forall a. (a -> a -> a) -> Ctx a a -> a)
-> (forall a. (a -> a -> a) -> Ctx a a -> a)
-> (forall a. Ctx a a -> [a])
-> (forall a. Ctx a a -> Bool)
-> (forall a. Ctx a a -> Int)
-> (forall a. Eq a => a -> Ctx a a -> Bool)
-> (forall a. Ord a => Ctx a a -> a)
-> (forall a. Ord a => Ctx a a -> a)
-> (forall a. Num a => Ctx a a -> a)
-> (forall a. Num a => Ctx a a -> a)
-> Foldable (Ctx a)
forall a. Eq a => a -> Ctx a a -> Bool
forall a. Num a => Ctx a a -> a
forall a. Ord a => Ctx a a -> a
forall m. Monoid m => Ctx a m -> m
forall a. Ctx a a -> Bool
forall a. Ctx a a -> Int
forall a. Ctx a a -> [a]
forall a. (a -> a -> a) -> Ctx a a -> a
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 b a. (b -> a -> b) -> b -> Ctx a a -> b
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 :: Ctx a a -> a
$cproduct :: forall a a. Num a => Ctx a a -> a
sum :: Ctx a a -> a
$csum :: forall a a. Num a => Ctx a a -> a
minimum :: Ctx a a -> a
$cminimum :: forall a a. Ord a => Ctx a a -> a
maximum :: Ctx a a -> a
$cmaximum :: forall a a. Ord a => Ctx a a -> a
elem :: a -> Ctx a a -> Bool
$celem :: forall a a. Eq a => a -> Ctx a a -> Bool
length :: Ctx a a -> Int
$clength :: forall a a. Ctx a a -> Int
null :: Ctx a a -> Bool
$cnull :: forall a a. Ctx a a -> Bool
toList :: Ctx a a -> [a]
$ctoList :: forall a a. Ctx a a -> [a]
foldl1 :: (a -> a -> a) -> Ctx a a -> a
$cfoldl1 :: forall a a. (a -> a -> a) -> Ctx a a -> a
foldr1 :: (a -> a -> a) -> Ctx a a -> a
$cfoldr1 :: forall a a. (a -> a -> a) -> Ctx a a -> a
foldl' :: (b -> a -> b) -> b -> Ctx a a -> b
$cfoldl' :: forall a b a. (b -> a -> b) -> b -> Ctx a a -> b
foldl :: (b -> a -> b) -> b -> Ctx a a -> b
$cfoldl :: forall a b a. (b -> a -> b) -> b -> Ctx a a -> b
foldr' :: (a -> b -> b) -> b -> Ctx a a -> b
$cfoldr' :: forall a a b. (a -> b -> b) -> b -> Ctx a a -> b
foldr :: (a -> b -> b) -> b -> Ctx a a -> b
$cfoldr :: forall a a b. (a -> b -> b) -> b -> Ctx a a -> b
foldMap' :: (a -> m) -> Ctx a a -> m
$cfoldMap' :: forall a m a. Monoid m => (a -> m) -> Ctx a a -> m
foldMap :: (a -> m) -> Ctx a a -> m
$cfoldMap :: forall a m a. Monoid m => (a -> m) -> Ctx a a -> m
fold :: Ctx a m -> m
$cfold :: forall a m. Monoid m => Ctx a m -> m
Foldable, Functor (Ctx a)
Foldable (Ctx a)
Functor (Ctx a)
-> Foldable (Ctx a)
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> Ctx a a -> f (Ctx a b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    Ctx a (f a) -> f (Ctx a a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> Ctx a a -> m (Ctx a b))
-> (forall (m :: * -> *) a. Monad m => Ctx a (m a) -> m (Ctx a a))
-> Traversable (Ctx a)
(a -> f b) -> Ctx a a -> f (Ctx a b)
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 (m :: * -> *) a. Monad m => Ctx a (m a) -> m (Ctx a a)
forall (f :: * -> *) a. Applicative f => Ctx a (f a) -> f (Ctx a a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Ctx a a -> m (Ctx a b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Ctx a a -> f (Ctx a b)
sequence :: Ctx a (m a) -> m (Ctx a a)
$csequence :: forall a (m :: * -> *) a. Monad m => Ctx a (m a) -> m (Ctx a a)
mapM :: (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 :: Ctx a (f a) -> f (Ctx a a)
$csequenceA :: forall a (f :: * -> *) a.
Applicative f =>
Ctx a (f a) -> f (Ctx a a)
traverse :: (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)
$cp2Traversable :: forall a. Foldable (Ctx a)
$cp1Traversable :: forall a. Functor (Ctx a)
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
(<>) = 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 = Ctx a b
forall a b. Ctx a b
emptyCtx
  mappend :: Ctx a b -> Ctx a b -> Ctx a b
mappend = Ctx a b -> Ctx a b -> Ctx a b
forall a. Semigroup a => a -> a -> a
(<>)

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

-- | The empty context.
emptyCtx :: Ctx a b
emptyCtx :: Ctx a b
emptyCtx = Map NameProvenance (Map (Name a) b) -> Ctx a b
forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx Map NameProvenance (Map (Name a) b)
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 :: QName a -> b -> Ctx a b
singleCtx (QName NameProvenance
p Name a
n) = Map NameProvenance (Map (Name a) b) -> Ctx a b
forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx (Map NameProvenance (Map (Name a) b) -> Ctx a b)
-> (b -> Map NameProvenance (Map (Name a) b)) -> b -> Ctx a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NameProvenance
-> Map (Name a) b -> Map NameProvenance (Map (Name a) b)
forall k a. k -> a -> Map k a
M.singleton NameProvenance
p (Map (Name a) b -> Map NameProvenance (Map (Name a) b))
-> (b -> Map (Name a) b)
-> b
-> Map NameProvenance (Map (Name a) b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name a -> b -> Map (Name a) b
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 :: [(QName a, b)] -> Ctx a b
fromList = Map NameProvenance (Map (Name a) b) -> Ctx a b
forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx (Map NameProvenance (Map (Name a) b) -> Ctx a b)
-> ([(QName a, b)] -> Map NameProvenance (Map (Name a) b))
-> [(QName a, b)]
-> Ctx a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map (Name a) b -> Map (Name a) b -> Map (Name a) b)
-> [(NameProvenance, Map (Name a) b)]
-> Map NameProvenance (Map (Name a) b)
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Map (Name a) b -> Map (Name a) b -> Map (Name a) b
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union ([(NameProvenance, Map (Name a) b)]
 -> Map NameProvenance (Map (Name a) b))
-> ([(QName a, b)] -> [(NameProvenance, Map (Name a) b)])
-> [(QName a, b)]
-> Map NameProvenance (Map (Name a) b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((QName a, b) -> (NameProvenance, Map (Name a) b))
-> [(QName a, b)] -> [(NameProvenance, Map (Name a) b)]
forall a b. (a -> b) -> [a] -> [b]
map (\(QName NameProvenance
p Name a
n, b
b) -> (NameProvenance
p, Name a -> b -> Map (Name a) b
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 :: ModuleName -> [(Name a, b)] -> Ctx a b
ctxForModule ModuleName
m = Map NameProvenance (Map (Name a) b) -> Ctx a b
forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx (Map NameProvenance (Map (Name a) b) -> Ctx a b)
-> ([(Name a, b)] -> Map NameProvenance (Map (Name a) b))
-> [(Name a, b)]
-> Ctx a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NameProvenance
-> Map (Name a) b -> Map NameProvenance (Map (Name a) b)
forall k a. k -> a -> Map k a
M.singleton (ModuleName -> NameProvenance
QualifiedName ModuleName
m) (Map (Name a) b -> Map NameProvenance (Map (Name a) b))
-> ([(Name a, b)] -> Map (Name a) b)
-> [(Name a, b)]
-> Map NameProvenance (Map (Name a) b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name a, b)] -> Map (Name a) b
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 :: [(Name a, b)] -> Ctx a b
localCtx = Map NameProvenance (Map (Name a) b) -> Ctx a b
forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx (Map NameProvenance (Map (Name a) b) -> Ctx a b)
-> ([(Name a, b)] -> Map NameProvenance (Map (Name a) b))
-> [(Name a, b)]
-> Ctx a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NameProvenance
-> Map (Name a) b -> Map NameProvenance (Map (Name a) b)
forall k a. k -> a -> Map k a
M.singleton NameProvenance
LocalName (Map (Name a) b -> Map NameProvenance (Map (Name a) b))
-> ([(Name a, b)] -> Map (Name a) b)
-> [(Name a, b)]
-> Map NameProvenance (Map (Name a) b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name a, b)] -> Map (Name a) b
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 :: QName a -> b -> Ctx a b -> Ctx a b
insert (QName NameProvenance
p Name a
n) b
b = Map NameProvenance (Map (Name a) b) -> Ctx a b
forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx (Map NameProvenance (Map (Name a) b) -> Ctx a b)
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> Ctx a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map (Name a) b -> Map (Name a) b -> Map (Name a) b)
-> NameProvenance
-> Map (Name a) b
-> Map NameProvenance (Map (Name a) b)
-> Map NameProvenance (Map (Name a) b)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Map (Name a) b -> Map (Name a) b -> Map (Name a) b
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union NameProvenance
p (Name a -> b -> Map (Name a) b
forall k a. k -> a -> Map k a
M.singleton Name a
n b
b) (Map NameProvenance (Map (Name a) b)
 -> Map NameProvenance (Map (Name a) b))
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> Map NameProvenance (Map (Name a) b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Map NameProvenance (Map (Name a) b)
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 :: QName a -> b -> Sem r c -> Sem r c
extend QName a
qn b
b = (Ctx a b -> Ctx a b) -> Sem r c -> Sem r c
forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local (QName a -> b -> Ctx a b -> Ctx a b
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 :: Ctx a b -> Sem r c -> Sem r c
extends = (Ctx a b -> Ctx a b) -> Sem r c -> Sem r c
forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local ((Ctx a b -> Ctx a b) -> Sem r c -> Sem r c)
-> (Ctx a b -> Ctx a b -> Ctx a b) -> Ctx a b -> Sem r c -> Sem r c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Ctx a b -> Ctx a b
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 :: Ctx a b -> Bool
null = (Map (Name a) b -> Bool)
-> Map NameProvenance (Map (Name a) b) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Map (Name a) b -> Bool
forall k a. Map k a -> Bool
M.null (Map NameProvenance (Map (Name a) b) -> Bool)
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Map NameProvenance (Map (Name a) b)
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 :: QName a -> Sem r (Maybe b)
lookup QName a
x = QName a -> Ctx a b -> Maybe b
forall a b. QName a -> Ctx a b -> Maybe b
lookup' QName a
x (Ctx a b -> Maybe b) -> Sem r (Ctx a b) -> Sem r (Maybe b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sem r (Ctx a 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' :: QName a -> Ctx a b -> Maybe b
lookup' (QName NameProvenance
p Name a
n) = (Name a -> Map (Name a) b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name a
n (Map (Name a) b -> Maybe b)
-> (Map NameProvenance (Map (Name a) b) -> Maybe (Map (Name a) b))
-> Map NameProvenance (Map (Name a) b)
-> Maybe b
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< NameProvenance
-> Map NameProvenance (Map (Name a) b) -> Maybe (Map (Name a) b)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup NameProvenance
p) (Map NameProvenance (Map (Name a) b) -> Maybe b)
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> Maybe b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Map NameProvenance (Map (Name a) b)
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 :: Name a -> Sem r [(ModuleName, b)]
lookupNonLocal Name a
n = Name a -> Ctx a b -> [(ModuleName, b)]
forall a b. Name a -> Ctx a b -> [(ModuleName, b)]
lookupNonLocal' Name a
n (Ctx a b -> [(ModuleName, b)])
-> Sem r (Ctx a b) -> Sem r [(ModuleName, b)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sem r (Ctx a 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' :: Name a -> Ctx a b -> [(ModuleName, b)]
lookupNonLocal' Name a
n = [(QName a, b)] -> [(ModuleName, b)]
forall a b. [(QName a, b)] -> [(ModuleName, b)]
nonLocal ([(QName a, b)] -> [(ModuleName, b)])
-> (Ctx a b -> [(QName a, b)]) -> Ctx a b -> [(ModuleName, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name a -> Ctx a b -> [(QName a, b)]
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 :: Name a -> Sem r [(QName a, b)]
lookupAll Name a
n = Name a -> Ctx a b -> [(QName a, b)]
forall a b. Name a -> Ctx a b -> [(QName a, b)]
lookupAll' Name a
n (Ctx a b -> [(QName a, b)])
-> Sem r (Ctx a b) -> Sem r [(QName a, b)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sem r (Ctx a 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' :: Name a -> Ctx a b -> [(QName a, b)]
lookupAll' Name a
n = ((NameProvenance, b) -> (QName a, b))
-> [(NameProvenance, b)] -> [(QName a, b)]
forall a b. (a -> b) -> [a] -> [b]
map ((NameProvenance -> QName a) -> (NameProvenance, b) -> (QName a, b)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (NameProvenance -> Name a -> QName a
forall a. NameProvenance -> Name a -> QName a
`QName` Name a
n)) ([(NameProvenance, b)] -> [(QName a, b)])
-> (Ctx a b -> [(NameProvenance, b)]) -> Ctx a b -> [(QName a, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map NameProvenance b -> [(NameProvenance, b)]
forall k a. Map k a -> [(k, a)]
M.assocs (Map NameProvenance b -> [(NameProvenance, b)])
-> (Ctx a b -> Map NameProvenance b)
-> Ctx a b
-> [(NameProvenance, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map (Name a) b -> Maybe b)
-> Map NameProvenance (Map (Name a) b) -> Map NameProvenance b
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe (Name a -> Map (Name a) b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name a
n) (Map NameProvenance (Map (Name a) b) -> Map NameProvenance b)
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> Map NameProvenance b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Map NameProvenance (Map (Name a) b)
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 :: Ctx a b -> [Name a]
names = (Map (Name a) b -> [Name a]) -> [Map (Name a) b] -> [Name a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Map (Name a) b -> [Name a]
forall k a. Map k a -> [k]
M.keys ([Map (Name a) b] -> [Name a])
-> (Ctx a b -> [Map (Name a) b]) -> Ctx a b -> [Name a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map NameProvenance (Map (Name a) b) -> [Map (Name a) b]
forall k a. Map k a -> [a]
M.elems (Map NameProvenance (Map (Name a) b) -> [Map (Name a) b])
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> [Map (Name a) b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Map NameProvenance (Map (Name a) b)
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 :: Ctx a b -> [b]
elems = (Map (Name a) b -> [b]) -> [Map (Name a) b] -> [b]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Map (Name a) b -> [b]
forall k a. Map k a -> [a]
M.elems ([Map (Name a) b] -> [b])
-> (Ctx a b -> [Map (Name a) b]) -> Ctx a b -> [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map NameProvenance (Map (Name a) b) -> [Map (Name a) b]
forall k a. Map k a -> [a]
M.elems (Map NameProvenance (Map (Name a) b) -> [Map (Name a) b])
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> [Map (Name a) b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Map NameProvenance (Map (Name a) b)
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 :: Ctx a b -> [(QName a, b)]
assocs = ((NameProvenance, Map (Name a) b) -> [(QName a, b)])
-> [(NameProvenance, Map (Name a) b)] -> [(QName a, b)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((NameProvenance -> Map (Name a) b -> [(QName a, b)])
-> (NameProvenance, Map (Name a) b) -> [(QName a, b)]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry NameProvenance -> Map (Name a) b -> [(QName a, b)]
forall a b. NameProvenance -> Map (Name a) b -> [(QName a, b)]
modAssocs) ([(NameProvenance, Map (Name a) b)] -> [(QName a, b)])
-> (Ctx a b -> [(NameProvenance, Map (Name a) b)])
-> Ctx a b
-> [(QName a, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map NameProvenance (Map (Name a) b)
-> [(NameProvenance, Map (Name a) b)]
forall k a. Map k a -> [(k, a)]
M.assocs (Map NameProvenance (Map (Name a) b)
 -> [(NameProvenance, Map (Name a) b)])
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> [(NameProvenance, Map (Name a) b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Map NameProvenance (Map (Name a) b)
forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx
  where
    modAssocs :: NameProvenance -> Map (Name a) b -> [(QName a, b)]
    modAssocs :: NameProvenance -> Map (Name a) b -> [(QName a, b)]
modAssocs NameProvenance
p = ((Name a, b) -> (QName a, b)) -> [(Name a, b)] -> [(QName a, b)]
forall a b. (a -> b) -> [a] -> [b]
map ((Name a -> QName a) -> (Name a, b) -> (QName a, b)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (NameProvenance -> Name a -> QName a
forall a. NameProvenance -> Name a -> QName a
QName NameProvenance
p)) ([(Name a, b)] -> [(QName a, b)])
-> (Map (Name a) b -> [(Name a, b)])
-> Map (Name a) b
-> [(QName a, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map (Name a) b -> [(Name a, b)]
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 :: Ctx a b -> Set (QName a)
keysSet = [Set (QName a)] -> Set (QName a)
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ([Set (QName a)] -> Set (QName a))
-> (Ctx a b -> [Set (QName a)]) -> Ctx a b -> Set (QName a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((NameProvenance, Map (Name a) b) -> Set (QName a))
-> [(NameProvenance, Map (Name a) b)] -> [Set (QName a)]
forall a b. (a -> b) -> [a] -> [b]
map ((NameProvenance -> Set (Name a) -> Set (QName a))
-> (NameProvenance, Set (Name a)) -> Set (QName a)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Name a -> QName a) -> Set (Name a) -> Set (QName a)
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map ((Name a -> QName a) -> Set (Name a) -> Set (QName a))
-> (NameProvenance -> Name a -> QName a)
-> NameProvenance
-> Set (Name a)
-> Set (QName a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NameProvenance -> Name a -> QName a
forall a. NameProvenance -> Name a -> QName a
QName) ((NameProvenance, Set (Name a)) -> Set (QName a))
-> ((NameProvenance, Map (Name a) b)
    -> (NameProvenance, Set (Name a)))
-> (NameProvenance, Map (Name a) b)
-> Set (QName a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map (Name a) b -> Set (Name a))
-> (NameProvenance, Map (Name a) b)
-> (NameProvenance, Set (Name a))
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Map (Name a) b -> Set (Name a)
forall k a. Map k a -> Set k
M.keysSet) ([(NameProvenance, Map (Name a) b)] -> [Set (QName a)])
-> (Ctx a b -> [(NameProvenance, Map (Name a) b)])
-> Ctx a b
-> [Set (QName a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map NameProvenance (Map (Name a) b)
-> [(NameProvenance, Map (Name a) b)]
forall k a. Map k a -> [(k, a)]
M.assocs (Map NameProvenance (Map (Name a) b)
 -> [(NameProvenance, Map (Name a) b)])
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> [(NameProvenance, Map (Name a) b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Map NameProvenance (Map (Name a) b)
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 :: Ctx a1 b -> Ctx a2 b
coerceKeys = Map NameProvenance (Map (Name a2) b) -> Ctx a2 b
forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx (Map NameProvenance (Map (Name a2) b) -> Ctx a2 b)
-> (Ctx a1 b -> Map NameProvenance (Map (Name a2) b))
-> Ctx a1 b
-> Ctx a2 b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map (Name a1) b -> Map (Name a2) b)
-> Map NameProvenance (Map (Name a1) b)
-> Map NameProvenance (Map (Name a2) b)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((Name a1 -> Name a2) -> Map (Name a1) b -> Map (Name a2) b
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys Name a1 -> Name a2
coerce) (Map NameProvenance (Map (Name a1) b)
 -> Map NameProvenance (Map (Name a2) b))
-> (Ctx a1 b -> Map NameProvenance (Map (Name a1) b))
-> Ctx a1 b
-> Map NameProvenance (Map (Name a2) b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a1 b -> Map NameProvenance (Map (Name a1) b)
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 :: Ctx a b -> Set (QName a) -> Ctx a b
restrictKeys Ctx a b
ctx Set (QName a)
xs = Map NameProvenance (Map (Name a) b) -> Ctx a b
forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx (Map NameProvenance (Map (Name a) b) -> Ctx a b)
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> Ctx a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map NameProvenance (Set (Name a))
-> Map NameProvenance (Map (Name a) b)
-> Map NameProvenance (Map (Name a) b)
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 (Map NameProvenance (Map (Name a) b)
 -> Map NameProvenance (Map (Name a) b))
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> Map NameProvenance (Map (Name a) b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Map NameProvenance (Map (Name a) b)
forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx (Ctx a b -> Ctx a b) -> Ctx a b -> Ctx a b
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 = SimpleWhenMissing NameProvenance (Set (Name a)) (Map (Name a) a)
-> SimpleWhenMissing
     NameProvenance (Map (Name a) a) (Map (Name a) a)
-> SimpleWhenMatched
     NameProvenance (Set (Name a)) (Map (Name a) a) (Map (Name a) a)
-> Map NameProvenance (Set (Name a))
-> Map NameProvenance (Map (Name a) a)
-> Map NameProvenance (Map (Name a) a)
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 SimpleWhenMissing NameProvenance (Set (Name a)) (Map (Name a) a)
forall (f :: * -> *) k x y. Applicative f => WhenMissing f k x y
MM.dropMissing SimpleWhenMissing NameProvenance (Map (Name a) a) (Map (Name a) a)
forall (f :: * -> *) k x y. Applicative f => WhenMissing f k x y
MM.dropMissing ((NameProvenance
 -> Set (Name a) -> Map (Name a) a -> Map (Name a) a)
-> SimpleWhenMatched
     NameProvenance (Set (Name a)) (Map (Name a) a) (Map (Name a) a)
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' -> Map (Name a) a -> Set (Name a) -> Map (Name a) a
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 = (Set (Name a) -> Set (Name a) -> Set (Name a))
-> [(NameProvenance, Set (Name a))]
-> Map NameProvenance (Set (Name a))
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Set (Name a) -> Set (Name a) -> Set (Name a)
forall a. Ord a => Set a -> Set a -> Set a
S.union ([(NameProvenance, Set (Name a))]
 -> Map NameProvenance (Set (Name a)))
-> (Set (QName a) -> [(NameProvenance, Set (Name a))])
-> Set (QName a)
-> Map NameProvenance (Set (Name a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (QName a -> (NameProvenance, Set (Name a)))
-> [QName a] -> [(NameProvenance, Set (Name a))]
forall a b. (a -> b) -> [a] -> [b]
map (\(QName NameProvenance
p Name a
n) -> (NameProvenance
p, Name a -> Set (Name a)
forall a. a -> Set a
S.singleton Name a
n)) ([QName a] -> [(NameProvenance, Set (Name a))])
-> (Set (QName a) -> [QName a])
-> Set (QName a)
-> [(NameProvenance, Set (Name a))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set (QName a) -> [QName a]
forall a. Set a -> [a]
S.toList (Set (QName a) -> Map NameProvenance (Set (Name a)))
-> Set (QName a) -> Map NameProvenance (Set (Name a))
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 :: Ctx a b -> Ctx a b -> Ctx a b
joinCtx Ctx a b
a Ctx a b
b = [Ctx a b] -> Ctx a 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 :: [Ctx a b] -> Ctx a b
joinCtxs = Map NameProvenance (Map (Name a) b) -> Ctx a b
forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx (Map NameProvenance (Map (Name a) b) -> Ctx a b)
-> ([Ctx a b] -> Map NameProvenance (Map (Name a) b))
-> [Ctx a b]
-> Ctx a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map (Name a) b -> Map (Name a) b -> Map (Name a) b)
-> [Map NameProvenance (Map (Name a) b)]
-> Map NameProvenance (Map (Name a) b)
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
M.unionsWith Map (Name a) b -> Map (Name a) b -> Map (Name a) b
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union ([Map NameProvenance (Map (Name a) b)]
 -> Map NameProvenance (Map (Name a) b))
-> ([Ctx a b] -> [Map NameProvenance (Map (Name a) b)])
-> [Ctx a b]
-> Map NameProvenance (Map (Name a) b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> [Ctx a b] -> [Map NameProvenance (Map (Name a) b)]
forall a b. (a -> b) -> [a] -> [b]
map Ctx a b -> Map NameProvenance (Map (Name a) b)
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 :: (b -> Bool) -> Ctx a b -> Ctx a b
filter b -> Bool
p = Map NameProvenance (Map (Name a) b) -> Ctx a b
forall a b. Map NameProvenance (Map (Name a) b) -> Ctx a b
Ctx (Map NameProvenance (Map (Name a) b) -> Ctx a b)
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> Ctx a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map (Name a) b -> Map (Name a) b)
-> Map NameProvenance (Map (Name a) b)
-> Map NameProvenance (Map (Name a) b)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((b -> Bool) -> Map (Name a) b -> Map (Name a) b
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter b -> Bool
p) (Map NameProvenance (Map (Name a) b)
 -> Map NameProvenance (Map (Name a) b))
-> (Ctx a b -> Map NameProvenance (Map (Name a) b))
-> Ctx a b
-> Map NameProvenance (Map (Name a) b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ctx a b -> Map NameProvenance (Map (Name a) b)
forall a b. Ctx a b -> Map NameProvenance (Map (Name a) b)
getCtx