module Unison.PrettyPrintEnv.MonadPretty
  ( MonadPretty,
    Env (..),
    runPretty,
    addTypeVars,
    willCaptureType,
    withBoundTerm,
    withBoundTerms,
  )
where

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

type MonadPretty v m = (Var v, MonadReader (Env v) m)

-- See Note [Bound and free term variables] for an explanation of boundTerms/freeTerms
data Env v = Env
  { forall v. Env v -> Set v
boundTerms :: !(Set v),
    forall v. Env v -> Set v
boundTypes :: !(Set v),
    forall v. Env v -> Set v
freeTerms :: !(Set v),
    forall v. Env v -> PrettyPrintEnv
ppe :: !PrettyPrintEnv
  }
  deriving stock ((forall x. Env v -> Rep (Env v) x)
-> (forall x. Rep (Env v) x -> Env v) -> Generic (Env v)
forall x. Rep (Env v) x -> Env v
forall x. Env v -> Rep (Env v) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall v x. Rep (Env v) x -> Env v
forall v x. Env v -> Rep (Env v) x
$cfrom :: forall v x. Env v -> Rep (Env v) x
from :: forall x. Env v -> Rep (Env v) x
$cto :: forall v x. Rep (Env v) x -> Env v
to :: forall x. Rep (Env v) x -> Env v
Generic)

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 = (Env v -> Env v) -> m a -> m a
forall a. (Env v -> Env v) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env v -> Env v) -> m a -> m a)
-> ((Set v -> Set v) -> Env v -> Env v)
-> (Set v -> Set v)
-> m a
-> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASetter (Env v) (Env v) (Set v) (Set v)
-> (Set v -> Set v) -> Env v -> Env v
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (Env v) (Env v) (Set v) (Set v)
#boundTypes

-- | 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
willCaptureType :: (MonadPretty v m) => [v] -> m Bool
willCaptureType :: forall v (m :: * -> *). MonadPretty v m => [v] -> m Bool
willCaptureType [v]
vs = LensLike' (Const Bool) (Env 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) (Env v) (Set v)
#boundTypes (Set v -> Set v -> Bool
forall a. Ord a => Set a -> Set a -> Bool
Set.intersects ([v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v]
vs))

withBoundTerm :: (MonadPretty v m) => v -> m a -> m a
withBoundTerm :: forall v (m :: * -> *) a. MonadPretty v m => v -> m a -> m a
withBoundTerm v
v =
  (Env v -> Env v) -> m a -> m a
forall a. (Env v -> Env v) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (ASetter (Env v) (Env v) (Set v) (Set v)
-> (Set v -> Set v) -> Env v -> Env v
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (Env v) (Env v) (Set v) (Set v)
#boundTerms (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert (v -> v
forall v. Var v => v -> v
Var.reset v
v)))

withBoundTerms :: (MonadPretty v m) => [v] -> m a -> m a
withBoundTerms :: forall v (m :: * -> *) a. MonadPretty v m => [v] -> m a -> m a
withBoundTerms [v]
vs =
  (Env v -> Env v) -> m a -> m a
forall a. (Env v -> Env v) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (ASetter (Env v) (Env v) (Set v) (Set v)
-> (Set v -> Set v) -> Env v -> Env v
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (Env v) (Env v) (Set v) (Set v)
#boundTerms (Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
Set.union ((v -> Set v) -> [v] -> Set v
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (v -> Set v
forall a. a -> Set a
Set.singleton (v -> Set v) -> (v -> v) -> v -> Set v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> v
forall v. Var v => v -> v
Var.reset) [v]
vs)))

runPretty :: (Var v) => PrettyPrintEnv -> Reader (Env v) a -> a
runPretty :: forall v a. Var v => PrettyPrintEnv -> Reader (Env v) a -> a
runPretty PrettyPrintEnv
ppe Reader (Env v) a
m =
  Reader (Env v) a -> Env v -> a
forall r a. Reader r a -> r -> a
runReader
    Reader (Env v) a
m
    Env
      { $sel:boundTerms:Env :: Set v
boundTerms = Set v
forall a. Set a
Set.empty,
        $sel:boundTypes:Env :: Set v
boundTypes = Set v
forall a. Set a
Set.empty,
        $sel:freeTerms:Env :: Set v
freeTerms = Set v
forall a. Set a
Set.empty,
        PrettyPrintEnv
$sel:ppe:Env :: PrettyPrintEnv
ppe :: PrettyPrintEnv
ppe
      }

-- Note [Bound and free term variables]
--
-- When rendering a Unison file, we render top-level bindings independently, which may end up referring to each
-- other after all are parsed together. Any individual term, therefore, may have free variables. For example,
--
--     foo = ... bar ...
--               ^^^
--               this "bar" variable is free in foo
--
--     bar = ...
--     ^^^
--     it's ultimately bound by a different top-level term rendering
--
-- Therefore, we pass down all free variables of a top-level term binding, so that we can decide, when rendering one of
-- them, whether to add a leading dot.
--
-- Now, when do we need to add a leading dot? Basically, any time a binder introduces a var that, when rendered reset,
-- clashes with the free var.
--
-- Here are a few examples using a made-up Unison syntax in which we can see whether a let is recursive or
-- non-recursive, and using "%" to separate a var name from its unique id.
--
-- Example 1
--
--   Made-up syntax                      Actual syntax
--   --------------                      ----------------
--   foo%0 =                             foo =
--     let bar%0 = bar%0                   bar = #someref -- rendered as ".bar", then parsed as var "bar"
--     in 5                                5
--
--   bar%0 = ...                         bar = ...
--
-- In this example, we have a non-recursive let that binds a local variable called bar%0. The body of the bar%0 binding
-- can itself refer to a different bar%0, which isn't captured, since a non-recursive let binding body can't refer to
-- the binding.
--
-- So, when rendering the free bar%0 in the right-hand side, we ask: should we add a leading dot? And the answer is: is
-- the var bar%0 in the set of all reset locally-bound vars {bar%0}? Yes? Then yes.
--
-- Example 2
--
--   Made-up syntax                      Actual syntax
--   --------------                      ----------------
--   foo%0 =                             foo =
--     letrec bar%1 = do bar%0 hey%0       bar = do #someref hey -- rendered as ".bar", then parsed as var "bar"
--            hey%0 = do bar%1             hey = do bar
--     in 5                                5
--
--   bar%0 = ...                         bar = ...
--
-- In this example, we have a recursive let that binds a bar%1 variable, and refers to bar%0 from inside. Same
-- situation, but variable resetting is relevant, because when walking underneath binder bar%1, we want to add its reset
-- form (bar%0) to the set of bound variables to check against, when rendering a free var (which we assume to have
-- unique id 0).