{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TemplateHaskell #-}

-- |
-- Module      :  Disco.Effects.LFresh
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Polysemy effect for local fresh name generation, compatible with
-- the unbound-generics library.
module Disco.Effects.LFresh where

import Data.Set (Set)
import qualified Data.Set as S
import Data.Typeable (Typeable)
import Disco.Util (gate, iterUntil)
import Polysemy
import Polysemy.ConstraintAbsorber
import Polysemy.Reader
import qualified Unbound.Generics.LocallyNameless as U
import Unbound.Generics.LocallyNameless.Name

-- | Local fresh name generation effect.
data LFresh m a where
  Lfresh :: Typeable a => Name a -> LFresh m (Name a)
  Avoid :: [AnyName] -> m a -> LFresh m a
  GetAvoids :: LFresh m (Set AnyName)

makeSem ''LFresh

-- | Dispatch an 'LFresh' effect via a 'Reader' effect to keep track
--   of a set of in-scope names.
runLFresh :: Sem (LFresh ': r) a -> Sem r a
runLFresh :: forall (r :: EffectRow) a. Sem (LFresh : r) a -> Sem r a
runLFresh = Set AnyName -> Sem (Reader (Set AnyName) : r) a -> Sem r a
forall i (r :: EffectRow) a. i -> Sem (Reader i : r) a -> Sem r a
runReader Set AnyName
forall a. Set a
S.empty (Sem (Reader (Set AnyName) : r) a -> Sem r a)
-> (Sem (LFresh : r) a -> Sem (Reader (Set AnyName) : r) a)
-> Sem (LFresh : r) a
-> Sem r a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sem (LFresh : r) a -> Sem (Reader (Set AnyName) : r) a
forall (r :: EffectRow) a.
Sem (LFresh : r) a -> Sem (Reader (Set AnyName) : r) a
runLFresh'

runLFresh' :: Sem (LFresh ': r) a -> Sem (Reader (Set AnyName) ': r) a
runLFresh' :: forall (r :: EffectRow) a.
Sem (LFresh : r) a -> Sem (Reader (Set AnyName) : r) a
runLFresh' =
  forall (e1 :: Effect) (e2 :: Effect) (r :: EffectRow) a.
(forall (rInitial :: EffectRow) x.
 e1 (Sem rInitial) x -> Tactical e1 (Sem rInitial) (e2 : r) x)
-> Sem (e1 : r) a -> Sem (e2 : r) a
reinterpretH @_ @(Reader (Set AnyName)) \case
    Lfresh Name a
nm -> do
      let s :: String
s = Name a -> String
forall a. Name a -> String
name2String Name a
nm
      Set AnyName
used <- Sem
  (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r))
  (Set AnyName)
forall i (r :: EffectRow). Member (Reader i) r => Sem r i
ask
      let ok :: Name a -> Bool
ok Name a
n = AnyName -> Set AnyName -> Bool
forall a. Ord a => a -> Set a -> Bool
S.notMember (Name a -> AnyName
forall a. Typeable a => Name a -> AnyName
AnyName Name a
n) Set AnyName
used
      x
-> Sem
     (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r))
     (f x)
forall (f :: * -> *) a (e :: Effect) (m :: * -> *)
       (r :: EffectRow).
Functor f =>
a -> Sem (WithTactics e f m r) (f a)
pureT (x
 -> Sem
      (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r))
      (f x))
-> x
-> Sem
     (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r))
     (f x)
forall a b. (a -> b) -> a -> b
$ (Integer -> Integer) -> (Integer -> Maybe x) -> Integer -> x
forall a b. (a -> a) -> (a -> Maybe b) -> a -> b
iterUntil (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) ((x -> Bool) -> x -> Maybe x
forall (f :: * -> *) a. Alternative f => (a -> Bool) -> a -> f a
gate x -> Bool
Name a -> Bool
ok (x -> Maybe x) -> (Integer -> x) -> Integer -> Maybe x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Integer -> Name a
forall a. String -> Integer -> Name a
makeName String
s) Integer
0
    Avoid [AnyName]
names Sem rInitial x
m -> do
      Sem (LFresh : Reader (Set AnyName) : r) (f x)
m' <- Sem rInitial x
-> Sem
     (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r))
     (Sem (LFresh : Reader (Set AnyName) : r) (f x))
forall (m :: * -> *) a (e :: Effect) (f :: * -> *)
       (r :: EffectRow).
m a -> Sem (WithTactics e f m r) (Sem (e : r) (f a))
runT Sem rInitial x
m
      Sem (Reader (Set AnyName) : r) (f x)
-> Sem
     (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r))
     (f x)
forall (e :: Effect) (r :: EffectRow) a. Sem r a -> Sem (e : r) a
raise (Sem (Reader (Set AnyName) : Reader (Set AnyName) : r) (f x)
-> Sem (Reader (Set AnyName) : r) (f x)
forall (e :: Effect) (r :: EffectRow) a.
Member e r =>
Sem (e : r) a -> Sem r a
subsume (Sem (LFresh : Reader (Set AnyName) : r) (f x)
-> Sem (Reader (Set AnyName) : Reader (Set AnyName) : r) (f x)
forall (r :: EffectRow) a.
Sem (LFresh : r) a -> Sem (Reader (Set AnyName) : r) a
runLFresh' ((Set AnyName -> Set AnyName)
-> Sem (LFresh : Reader (Set AnyName) : r) (f x)
-> Sem (LFresh : Reader (Set AnyName) : r) (f x)
forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local (Set AnyName -> Set AnyName -> Set AnyName
forall a. Ord a => Set a -> Set a -> Set a
S.union ([AnyName] -> Set AnyName
forall a. Ord a => [a] -> Set a
S.fromList [AnyName]
names)) Sem (LFresh : Reader (Set AnyName) : r) (f x)
m')))
    LFresh (Sem rInitial) x
GetAvoids -> Sem
  (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r)) x
forall i (r :: EffectRow). Member (Reader i) r => Sem r i
ask Sem
  (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r)) x
-> (x
    -> Sem
         (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r))
         (f x))
-> Sem
     (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r))
     (f x)
forall a b.
Sem
  (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r)) a
-> (a
    -> Sem
         (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r)) b)
-> Sem
     (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r)) b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= x
-> Sem
     (WithTactics LFresh f (Sem rInitial) (Reader (Set AnyName) : r))
     (f x)
forall (f :: * -> *) a (e :: Effect) (m :: * -> *)
       (r :: EffectRow).
Functor f =>
a -> Sem (WithTactics e f m r) (f a)
pureT

-- Much of the above code copied from
-- https://hackage.haskell.org/package/unbound-generics-0.4.1/docs/src/Unbound.Generics.LocallyNameless.LFresh.html
-- (see instance Monad m => LFresh (LFreshMT m))

-- It turns out to make things much simpler to reimplement the
-- LFresh effect ourselves in terms of a reader effect, since then
-- we can immediately dispatch it as above.  The alternative would
-- be to implement it in terms of (Final U.LFreshM) (see the
-- commented code at the bottom of this file), but then we are stuck
-- with that constraint.  Given the constraint-absorbing machinery
-- below, just impementing the 'LFresh' effect itself means we can
-- then reuse other things from unbound-generics that depend on a
-- Fresh constraint, such as the 'lunbind' function below.

-- NOTE: originally, there was a single function runLFresh which
-- called reinterpretH and then immediately dispatched the Reader
-- (Set AnyName) effect.  However, since runLFresh is recursive,
-- this means that the recursive calls were running with a
-- completely *separate* Reader effect that started over from the
-- empty set! This meant that LFresh basically never changed any
-- names, leading to all sorts of name clashes and crashes.
--
-- Instead, we need to organize things as above: runLFresh' is
-- recursive, and keeps the Reader effect (using 'subsume' to squash
-- the duplicated Reader effects together).  Then a top-level
-- runLFresh function finally runs the Reader effect.

--------------------------------------------------
-- Other functions

-- | Open a binder, automatically freshening the names of the bound
--   variables, and providing the opened pattern and term to the
--   provided continuation.  The bound variables are also added to the
--   set of in-scope variables within in the continuation.
lunbind ::
  (Member LFresh r, U.Alpha p, U.Alpha t) =>
  U.Bind p t ->
  ((p, t) -> Sem r c) ->
  Sem r c
lunbind :: forall (r :: EffectRow) p t c.
(Member LFresh r, Alpha p, Alpha t) =>
Bind p t -> ((p, t) -> Sem r c) -> Sem r c
lunbind Bind p t
b (p, t) -> Sem r c
k = (LFresh (Sem r) => Sem r c) -> Sem r c
forall (r :: EffectRow) a.
Member LFresh r =>
(LFresh (Sem r) => Sem r a) -> Sem r a
absorbLFresh (Bind p t -> ((p, t) -> Sem r c) -> Sem r c
forall (m :: * -> *) p t c.
(LFresh m, Alpha p, Alpha t) =>
Bind p t -> ((p, t) -> m c) -> m c
U.lunbind Bind p t
b (p, t) -> Sem r c
k)

------------------------------------------------------------
-- Machinery for absorbing MTL-style constraint.
-- See https://hackage.haskell.org/package/polysemy-zoo-0.7.0.1/docs/Polysemy-ConstraintAbsorber.html
-- Used https://hackage.haskell.org/package/polysemy-zoo-0.7.0.1/docs/src/Polysemy.ConstraintAbsorber.MonadState.html#absorbState as a template.

absorbLFresh :: Member LFresh r => (U.LFresh (Sem r) => Sem r a) -> Sem r a
absorbLFresh :: forall (r :: EffectRow) a.
Member LFresh r =>
(LFresh (Sem r) => Sem r a) -> Sem r a
absorbLFresh = forall (p :: (* -> *) -> Constraint) (x :: (* -> *) -> * -> * -> *)
       d (r :: EffectRow) a.
d
-> (forall s. Reifies s d :- p (x (Sem r) s))
-> (p (Sem r) => Sem r a)
-> Sem r a
absorbWithSem @U.LFresh @Action ((forall a. Typeable a => Name a -> Sem r (Name a))
-> (forall a. [AnyName] -> Sem r a -> Sem r a)
-> Sem r (Set AnyName)
-> LFreshDict (Sem r)
forall (m :: * -> *).
(forall a. Typeable a => Name a -> m (Name a))
-> (forall a. [AnyName] -> m a -> m a)
-> m (Set AnyName)
-> LFreshDict m
LFreshDict Name a -> Sem r (Name a)
forall (r :: EffectRow) a.
(Member LFresh r, Typeable a) =>
Name a -> Sem r (Name a)
forall a. Typeable a => Name a -> Sem r (Name a)
lfresh [AnyName] -> Sem r a -> Sem r a
forall (r :: EffectRow) a.
Member LFresh r =>
[AnyName] -> Sem r a -> Sem r a
forall a. [AnyName] -> Sem r a -> Sem r a
avoid Sem r (Set AnyName)
forall (r :: EffectRow). Member LFresh r => Sem r (Set AnyName)
getAvoids) ((Reifies s (LFreshDict (Sem r)) =>
 Dict (LFresh (Action (Sem r) s)))
-> Reifies s (LFreshDict (Sem r)) :- LFresh (Action (Sem r) s)
forall (a :: Constraint) (b :: Constraint). (a => Dict b) -> a :- b
Sub Dict (LFresh (Action (Sem r) s))
Reifies s (LFreshDict (Sem r)) => Dict (LFresh (Action (Sem r) s))
forall (a :: Constraint). a => Dict a
Dict)
{-# INLINEABLE absorbLFresh #-}

data LFreshDict m = LFreshDict
  { forall (m :: * -> *).
LFreshDict m -> forall a. Typeable a => Name a -> m (Name a)
lfresh_ :: forall a. Typeable a => Name a -> m (Name a)
  , forall (m :: * -> *).
LFreshDict m -> forall a. [AnyName] -> m a -> m a
avoid_ :: forall a. [AnyName] -> m a -> m a
  , forall (m :: * -> *). LFreshDict m -> m (Set AnyName)
getAvoids_ :: m (Set AnyName)
  }

-- | Wrapper for a monadic action with phantom type parameter for reflection.
--   Locally defined so that the instance we are going to build with reflection
--   must be coherent, that is there cannot be orphans.
newtype Action m s' a = Action (m a)
  deriving ((forall a b. (a -> b) -> Action m s' a -> Action m s' b)
-> (forall a b. a -> Action m s' b -> Action m s' a)
-> Functor (Action m s')
forall a b. a -> Action m s' b -> Action m s' a
forall a b. (a -> b) -> Action m s' a -> Action m s' b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (m :: * -> *) k (s' :: k) a b.
Functor m =>
a -> Action m s' b -> Action m s' a
forall (m :: * -> *) k (s' :: k) a b.
Functor m =>
(a -> b) -> Action m s' a -> Action m s' b
$cfmap :: forall (m :: * -> *) k (s' :: k) a b.
Functor m =>
(a -> b) -> Action m s' a -> Action m s' b
fmap :: forall a b. (a -> b) -> Action m s' a -> Action m s' b
$c<$ :: forall (m :: * -> *) k (s' :: k) a b.
Functor m =>
a -> Action m s' b -> Action m s' a
<$ :: forall a b. a -> Action m s' b -> Action m s' a
Functor, Functor (Action m s')
Functor (Action m s') =>
(forall a. a -> Action m s' a)
-> (forall a b.
    Action m s' (a -> b) -> Action m s' a -> Action m s' b)
-> (forall a b c.
    (a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c)
-> (forall a b. Action m s' a -> Action m s' b -> Action m s' b)
-> (forall a b. Action m s' a -> Action m s' b -> Action m s' a)
-> Applicative (Action m s')
forall a. a -> Action m s' a
forall a b. Action m s' a -> Action m s' b -> Action m s' a
forall a b. Action m s' a -> Action m s' b -> Action m s' b
forall a b. Action m s' (a -> b) -> Action m s' a -> Action m s' b
forall a b c.
(a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall (m :: * -> *) k (s' :: k).
Applicative m =>
Functor (Action m s')
forall (m :: * -> *) k (s' :: k) a.
Applicative m =>
a -> Action m s' a
forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' a -> Action m s' b -> Action m s' a
forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' a -> Action m s' b -> Action m s' b
forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' (a -> b) -> Action m s' a -> Action m s' b
forall (m :: * -> *) k (s' :: k) a b c.
Applicative m =>
(a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c
$cpure :: forall (m :: * -> *) k (s' :: k) a.
Applicative m =>
a -> Action m s' a
pure :: forall a. a -> Action m s' a
$c<*> :: forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' (a -> b) -> Action m s' a -> Action m s' b
<*> :: forall a b. Action m s' (a -> b) -> Action m s' a -> Action m s' b
$cliftA2 :: forall (m :: * -> *) k (s' :: k) a b c.
Applicative m =>
(a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c
liftA2 :: forall a b c.
(a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c
$c*> :: forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' a -> Action m s' b -> Action m s' b
*> :: forall a b. Action m s' a -> Action m s' b -> Action m s' b
$c<* :: forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' a -> Action m s' b -> Action m s' a
<* :: forall a b. Action m s' a -> Action m s' b -> Action m s' a
Applicative, Applicative (Action m s')
Applicative (Action m s') =>
(forall a b.
 Action m s' a -> (a -> Action m s' b) -> Action m s' b)
-> (forall a b. Action m s' a -> Action m s' b -> Action m s' b)
-> (forall a. a -> Action m s' a)
-> Monad (Action m s')
forall a. a -> Action m s' a
forall a b. Action m s' a -> Action m s' b -> Action m s' b
forall a b. Action m s' a -> (a -> Action m s' b) -> Action m s' b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall (m :: * -> *) k (s' :: k).
Monad m =>
Applicative (Action m s')
forall (m :: * -> *) k (s' :: k) a. Monad m => a -> Action m s' a
forall (m :: * -> *) k (s' :: k) a b.
Monad m =>
Action m s' a -> Action m s' b -> Action m s' b
forall (m :: * -> *) k (s' :: k) a b.
Monad m =>
Action m s' a -> (a -> Action m s' b) -> Action m s' b
$c>>= :: forall (m :: * -> *) k (s' :: k) a b.
Monad m =>
Action m s' a -> (a -> Action m s' b) -> Action m s' b
>>= :: forall a b. Action m s' a -> (a -> Action m s' b) -> Action m s' b
$c>> :: forall (m :: * -> *) k (s' :: k) a b.
Monad m =>
Action m s' a -> Action m s' b -> Action m s' b
>> :: forall a b. Action m s' a -> Action m s' b -> Action m s' b
$creturn :: forall (m :: * -> *) k (s' :: k) a. Monad m => a -> Action m s' a
return :: forall a. a -> Action m s' a
Monad)

instance
  ( Monad m
  , Reifies s' (LFreshDict m)
  ) =>
  U.LFresh (Action m s')
  where
  lfresh :: forall a. Typeable a => Name a -> Action m s' (Name a)
lfresh Name a
x = m (Name a) -> Action m s' (Name a)
forall {k} {k} (m :: k -> *) (s' :: k) (a :: k).
m a -> Action m s' a
Action (m (Name a) -> Action m s' (Name a))
-> m (Name a) -> Action m s' (Name a)
forall a b. (a -> b) -> a -> b
$ LFreshDict m -> forall a. Typeable a => Name a -> m (Name a)
forall (m :: * -> *).
LFreshDict m -> forall a. Typeable a => Name a -> m (Name a)
lfresh_ (Proxy s' -> LFreshDict m
forall {k} (s :: k) a (proxy :: k -> *).
Reifies s a =>
proxy s -> a
forall (proxy :: k -> *). proxy s' -> LFreshDict m
reflect (Proxy s' -> LFreshDict m) -> Proxy s' -> LFreshDict m
forall a b. (a -> b) -> a -> b
$ forall (t :: k). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s') Name a
x
  {-# INLINEABLE lfresh #-}
  avoid :: forall a. [AnyName] -> Action m s' a -> Action m s' a
avoid [AnyName]
xs (Action m a
m) = m a -> Action m s' a
forall {k} {k} (m :: k -> *) (s' :: k) (a :: k).
m a -> Action m s' a
Action (m a -> Action m s' a) -> m a -> Action m s' a
forall a b. (a -> b) -> a -> b
$ LFreshDict m -> forall a. [AnyName] -> m a -> m a
forall (m :: * -> *).
LFreshDict m -> forall a. [AnyName] -> m a -> m a
avoid_ (Proxy s' -> LFreshDict m
forall {k} (s :: k) a (proxy :: k -> *).
Reifies s a =>
proxy s -> a
forall (proxy :: k -> *). proxy s' -> LFreshDict m
reflect (Proxy s' -> LFreshDict m) -> Proxy s' -> LFreshDict m
forall a b. (a -> b) -> a -> b
$ forall (t :: k). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s') [AnyName]
xs m a
m
  {-# INLINEABLE avoid #-}
  getAvoids :: Action m s' (Set AnyName)
getAvoids = m (Set AnyName) -> Action m s' (Set AnyName)
forall {k} {k} (m :: k -> *) (s' :: k) (a :: k).
m a -> Action m s' a
Action (m (Set AnyName) -> Action m s' (Set AnyName))
-> m (Set AnyName) -> Action m s' (Set AnyName)
forall a b. (a -> b) -> a -> b
$ LFreshDict m -> m (Set AnyName)
forall (m :: * -> *). LFreshDict m -> m (Set AnyName)
getAvoids_ (Proxy s' -> LFreshDict m
forall {k} (s :: k) a (proxy :: k -> *).
Reifies s a =>
proxy s -> a
forall (proxy :: k -> *). proxy s' -> LFreshDict m
reflect (Proxy s' -> LFreshDict m) -> Proxy s' -> LFreshDict m
forall a b. (a -> b) -> a -> b
$ forall (t :: k). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s')
  {-# INLINEABLE getAvoids #-}

----------------------------------------------------------------------
-- Old code I don't want to delete because I spent so much time
-- banging my head against it.  It wasn't wasted, though, since I used
-- some of my hard-earned knowledge to write runLFresh' above.

-- -- | Dispatch the local fresh name generation effect in an effect stack
-- --   containing the 'LFreshM' monad from @unbound-generics@.
-- runLFreshR :: Member (Final U.LFreshM) r => Sem (LFresh ': r) a -> Sem r a
-- runLFreshR = interpretFinal @U.LFreshM $ \case
--   Avoid xs m  -> do
--     m' <- runS m
--     pure (U.avoid xs m')
--   Lunbind b k -> do
--     s <- getInitialStateS
--     k' <- bindS k
--     pure (U.lunbind b (k' . (<$ s)))

-- -- The above code took me a long time to figure out how to write.
-- -- lunbind is a higher-order effect, so we have to use more
-- -- complicated machinery.  See my Stack Overflow question,
-- -- https://stackoverflow.com/questions/68384508/how-to-incorporate-mtl-style-cps-style-higher-order-effect-into-polysemy/68397358#68397358

-- -- | Run a computation requiring only fresh name generation.
-- runLFresh :: Sem '[LFresh, Final U.LFreshM] a -> a
-- runLFresh = U.runLFreshM . runFinal . runLFreshR