Context.hs

{-# 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 {getCtx :: M.Map NameProvenance (M.Map (Name a) b)}
  deriving (Eq, Show, Functor, Foldable, 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
  (<>) = joinCtx

instance Monoid (Ctx a b) where
  mempty = emptyCtx
  mappend = (<>)

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

-- | The empty context.
emptyCtx :: Ctx a b
emptyCtx = Ctx M.empty

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

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

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

-- | Create a context with local bindings.
localCtx :: [(Name a, b)] -> Ctx a b
localCtx = Ctx . M.singleton LocalName . 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 p n) b = Ctx . M.insertWith M.union p (M.singleton n 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 qn b = local (insert qn 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 = local . joinCtx

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

-- | Check if a context is empty.
null :: Ctx a b -> Bool
null = all M.null . getCtx

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

-- | Look up a qualified name in a context.
lookup' :: QName a -> Ctx a b -> Maybe b
lookup' (QName p n) = (M.lookup n <=< M.lookup p) . 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 n = lookupNonLocal' n <$> ask

-- | Look up all the non-local bindings of a name in a context.
lookupNonLocal' :: Name a -> Ctx a b -> [(ModuleName, b)]
lookupNonLocal' n = nonLocal . lookupAll' n
 where
  nonLocal bs = [(m, b) | (QName (QualifiedName m) _, 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 n = lookupAll' n <$> ask

-- | Look up all the bindings of an (unqualified) name in a context.
lookupAll' :: Name a -> Ctx a b -> [(QName a, b)]
lookupAll' n = map (first (`QName` n)) . M.assocs . M.mapMaybe (M.lookup n) . getCtx

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

-- | Return a list of the names defined by the context.
names :: Ctx a b -> [Name a]
names = concatMap M.keys . M.elems . getCtx

-- | Return a list of all the values bound by the context.
elems :: Ctx a b -> [b]
elems = concatMap M.elems . M.elems . getCtx

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

-- | Return a set of all qualified names in the context.
keysSet :: Ctx a b -> Set (QName a)
keysSet = S.unions . map (uncurry (S.map . QName) . second M.keysSet) . M.assocs . getCtx

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

-- | Coerce the type of the qualified name keys in a context.
coerceKeys :: Ctx a1 b -> Ctx a2 b
coerceKeys = Ctx . M.map (M.mapKeys coerce) . getCtx

-- | Restrict a context to only the keys in the given set.
restrictKeys :: Ctx a b -> Set (QName a) -> Ctx a b
restrictKeys ctx xs = Ctx . restrict m . getCtx $ ctx
 where
  restrict = MM.merge MM.dropMissing MM.dropMissing (MM.zipWithMatched (\_ ns m' -> M.restrictKeys m' ns))
  m = M.fromListWith S.union . map (\(QName p n) -> (p, S.singleton n)) . S.toList $ 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 a b = joinCtxs [a, b]

-- | Join a list of contexts (left-biased).
joinCtxs :: [Ctx a b] -> Ctx a b
joinCtxs = Ctx . M.unionsWith M.union . map getCtx

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

-- | Filter a context using a predicate.
filter :: (b -> Bool) -> Ctx a b -> Ctx a b
filter p = Ctx . M.map (M.filter p) . getCtx