{-# LANGUAGE BlockArguments             #-}
{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE QuantifiedConstraints      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}
-- | Main definitions of the foil that can be
-- reused for specific implementations.
-- This is an internal module, so it also contains implementation details of the foil.
--
-- The original description of this approach
-- is described in the IFL 2022 paper by Maclaurin, Radul, and Paszke
-- [«The Foil: Capture-Avoiding Substitution With No Sharp Edges»](https://doi.org/10.1145/3587216.3587224).
-- This module also introduces 'CoSinkable' class,
-- generalizing handling of patterns, as described in
-- [«Free Foil: Generating Efficient and Scope-Safe Abstract Syntax»](https://arxiv.org/abs/2405.16384).
--
-- Since the representation of scopes and substitutions
-- is either 'IntMap' or 'IntSet', many of the operations
-- have a worst-case complexity of \(O(\min(n,W))\).
-- This means that the operation can become linear in the size of the scope \(n\) with a maximum of \(W\)
-- — the number of bits in an 'Int' (32 or 64).
module Control.Monad.Foil.Internal where

import           Control.DeepSeq (NFData (..))
import           Data.IntMap
import qualified Data.IntMap     as IntMap
import           Data.IntSet
import qualified Data.IntSet     as IntSet
import           Data.Kind       (Type)
import           Unsafe.Coerce

-- * Safe types and operations

-- | 'S' is a data kind of scope indices.
data S
  = VoidS -- ^ 'VoidS' is the only explicit scope available to the users, representing an empty scope.
          -- All other scopes are represented with type variables,
          -- bound in rank-2 polymophic functions like 'withFreshBinder'.

-- | A safe scope, indexed by a type-level scope index 'n'.
newtype Scope (n :: S) = UnsafeScope RawScope
  deriving newtype Scope n -> ()
(Scope n -> ()) -> NFData (Scope n)
forall a. (a -> ()) -> NFData a
forall (n :: S). Scope n -> ()
$crnf :: forall (n :: S). Scope n -> ()
rnf :: Scope n -> ()
NFData

-- | A name in a safe scope, indexed by a type-level scope index 'n'.
newtype Name (n :: S) = UnsafeName RawName
  deriving newtype (Name n -> ()
(Name n -> ()) -> NFData (Name n)
forall a. (a -> ()) -> NFData a
forall (n :: S). Name n -> ()
$crnf :: forall (n :: S). Name n -> ()
rnf :: Name n -> ()
NFData, Name n -> Name n -> Bool
(Name n -> Name n -> Bool)
-> (Name n -> Name n -> Bool) -> Eq (Name n)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (n :: S). Name n -> Name n -> Bool
$c== :: forall (n :: S). Name n -> Name n -> Bool
== :: Name n -> Name n -> Bool
$c/= :: forall (n :: S). Name n -> Name n -> Bool
/= :: Name n -> Name n -> Bool
Eq, Eq (Name n)
Eq (Name n) =>
(Name n -> Name n -> Ordering)
-> (Name n -> Name n -> Bool)
-> (Name n -> Name n -> Bool)
-> (Name n -> Name n -> Bool)
-> (Name n -> Name n -> Bool)
-> (Name n -> Name n -> Name n)
-> (Name n -> Name n -> Name n)
-> Ord (Name n)
Name n -> Name n -> Bool
Name n -> Name n -> Ordering
Name n -> Name n -> Name n
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 (n :: S). Eq (Name n)
forall (n :: S). Name n -> Name n -> Bool
forall (n :: S). Name n -> Name n -> Ordering
forall (n :: S). Name n -> Name n -> Name n
$ccompare :: forall (n :: S). Name n -> Name n -> Ordering
compare :: Name n -> Name n -> Ordering
$c< :: forall (n :: S). Name n -> Name n -> Bool
< :: Name n -> Name n -> Bool
$c<= :: forall (n :: S). Name n -> Name n -> Bool
<= :: Name n -> Name n -> Bool
$c> :: forall (n :: S). Name n -> Name n -> Bool
> :: Name n -> Name n -> Bool
$c>= :: forall (n :: S). Name n -> Name n -> Bool
>= :: Name n -> Name n -> Bool
$cmax :: forall (n :: S). Name n -> Name n -> Name n
max :: Name n -> Name n -> Name n
$cmin :: forall (n :: S). Name n -> Name n -> Name n
min :: Name n -> Name n -> Name n
Ord, RawName -> Name n -> ShowS
[Name n] -> ShowS
Name n -> String
(RawName -> Name n -> ShowS)
-> (Name n -> String) -> ([Name n] -> ShowS) -> Show (Name n)
forall a.
(RawName -> a -> ShowS)
-> (a -> String) -> ([a] -> ShowS) -> Show a
forall (n :: S). RawName -> Name n -> ShowS
forall (n :: S). [Name n] -> ShowS
forall (n :: S). Name n -> String
$cshowsPrec :: forall (n :: S). RawName -> Name n -> ShowS
showsPrec :: RawName -> Name n -> ShowS
$cshow :: forall (n :: S). Name n -> String
show :: Name n -> String
$cshowList :: forall (n :: S). [Name n] -> ShowS
showList :: [Name n] -> ShowS
Show)

-- | A name binder is a name that extends scope @n@ to a (larger) scope @l@.
newtype NameBinder (n :: S) (l :: S) =
  UnsafeNameBinder (Name l)
    deriving newtype (NameBinder n l -> ()
(NameBinder n l -> ()) -> NFData (NameBinder n l)
forall a. (a -> ()) -> NFData a
forall (n :: S) (l :: S). NameBinder n l -> ()
$crnf :: forall (n :: S) (l :: S). NameBinder n l -> ()
rnf :: NameBinder n l -> ()
NFData, NameBinder n l -> NameBinder n l -> Bool
(NameBinder n l -> NameBinder n l -> Bool)
-> (NameBinder n l -> NameBinder n l -> Bool)
-> Eq (NameBinder n l)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (n :: S) (l :: S). NameBinder n l -> NameBinder n l -> Bool
$c== :: forall (n :: S) (l :: S). NameBinder n l -> NameBinder n l -> Bool
== :: NameBinder n l -> NameBinder n l -> Bool
$c/= :: forall (n :: S) (l :: S). NameBinder n l -> NameBinder n l -> Bool
/= :: NameBinder n l -> NameBinder n l -> Bool
Eq, Eq (NameBinder n l)
Eq (NameBinder n l) =>
(NameBinder n l -> NameBinder n l -> Ordering)
-> (NameBinder n l -> NameBinder n l -> Bool)
-> (NameBinder n l -> NameBinder n l -> Bool)
-> (NameBinder n l -> NameBinder n l -> Bool)
-> (NameBinder n l -> NameBinder n l -> Bool)
-> (NameBinder n l -> NameBinder n l -> NameBinder n l)
-> (NameBinder n l -> NameBinder n l -> NameBinder n l)
-> Ord (NameBinder n l)
NameBinder n l -> NameBinder n l -> Bool
NameBinder n l -> NameBinder n l -> Ordering
NameBinder n l -> NameBinder n l -> NameBinder n l
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 (n :: S) (l :: S). Eq (NameBinder n l)
forall (n :: S) (l :: S). NameBinder n l -> NameBinder n l -> Bool
forall (n :: S) (l :: S).
NameBinder n l -> NameBinder n l -> Ordering
forall (n :: S) (l :: S).
NameBinder n l -> NameBinder n l -> NameBinder n l
$ccompare :: forall (n :: S) (l :: S).
NameBinder n l -> NameBinder n l -> Ordering
compare :: NameBinder n l -> NameBinder n l -> Ordering
$c< :: forall (n :: S) (l :: S). NameBinder n l -> NameBinder n l -> Bool
< :: NameBinder n l -> NameBinder n l -> Bool
$c<= :: forall (n :: S) (l :: S). NameBinder n l -> NameBinder n l -> Bool
<= :: NameBinder n l -> NameBinder n l -> Bool
$c> :: forall (n :: S) (l :: S). NameBinder n l -> NameBinder n l -> Bool
> :: NameBinder n l -> NameBinder n l -> Bool
$c>= :: forall (n :: S) (l :: S). NameBinder n l -> NameBinder n l -> Bool
>= :: NameBinder n l -> NameBinder n l -> Bool
$cmax :: forall (n :: S) (l :: S).
NameBinder n l -> NameBinder n l -> NameBinder n l
max :: NameBinder n l -> NameBinder n l -> NameBinder n l
$cmin :: forall (n :: S) (l :: S).
NameBinder n l -> NameBinder n l -> NameBinder n l
min :: NameBinder n l -> NameBinder n l -> NameBinder n l
Ord, RawName -> NameBinder n l -> ShowS
[NameBinder n l] -> ShowS
NameBinder n l -> String
(RawName -> NameBinder n l -> ShowS)
-> (NameBinder n l -> String)
-> ([NameBinder n l] -> ShowS)
-> Show (NameBinder n l)
forall a.
(RawName -> a -> ShowS)
-> (a -> String) -> ([a] -> ShowS) -> Show a
forall (n :: S) (l :: S). RawName -> NameBinder n l -> ShowS
forall (n :: S) (l :: S). [NameBinder n l] -> ShowS
forall (n :: S) (l :: S). NameBinder n l -> String
$cshowsPrec :: forall (n :: S) (l :: S). RawName -> NameBinder n l -> ShowS
showsPrec :: RawName -> NameBinder n l -> ShowS
$cshow :: forall (n :: S) (l :: S). NameBinder n l -> String
show :: NameBinder n l -> String
$cshowList :: forall (n :: S) (l :: S). [NameBinder n l] -> ShowS
showList :: [NameBinder n l] -> ShowS
Show)

-- | An empty scope (without any names).
emptyScope :: Scope VoidS
emptyScope :: Scope 'VoidS
emptyScope = RawScope -> Scope 'VoidS
forall (n :: S). RawScope -> Scope n
UnsafeScope RawScope
IntSet.empty

-- | \(O(\min(n,W))\).
-- Extend a scope with one name (safely).
-- Note that as long as the foil is used as intended,
-- the name binder is guaranteed to introduce a name
-- that does not appear in the initial scope.
extendScope :: NameBinder n l -> Scope n -> Scope l
extendScope :: forall (n :: S) (l :: S). NameBinder n l -> Scope n -> Scope l
extendScope (UnsafeNameBinder (UnsafeName RawName
name)) (UnsafeScope RawScope
scope) =
  RawScope -> Scope l
forall (n :: S). RawScope -> Scope n
UnsafeScope (RawName -> RawScope -> RawScope
IntSet.insert RawName
name RawScope
scope)

-- | A runtime check for potential name capture.
member :: Name l -> Scope n -> Bool
member :: forall (l :: S) (n :: S). Name l -> Scope n -> Bool
member (UnsafeName RawName
name) (UnsafeScope RawScope
s) = RawName -> RawScope -> Bool
rawMember RawName
name RawScope
s

-- | Extract name from a name binder.
nameOf :: NameBinder n l -> Name l
nameOf :: forall (n :: S) (l :: S). NameBinder n l -> Name l
nameOf (UnsafeNameBinder Name l
name) = Name l
name

-- | Convert 'Name' into an identifier.
-- This may be useful for printing and debugging.
nameId :: Name l -> Id
nameId :: forall (l :: S). Name l -> RawName
nameId (UnsafeName RawName
i) = RawName
i

-- | Allocate a fresh binder for a given scope.
withFreshBinder
  :: Scope n
  -> (forall l. NameBinder n l -> r) -> r
withFreshBinder :: forall (n :: S) r.
Scope n -> (forall (l :: S). NameBinder n l -> r) -> r
withFreshBinder (UnsafeScope RawScope
scope) forall (l :: S). NameBinder n l -> r
cont =
  NameBinder n Any -> r
forall (l :: S). NameBinder n l -> r
cont NameBinder n Any
binder
  where
    binder :: NameBinder n Any
binder = Name Any -> NameBinder n Any
forall (n :: S) (l :: S). Name l -> NameBinder n l
UnsafeNameBinder (RawName -> Name Any
forall (n :: S). RawName -> Name n
UnsafeName (RawScope -> RawName
rawFreshName RawScope
scope))

-- | Evidence that scope @n@ contains distinct names.
data DistinctEvidence (n :: S) where
  Distinct :: Distinct n => DistinctEvidence n

-- | Unsafely declare that scope @n@ is distinct.
-- Used in 'unsafeAssertFresh'.
unsafeDistinct :: DistinctEvidence n
unsafeDistinct :: forall (n :: S). DistinctEvidence n
unsafeDistinct = DistinctEvidence 'VoidS -> DistinctEvidence n
forall a b. a -> b
unsafeCoerce (DistinctEvidence 'VoidS
forall (n :: S). Distinct n => DistinctEvidence n
Distinct :: DistinctEvidence VoidS)

-- | Evidence that scope @l@ extends scope @n@.
data ExtEvidence (n :: S) (l :: S) where
  Ext :: Ext n l => ExtEvidence n l

-- | Unsafely declare that scope @l@ extends scope @n@.
-- Used in 'unsafeAssertFresh'.
unsafeExt :: ExtEvidence n l
unsafeExt :: forall (n :: S) (l :: S). ExtEvidence n l
unsafeExt = ExtEvidence 'VoidS 'VoidS -> ExtEvidence n l
forall a b. a -> b
unsafeCoerce (ExtEvidence 'VoidS 'VoidS
forall (n :: S) (l :: S). Ext n l => ExtEvidence n l
Ext :: ExtEvidence VoidS VoidS)

-- | Safely produce a fresh name binder with respect to a given scope.
withFresh
  :: Distinct n => Scope n
  -> (forall l. DExt n l => NameBinder n l -> r) -> r
withFresh :: forall (n :: S) r.
Distinct n =>
Scope n -> (forall (l :: S). DExt n l => NameBinder n l -> r) -> r
withFresh Scope n
scope forall (l :: S). DExt n l => NameBinder n l -> r
cont = Scope n -> (forall (l :: S). NameBinder n l -> r) -> r
forall (n :: S) r.
Scope n -> (forall (l :: S). NameBinder n l -> r) -> r
withFreshBinder Scope n
scope (NameBinder n l -> (DExt n Any => NameBinder n Any -> r) -> r
forall (n :: S) (l :: S) (n' :: S) (l' :: S) r.
NameBinder n l -> (DExt n' l' => NameBinder n' l' -> r) -> r
`unsafeAssertFresh` DExt n Any => NameBinder n Any -> r
NameBinder n Any -> r
forall (l :: S). DExt n l => NameBinder n l -> r
cont)

-- | Unsafely declare that a given name (binder)
-- is already fresh in any scope @n'@.
unsafeAssertFresh :: forall n l n' l' r. NameBinder n l
  -> (DExt n' l' => NameBinder n' l' -> r) -> r
unsafeAssertFresh :: forall (n :: S) (l :: S) (n' :: S) (l' :: S) r.
NameBinder n l -> (DExt n' l' => NameBinder n' l' -> r) -> r
unsafeAssertFresh NameBinder n l
binder DExt n' l' => NameBinder n' l' -> r
cont =
  case forall (n :: S). DistinctEvidence n
unsafeDistinct @l' of
    DistinctEvidence l'
Distinct -> case forall (n :: S) (l :: S). ExtEvidence n l
unsafeExt @n' @l' of
      ExtEvidence n' l'
Ext -> DExt n' l' => NameBinder n' l' -> r
NameBinder n' l' -> r
cont (NameBinder n l -> NameBinder n' l'
forall a b. a -> b
unsafeCoerce NameBinder n l
binder)

-- | A distinct scope extended with a 'NameBinder' is also distinct.
assertDistinct :: Distinct n => NameBinder n l -> DistinctEvidence l
assertDistinct :: forall (n :: S) (l :: S).
Distinct n =>
NameBinder n l -> DistinctEvidence l
assertDistinct NameBinder n l
_ = DistinctEvidence l
forall (n :: S). DistinctEvidence n
unsafeDistinct

-- | Safely rename (if necessary) a given name to extend a given scope.
-- This is similar to 'withFresh', except if the name does not clash with
-- the scope, it can be used immediately, without renaming.
withRefreshed
  :: Distinct o
  => Scope o    -- ^ Ambient scope.
  -> Name i     -- ^ Name to refresh (if it clashes with the ambient scope).
  -> (forall o'. DExt o o' => NameBinder o o' -> r)
  -- ^ Continuation, accepting the refreshed name.
  -> r
withRefreshed :: forall (o :: S) (i :: S) r.
Distinct o =>
Scope o
-> Name i
-> (forall (o' :: S). DExt o o' => NameBinder o o' -> r)
-> r
withRefreshed scope :: Scope o
scope@(UnsafeScope RawScope
rawScope) name :: Name i
name@(UnsafeName RawName
rawName) forall (o' :: S). DExt o o' => NameBinder o o' -> r
cont
  | RawName -> RawScope -> Bool
IntSet.member RawName
rawName RawScope
rawScope = Scope o
-> (forall (o' :: S). DExt o o' => NameBinder o o' -> r) -> r
forall (n :: S) r.
Distinct n =>
Scope n -> (forall (l :: S). DExt n l => NameBinder n l -> r) -> r
withFresh Scope o
scope NameBinder o l -> r
forall (o' :: S). DExt o o' => NameBinder o o' -> r
cont
  | Bool
otherwise = NameBinder Any i -> (DExt o Any => NameBinder o Any -> r) -> r
forall (n :: S) (l :: S) (n' :: S) (l' :: S) r.
NameBinder n l -> (DExt n' l' => NameBinder n' l' -> r) -> r
unsafeAssertFresh (Name i -> NameBinder Any i
forall (n :: S) (l :: S). Name l -> NameBinder n l
UnsafeNameBinder Name i
name) DExt o Any => NameBinder o Any -> r
NameBinder o Any -> r
forall (o' :: S). DExt o o' => NameBinder o o' -> r
cont

-- | Try coercing the name back to the (smaller) scope,
-- given a binder that extends that scope.
unsinkName :: NameBinder n l -> Name l -> Maybe (Name n)
unsinkName :: forall (n :: S) (l :: S).
NameBinder n l -> Name l -> Maybe (Name n)
unsinkName NameBinder n l
binder name :: Name l
name@(UnsafeName RawName
raw)
  | NameBinder n l -> Name l
forall (n :: S) (l :: S). NameBinder n l -> Name l
nameOf NameBinder n l
binder Name l -> Name l -> Bool
forall a. Eq a => a -> a -> Bool
== Name l
name = Maybe (Name n)
forall a. Maybe a
Nothing
  | Bool
otherwise = Name n -> Maybe (Name n)
forall a. a -> Maybe a
Just (RawName -> Name n
forall (n :: S). RawName -> Name n
UnsafeName RawName
raw)

-- * Safe sinking

-- | Sinking an expression from scope @n@ into a (usualy extended) scope @l@,
-- given the renaming (injection from scope @n@ to scope @l@).
class Sinkable (e :: S -> Type) where
  -- | An implementation of this method that typechecks
  -- proves to the compiler that the expression is indeed
  -- 'Sinkable'. However, instead of this implementation, 'sink'
  -- should be used at all call sites for efficiency.
  sinkabilityProof
    :: (Name n -> Name l)   -- ^ Map names from scope @n@ to a (possibly larger) scope @l@.
    -> e n                  -- ^ Expression with free variables in scope @n@.
    -> e l

-- | Sinking a 'Name' is as simple as applying the renaming.
instance Sinkable Name where
  sinkabilityProof :: forall (n :: S) (l :: S). (Name n -> Name l) -> Name n -> Name l
sinkabilityProof Name n -> Name l
rename = Name n -> Name l
rename

-- | Efficient version of 'sinkabilityProof'.
-- In fact, once 'sinkabilityProof' typechecks,
-- it is safe to 'sink' by coercion.
-- See Section 3.5 in [«The Foil: Capture-Avoiding Substitution With No Sharp Edges»](https://doi.org/10.1145/3587216.3587224) for the details.
sink :: (Sinkable e, DExt n l) => e n -> e l
sink :: forall (e :: S -> *) (n :: S) (l :: S).
(Sinkable e, DExt n l) =>
e n -> e l
sink = e n -> e l
forall a b. a -> b
unsafeCoerce

-- | Extend renaming when going under a 'CoSinkable' pattern (generalized binder).
-- Note that the scope under pattern is independent of the codomain of the renaming.
--
-- This function is used to go under binders when implementing 'sinkabilityProof'
-- and is both a generalization of 'extendRenamingNameBinder' and an efficient implementation of 'coSinkabilityProof'.
extendRenaming
  :: CoSinkable pattern
  => (Name n -> Name n')  -- ^ Map names from scope @n@ to a (possibly larger) scope @n'@.
  -> pattern n l          -- ^ A pattern that extends scope @n@ to another scope @l@.
  -> (forall l'. (Name l -> Name l') -> pattern n' l' -> r )
  -- ^ A continuation, accepting an extended renaming from @l@ to @l'@ (which itself extends @n'@)
  -- and a (possibly refreshed) pattern that extends @n'@ to @l'@.
  -> r
extendRenaming :: forall (pattern :: S -> S -> *) (n :: S) (n' :: S) (l :: S) r.
CoSinkable pattern =>
(Name n -> Name n')
-> pattern n l
-> (forall (l' :: S). (Name l -> Name l') -> pattern n' l' -> r)
-> r
extendRenaming Name n -> Name n'
_ pattern n l
pattern forall (l' :: S). (Name l -> Name l') -> pattern n' l' -> r
cont =
  (Name l -> Name Any) -> pattern n' Any -> r
forall (l' :: S). (Name l -> Name l') -> pattern n' l' -> r
cont Name l -> Name Any
forall a b. a -> b
unsafeCoerce (pattern n l -> pattern n' Any
forall a b. a -> b
unsafeCoerce pattern n l
pattern)

-- | Extend renaming when going under a 'NameBinder'.
-- Note that the scope under binder is independent of the codomain of the renaming.
--
-- Semantically, this function may need to rename the binder (resulting in the new scope @l'@),
-- to make sure it does not clash with scope @n'@.
-- However, as it turns out, the foil makes it safe
-- to implement this function as a coercion.
-- See Appendix A in [«The Foil: Capture-Avoiding Substitution With No Sharp Edges»](https://doi.org/10.1145/3587216.3587224) for the details.
--
-- This function is used to go under binders when implementing 'sinkabilityProof'.
-- A generalization of this function is 'extendRenaming' (which is an efficient version of 'coSinkabilityProof').
extendRenamingNameBinder
  :: (Name n -> Name n')  -- ^ Map names from scope @n@ to a (possibly larger) scope @n'@.
  -> NameBinder n l       -- ^ A name binder that extends scope @n@ to another scope @l@.
  -> (forall l'. (Name l -> Name l') -> NameBinder n' l' -> r )
  -- ^ A continuation, accepting an extended renaming from @l@ to @l'@ (which itself extends @n'@)
  -- and a (possibly refreshed) binder that extends @n'@ to @l'@.
  -> r
extendRenamingNameBinder :: forall (n :: S) (n' :: S) (l :: S) r.
(Name n -> Name n')
-> NameBinder n l
-> (forall (l' :: S). (Name l -> Name l') -> NameBinder n' l' -> r)
-> r
extendRenamingNameBinder Name n -> Name n'
_ (UnsafeNameBinder Name l
name) forall (l' :: S). (Name l -> Name l') -> NameBinder n' l' -> r
cont =
  (Name l -> Name l) -> NameBinder n' l -> r
forall (l' :: S). (Name l -> Name l') -> NameBinder n' l' -> r
cont Name l -> Name l
forall a b. a -> b
unsafeCoerce (Name l -> NameBinder n' l
forall (n :: S) (l :: S). Name l -> NameBinder n l
UnsafeNameBinder Name l
name)

-- | 'CoSinkable' is to patterns (generalized binders)
-- what 'Sinkable' is to expressions.
--
-- See Section 2.3 of [«Free Foil: Generating Efficient and Scope-Safe Abstract Syntax»](https://arxiv.org/abs/2405.16384) for more details.
class CoSinkable (pattern :: S -> S -> Type) where
  -- | An implementation of this method that typechecks
  -- proves to the compiler that the pattern is indeed
  -- 'CoSinkable'. However, instead of this implementation,
  -- 'extendRenaming' should be used at all call sites for efficiency.
  coSinkabilityProof
    :: (Name n -> Name n')  -- ^ Map names from scope @n@ to a (possibly larger) scope @n'@.
    -> pattern n l          -- ^ A pattern that extends scope @n@ to another scope @l@.
    -> (forall l'. (Name l -> Name l') -> pattern n' l' -> r)
    -- ^ A continuation, accepting an extended renaming from @l@ to @l'@ (which itself extends @n'@)
    -- and a (possibly refreshed) pattern that extends @n'@ to @l'@.
    -> r

instance CoSinkable NameBinder where
  coSinkabilityProof :: forall (n :: S) (n' :: S) (l :: S) r.
(Name n -> Name n')
-> NameBinder n l
-> (forall (l' :: S). (Name l -> Name l') -> NameBinder n' l' -> r)
-> r
coSinkabilityProof Name n -> Name n'
_rename (UnsafeNameBinder Name l
name) forall (l' :: S). (Name l -> Name l') -> NameBinder n' l' -> r
cont =
    (Name l -> Name l) -> NameBinder n' l -> r
forall (l' :: S). (Name l -> Name l') -> NameBinder n' l' -> r
cont Name l -> Name l
forall a b. a -> b
unsafeCoerce (Name l -> NameBinder n' l
forall (n :: S) (l :: S). Name l -> NameBinder n l
UnsafeNameBinder Name l
name)

-- * Safe substitions

-- | A substitution is a mapping from names in scope @i@
-- to expressions @e o@ in scope @o@.
newtype Substitution (e :: S -> Type) (i :: S) (o :: S) =
  UnsafeSubstitution (IntMap (e o))

-- | Apply substitution to a given name.
lookupSubst :: InjectName e => Substitution e i o -> Name i -> e o
lookupSubst :: forall (e :: S -> *) (i :: S) (o :: S).
InjectName e =>
Substitution e i o -> Name i -> e o
lookupSubst (UnsafeSubstitution IntMap (e o)
env) (UnsafeName RawName
name) =
    case RawName -> IntMap (e o) -> Maybe (e o)
forall a. RawName -> IntMap a -> Maybe a
IntMap.lookup RawName
name IntMap (e o)
env of
        Just e o
ex -> e o
ex
        Maybe (e o)
Nothing -> Name o -> e o
forall (n :: S). Name n -> e n
forall (e :: S -> *) (n :: S). InjectName e => Name n -> e n
injectName (RawName -> Name o
forall (n :: S). RawName -> Name n
UnsafeName RawName
name)

-- | Identity substitution maps all names to expresion-variables.
identitySubst
  :: InjectName e => Substitution e i i
identitySubst :: forall (e :: S -> *) (i :: S). InjectName e => Substitution e i i
identitySubst = IntMap (e i) -> Substitution e i i
forall (e :: S -> *) (i :: S) (o :: S).
IntMap (e o) -> Substitution e i o
UnsafeSubstitution IntMap (e i)
forall a. IntMap a
IntMap.empty

-- | Extend substitution with a particular mapping.
addSubst
  :: Substitution e i o
  -> NameBinder i i'
  -> e o
  -> Substitution e i' o
addSubst :: forall (e :: S -> *) (i :: S) (o :: S) (i' :: S).
Substitution e i o -> NameBinder i i' -> e o -> Substitution e i' o
addSubst (UnsafeSubstitution IntMap (e o)
env) (UnsafeNameBinder (UnsafeName RawName
name)) e o
ex
  = IntMap (e o) -> Substitution e i' o
forall (e :: S -> *) (i :: S) (o :: S).
IntMap (e o) -> Substitution e i o
UnsafeSubstitution (RawName -> e o -> IntMap (e o) -> IntMap (e o)
forall a. RawName -> a -> IntMap a -> IntMap a
IntMap.insert RawName
name e o
ex IntMap (e o)
env)

-- | Add variable renaming to a substitution.
-- This includes the performance optimization of eliding names mapped to themselves.
addRename :: InjectName e => Substitution e i o -> NameBinder i i' -> Name o -> Substitution e i' o
addRename :: forall (e :: S -> *) (i :: S) (o :: S) (i' :: S).
InjectName e =>
Substitution e i o
-> NameBinder i i' -> Name o -> Substitution e i' o
addRename s :: Substitution e i o
s@(UnsafeSubstitution IntMap (e o)
env) b :: NameBinder i i'
b@(UnsafeNameBinder (UnsafeName RawName
name1)) n :: Name o
n@(UnsafeName RawName
name2)
    | RawName
name1 RawName -> RawName -> Bool
forall a. Eq a => a -> a -> Bool
== RawName
name2 = IntMap (e o) -> Substitution e i' o
forall (e :: S -> *) (i :: S) (o :: S).
IntMap (e o) -> Substitution e i o
UnsafeSubstitution (RawName -> IntMap (e o) -> IntMap (e o)
forall a. RawName -> IntMap a -> IntMap a
IntMap.delete RawName
name1 IntMap (e o)
env)
    | Bool
otherwise = Substitution e i o -> NameBinder i i' -> e o -> Substitution e i' o
forall (e :: S -> *) (i :: S) (o :: S) (i' :: S).
Substitution e i o -> NameBinder i i' -> e o -> Substitution e i' o
addSubst Substitution e i o
s NameBinder i i'
b (Name o -> e o
forall (n :: S). Name n -> e n
forall (e :: S -> *) (n :: S). InjectName e => Name n -> e n
injectName Name o
n)

-- | Substitutions are sinkable as long as corresponding expressions are.
instance (Sinkable e) => Sinkable (Substitution e i) where
  sinkabilityProof :: forall (n :: S) (l :: S).
(Name n -> Name l) -> Substitution e i n -> Substitution e i l
sinkabilityProof Name n -> Name l
rename (UnsafeSubstitution IntMap (e n)
env) =
    IntMap (e l) -> Substitution e i l
forall (e :: S -> *) (i :: S) (o :: S).
IntMap (e o) -> Substitution e i o
UnsafeSubstitution ((e n -> e l) -> IntMap (e n) -> IntMap (e l)
forall a b. (a -> b) -> IntMap a -> IntMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Name n -> Name l) -> e n -> e l
forall (n :: S) (l :: S). (Name n -> Name l) -> e n -> e l
forall (e :: S -> *) (n :: S) (l :: S).
Sinkable e =>
(Name n -> Name l) -> e n -> e l
sinkabilityProof Name n -> Name l
rename) IntMap (e n)
env)

-- * 'Name' maps

-- | A /total/ map from names in scope @n@ to elements of type @a@.
newtype NameMap (n :: S) a = NameMap { forall (n :: S) a. NameMap n a -> IntMap a
getNameMap :: IntMap a }

-- | An empty map belongs in the empty scope.
emptyNameMap :: NameMap VoidS a
emptyNameMap :: forall a. NameMap 'VoidS a
emptyNameMap = IntMap a -> NameMap 'VoidS a
forall (n :: S) a. IntMap a -> NameMap n a
NameMap IntMap a
forall a. IntMap a
IntMap.empty

-- | Looking up a name should always succeed.
--
-- Note that since 'Name' is 'Sinkable', you can lookup a name from scope @n@ in a 'NameMap' for scope @l@ whenever @l@ extends @n@.
lookupName :: Name n -> NameMap n a -> a
lookupName :: forall (n :: S) a. Name n -> NameMap n a -> a
lookupName Name n
name (NameMap IntMap a
m) =
  case RawName -> IntMap a -> Maybe a
forall a. RawName -> IntMap a -> Maybe a
IntMap.lookup (Name n -> RawName
forall (l :: S). Name l -> RawName
nameId Name n
name) IntMap a
m of
    Maybe a
Nothing -> String -> a
forall a. HasCallStack => String -> a
error String
"impossible: unknown name in a NameMap"
    Just a
x  -> a
x

-- | Extending a map with a single mapping.
--
-- Note that the scope parameter of the result differs from the initial map.
addNameBinder :: NameBinder n l -> a -> NameMap n a -> NameMap l a
addNameBinder :: forall (n :: S) (l :: S) a.
NameBinder n l -> a -> NameMap n a -> NameMap l a
addNameBinder NameBinder n l
name a
x (NameMap IntMap a
m) = IntMap a -> NameMap l a
forall (n :: S) a. IntMap a -> NameMap n a
NameMap (RawName -> a -> IntMap a -> IntMap a
forall a. RawName -> a -> IntMap a -> IntMap a
IntMap.insert (Name l -> RawName
forall (l :: S). Name l -> RawName
nameId (NameBinder n l -> Name l
forall (n :: S) (l :: S). NameBinder n l -> Name l
nameOf NameBinder n l
name)) a
x IntMap a
m)

-- * Raw types and operations

-- | We will use 'Int' for efficient representation of identifiers.
type Id = Int

-- | Raw name is simply an identifier.
type RawName = Id

-- | A raw scope is a set of raw names.
type RawScope = IntSet

-- | \(O(\min(n, W))\).
-- Generate a fresh raw name that
-- does not appear in a given raw scope.
rawFreshName :: RawScope -> RawName
rawFreshName :: RawScope -> RawName
rawFreshName RawScope
scope | RawScope -> Bool
IntSet.null RawScope
scope = RawName
0
                   | Bool
otherwise = RawScope -> RawName
IntSet.findMax RawScope
scope RawName -> RawName -> RawName
forall a. Num a => a -> a -> a
+ RawName
1

-- | Check if a raw name is contained in a raw scope.
rawMember :: RawName -> RawScope -> Bool
rawMember :: RawName -> RawScope -> Bool
rawMember = RawName -> RawScope -> Bool
IntSet.member

-- * Constraints

-- | Every scope is a (trivial) extension of itself.
--
-- __Important__: this class exists to assist tracking scope extensions
-- for type variables of kind 'S'.
-- Users of the foil are not supposed to implement any instances of 'ExtEndo'.
class ExtEndo (n :: S)

-- | Some scopes are extensions of other scopes.
--
-- __Important__: this class exists to assist tracking scope extensions
-- for type variables of kind 'S'.
-- Users of the foil are not supposed to implement any instances of 'Ext'.
class (ExtEndo n => ExtEndo l ) => Ext (n :: S) (l :: S)
instance ( ExtEndo n => ExtEndo l ) => Ext n l

-- | Scopes with distinct names.
--
-- __Important__: this class exists to explicitly
-- mark scopes with distinct names.
-- Users of the foil are not supposed to implement any instances of 'Distinct'.
class Distinct (n :: S)
instance Distinct VoidS

-- | Scope extensions with distinct names.
type DExt n l = (Distinct l, Ext n l)

-- | Instances of this typeclass possess the ability to inject names.
-- Usually, this is a variable data constructor.
class InjectName (e :: S -> Type) where
  -- | Inject names into expressions.
  injectName :: Name n -> e n