{-# LANGUAGE OverloadedStrings #-}

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

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

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 { Substitution a -> Map (Name a) a
getSubst :: Map (Name a) a }
  deriving (Substitution a -> Substitution a -> Bool
(Substitution a -> Substitution a -> Bool)
-> (Substitution a -> Substitution a -> Bool)
-> Eq (Substitution a)
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, Eq (Substitution a)
Eq (Substitution a)
-> (Substitution a -> Substitution a -> Ordering)
-> (Substitution a -> Substitution a -> Bool)
-> (Substitution a -> Substitution a -> Bool)
-> (Substitution a -> Substitution a -> Bool)
-> (Substitution a -> Substitution a -> Bool)
-> (Substitution a -> Substitution a -> Substitution a)
-> (Substitution a -> Substitution a -> Substitution a)
-> Ord (Substitution a)
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
$cp1Ord :: forall a. Ord a => Eq (Substitution a)
Ord, Int -> Substitution a -> ShowS
[Substitution a] -> ShowS
Substitution a -> String
(Int -> Substitution a -> ShowS)
-> (Substitution a -> String)
-> ([Substitution a] -> ShowS)
-> Show (Substitution a)
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 :: (a -> b) -> Substitution a -> Substitution b
fmap a -> b
f (Substitution Map (Name a) a
m) = Map (Name b) b -> Substitution b
forall a. Map (Name a) a -> Substitution a
Substitution ((Name a -> Name b) -> Map (Name a) b -> Map (Name b) b
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys Name a -> Name b
coerce (Map (Name a) b -> Map (Name b) b)
-> (Map (Name a) a -> Map (Name a) b)
-> Map (Name a) a
-> Map (Name b) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> Map (Name a) a -> Map (Name a) b
forall a b k. (a -> b) -> Map k a -> Map k b
M.map a -> b
f (Map (Name a) a -> Map (Name b) b)
-> Map (Name a) a -> Map (Name b) b
forall a b. (a -> b) -> a -> b
$ Map (Name a) a
m)

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

prettyMapping :: (Pretty a, Members '[Reader PA, LFresh] r) => Name a -> a -> Sem r Doc
prettyMapping :: Name a -> a -> Sem r Doc
prettyMapping Name a
x a
a = Name a -> Sem r Doc
forall t (r :: EffectRow).
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r Doc
pretty Name a
x Sem r Doc -> Sem r Doc -> Sem r Doc
forall (f :: * -> *). Applicative f => f Doc -> f Doc -> f Doc
<+> Sem r Doc
"->" Sem r Doc -> Sem r Doc -> Sem r Doc
forall (f :: * -> *). Applicative f => f Doc -> f Doc -> f Doc
<+> a -> Sem r Doc
forall t (r :: EffectRow).
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r Doc
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 :: Substitution a -> Set (Name a)
dom = Map (Name a) a -> Set (Name a)
forall k a. Map k a -> Set k
M.keysSet (Map (Name a) a -> Set (Name a))
-> (Substitution a -> Map (Name a) a)
-> Substitution a
-> Set (Name a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Substitution a -> Map (Name a) a
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 :: Substitution a
idS = Map (Name a) a -> Substitution a
forall a. Map (Name a) a -> Substitution a
Substitution Map (Name a) a
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 |-> :: Name a -> a -> Substitution a
|-> a
t = Map (Name a) a -> Substitution a
forall a. Map (Name a) a -> Substitution a
Substitution (Name a -> a -> Map (Name a) a
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) @@ :: Substitution a -> Substitution a -> Substitution a
@@ (Substitution Map (Name a) a
s2) = Map (Name a) a -> Substitution a
forall a. Map (Name a) a -> Substitution a
Substitution (((a -> a) -> Map (Name a) a -> Map (Name a) a
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Substitution a -> a -> a
forall b a. Subst b a => Substitution b -> a -> a
applySubst (Map (Name a) a -> Substitution a
forall a. Map (Name a) a -> Substitution a
Substitution Map (Name a) a
s1))) Map (Name a) a
s2 Map (Name a) a -> Map (Name a) a -> Map (Name a) a
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 :: t (Substitution a) -> Substitution a
compose = (Substitution a -> Substitution a -> Substitution a)
-> Substitution a -> t (Substitution a) -> Substitution a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Substitution a -> Substitution a -> Substitution a
forall a.
Subst a a =>
Substitution 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 :: Substitution b -> a -> a
applySubst (Substitution Map (Name b) b
s) = [(Name b, b)] -> a -> a
forall b a. Subst b a => [(Name b, b)] -> a -> a
substs (Map (Name b) b -> [(Name b, b)]
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 :: [(Name a, a)] -> Substitution a
fromList = Map (Name a) a -> Substitution a
forall a. Map (Name a) a -> Substitution a
Substitution (Map (Name a) a -> Substitution a)
-> ([(Name a, a)] -> Map (Name a) a)
-> [(Name a, a)]
-> Substitution a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name a, a)] -> Map (Name a) a
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 :: Substitution a -> [(Name a, a)]
toList (Substitution Map (Name a) a
m) = Map (Name a) a -> [(Name a, a)]
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 :: Name a -> Substitution a -> Maybe a
lookup Name a
x (Substitution Map (Name a) a
m) = Name a -> Map (Name a) a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name a
x Map (Name a) a
m