{-# LANGUAGE OverloadedStrings #-}

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

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

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

-- |
-- Module      :  Disco.Subst
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- The "Disco.Subst" module defines a generic type of substitutions
-- that map variable names to values.
module Disco.Subst (
  -- * Substitutions
  Substitution (..),
  dom,

  -- ** Constructing/destructing substitutions
  idS,
  (|->),
  fromList,
  toList,

  -- ** Substitution operations
  (@@),
  compose,
  applySubst,
  lookup,
)
where

import Prelude hiding (lookup)

import Unbound.Generics.LocallyNameless (Name, Subst, substs)

import Data.Coerce

import Data.Map (Map)
import qualified Data.Map as M
import Data.Set (Set)

import Disco.Effects.LFresh
import Disco.Pretty
import Polysemy
import Polysemy.Reader

-- | A value of type @Substitution a@ is a substitution which maps some set of
--   names (the /domain/, see 'dom') to values of type @a@.
--   Substitutions can be /applied/ to certain terms (see
--   'applySubst'), replacing any free occurrences of names in the
--   domain with their corresponding values.  Thus, substitutions can
--   be thought of as functions of type @Term -> Term@ (for suitable
--   @Term@s that contain names and values of the right type).
--
--   Concretely, substitutions are stored using a @Map@.
--
--   See also "Disco.Types", which defines 'S' as an alias for
--   substitutions on types (the most common kind in the disco
--   codebase).
newtype Substitution a = Substitution {forall a. Substitution a -> Map (Name a) a
getSubst :: Map (Name a) a}
  deriving (Substitution a -> Substitution a -> Bool
forall a. Eq a => Substitution a -> Substitution a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Substitution a -> Substitution a -> Bool
$c/= :: forall a. Eq a => Substitution a -> Substitution a -> Bool
== :: Substitution a -> Substitution a -> Bool
$c== :: forall a. Eq a => Substitution a -> Substitution a -> Bool
Eq, Substitution a -> Substitution a -> Bool
Substitution a -> Substitution a -> Ordering
Substitution a -> Substitution a -> Substitution a
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {a}. Ord a => Eq (Substitution a)
forall a. Ord a => Substitution a -> Substitution a -> Bool
forall a. Ord a => Substitution a -> Substitution a -> Ordering
forall a.
Ord a =>
Substitution a -> Substitution a -> Substitution a
min :: Substitution a -> Substitution a -> Substitution a
$cmin :: forall a.
Ord a =>
Substitution a -> Substitution a -> Substitution a
max :: Substitution a -> Substitution a -> Substitution a
$cmax :: forall a.
Ord a =>
Substitution a -> Substitution a -> Substitution a
>= :: Substitution a -> Substitution a -> Bool
$c>= :: forall a. Ord a => Substitution a -> Substitution a -> Bool
> :: Substitution a -> Substitution a -> Bool
$c> :: forall a. Ord a => Substitution a -> Substitution a -> Bool
<= :: Substitution a -> Substitution a -> Bool
$c<= :: forall a. Ord a => Substitution a -> Substitution a -> Bool
< :: Substitution a -> Substitution a -> Bool
$c< :: forall a. Ord a => Substitution a -> Substitution a -> Bool
compare :: Substitution a -> Substitution a -> Ordering
$ccompare :: forall a. Ord a => Substitution a -> Substitution a -> Ordering
Ord, Int -> Substitution a -> ShowS
forall a. Show a => Int -> Substitution a -> ShowS
forall a. Show a => [Substitution a] -> ShowS
forall a. Show a => Substitution a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Substitution a] -> ShowS
$cshowList :: forall a. Show a => [Substitution a] -> ShowS
show :: Substitution a -> String
$cshow :: forall a. Show a => Substitution a -> String
showsPrec :: Int -> Substitution a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Substitution a -> ShowS
Show)

instance Functor Substitution where
  fmap :: forall a b. (a -> b) -> Substitution a -> Substitution b
fmap a -> b
f (Substitution Map (Name a) a
m) = forall a. Map (Name a) a -> Substitution a
Substitution (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 k. (a -> b) -> Map k a -> Map k b
M.map a -> b
f forall a b. (a -> b) -> a -> b
$ Map (Name a) a
m)

instance Pretty a => Pretty (Substitution a) where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
Substitution a -> Sem r (Doc ann)
pretty (Substitution Map (Name a) a
s) = do
    let es :: [Sem r (Doc ann)]
es = forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a (r :: EffectRow) ann.
(Pretty a, Members '[Reader PA, LFresh] r) =>
Name a -> a -> Sem r (Doc ann)
prettyMapping) (forall k a. Map k a -> [(k, a)]
M.assocs Map (Name a) a
s)
    [Sem r (Doc ann)]
ds <- forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> [f (Doc ann)] -> f [f (Doc ann)]
punctuate Sem r (Doc ann)
"," [Sem r (Doc ann)]
es
    forall (f :: * -> *) ann. Functor f => f (Doc ann) -> f (Doc ann)
braces (forall (f :: * -> *) ann.
Applicative f =>
[f (Doc ann)] -> f (Doc ann)
hsep [Sem r (Doc ann)]
ds)

prettyMapping :: (Pretty a, Members '[Reader PA, LFresh] r) => Name a -> a -> Sem r (Doc ann)
prettyMapping :: forall a (r :: EffectRow) ann.
(Pretty a, Members '[Reader PA, LFresh] r) =>
Name a -> a -> Sem r (Doc ann)
prettyMapping Name a
x a
a = forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty Name a
x forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<+> Sem r (Doc ann)
"->" forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<+> forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty a
a

-- | The domain of a substitution is the set of names for which the
--   substitution is defined.
dom :: Substitution a -> Set (Name a)
dom :: forall a. Substitution a -> Set (Name a)
dom = forall k a. Map k a -> Set k
M.keysSet forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Substitution a -> Map (Name a) a
getSubst

-- | The identity substitution, /i.e./ the unique substitution with an
--   empty domain, which acts as the identity function on terms.
idS :: Substitution a
idS :: forall a. Substitution a
idS = forall a. Map (Name a) a -> Substitution a
Substitution forall k a. Map k a
M.empty

-- | Construct a singleton substitution, which maps the given name to
--   the given value.
(|->) :: Name a -> a -> Substitution a
Name a
x |-> :: forall a. Name a -> a -> Substitution a
|-> a
t = forall a. Map (Name a) a -> Substitution a
Substitution (forall k a. k -> a -> Map k a
M.singleton Name a
x a
t)

-- | Compose two substitutions.  Applying @s1 \@\@ s2@ is the same as
--   applying first @s2@, then @s1@; that is, semantically,
--   composition of substitutions corresponds exactly to function
--   composition when they are considered as functions on terms.
--
--   As one would expect, composition is associative and has 'idS' as
--   its identity.
(@@) :: Subst a a => Substitution a -> Substitution a -> Substitution a
(Substitution Map (Name a) a
s1) @@ :: forall a.
Subst a a =>
Substitution a -> Substitution a -> Substitution a
@@ (Substitution Map (Name a) a
s2) = forall a. Map (Name a) a -> Substitution a
Substitution ((forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall b a. Subst b a => Substitution b -> a -> a
applySubst (forall a. Map (Name a) a -> Substitution a
Substitution Map (Name a) a
s1))) Map (Name a) a
s2 forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map (Name a) a
s1)

-- | Compose a whole container of substitutions.  For example,
--   @compose [s1, s2, s3] = s1 \@\@ s2 \@\@ s3@.
compose :: (Subst a a, Foldable t) => t (Substitution a) -> Substitution a
compose :: forall a (t :: * -> *).
(Subst a a, Foldable t) =>
t (Substitution a) -> Substitution a
compose = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall a.
Subst a a =>
Substitution a -> Substitution a -> Substitution a
(@@) forall a. Substitution a
idS

-- | Apply a substitution to a term, resulting in a new term in which
--   any free variables in the domain of the substitution have been
--   replaced by their corresponding values.  Note this requires a
--   @Subst b a@ constraint, which intuitively means that values of
--   type @a@ contain variables of type @b@ we can substitute for.
applySubst :: Subst b a => Substitution b -> a -> a
applySubst :: forall b a. Subst b a => Substitution b -> a -> a
applySubst (Substitution Map (Name b) b
s) = forall b a. Subst b a => [(Name b, b)] -> a -> a
substs (forall k a. Map k a -> [(k, a)]
M.assocs Map (Name b) b
s)

-- | Create a substitution from an association list of names and
--   values.
fromList :: [(Name a, a)] -> Substitution a
fromList :: forall a. [(Name a, a)] -> Substitution a
fromList = forall a. Map (Name a) a -> Substitution a
Substitution forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList

-- | Convert a substitution into an association list.
toList :: Substitution a -> [(Name a, a)]
toList :: forall a. Substitution a -> [(Name a, a)]
toList (Substitution Map (Name a) a
m) = forall k a. Map k a -> [(k, a)]
M.assocs Map (Name a) a
m

-- | Look up the value a particular name maps to under the given
--   substitution; or return @Nothing@ if the name being looked up is
--   not in the domain.
lookup :: Name a -> Substitution a -> Maybe a
lookup :: forall a. Name a -> Substitution a -> Maybe a
lookup Name a
x (Substitution Map (Name a) a
m) = forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name a
x Map (Name a) a
m