Pretty.hs

{-# OPTIONS_GHC -Wno-orphans #-}

module Zwirn.Language.Pretty where

{-
    Pretty.hs - prettyprinter for zwirn
    Copyright (C) 2025, Martin Gius

    This library is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this library.  If not, see <http://www.gnu.org/licenses/>.
-}

import qualified Data.Text as T
import Prettyprinter
import Prettyprinter.Render.Text (renderStrict)
import Text.Read (readMaybe)
import Zwirn.Language.Location (Located (..), RealSrcLoc (..), SrcLoc (..), noLoc)
import Zwirn.Language.Syntax
import Zwirn.Language.TypeCheck.Constraint (TypeError (..))
import Zwirn.Language.TypeCheck.Types

instance (Pretty a) => Pretty (Located a) where
  pretty (Located _ x) = pretty x

instance Pretty Term where
  pretty (TVar x) = pretty x
  pretty (TNum x) = case (readMaybe (T.unpack x) :: Maybe Int) of
    Just i -> pretty i
    Nothing -> pretty x
  pretty (TText x) = pretty x
  pretty (TMacro x) = "!" <> pretty x
  pretty (TBracket x) = parens $ pretty x
  pretty TRest = "~"
  pretty (TRepeat x Nothing) = pretty x <> "!"
  pretty (TRepeat x (Just i)) = pretty x <> sep (replicate i "!")
  pretty (TSeq [t]) = pretty t
  pretty (TSeq ts) = group $ brackets $ sep $ map pretty ts
  pretty (TStack ts) = alignedBetween (map pretty ts) lbracket rbracket comma
  pretty (TAlt ts) = group $ alignedBetween (map pretty ts) langle rangle space
  pretty (TChoice _ ts) = group $ alignedBetween (map pretty ts) lbracket rbracket pipe
  pretty (TPoly x y) = pretty x <> "%" <> pretty y
  pretty (TApp x y) = pretty x <+> pretty y
  pretty (TLambda vs x) = "\\" <> hcat (punctuate space $ map (pretty . T.unpack) vs) <+> "->" <+> pretty x
  pretty (TIfThenElse x y (Just z)) = "if" <+> pretty x <+> "then" <+> pretty y <+> "else" <+> pretty z
  pretty (TIfThenElse x y Nothing) = "if" <+> pretty x <+> "then" <+> pretty y
  pretty (TSectionL t n) = pretty t <+> pretty (unpackOp $ lValue n)
  pretty (TSectionR n t) = pretty (unpackOp $ lValue n) <+> pretty t
  pretty (TEnum Run x y) = brackets (pretty x <+> ".." <+> pretty y)
  pretty (TEnumThen Run x y z) = brackets (pretty x <+> pretty y <+> ".." <+> pretty z)
  pretty (TEnum Cord x y) = brackets (pretty x <+> ", .." <+> pretty y)
  pretty (TEnumThen Cord x y z) = brackets (pretty x <> comma <+> pretty y <+> ".." <+> pretty z)
  pretty (TEnum Choice x y) = brackets (pretty x <+> "| .." <+> pretty y)
  pretty (TEnumThen Choice x y z) = brackets (pretty x <+> pipe <+> pretty y <+> ".." <+> pretty z)
  pretty (TEnum Alt x y) = angles (pretty x <+> ".." <+> pretty y)
  pretty (TEnumThen Alt x y z) = angles (pretty x <+> pretty y <+> ".." <+> pretty z)
  pretty inf@(TInfix {}) = startThenAlign (pretty x) (map (\(op, y) -> pretty (unpackOp $ lValue op) <+> pretty y) xs)
    where
      (x, xs) = infixChain (noLoc inf)

alignedBetween :: [Doc a] -> Doc a -> Doc a -> Doc a -> Doc a
alignedBetween [] _ _ _ = mempty
alignedBetween (x : xs) l r s = align $ vcat $ (l <> x) : map (s <>) xs ++ [r]

startThenAlign :: Doc ann -> [Doc ann] -> Doc ann
startThenAlign x xs = x <+> align (vsep xs)

unpackOp :: T.Text -> String
unpackOp t = filter (\c -> c /= '(' && c /= ')') $ T.unpack t

infixChain :: LocTerm -> (LocTerm, [(LocVar, LocTerm)])
infixChain (Located _ (TInfix x op y)) = let (t, cs) = infixChain y in (x, (op, t) : cs)
infixChain t = (t, [])

parensIf :: Bool -> Doc a -> Doc a
parensIf True = parens
parensIf False = id

instance Pretty Type where
  pretty (TypeArr a b) = parensIf (isArrow a) (pretty a) <+> "->" <+> pretty b
    where
      isArrow (Located _ TypeArr {}) = True
      isArrow _ = False
  pretty (TypeVar a) = pretty a
  pretty (TypeCon a) = pretty a

instance Pretty Predicate where
  pretty (IsIn c t) = pretty c <+> pretty t

prettyPredicates :: [Predicate] -> Doc a
prettyPredicates ps = parensIf (length ps > 1) (hcat (punctuate comma (map pretty ps)))

instance (Pretty a) => Pretty (Qualified a) where
  pretty (Qual [] _ t) = pretty t
  pretty (Qual ps _ t) = prettyPredicates ps <+> "=>" <+> pretty t

instance Pretty Scheme where
  pretty (Forall _ t) = pretty t

instance Pretty SrcLoc where
  pretty NoLoc = "NoLoc"
  pretty (SrcLoc (RealSrcLoc _ lst cst len cen)) = parens $ vcat $ punctuate comma [pretty lst, pretty cst, pretty len, pretty cen]

renderDoc :: Doc a -> T.Text
renderDoc = renderStrict . layoutPretty defaultLayoutOptions

render :: (Pretty a) => a -> T.Text
render = renderDoc . pretty

pptype :: Type -> T.Text
pptype = render

ppscheme :: Scheme -> T.Text
ppscheme = render

ppterm :: Term -> T.Text
ppterm = render

ppTermHasType :: (LocTerm, Scheme) -> T.Text
ppTermHasType (t, s) = renderDoc $ pretty t <+> "::" <+> pretty s

instance Pretty TypeError where
  pretty (UnificationFail (Located _ (a, b))) = "Cannot unify types:" <+> pretty a <+> "~" <+> pretty b
  pretty (InfiniteType a b) = "Cannot construct the infinite type:" <+> pretty a <+> "=" <+> pretty b
  pretty (Ambigious cs) = vsep ["Cannot not match expected type: '" <> pretty a <> "' with actual type: '" <> pretty b <> "'\n" | Located _ (a, b) <- cs]
  pretty (UnboundVariable a) = "Variable not in scope:" <+> pretty a
  pretty (NoInstance (Located _ (IsIn c x))) = "No instance for" <+> pretty c <+> pretty x
  pretty NotImplemented = "Case not implemented in type-checker."

instance Pretty Command where
  pretty (TypeCommand t) = ":t" <+> pretty t
  pretty (ShowCommand t) = ":show" <+> pretty t
  pretty (InfoCommand t) = ":info" <+> pretty t
  pretty (SetCommand t) = ":set" <+> pretty t
  pretty (UnsetCommand t) = ":unset" <+> pretty t
  pretty (LoadCommand t) = ":load" <+> pretty t
  pretty ResetConfigCommand = ":resetconfig"
  pretty ShowConfigPathCommand = ":showconfig"
  pretty ResetEnvCommand = ":reset"
  pretty StatusCommand = ":status"
  pretty EnvCommand = ":env"

instance Pretty Definition where
  pretty (Definition x xs t) = pretty x <+> vsep (map pretty xs) <+> "=" <+> pretty t

instance Pretty DynamicDefinition where
  pretty (DynamicDefinition x t) = pretty x <+> "<-" <+> pretty t

instance Pretty MacroDefinition where
  pretty (MacroDefinition x t) = "!" <> pretty x <+> "=" <+> pretty t

instance Pretty Syntax where
  pretty (Exec t) = pretty t
  pretty (Def d) = pretty d
  pretty (DynDef d) = pretty d
  pretty (MacroDef d) = pretty d
  pretty (Command c) = pretty c