module Covenant.Internal.PrettyPrint
  ( ScopeBoundary (..),
    PrettyContext (..),
    PrettyM,
    runPrettyM,
    bindVars,
    mkForall,
    lookupAbstraction,
  )
where

import Control.Monad.Reader
  ( MonadReader (local),
    Reader,
    asks,
    runReader,
  )
import Covenant.Index
  ( Count,
    Index,
    intCount,
    intIndex,
  )
import Data.Kind (Type)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import GHC.Exts (fromListN)
import Optics.At ()
import Optics.Core
  ( A_Lens,
    LabelOptic (labelOptic),
    ix,
    lens,
    over,
    preview,
    review,
    set,
    view,
    (%),
  )
import Prettyprinter
  ( Doc,
    Pretty (pretty),
    hsep,
    (<+>),
  )

newtype ScopeBoundary = ScopeBoundary Int
  deriving (Int -> ScopeBoundary -> ShowS
[ScopeBoundary] -> ShowS
ScopeBoundary -> String
(Int -> ScopeBoundary -> ShowS)
-> (ScopeBoundary -> String)
-> ([ScopeBoundary] -> ShowS)
-> Show ScopeBoundary
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ScopeBoundary -> ShowS
showsPrec :: Int -> ScopeBoundary -> ShowS
$cshow :: ScopeBoundary -> String
show :: ScopeBoundary -> String
$cshowList :: [ScopeBoundary] -> ShowS
showList :: [ScopeBoundary] -> ShowS
Show, ScopeBoundary -> ScopeBoundary -> Bool
(ScopeBoundary -> ScopeBoundary -> Bool)
-> (ScopeBoundary -> ScopeBoundary -> Bool) -> Eq ScopeBoundary
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ScopeBoundary -> ScopeBoundary -> Bool
== :: ScopeBoundary -> ScopeBoundary -> Bool
$c/= :: ScopeBoundary -> ScopeBoundary -> Bool
/= :: ScopeBoundary -> ScopeBoundary -> Bool
Eq, Eq ScopeBoundary
Eq ScopeBoundary =>
(ScopeBoundary -> ScopeBoundary -> Ordering)
-> (ScopeBoundary -> ScopeBoundary -> Bool)
-> (ScopeBoundary -> ScopeBoundary -> Bool)
-> (ScopeBoundary -> ScopeBoundary -> Bool)
-> (ScopeBoundary -> ScopeBoundary -> Bool)
-> (ScopeBoundary -> ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary -> ScopeBoundary -> ScopeBoundary)
-> Ord ScopeBoundary
ScopeBoundary -> ScopeBoundary -> Bool
ScopeBoundary -> ScopeBoundary -> Ordering
ScopeBoundary -> ScopeBoundary -> ScopeBoundary
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
$ccompare :: ScopeBoundary -> ScopeBoundary -> Ordering
compare :: ScopeBoundary -> ScopeBoundary -> Ordering
$c< :: ScopeBoundary -> ScopeBoundary -> Bool
< :: ScopeBoundary -> ScopeBoundary -> Bool
$c<= :: ScopeBoundary -> ScopeBoundary -> Bool
<= :: ScopeBoundary -> ScopeBoundary -> Bool
$c> :: ScopeBoundary -> ScopeBoundary -> Bool
> :: ScopeBoundary -> ScopeBoundary -> Bool
$c>= :: ScopeBoundary -> ScopeBoundary -> Bool
>= :: ScopeBoundary -> ScopeBoundary -> Bool
$cmax :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
max :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
$cmin :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
min :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
Ord, Integer -> ScopeBoundary
ScopeBoundary -> ScopeBoundary
ScopeBoundary -> ScopeBoundary -> ScopeBoundary
(ScopeBoundary -> ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary -> ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary -> ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary -> ScopeBoundary)
-> (Integer -> ScopeBoundary)
-> Num ScopeBoundary
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
$c+ :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
+ :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
$c- :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
- :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
$c* :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
* :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
$cnegate :: ScopeBoundary -> ScopeBoundary
negate :: ScopeBoundary -> ScopeBoundary
$cabs :: ScopeBoundary -> ScopeBoundary
abs :: ScopeBoundary -> ScopeBoundary
$csignum :: ScopeBoundary -> ScopeBoundary
signum :: ScopeBoundary -> ScopeBoundary
$cfromInteger :: Integer -> ScopeBoundary
fromInteger :: Integer -> ScopeBoundary
Num, Num ScopeBoundary
Ord ScopeBoundary
(Num ScopeBoundary, Ord ScopeBoundary) =>
(ScopeBoundary -> Rational) -> Real ScopeBoundary
ScopeBoundary -> Rational
forall a. (Num a, Ord a) => (a -> Rational) -> Real a
$ctoRational :: ScopeBoundary -> Rational
toRational :: ScopeBoundary -> Rational
Real, Int -> ScopeBoundary
ScopeBoundary -> Int
ScopeBoundary -> [ScopeBoundary]
ScopeBoundary -> ScopeBoundary
ScopeBoundary -> ScopeBoundary -> [ScopeBoundary]
ScopeBoundary -> ScopeBoundary -> ScopeBoundary -> [ScopeBoundary]
(ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary -> ScopeBoundary)
-> (Int -> ScopeBoundary)
-> (ScopeBoundary -> Int)
-> (ScopeBoundary -> [ScopeBoundary])
-> (ScopeBoundary -> ScopeBoundary -> [ScopeBoundary])
-> (ScopeBoundary -> ScopeBoundary -> [ScopeBoundary])
-> (ScopeBoundary
    -> ScopeBoundary -> ScopeBoundary -> [ScopeBoundary])
-> Enum ScopeBoundary
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: ScopeBoundary -> ScopeBoundary
succ :: ScopeBoundary -> ScopeBoundary
$cpred :: ScopeBoundary -> ScopeBoundary
pred :: ScopeBoundary -> ScopeBoundary
$ctoEnum :: Int -> ScopeBoundary
toEnum :: Int -> ScopeBoundary
$cfromEnum :: ScopeBoundary -> Int
fromEnum :: ScopeBoundary -> Int
$cenumFrom :: ScopeBoundary -> [ScopeBoundary]
enumFrom :: ScopeBoundary -> [ScopeBoundary]
$cenumFromThen :: ScopeBoundary -> ScopeBoundary -> [ScopeBoundary]
enumFromThen :: ScopeBoundary -> ScopeBoundary -> [ScopeBoundary]
$cenumFromTo :: ScopeBoundary -> ScopeBoundary -> [ScopeBoundary]
enumFromTo :: ScopeBoundary -> ScopeBoundary -> [ScopeBoundary]
$cenumFromThenTo :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary -> [ScopeBoundary]
enumFromThenTo :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary -> [ScopeBoundary]
Enum, Enum ScopeBoundary
Real ScopeBoundary
(Real ScopeBoundary, Enum ScopeBoundary) =>
(ScopeBoundary -> ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary -> ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary -> ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary -> ScopeBoundary -> ScopeBoundary)
-> (ScopeBoundary
    -> ScopeBoundary -> (ScopeBoundary, ScopeBoundary))
-> (ScopeBoundary
    -> ScopeBoundary -> (ScopeBoundary, ScopeBoundary))
-> (ScopeBoundary -> Integer)
-> Integral ScopeBoundary
ScopeBoundary -> Integer
ScopeBoundary -> ScopeBoundary -> (ScopeBoundary, ScopeBoundary)
ScopeBoundary -> ScopeBoundary -> ScopeBoundary
forall a.
(Real a, Enum a) =>
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> (a, a))
-> (a -> a -> (a, a))
-> (a -> Integer)
-> Integral a
$cquot :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
quot :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
$crem :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
rem :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
$cdiv :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
div :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
$cmod :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
mod :: ScopeBoundary -> ScopeBoundary -> ScopeBoundary
$cquotRem :: ScopeBoundary -> ScopeBoundary -> (ScopeBoundary, ScopeBoundary)
quotRem :: ScopeBoundary -> ScopeBoundary -> (ScopeBoundary, ScopeBoundary)
$cdivMod :: ScopeBoundary -> ScopeBoundary -> (ScopeBoundary, ScopeBoundary)
divMod :: ScopeBoundary -> ScopeBoundary -> (ScopeBoundary, ScopeBoundary)
$ctoInteger :: ScopeBoundary -> Integer
toInteger :: ScopeBoundary -> Integer
Integral) via Int

-- Keeping the field names for clarity even if we don't use them
data PrettyContext (ann :: Type)
  = PrettyContext
  { forall ann.
PrettyContext ann -> Map ScopeBoundary (Vector (Doc ann))
_boundIdents :: Map ScopeBoundary (Vector (Doc ann)),
    forall ann. PrettyContext ann -> ScopeBoundary
_currentScope :: ScopeBoundary,
    forall ann. PrettyContext ann -> [Doc ann]
_varStream :: [Doc ann]
  }

instance
  (k ~ A_Lens, a ~ Map ScopeBoundary (Vector (Doc ann)), b ~ Map ScopeBoundary (Vector (Doc ann))) =>
  LabelOptic "boundIdents" k (PrettyContext ann) (PrettyContext ann) a b
  where
  {-# INLINEABLE labelOptic #-}
  labelOptic :: Optic k NoIx (PrettyContext ann) (PrettyContext ann) a b
labelOptic =
    (PrettyContext ann -> a)
-> (PrettyContext ann -> b -> PrettyContext ann)
-> Lens (PrettyContext ann) (PrettyContext ann) a b
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens
      (\(PrettyContext Map ScopeBoundary (Vector (Doc ann))
x ScopeBoundary
_ [Doc ann]
_) -> a
Map ScopeBoundary (Vector (Doc ann))
x)
      (\(PrettyContext Map ScopeBoundary (Vector (Doc ann))
_ ScopeBoundary
y [Doc ann]
z) b
x -> Map ScopeBoundary (Vector (Doc ann))
-> ScopeBoundary -> [Doc ann] -> PrettyContext ann
forall ann.
Map ScopeBoundary (Vector (Doc ann))
-> ScopeBoundary -> [Doc ann] -> PrettyContext ann
PrettyContext b
Map ScopeBoundary (Vector (Doc ann))
x ScopeBoundary
y [Doc ann]
z)

instance
  (k ~ A_Lens, a ~ ScopeBoundary, b ~ ScopeBoundary) =>
  LabelOptic "currentScope" k (PrettyContext ann) (PrettyContext ann) a b
  where
  {-# INLINEABLE labelOptic #-}
  labelOptic :: Optic k NoIx (PrettyContext ann) (PrettyContext ann) a b
labelOptic =
    (PrettyContext ann -> a)
-> (PrettyContext ann -> b -> PrettyContext ann)
-> Lens (PrettyContext ann) (PrettyContext ann) a b
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens
      (\(PrettyContext Map ScopeBoundary (Vector (Doc ann))
_ ScopeBoundary
x [Doc ann]
_) -> a
ScopeBoundary
x)
      (\(PrettyContext Map ScopeBoundary (Vector (Doc ann))
x ScopeBoundary
_ [Doc ann]
z) b
y -> Map ScopeBoundary (Vector (Doc ann))
-> ScopeBoundary -> [Doc ann] -> PrettyContext ann
forall ann.
Map ScopeBoundary (Vector (Doc ann))
-> ScopeBoundary -> [Doc ann] -> PrettyContext ann
PrettyContext Map ScopeBoundary (Vector (Doc ann))
x b
ScopeBoundary
y [Doc ann]
z)

instance
  (k ~ A_Lens, a ~ [Doc ann], b ~ [Doc ann]) =>
  LabelOptic "varStream" k (PrettyContext ann) (PrettyContext ann) a b
  where
  {-# INLINEABLE labelOptic #-}
  labelOptic :: Optic k NoIx (PrettyContext ann) (PrettyContext ann) a b
labelOptic =
    (PrettyContext ann -> a)
-> (PrettyContext ann -> b -> PrettyContext ann)
-> Lens (PrettyContext ann) (PrettyContext ann) a b
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens
      (\(PrettyContext Map ScopeBoundary (Vector (Doc ann))
_ ScopeBoundary
_ [Doc ann]
x) -> a
[Doc ann]
x)
      (\(PrettyContext Map ScopeBoundary (Vector (Doc ann))
x ScopeBoundary
y [Doc ann]
_) b
z -> Map ScopeBoundary (Vector (Doc ann))
-> ScopeBoundary -> [Doc ann] -> PrettyContext ann
forall ann.
Map ScopeBoundary (Vector (Doc ann))
-> ScopeBoundary -> [Doc ann] -> PrettyContext ann
PrettyContext Map ScopeBoundary (Vector (Doc ann))
x ScopeBoundary
y b
[Doc ann]
z)

-- Maybe make a newtype with error reporting since this can fail, but do later since *should't* fail
newtype PrettyM (ann :: Type) (a :: Type) = PrettyM (Reader (PrettyContext ann) a)
  deriving
    ( (forall a b. (a -> b) -> PrettyM ann a -> PrettyM ann b)
-> (forall a b. a -> PrettyM ann b -> PrettyM ann a)
-> Functor (PrettyM ann)
forall a b. a -> PrettyM ann b -> PrettyM ann a
forall a b. (a -> b) -> PrettyM ann a -> PrettyM ann b
forall ann a b. a -> PrettyM ann b -> PrettyM ann a
forall ann a b. (a -> b) -> PrettyM ann a -> PrettyM ann b
forall (f :: Type -> Type).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall ann a b. (a -> b) -> PrettyM ann a -> PrettyM ann b
fmap :: forall a b. (a -> b) -> PrettyM ann a -> PrettyM ann b
$c<$ :: forall ann a b. a -> PrettyM ann b -> PrettyM ann a
<$ :: forall a b. a -> PrettyM ann b -> PrettyM ann a
Functor,
      Functor (PrettyM ann)
Functor (PrettyM ann) =>
(forall a. a -> PrettyM ann a)
-> (forall a b.
    PrettyM ann (a -> b) -> PrettyM ann a -> PrettyM ann b)
-> (forall a b c.
    (a -> b -> c) -> PrettyM ann a -> PrettyM ann b -> PrettyM ann c)
-> (forall a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann b)
-> (forall a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann a)
-> Applicative (PrettyM ann)
forall ann. Functor (PrettyM ann)
forall a. a -> PrettyM ann a
forall ann a. a -> PrettyM ann a
forall a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann a
forall a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann b
forall a b. PrettyM ann (a -> b) -> PrettyM ann a -> PrettyM ann b
forall ann a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann a
forall ann a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann b
forall ann a b.
PrettyM ann (a -> b) -> PrettyM ann a -> PrettyM ann b
forall a b c.
(a -> b -> c) -> PrettyM ann a -> PrettyM ann b -> PrettyM ann c
forall ann a b c.
(a -> b -> c) -> PrettyM ann a -> PrettyM ann b -> PrettyM ann c
forall (f :: Type -> Type).
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
$cpure :: forall ann a. a -> PrettyM ann a
pure :: forall a. a -> PrettyM ann a
$c<*> :: forall ann a b.
PrettyM ann (a -> b) -> PrettyM ann a -> PrettyM ann b
<*> :: forall a b. PrettyM ann (a -> b) -> PrettyM ann a -> PrettyM ann b
$cliftA2 :: forall ann a b c.
(a -> b -> c) -> PrettyM ann a -> PrettyM ann b -> PrettyM ann c
liftA2 :: forall a b c.
(a -> b -> c) -> PrettyM ann a -> PrettyM ann b -> PrettyM ann c
$c*> :: forall ann a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann b
*> :: forall a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann b
$c<* :: forall ann a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann a
<* :: forall a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann a
Applicative,
      Applicative (PrettyM ann)
Applicative (PrettyM ann) =>
(forall a b.
 PrettyM ann a -> (a -> PrettyM ann b) -> PrettyM ann b)
-> (forall a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann b)
-> (forall a. a -> PrettyM ann a)
-> Monad (PrettyM ann)
forall ann. Applicative (PrettyM ann)
forall a. a -> PrettyM ann a
forall ann a. a -> PrettyM ann a
forall a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann b
forall a b. PrettyM ann a -> (a -> PrettyM ann b) -> PrettyM ann b
forall ann a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann b
forall ann a b.
PrettyM ann a -> (a -> PrettyM ann b) -> PrettyM ann b
forall (m :: Type -> Type).
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
$c>>= :: forall ann a b.
PrettyM ann a -> (a -> PrettyM ann b) -> PrettyM ann b
>>= :: forall a b. PrettyM ann a -> (a -> PrettyM ann b) -> PrettyM ann b
$c>> :: forall ann a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann b
>> :: forall a b. PrettyM ann a -> PrettyM ann b -> PrettyM ann b
$creturn :: forall ann a. a -> PrettyM ann a
return :: forall a. a -> PrettyM ann a
Monad,
      MonadReader (PrettyContext ann)
    )
    via (Reader (PrettyContext ann))

runPrettyM :: forall (ann :: Type) (a :: Type). PrettyM ann a -> a
runPrettyM :: forall ann a. PrettyM ann a -> a
runPrettyM (PrettyM Reader (PrettyContext ann) a
ma) = Reader (PrettyContext ann) a -> PrettyContext ann -> a
forall r a. Reader r a -> r -> a
runReader Reader (PrettyContext ann) a
ma (Map ScopeBoundary (Vector (Doc ann))
-> ScopeBoundary -> [Doc ann] -> PrettyContext ann
forall ann.
Map ScopeBoundary (Vector (Doc ann))
-> ScopeBoundary -> [Doc ann] -> PrettyContext ann
PrettyContext Map ScopeBoundary (Vector (Doc ann))
forall a. Monoid a => a
mempty ScopeBoundary
0 [Doc ann]
infiniteVars)
  where
    -- Lazily generated infinite list of variables. Will start with a, b, c...
    -- and cycle around to a1, b2, c3 etc.
    -- We could do something more sophisticated but this should work.
    infiniteVars :: [Doc ann]
    infiniteVars :: [Doc ann]
infiniteVars =
      let aToZ :: String
aToZ = [Char
'a' .. Char
'z']
          intStrings :: [String]
intStrings = (String
"" String -> String -> [String]
forall a b. a -> [b] -> [a]
forall (f :: Type -> Type) a b. Functor f => a -> f b -> f a
<$ String
aToZ) [String] -> [String] -> [String]
forall a. Semigroup a => a -> a -> a
<> (Integer -> String) -> [Integer] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (forall a. Show a => a -> String
show @Integer) [Integer
0 ..]
       in (Char -> String -> Doc ann) -> String -> [String] -> [Doc ann]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Char
x String
xs -> String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Char
x Char -> ShowS
forall a. a -> [a] -> [a]
: String
xs)) String
aToZ [String]
intStrings

bindVars ::
  forall (ann :: Type) (a :: Type).
  Count "tyvar" ->
  (Vector (Doc ann) -> PrettyM ann a) ->
  PrettyM ann a
bindVars :: forall ann a.
Count "tyvar"
-> (Vector (Doc ann) -> PrettyM ann a) -> PrettyM ann a
bindVars Count "tyvar"
count' Vector (Doc ann) -> PrettyM ann a
act
  | Int
count Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = PrettyM ann a -> PrettyM ann a
crossBoundary (Vector (Doc ann) -> PrettyM ann a
act Vector (Doc ann)
forall a. Vector a
Vector.empty)
  | Bool
otherwise = PrettyM ann a -> PrettyM ann a
crossBoundary (PrettyM ann a -> PrettyM ann a) -> PrettyM ann a -> PrettyM ann a
forall a b. (a -> b) -> a -> b
$ do
      ScopeBoundary
here <- (PrettyContext ann -> ScopeBoundary) -> PrettyM ann ScopeBoundary
forall r (m :: Type -> Type) a. MonadReader r m => (r -> a) -> m a
asks (Optic
  A_Lens
  NoIx
  (PrettyContext ann)
  (PrettyContext ann)
  ScopeBoundary
  ScopeBoundary
-> PrettyContext ann -> ScopeBoundary
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
view Optic
  A_Lens
  NoIx
  (PrettyContext ann)
  (PrettyContext ann)
  ScopeBoundary
  ScopeBoundary
#currentScope)
      Int -> (Vector (Doc ann) -> PrettyM ann a) -> PrettyM ann a
forall ann a.
Int -> (Vector (Doc ann) -> PrettyM ann a) -> PrettyM ann a
withFreshVarNames Int
count ((Vector (Doc ann) -> PrettyM ann a) -> PrettyM ann a)
-> (Vector (Doc ann) -> PrettyM ann a) -> PrettyM ann a
forall a b. (a -> b) -> a -> b
$ \Vector (Doc ann)
newBoundVars ->
        (PrettyContext ann -> PrettyContext ann)
-> PrettyM ann a -> PrettyM ann a
forall a.
(PrettyContext ann -> PrettyContext ann)
-> PrettyM ann a -> PrettyM ann a
forall r (m :: Type -> Type) a.
MonadReader r m =>
(r -> r) -> m a -> m a
local (Optic
  A_Lens
  NoIx
  (PrettyContext ann)
  (PrettyContext ann)
  (Map ScopeBoundary (Vector (Doc ann)))
  (Map ScopeBoundary (Vector (Doc ann)))
-> (Map ScopeBoundary (Vector (Doc ann))
    -> Map ScopeBoundary (Vector (Doc ann)))
-> PrettyContext ann
-> PrettyContext ann
forall k (is :: IxList) s t a b.
Is k A_Setter =>
Optic k is s t a b -> (a -> b) -> s -> t
over Optic
  A_Lens
  NoIx
  (PrettyContext ann)
  (PrettyContext ann)
  (Map ScopeBoundary (Vector (Doc ann)))
  (Map ScopeBoundary (Vector (Doc ann)))
#boundIdents (ScopeBoundary
-> Vector (Doc ann)
-> Map ScopeBoundary (Vector (Doc ann))
-> Map ScopeBoundary (Vector (Doc ann))
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ScopeBoundary
here Vector (Doc ann)
newBoundVars)) (Vector (Doc ann) -> PrettyM ann a
act Vector (Doc ann)
newBoundVars)
  where
    -- Increment the current scope
    crossBoundary :: PrettyM ann a -> PrettyM ann a
    crossBoundary :: PrettyM ann a -> PrettyM ann a
crossBoundary = (PrettyContext ann -> PrettyContext ann)
-> PrettyM ann a -> PrettyM ann a
forall a.
(PrettyContext ann -> PrettyContext ann)
-> PrettyM ann a -> PrettyM ann a
forall r (m :: Type -> Type) a.
MonadReader r m =>
(r -> r) -> m a -> m a
local (Optic
  A_Lens
  NoIx
  (PrettyContext ann)
  (PrettyContext ann)
  ScopeBoundary
  ScopeBoundary
-> (ScopeBoundary -> ScopeBoundary)
-> PrettyContext ann
-> PrettyContext ann
forall k (is :: IxList) s t a b.
Is k A_Setter =>
Optic k is s t a b -> (a -> b) -> s -> t
over Optic
  A_Lens
  NoIx
  (PrettyContext ann)
  (PrettyContext ann)
  ScopeBoundary
  ScopeBoundary
#currentScope (ScopeBoundary -> ScopeBoundary -> ScopeBoundary
forall a. Num a => a -> a -> a
+ ScopeBoundary
1))
    count :: Int
    count :: Int
count = Optic' A_Prism NoIx Int (Count "tyvar") -> Count "tyvar" -> Int
forall k (is :: IxList) t b.
Is k A_Review =>
Optic' k is t b -> b -> t
review Optic' A_Prism NoIx Int (Count "tyvar")
forall (ofWhat :: Symbol). Prism' Int (Count ofWhat)
intCount Count "tyvar"
count'

mkForall ::
  forall (ann :: Type).
  Vector (Doc ann) ->
  Doc ann ->
  Doc ann
mkForall :: forall ann. Vector (Doc ann) -> Doc ann -> Doc ann
mkForall Vector (Doc ann)
tvars Doc ann
funTyBody =
  if Vector (Doc ann) -> Bool
forall a. Vector a -> Bool
Vector.null Vector (Doc ann)
tvars
    then Doc ann
funTyBody
    else Doc ann
"forall" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
hsep (Vector (Doc ann) -> [Doc ann]
forall a. Vector a -> [a]
Vector.toList Vector (Doc ann)
tvars) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"." Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
funTyBody

lookupAbstraction :: forall (ann :: Type). Int -> Index "tyvar" -> PrettyM ann (Doc ann)
lookupAbstraction :: forall ann. Int -> Index "tyvar" -> PrettyM ann (Doc ann)
lookupAbstraction Int
offset Index "tyvar"
argIndex = do
  let scopeOffset :: ScopeBoundary
scopeOffset = Int -> ScopeBoundary
ScopeBoundary Int
offset
  let argIndex' :: Int
argIndex' = Optic' A_Prism NoIx Int (Index "tyvar") -> Index "tyvar" -> Int
forall k (is :: IxList) t b.
Is k A_Review =>
Optic' k is t b -> b -> t
review Optic' A_Prism NoIx Int (Index "tyvar")
forall (ofWhat :: Symbol). Prism' Int (Index ofWhat)
intIndex Index "tyvar"
argIndex
  ScopeBoundary
here <- (PrettyContext ann -> ScopeBoundary) -> PrettyM ann ScopeBoundary
forall r (m :: Type -> Type) a. MonadReader r m => (r -> a) -> m a
asks (Optic' A_Lens NoIx (PrettyContext ann) ScopeBoundary
-> PrettyContext ann -> ScopeBoundary
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
view Optic' A_Lens NoIx (PrettyContext ann) ScopeBoundary
#currentScope)
  (PrettyContext ann -> Maybe (Doc ann))
-> PrettyM ann (Maybe (Doc ann))
forall r (m :: Type -> Type) a. MonadReader r m => (r -> a) -> m a
asks (Optic' An_AffineTraversal NoIx (PrettyContext ann) (Doc ann)
-> PrettyContext ann -> Maybe (Doc ann)
forall k (is :: IxList) s a.
Is k An_AffineFold =>
Optic' k is s a -> s -> Maybe a
preview (Optic
  A_Lens
  NoIx
  (PrettyContext ann)
  (PrettyContext ann)
  (Map ScopeBoundary (Vector (Doc ann)))
  (Map ScopeBoundary (Vector (Doc ann)))
#boundIdents Optic
  A_Lens
  NoIx
  (PrettyContext ann)
  (PrettyContext ann)
  (Map ScopeBoundary (Vector (Doc ann)))
  (Map ScopeBoundary (Vector (Doc ann)))
-> Optic
     (IxKind (Map ScopeBoundary (Vector (Doc ann))))
     NoIx
     (Map ScopeBoundary (Vector (Doc ann)))
     (Map ScopeBoundary (Vector (Doc ann)))
     (IxValue (Map ScopeBoundary (Vector (Doc ann))))
     (IxValue (Map ScopeBoundary (Vector (Doc ann))))
-> Optic
     An_AffineTraversal
     NoIx
     (PrettyContext ann)
     (PrettyContext ann)
     (IxValue (Map ScopeBoundary (Vector (Doc ann))))
     (IxValue (Map ScopeBoundary (Vector (Doc ann))))
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% Index (Map ScopeBoundary (Vector (Doc ann)))
-> Optic
     (IxKind (Map ScopeBoundary (Vector (Doc ann))))
     NoIx
     (Map ScopeBoundary (Vector (Doc ann)))
     (Map ScopeBoundary (Vector (Doc ann)))
     (IxValue (Map ScopeBoundary (Vector (Doc ann))))
     (IxValue (Map ScopeBoundary (Vector (Doc ann))))
forall m. Ixed m => Index m -> Optic' (IxKind m) NoIx m (IxValue m)
ix (ScopeBoundary
here ScopeBoundary -> ScopeBoundary -> ScopeBoundary
forall a. Num a => a -> a -> a
+ ScopeBoundary
scopeOffset) Optic
  An_AffineTraversal
  NoIx
  (PrettyContext ann)
  (PrettyContext ann)
  (IxValue (Map ScopeBoundary (Vector (Doc ann))))
  (IxValue (Map ScopeBoundary (Vector (Doc ann))))
-> Optic
     (IxKind (IxValue (Map ScopeBoundary (Vector (Doc ann)))))
     NoIx
     (IxValue (Map ScopeBoundary (Vector (Doc ann))))
     (IxValue (Map ScopeBoundary (Vector (Doc ann))))
     (Doc ann)
     (Doc ann)
-> Optic' An_AffineTraversal NoIx (PrettyContext ann) (Doc ann)
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% Index (IxValue (Map ScopeBoundary (Vector (Doc ann))))
-> Optic'
     (IxKind (IxValue (Map ScopeBoundary (Vector (Doc ann)))))
     NoIx
     (IxValue (Map ScopeBoundary (Vector (Doc ann))))
     (IxValue (IxValue (Map ScopeBoundary (Vector (Doc ann)))))
forall m. Ixed m => Index m -> Optic' (IxKind m) NoIx m (IxValue m)
ix Int
Index (IxValue (Map ScopeBoundary (Vector (Doc ann))))
argIndex')) PrettyM ann (Maybe (Doc ann))
-> (Maybe (Doc ann) -> PrettyM ann (Doc ann))
-> PrettyM ann (Doc ann)
forall a b. PrettyM ann a -> (a -> PrettyM ann b) -> PrettyM ann b
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe (Doc ann)
Nothing ->
      -- TODO: actual error reporting
      String -> PrettyM ann (Doc ann)
forall a. HasCallStack => String -> a
error (String -> PrettyM ann (Doc ann))
-> String -> PrettyM ann (Doc ann)
forall a b. (a -> b) -> a -> b
$
        String
"Internal error: The encountered a variable at arg index "
          String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
argIndex'
          String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" with true level "
          String -> ShowS
forall a. Semigroup a => a -> a -> a
<> ScopeBoundary -> String
forall a. Show a => a -> String
show ScopeBoundary
scopeOffset
          String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" but could not locate the corresponding pretty form at scope level "
          String -> ShowS
forall a. Semigroup a => a -> a -> a
<> ScopeBoundary -> String
forall a. Show a => a -> String
show ScopeBoundary
here
    Just Doc ann
res' -> Doc ann -> PrettyM ann (Doc ann)
forall a. a -> PrettyM ann a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Doc ann
res'

-- Helpers

-- Generate N fresh var names and use the supplied monadic function to do something with them.
withFreshVarNames ::
  forall (ann :: Type) (a :: Type).
  Int ->
  (Vector (Doc ann) -> PrettyM ann a) ->
  PrettyM ann a
withFreshVarNames :: forall ann a.
Int -> (Vector (Doc ann) -> PrettyM ann a) -> PrettyM ann a
withFreshVarNames Int
n Vector (Doc ann) -> PrettyM ann a
act = do
  [Doc ann]
stream <- (PrettyContext ann -> [Doc ann]) -> PrettyM ann [Doc ann]
forall r (m :: Type -> Type) a. MonadReader r m => (r -> a) -> m a
asks (Optic' A_Lens NoIx (PrettyContext ann) [Doc ann]
-> PrettyContext ann -> [Doc ann]
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
view Optic' A_Lens NoIx (PrettyContext ann) [Doc ann]
#varStream)
  let ([Doc ann]
used, [Doc ann]
rest) = Int -> [Doc ann] -> ([Doc ann], [Doc ann])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [Doc ann]
stream
  (PrettyContext ann -> PrettyContext ann)
-> PrettyM ann a -> PrettyM ann a
forall a.
(PrettyContext ann -> PrettyContext ann)
-> PrettyM ann a -> PrettyM ann a
forall r (m :: Type -> Type) a.
MonadReader r m =>
(r -> r) -> m a -> m a
local (Optic' A_Lens NoIx (PrettyContext ann) [Doc ann]
-> [Doc ann] -> PrettyContext ann -> PrettyContext ann
forall k (is :: IxList) s t a b.
Is k A_Setter =>
Optic k is s t a b -> b -> s -> t
set Optic' A_Lens NoIx (PrettyContext ann) [Doc ann]
#varStream [Doc ann]
rest) (PrettyM ann a -> PrettyM ann a)
-> ([Doc ann] -> PrettyM ann a) -> [Doc ann] -> PrettyM ann a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Doc ann) -> PrettyM ann a
act (Vector (Doc ann) -> PrettyM ann a)
-> ([Doc ann] -> Vector (Doc ann)) -> [Doc ann] -> PrettyM ann a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Item (Vector (Doc ann))] -> Vector (Doc ann)
forall l. IsList l => Int -> [Item l] -> l
fromListN Int
n ([Doc ann] -> PrettyM ann a) -> [Doc ann] -> PrettyM ann a
forall a b. (a -> b) -> a -> b
$ [Doc ann]
used