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
import Zwirn.Stream.Types (Identifier (..))

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 Identifier where
  pretty (TextID t) = pretty t
  pretty (NumID i) = pretty i

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

ppID :: Identifier -> T.Text
ppID = render

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