{-# LANGUAGE ConstraintKinds #-}

module Unison.PrettyPrintEnv.MonadPretty where

import Control.Lens (views, _1, _2)
import Control.Monad.Reader (MonadReader, Reader, local, runReader)
import Data.Set qualified as Set
import Unison.Prelude
import Unison.PrettyPrintEnv (PrettyPrintEnv)
import Unison.Var (Var)

type MonadPretty v m = (Var v, MonadReader (PrettyPrintEnv, Set v) m)

getPPE :: (MonadPretty v m) => m PrettyPrintEnv
getPPE :: forall v (m :: * -> *). MonadPretty v m => m PrettyPrintEnv
getPPE = Getting PrettyPrintEnv (PrettyPrintEnv, Set v) PrettyPrintEnv
-> m PrettyPrintEnv
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting PrettyPrintEnv (PrettyPrintEnv, Set v) PrettyPrintEnv
forall s t a b. Field1 s t a b => Lens s t a b
Lens
  (PrettyPrintEnv, Set v)
  (PrettyPrintEnv, Set v)
  PrettyPrintEnv
  PrettyPrintEnv
_1

-- | Run a computation with a modified PrettyPrintEnv, restoring the original
withPPE :: (MonadPretty v m) => PrettyPrintEnv -> m a -> m a
withPPE :: forall v (m :: * -> *) a.
MonadPretty v m =>
PrettyPrintEnv -> m a -> m a
withPPE PrettyPrintEnv
p = ((PrettyPrintEnv, Set v) -> (PrettyPrintEnv, Set v)) -> m a -> m a
forall a.
((PrettyPrintEnv, Set v) -> (PrettyPrintEnv, Set v)) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (ASetter
  (PrettyPrintEnv, Set v)
  (PrettyPrintEnv, Set v)
  PrettyPrintEnv
  PrettyPrintEnv
-> PrettyPrintEnv
-> (PrettyPrintEnv, Set v)
-> (PrettyPrintEnv, Set v)
forall s t a b. ASetter s t a b -> b -> s -> t
set ASetter
  (PrettyPrintEnv, Set v)
  (PrettyPrintEnv, Set v)
  PrettyPrintEnv
  PrettyPrintEnv
forall s t a b. Field1 s t a b => Lens s t a b
Lens
  (PrettyPrintEnv, Set v)
  (PrettyPrintEnv, Set v)
  PrettyPrintEnv
  PrettyPrintEnv
_1 PrettyPrintEnv
p)

applyPPE :: (MonadPretty v m) => (PrettyPrintEnv -> a) -> m a
applyPPE :: forall v (m :: * -> *) a.
MonadPretty v m =>
(PrettyPrintEnv -> a) -> m a
applyPPE = LensLike' (Const a) (PrettyPrintEnv, Set v) PrettyPrintEnv
-> (PrettyPrintEnv -> a) -> m a
forall s (m :: * -> *) r a.
MonadReader s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
views LensLike' (Const a) (PrettyPrintEnv, Set v) PrettyPrintEnv
forall s t a b. Field1 s t a b => Lens s t a b
Lens
  (PrettyPrintEnv, Set v)
  (PrettyPrintEnv, Set v)
  PrettyPrintEnv
  PrettyPrintEnv
_1

applyPPE2 :: (MonadPretty v m) => (PrettyPrintEnv -> a -> b) -> a -> m b
applyPPE2 :: forall v (m :: * -> *) a b.
MonadPretty v m =>
(PrettyPrintEnv -> a -> b) -> a -> m b
applyPPE2 PrettyPrintEnv -> a -> b
f a
a = LensLike' (Const b) (PrettyPrintEnv, Set v) PrettyPrintEnv
-> (PrettyPrintEnv -> b) -> m b
forall s (m :: * -> *) r a.
MonadReader s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
views LensLike' (Const b) (PrettyPrintEnv, Set v) PrettyPrintEnv
forall s t a b. Field1 s t a b => Lens s t a b
Lens
  (PrettyPrintEnv, Set v)
  (PrettyPrintEnv, Set v)
  PrettyPrintEnv
  PrettyPrintEnv
_1 (PrettyPrintEnv -> a -> b
`f` a
a)

applyPPE3 :: (MonadPretty v m) => (PrettyPrintEnv -> a -> b -> c) -> a -> b -> m c
applyPPE3 :: forall v (m :: * -> *) a b c.
MonadPretty v m =>
(PrettyPrintEnv -> a -> b -> c) -> a -> b -> m c
applyPPE3 PrettyPrintEnv -> a -> b -> c
f a
a b
b = LensLike' (Const c) (PrettyPrintEnv, Set v) PrettyPrintEnv
-> (PrettyPrintEnv -> c) -> m c
forall s (m :: * -> *) r a.
MonadReader s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
views LensLike' (Const c) (PrettyPrintEnv, Set v) PrettyPrintEnv
forall s t a b. Field1 s t a b => Lens s t a b
Lens
  (PrettyPrintEnv, Set v)
  (PrettyPrintEnv, Set v)
  PrettyPrintEnv
  PrettyPrintEnv
_1 (\PrettyPrintEnv
ppe -> PrettyPrintEnv -> a -> b -> c
f PrettyPrintEnv
ppe a
a b
b)

-- | Run a computation with a modified PrettyPrintEnv, restoring the original
modifyPPE :: (MonadPretty v m) => (PrettyPrintEnv -> PrettyPrintEnv) -> m a -> m a
modifyPPE :: forall v (m :: * -> *) a.
MonadPretty v m =>
(PrettyPrintEnv -> PrettyPrintEnv) -> m a -> m a
modifyPPE = ((PrettyPrintEnv, Set v) -> (PrettyPrintEnv, Set v)) -> m a -> m a
forall a.
((PrettyPrintEnv, Set v) -> (PrettyPrintEnv, Set v)) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (((PrettyPrintEnv, Set v) -> (PrettyPrintEnv, Set v))
 -> m a -> m a)
-> ((PrettyPrintEnv -> PrettyPrintEnv)
    -> (PrettyPrintEnv, Set v) -> (PrettyPrintEnv, Set v))
-> (PrettyPrintEnv -> PrettyPrintEnv)
-> m a
-> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASetter
  (PrettyPrintEnv, Set v)
  (PrettyPrintEnv, Set v)
  PrettyPrintEnv
  PrettyPrintEnv
-> (PrettyPrintEnv -> PrettyPrintEnv)
-> (PrettyPrintEnv, Set v)
-> (PrettyPrintEnv, Set v)
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter
  (PrettyPrintEnv, Set v)
  (PrettyPrintEnv, Set v)
  PrettyPrintEnv
  PrettyPrintEnv
forall s t a b. Field1 s t a b => Lens s t a b
Lens
  (PrettyPrintEnv, Set v)
  (PrettyPrintEnv, Set v)
  PrettyPrintEnv
  PrettyPrintEnv
_1

modifyTypeVars :: (MonadPretty v m) => (Set v -> Set v) -> m a -> m a
modifyTypeVars :: forall v (m :: * -> *) a.
MonadPretty v m =>
(Set v -> Set v) -> m a -> m a
modifyTypeVars = ((PrettyPrintEnv, Set v) -> (PrettyPrintEnv, Set v)) -> m a -> m a
forall a.
((PrettyPrintEnv, Set v) -> (PrettyPrintEnv, Set v)) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (((PrettyPrintEnv, Set v) -> (PrettyPrintEnv, Set v))
 -> m a -> m a)
-> ((Set v -> Set v)
    -> (PrettyPrintEnv, Set v) -> (PrettyPrintEnv, Set v))
-> (Set v -> Set v)
-> m a
-> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASetter
  (PrettyPrintEnv, Set v) (PrettyPrintEnv, Set v) (Set v) (Set v)
-> (Set v -> Set v)
-> (PrettyPrintEnv, Set v)
-> (PrettyPrintEnv, Set v)
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter
  (PrettyPrintEnv, Set v) (PrettyPrintEnv, Set v) (Set v) (Set v)
forall s t a b. Field2 s t a b => Lens s t a b
Lens
  (PrettyPrintEnv, Set v) (PrettyPrintEnv, Set v) (Set v) (Set v)
_2

-- | Add type variables to the set of variables that need to be avoided
addTypeVars :: (MonadPretty v m) => [v] -> m a -> m a
addTypeVars :: forall v (m :: * -> *) a. MonadPretty v m => [v] -> m a -> m a
addTypeVars = (Set v -> Set v) -> m a -> m a
forall v (m :: * -> *) a.
MonadPretty v m =>
(Set v -> Set v) -> m a -> m a
modifyTypeVars ((Set v -> Set v) -> m a -> m a)
-> ([v] -> Set v -> Set v) -> [v] -> m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
Set.union (Set v -> Set v -> Set v)
-> ([v] -> Set v) -> [v] -> Set v -> Set v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList

-- | Check if a list of type variables contains any variables that need to be
-- avoided
willCapture :: (MonadPretty v m) => [v] -> m Bool
willCapture :: forall v (m :: * -> *). MonadPretty v m => [v] -> m Bool
willCapture [v]
vs = LensLike' (Const Bool) (PrettyPrintEnv, Set v) (Set v)
-> (Set v -> Bool) -> m Bool
forall s (m :: * -> *) r a.
MonadReader s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
views LensLike' (Const Bool) (PrettyPrintEnv, Set v) (Set v)
forall s t a b. Field2 s t a b => Lens s t a b
Lens
  (PrettyPrintEnv, Set v) (PrettyPrintEnv, Set v) (Set v) (Set v)
_2 (Bool -> Bool
not (Bool -> Bool) -> (Set v -> Bool) -> Set v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set v -> Bool
forall a. Set a -> Bool
Set.null (Set v -> Bool) -> (Set v -> Set v) -> Set v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
Set.intersection ([v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v]
vs))

runPretty :: (Var v) => PrettyPrintEnv -> Reader (PrettyPrintEnv, Set v) a -> a
runPretty :: forall v a.
Var v =>
PrettyPrintEnv -> Reader (PrettyPrintEnv, Set v) a -> a
runPretty PrettyPrintEnv
ppe Reader (PrettyPrintEnv, Set v) a
m = Reader (PrettyPrintEnv, Set v) a -> (PrettyPrintEnv, Set v) -> a
forall r a. Reader r a -> r -> a
runReader Reader (PrettyPrintEnv, Set v) a
m (PrettyPrintEnv
ppe, Set v
forall a. Monoid a => a
mempty)