amelia
/
cubical


{-# OPTIONS_GHC -Wno-orphans #-}

{-# LANGUAGE ViewPatterns #-}

module Syntax.Pretty where


import Control.Arrow (Arrow(first))


import qualified Data.Map.Strict as Map

import qualified Data.Text.Lazy as L

import qualified Data.Set as Set

import qualified Data.Text as T

import Data.Map.Strict (Map)

import Data.Text (Text)

import Data.Set (Set)

import Data.Generics


import Presyntax.Presyntax (Plicity(..))


import Prettyprinter.Render.Terminal

import Prettyprinter


import Syntax


instance Pretty Name where

 pretty = pretty . getNameText


prettyTm :: Term -> Doc AnsiStyle

prettyTm = prettyTm . everywhere (mkT beautify) where

 prettyTm (Ref v) =

 case T.uncons (getNameText v) of

 Just ('.', w) -> keyword (pretty w)

 _ -> pretty v

 prettyTm (Con v) = keyword (pretty v)

 prettyTm (Data v) = operator (pretty v)


 prettyTm (App Im f x) = parenFun f <+> braces (prettyTm x)

 prettyTm (App Ex f x) = parenFun f <+> parenArg x


 prettyTm (Pair x y) = parens $ prettyTm x <> operator comma <+> prettyTm y

 prettyTm (Proj1 x) = prettyTm x <> operator (pretty ".1")

 prettyTm (Proj2 x) = prettyTm x <> operator (pretty ".2")


 prettyTm l@(Lam _ _ _) = pretty '\\' <> hsep (map prettyArgList al) <+> pretty "->" <+> prettyTm bod where

 unwindLam (Lam p x y) = first ((p, x):) (unwindLam y)

 unwindLam (PathIntro _ _ _ (Lam p x y)) = first ((p, x):) (unwindLam y)

 unwindLam t = ([], t)


 (al, bod) = unwindLam l


 used = freeVars bod


 prettyArgList (Ex, v)

 | v `Set.member` used = pretty v

 | otherwise = pretty "_"

 prettyArgList (Im, v)

 | v `Set.member` used = braces $ pretty v

 | otherwise = pretty "_"


 prettyTm (Meta x) = keyword $ pretty '?' <> pretty (mvName x)

 prettyTm Type{} = keyword $ pretty "Type"

 prettyTm Typeω{} = keyword $ pretty "Typeω"

 prettyTm I{} = keyword $ pretty "I"

 prettyTm I0{} = keyword $ pretty "i0"

 prettyTm I1{} = keyword $ pretty "i1"


 prettyTm (Pi Ex (T.unpack . getNameText -> "_") d r) = prettyDom d <+> pretty "->" <+> prettyTm r

 prettyTm (Pi Im v d r) = group (braces (pretty v <+> colon <+> prettyTm d)) <> softline <> pretty "->" <+> prettyTm r

 prettyTm (Pi Ex v d r) = group (parens (pretty v <+> colon <+> prettyTm d)) <> softline <> pretty "->" <+> prettyTm r


 prettyTm (Sigma (T.unpack . getNameText -> "_") d r) = prettyDom d <+> pretty "*" <+> prettyTm r

 prettyTm (Sigma v d r) = parens (pretty v <+> colon <+> prettyTm d) <+> pretty "*" <+> prettyTm r


 prettyTm (IAnd x y) = parens $ prettyTm x <+> operator (pretty "&&") <+> prettyTm y

 prettyTm (IOr x y) = parens $ prettyTm x <+> operator (pretty "||") <+> prettyTm y

 prettyTm (INot x) = operator (pretty '~') <> prettyTm x


 prettyTm (System xs) = braces (mempty <+> hsep (punctuate comma (map go (Map.toList xs))) <+> mempty) where

 go (f, t) = prettyTm f <+> operator (pretty "=>") <+> prettyTm t


 prettyTm (Case x xs) = keyword (pretty "case") <+> prettyTm x <+> keyword (pretty "of") <+> braces (prettyCase xs)


 prettyTm x = error (show x)


 prettyCase = vsep . map go where

 go (x, xs) = prettyTm x <+> operator (pretty "=>") <+> prettyTm xs


 beautify (PathP l x y) = toFun "PathP" [l, x, y]

 beautify (IElim _ _ _ f i) = App Ex f i

 beautify (PathIntro _ _ _ f) = f


 beautify (IsOne phi) = toFun "IsOne" [phi]

 beautify ItIsOne = Ref (Bound (T.pack ".1=1") 0)


 beautify (Lam Ex v (App Ex f (Ref v')))

 | v == v', v `Set.notMember` freeVars f = f


 beautify (Partial phi a) = toFun "Partial" [phi, a]

 beautify (PartialP phi a) = toFun "PartialP" [phi, a]

 beautify (Comp a phi u a0) = toFun "comp" [a, phi, u, a0]

 beautify (Sub a phi u) = toFun "Sub" [a, phi, u]

 beautify (Inc _ _ u) = toFun "inS" [u]

 beautify (Ouc _ _ _ u) = toFun "outS" [u]


 beautify (GlueTy a I1 _ _) = a

 beautify (GlueTy a b c d) = toFun "Glue" [a,b,c,d]

 beautify (Glue a b c d e f) = toFun "glue" [a,b,c,d,e,f]

 beautify (Unglue a b c d e) = toFun "unglue" [a,b,c,d,e]

 beautify x = x


 toFun s a = foldl (App Ex) (Ref (Bound (T.pack ('.':s)) 0)) a


keyword :: Doc AnsiStyle -> Doc AnsiStyle

keyword = annotate (color Magenta)


operator :: Doc AnsiStyle -> Doc AnsiStyle

operator = annotate (color Yellow)


parenArg :: Term -> Doc AnsiStyle

parenArg x@App{} = parens (prettyTm x)

parenArg x@IElim{} = parens (prettyTm x)

parenArg x@IsOne{} = parens $ prettyTm x

parenArg x@Partial{} = parens $ prettyTm x

parenArg x@PartialP{} = parens $ prettyTm x


parenArg x@Sub{} = parens $ prettyTm x

parenArg x@Inc{} = parens $ prettyTm x

parenArg x@Ouc{} = parens $ prettyTm x


parenArg x@Comp{} = parens $ prettyTm x


parenArg x = prettyDom x


prettyDom :: Term -> Doc AnsiStyle

prettyDom x@Pi{} = parens (prettyTm x)

prettyDom x@Sigma{} = parens (prettyTm x)

prettyDom x = parenFun x


parenFun :: Term -> Doc AnsiStyle

parenFun x@Lam{} = parens $ prettyTm x

parenFun x@PathIntro{} = parens $ prettyTm x

parenFun x = prettyTm x


render :: Doc AnsiStyle -> Text

render = L.toStrict . renderLazy . layoutSmart defaultLayoutOptions


showValue :: Value -> String

showValue = L.unpack . renderLazy . layoutSmart defaultLayoutOptions . prettyTm . quote


showFace :: Map Head Bool -> Doc AnsiStyle

showFace = hsep . map go . Map.toList where

 go (h, b) = parens $ prettyTm (quote (VNe h mempty)) <+> operator (pretty "=") <+> pretty (if b then "i1" else "i0")


freeVars :: Term -> Set Name

freeVars (Ref v) = Set.singleton v

freeVars (App _ f x) = Set.union (freeVars f) (freeVars x)

freeVars (Pi _ n x y) = Set.union (freeVars x) (n `Set.delete` freeVars y)

freeVars (Lam _ n x) = n `Set.delete` freeVars x

freeVars (Let ns x) = Set.union (freeVars x `Set.difference` bound) freed where

 bound = Set.fromList (map (\(x, _, _) -> x) ns)

 freed = foldr (\(_, x, y) -> Set.union (Set.union (freeVars x) (freeVars y))) mempty ns

freeVars Meta{} = mempty

freeVars Con{} = mempty

freeVars Data{} = mempty

freeVars Type{} = mempty

freeVars Typeω{} = mempty

freeVars I{} = mempty

freeVars I0{} = mempty

freeVars I1{} = mempty

freeVars (Sigma n x y) = Set.union (freeVars x) (n `Set.delete` freeVars y)

freeVars (Pair x y) = Set.unions $ map freeVars [x, y]

freeVars (Proj1 x) = Set.unions $ map freeVars [x]

freeVars (Proj2 x) = Set.unions $ map freeVars [x]

freeVars (IAnd x y) = Set.unions $ map freeVars [x, y]

freeVars (IOr x y) = Set.unions $ map freeVars [x, y]

freeVars (INot x) = Set.unions $ map freeVars [x]

freeVars (PathP x y z) = Set.unions $ map freeVars [x, y, z]

freeVars (IElim x y z a b) = Set.unions $ map freeVars [x, y, z, a, b]

freeVars (PathIntro x y z a) = Set.unions $ map freeVars [x, y, z, a]

freeVars (IsOne a) = Set.unions $ map freeVars [a]

freeVars ItIsOne{} = mempty

freeVars (Partial x y) = Set.unions $ map freeVars [x, y]

freeVars (PartialP x y) = Set.unions $ map freeVars [x, y]

freeVars (System fs) = foldr (\(x, y) -> Set.union (Set.union (freeVars x) (freeVars y))) mempty (Map.toList fs)


freeVars (Sub x y z) = Set.unions $ map freeVars [x, y, z]

freeVars (Inc x y z) = Set.unions $ map freeVars [x, y, z]

freeVars (Ouc x y z a) = Set.unions $ map freeVars [x, y, z, a]

freeVars (Comp x y z a) = Set.unions $ map freeVars [x, y, z, a]

freeVars (GlueTy x y z a) = Set.unions $ map freeVars [x, y, z, a]

freeVars (Glue x y z a b c) = Set.unions $ map freeVars [x, y, z, a, b, c]

freeVars (Unglue x y z a c) = Set.unions $ map freeVars [x, y, z, a, c]

freeVars (Case x y) = freeVars x <> foldMap (freeVars . snd) y