quox/lib/Quox/Pretty.idr

module Quox.Pretty

import Quox.Loc
import Quox.Name

import Control.Monad.ST.Extra
import public Text.PrettyPrint.Bernardy
import public Text.PrettyPrint.Bernardy.Core.Decorate
import public Quox.EffExtra
import public Data.String
import Control.ANSI.SGR

import Data.DPair
import Data.SnocList
import Derive.Prelude

%default total
%language ElabReflection

%hide TT.Name


public export
data PPrec
= Outer
| Times   -- "_ × _"
| InTimes -- arguments of ×
| AnnL    -- left of "∷"
| Eq      -- "_ ≡ _ : _"
| InEq    -- arguments of ≡
  -- ...
| App  -- term/dimension application
| Arg  -- argument to nonfix function
%runElab derive "PPrec" [Eq, Ord, Show]


public export
data HL
= Delim
| Free | TVar | TVarErr
| Dim  | DVar | DVarErr
| Qty  | Universe
| Syntax
| Constant
%runElab derive "HL" [Eq, Ord, Show]


public export
data Flavor = Unicode | Ascii
%runElab derive "Flavor" [Eq, Ord, Show]

export %inline
noHighlight : HL -> Highlight
noHighlight _ = MkHighlight "" ""


public export
data EffTag = PREC | FLAVOR | HIGHLIGHT | INDENT

public export
Pretty : List (Type -> Type)
Pretty = [StateL PREC PPrec, ReaderL FLAVOR Flavor,
          ReaderL HIGHLIGHT (HL -> Highlight), ReaderL INDENT Nat]

export %inline
runPrettyWith : PPrec -> Flavor -> (HL -> Highlight) -> Nat ->
                Eff Pretty a -> a
runPrettyWith prec flavor highlight indent act =
  runST $ do
    runEff act $ with Union.(::)
      [handleStateSTRef !(newSTRef prec),
       handleReaderConst flavor,
       handleReaderConst highlight,
       handleReaderConst indent]


export %inline
toSGR : HL -> List SGR
toSGR Delim    = []
toSGR Free     = [SetForeground BrightBlue]
toSGR TVar     = [SetForeground BrightYellow]
toSGR TVarErr  = [SetForeground BrightYellow, SetStyle SingleUnderline]
toSGR Dim      = [SetForeground BrightGreen]
toSGR DVar     = [SetForeground BrightGreen]
toSGR DVarErr  = [SetForeground BrightGreen, SetStyle SingleUnderline]
toSGR Qty      = [SetForeground BrightMagenta]
toSGR Universe = [SetForeground BrightRed]
toSGR Syntax   = [SetForeground BrightCyan]
toSGR Constant = [SetForeground BrightRed]

export %inline
highlightSGR : HL -> Highlight
highlightSGR h = MkHighlight (escapeSGR $ toSGR h) (escapeSGR [Reset])


export %inline
runPretty : Eff Pretty a -> a
runPretty = runPrettyWith Outer Unicode noHighlight 2

export %inline
runPrettyColor : Eff Pretty a -> a
runPrettyColor = runPrettyWith Outer Unicode highlightSGR 2


export %inline
hl : {opts : LayoutOpts} -> HL -> Doc opts -> Eff Pretty (Doc opts)
hl h doc = asksAt HIGHLIGHT $ \f => decorate (f h) doc


export %inline
indentD : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
indentD doc = pure $ indent !(askAt INDENT) doc

export %inline
hangD : {opts : LayoutOpts} -> Doc opts -> Doc opts -> Eff Pretty (Doc opts)
hangD d1 d2 = pure $ hangSep !(askAt INDENT) d1 d2

export %inline
hangSingle : {opts : LayoutOpts} -> Nat -> Doc opts -> Doc opts -> Doc opts
hangSingle n d1 d2 = ifMultiline (d1 <++> d2) (vappend d1 (indent n d2))

export %inline
hangDSingle : {opts : LayoutOpts} -> Doc opts -> Doc opts ->
              Eff Pretty (Doc opts)
hangDSingle d1 d2 = pure $ hangSingle !(askAt INDENT) d1 d2


export
tightDelims : {opts : LayoutOpts} -> (l, r : String) -> (inner : Doc opts) ->
              Eff Pretty (Doc opts)
tightDelims l r inner = do
  l <- hl Delim $ text l
  r <- hl Delim $ text r
  pure $ hcat [l, inner, r]

export
looseDelims : {opts : LayoutOpts} -> (l, r : String) -> (inner : Doc opts) ->
              Eff Pretty (Doc opts)
looseDelims l r inner = do
  l <- hl Delim $ text l
  r <- hl Delim $ text r
  let short = hsep [l, inner, r]
      long  = vsep [l, !(indentD inner), r]
  pure $ ifMultiline short long

export %inline
parens : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
parens = tightDelims "(" ")"

export %inline
bracks : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
bracks = tightDelims "[" "]"

export %inline
braces : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
braces = looseDelims "{" "}"

export %inline
tightBraces : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
tightBraces = tightDelims "{" "}"

export %inline
parensIf : {opts : LayoutOpts} -> Bool -> Doc opts -> Eff Pretty (Doc opts)
parensIf True  = parens
parensIf False = pure


||| uses hsep only if the whole list fits on one line
export
sepSingle : {opts : LayoutOpts} -> List (Doc opts) -> Doc opts
sepSingle xs = ifMultiline (hsep xs) (vsep xs)

export
fillSep : {opts : LayoutOpts} -> List (Doc opts) -> Doc opts
fillSep []        = empty
fillSep (x :: xs) = foldl (\x, y => sep [x, y]) x xs

export
exceptLast : {opts : LayoutOpts} -> (Doc opts -> Doc opts) ->
             List (Doc opts) -> List (Doc opts)
exceptLast f []        = []
exceptLast f [x]       = [x]
exceptLast f (x :: xs) = f x :: exceptLast f xs

parameters {opts : LayoutOpts} {auto _ : Foldable t}
  export
  separateLoose : Doc opts -> t (Doc opts) -> Doc opts
  separateLoose d = sep . exceptLast (<++> d) . toList

  export
  separateTight : Doc opts -> t (Doc opts) -> Doc opts
  separateTight d = sep . exceptLast (<+> d) . toList

  export
  hseparateTight : Doc opts -> t (Doc opts) -> Doc opts
  hseparateTight d = hsep . exceptLast (<+> d) . toList

  export
  vseparateTight : Doc opts -> t (Doc opts) -> Doc opts
  vseparateTight d = vsep . exceptLast (<+> d) . toList

  export
  fillSeparateTight : Doc opts -> t (Doc opts) -> Doc opts
  fillSeparateTight d = fillSep . exceptLast (<+> d) . toList


export %inline
pshow : {opts : LayoutOpts} -> Show a => a -> Doc opts
pshow = text . show


export %inline
ifUnicode : (uni, asc : Lazy a) -> Eff Pretty a
ifUnicode uni asc =
  asksAt FLAVOR $ \case
    Unicode => uni
    Ascii   => asc

export %inline
parensIfM : {opts : LayoutOpts} -> PPrec -> Doc opts -> Eff Pretty (Doc opts)
parensIfM d doc = parensIf (!(getAt PREC) > d) doc

export %inline
withPrec : PPrec -> Eff Pretty a -> Eff Pretty a
withPrec = localAt_ PREC


export
prettyName : Name -> Doc opts
prettyName = text . toDots

export
prettyFree : {opts : LayoutOpts} -> Name -> Eff Pretty (Doc opts)
prettyFree = hl Free . prettyName

export
prettyBind' : BindName -> Doc opts
prettyBind' = text . baseStr . val

export
prettyTBind : {opts : LayoutOpts} -> BindName -> Eff Pretty (Doc opts)
prettyTBind = hl TVar . prettyBind'

export
prettyDBind : {opts : LayoutOpts} -> BindName -> Eff Pretty (Doc opts)
prettyDBind = hl DVar . prettyBind'



export %inline
typeD, ioStateD, arrowD, darrowD, timesD, lamD, eqndD, dlamD, annD, natD,
stringD, eqD, colonD, commaD, semiD, caseD, typecaseD, returnD,
ofD, dotD, zeroD, succD, coeD, compD, undD, cstD, pipeD, fstD, sndD :
  {opts : LayoutOpts} -> Eff Pretty (Doc opts)
typeD     = hl Syntax . text =<< ifUnicode "★" "Type"
ioStateD  = hl Syntax $ text "IOState"
arrowD    = hl Delim  . text =<< ifUnicode "→" "->"
darrowD   = hl Delim  . text =<< ifUnicode "⇒" "=>"
timesD    = hl Delim  . text =<< ifUnicode "×" "**"
lamD      = hl Syntax . text =<< ifUnicode "λ" "fun"
eqndD     = hl Delim  . text =<< ifUnicode "≡" "=="
dlamD     = hl Syntax . text =<< ifUnicode "δ" "dfun"
annD      = hl Delim  . text =<< ifUnicode "∷" "::"
natD      = hl Syntax . text =<< ifUnicode "ℕ" "Nat"
stringD   = hl Syntax   $ text "String"
eqD       = hl Syntax   $ text "Eq"
colonD    = hl Delim    $ text ":"
commaD    = hl Delim    $ text ","
semiD     = hl Delim    $ text ";"
caseD     = hl Syntax   $ text "case"
typecaseD = hl Syntax   $ text "type-case"
ofD       = hl Syntax   $ text "of"
returnD   = hl Syntax   $ text "return"
dotD      = hl Delim    $ text "."
zeroD     = hl Constant $ text "zero"
succD     = hl Constant $ text "succ"
coeD      = hl Syntax   $ text "coe"
compD     = hl Syntax   $ text "comp"
undD      = hl Syntax   $ text "_"
cstD      = hl Syntax   $ text "="
pipeD     = hl Syntax   $ text "|"
fstD      = hl Syntax   $ text "fst"
sndD      = hl Syntax   $ text "snd"


export
prettyApp : {opts : LayoutOpts} -> Nat -> Doc opts ->
            List (Doc opts) -> Doc opts
prettyApp ind f args =
      hsep (f :: args)
  <|> hsep [f, vsep args]
  <|> vsep (f :: map (indent ind) args)

export
prettyAppD : {opts : LayoutOpts} -> Doc opts -> List (Doc opts) ->
             Eff Pretty (Doc opts)
prettyAppD f args = pure $ prettyApp !(askAt INDENT) f args


export
escapeString : String -> String
escapeString = concatMap esc1 . unpack where
  esc1 : Char -> String
  esc1 '"'  = #"\""#
  esc1 '\\' = #"\\"#
  esc1 '\n' = #"\n"#
  esc1 c    = singleton c

export
quoteTag : String -> String
quoteTag tag =
  if isName tag then tag else
  "\"" ++ escapeString tag ++ "\""

export
prettyBounds : {opts : LayoutOpts} -> Bounds -> Eff Pretty (Doc opts)
prettyBounds (MkBounds l1 c1 l2 c2) =
  hcat <$> sequence
    [hl TVar $ text $ show l1, colonD,
     hl DVar $ text $ show c1, hl Delim "-",
     hl TVar $ text $ show l2, colonD,
     hl DVar $ text $ show c2, colonD]

export
prettyLoc : {opts : LayoutOpts} -> Loc -> Eff Pretty (Doc opts)
prettyLoc (L NoLoc) =
  hcat <$> sequence [hl TVarErr "no location", colonD]
prettyLoc (L (YesLoc file b)) =
  hcat <$> sequence [hl Free $ text file, colonD, prettyBounds b]

export
prettyTag : {opts : _} -> String -> Eff Pretty (Doc opts)
prettyTag tag = hl Constant $ text $ "'" ++ quoteTag tag

export
prettyStrLit : {opts : _} -> String -> Eff Pretty (Doc opts)
prettyStrLit s =
  let s = concatMap esc1 $ unpack s in
  hl Constant $ hcat ["\"", text s, "\""]
where
  esc1 : Char -> String
  esc1 '"' = "\""; esc1 '\\' = "\\"
  esc1 c   = singleton c