Move module hierarchy under CHAD.

1 files changed, 525 insertions, 0 deletions

diff --git a/src/CHAD/AST/Pretty.hs b/src/CHAD/AST/Pretty.hs
new file mode 100644
index 0000000..3f6a3af
--- /dev/null
+++ b/src/CHAD/AST/Pretty.hs
@@ -0,0 +1,525 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE EmptyCase #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeOperators #-}
+module CHAD.AST.Pretty (pprintExpr, ppExpr, ppSTy, ppSMTy, PrettyX(..)) where
+
+import Control.Monad (ap)
+import Data.List (intersperse, intercalate)
+import Data.Functor.Const
+import qualified Data.Functor.Product as Product
+import Data.String (fromString)
+import Prettyprinter
+import Prettyprinter.Render.String
+
+import qualified Data.Text.Lazy as TL
+import qualified Prettyprinter.Render.Terminal as PT
+import System.Console.ANSI (hSupportsANSI)
+import System.IO (stdout)
+import System.IO.Unsafe (unsafePerformIO)
+
+import CHAD.AST
+import CHAD.AST.Count
+import CHAD.AST.Sparse.Types
+import CHAD.Data
+import CHAD.Drev.Types
+
+
+class PrettyX x where
+  prettyX :: x t -> String
+
+  prettyXsuffix :: x t -> String
+  prettyXsuffix x = "<" ++ prettyX x ++ ">"
+
+instance PrettyX (Const ()) where
+  prettyX _ = ""
+  prettyXsuffix _ = ""
+
+
+type SVal = SList (Const String)
+
+newtype M a = M { runM :: Int -> (a, Int) }
+  deriving (Functor)
+instance Applicative M where { pure x = M (\i -> (x, i)) ; (<*>) = ap }
+instance Monad M where { M f >>= g = M (\i -> let (x, j) = f i in runM (g x) j) }
+
+genId :: M Int
+genId = M (\i -> (i, i + 1))
+
+nameBaseForType :: STy t -> String
+nameBaseForType STNil = "nil"
+nameBaseForType (STPair{}) = "p"
+nameBaseForType (STEither{}) = "e"
+nameBaseForType (STMaybe{}) = "m"
+nameBaseForType (STScal STI32) = "n"
+nameBaseForType (STScal STI64) = "n"
+nameBaseForType (STArr{}) = "a"
+nameBaseForType (STAccum{}) = "ac"
+nameBaseForType _ = "x"
+
+genName' :: String -> M String
+genName' prefix = (prefix ++) . show <$> genId
+
+genNameIfUsedIn' :: String -> STy a -> Idx env a -> Expr x env t -> M String
+genNameIfUsedIn' prefix ty idx ex
+  | occCount idx ex == mempty = case ty of STNil -> return "()"
+                                           _ -> return "_"
+  | otherwise                 = genName' prefix
+
+-- TODO: let this return a type-tagged thing so that name environments are more typed than Const
+genNameIfUsedIn :: STy a -> Idx env a -> Expr x env t -> M String
+genNameIfUsedIn = \t -> genNameIfUsedIn' (nameBaseForType t) t
+
+pprintExpr :: (KnownEnv env, PrettyX x) => Expr x env t -> IO ()
+pprintExpr = putStrLn . ppExpr knownEnv
+
+ppExpr :: PrettyX x => SList STy env -> Expr x env t -> String
+ppExpr senv e = render $ fst . flip runM 1 $ do
+  val <- mkVal senv
+  e' <- ppExpr' 0 val e
+  let lam = "λ" ++ intercalate " " (reverse (unSList (\(Product.Pair (Const name) ty) -> "(" ++ name ++ " : " ++ ppSTy 0 ty ++ ")") (slistZip val senv))) ++ "."
+  return $ group $ flatAlt
+    (hang 2 $
+       ppString lam
+       <> hardline <> e')
+    (ppString lam <+> e')
+  where
+    mkVal :: SList f env -> M (SVal env)
+    mkVal SNil = return SNil
+    mkVal (SCons _ v) = do
+      val <- mkVal v
+      name <- genName' "arg"
+      return (Const name `SCons` val)
+
+ppExpr' :: PrettyX x => Int -> SVal env -> Expr x env t -> M ADoc
+ppExpr' d val expr = case expr of
+  EVar _ _ i -> return $ ppString (getConst (slistIdx val i)) <> ppX expr
+
+  e@ELet{} -> ppExprLet d val e
+
+  EPair _ a b -> do
+    a' <- ppExpr' 0 val a
+    b' <- ppExpr' 0 val b
+    return $ group $ flatAlt (align $ ppString "(" <> a' <> hardline <> ppString "," <> b' <> ppString ")" <> ppX expr)
+                             (ppString "(" <> a' <> ppString "," <+> b' <> ppString ")" <> ppX expr)
+
+  EFst _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "fst" <> ppX expr <+> e'
+
+  ESnd _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "snd" <> ppX expr <+> e'
+
+  ENil _ -> return $ ppString "()"
+
+  EInl _ _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "Inl" <> ppX expr <+> e'
+
+  EInr _ _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "Inr" <> ppX expr <+> e'
+
+  ECase _ e a b -> do
+    e' <- ppExpr' 0 val e
+    let STEither t1 t2 = typeOf e
+    name1 <- genNameIfUsedIn t1 IZ a
+    a' <- ppExpr' 0 (Const name1 `SCons` val) a
+    name2 <- genNameIfUsedIn t2 IZ b
+    b' <- ppExpr' 0 (Const name2 `SCons` val) b
+    return $ ppParen (d > 0) $
+      hang 2 $
+        annotate AKey (ppString "case") <> ppX expr <+> e' <+> annotate AKey (ppString "of")
+        <> hardline <> ppString "Inl" <+> ppString name1 <+> ppString "->" <+> a'
+        <> hardline <> ppString "Inr" <+> ppString name2 <+> ppString "->" <+> b'
+
+  ENothing _ _ -> return $ ppString "Nothing"
+
+  EJust _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "Just" <> ppX expr <+> e'
+
+  EMaybe _ a b e -> do
+    let STMaybe t = typeOf e
+    e' <- ppExpr' 0 val e
+    a' <- ppExpr' 0 val a
+    name <- genNameIfUsedIn t IZ b
+    b' <- ppExpr' 0 (Const name `SCons` val) b
+    return $ ppParen (d > 0) $
+      align $
+        group (flatAlt
+          (annotate AKey (ppString "case") <> ppX expr <+> e'
+           <> hardline <> annotate AKey (ppString "of"))
+          (annotate AKey (ppString "case") <> ppX expr <+> e' <+> annotate AKey (ppString "of")))
+        <> hardline
+        <> indent 2
+          (ppString "Nothing" <+> ppString "->" <+> a'
+           <> hardline <> ppString "Just" <+> ppString name <+> ppString "->" <+> b')
+
+  ELNil _ _ _ -> return (ppString "LNil")
+
+  ELInl _ _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "LInl" <> ppX expr <+> e'
+
+  ELInr _ _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "LInr" <> ppX expr <+> e'
+
+  ELCase _ e a b c -> do
+    e' <- ppExpr' 0 val e
+    let STLEither t1 t2 = typeOf e
+    a' <- ppExpr' 11 val a
+    name1 <- genNameIfUsedIn t1 IZ b
+    b' <- ppExpr' 0 (Const name1 `SCons` val) b
+    name2 <- genNameIfUsedIn t2 IZ c
+    c' <- ppExpr' 0 (Const name2 `SCons` val) c
+    return $ ppParen (d > 0) $
+      hang 2 $
+        annotate AKey (ppString "lcase") <> ppX expr <+> e' <+> annotate AKey (ppString "of")
+        <> hardline <> ppString "LNil" <+> ppString "->" <+> a'
+        <> hardline <> ppString "LInl" <+> ppString name1 <+> ppString "->" <+> b'
+        <> hardline <> ppString "LInr" <+> ppString name2 <+> ppString "->" <+> c'
+
+  EConstArr _ _ ty v
+    | Dict <- scalRepIsShow ty -> return $ ppString (showsPrec d v "") <> ppX expr
+
+  EBuild _ n a b -> do
+    a' <- ppExpr' 11 val a
+    name <- genNameIfUsedIn' "i" (tTup (sreplicate n tIx)) IZ b
+    e' <- ppExpr' 0 (Const name `SCons` val) b
+    let primName = ppString ("build" ++ intSubscript (fromSNat n))
+    return $ ppParen (d > 0) $
+      group $ flatAlt
+        (hang 2 $
+           annotate AHighlight primName <> ppX expr <+> a'
+           <+> ppString "$" <+> ppString "\\" <> ppString name <+> ppString "->"
+           <> hardline <> e')
+        (ppApp (annotate AHighlight primName <> ppX expr) [a', ppLam [ppString name] e'])
+
+  EMap _ a b -> do
+    let STArr _ t1 = typeOf b
+    name <- genNameIfUsedIn t1 IZ a
+    a' <- ppExpr' 0 (Const name `SCons` val) a
+    b' <- ppExpr' 11 val b
+    return $ ppParen (d > 0) $
+      ppApp (annotate AHighlight (ppString "map") <> ppX expr) [ppLam [ppString name] a', b']
+
+  EFold1Inner _ cm a b c -> do
+    name <- genNameIfUsedIn (STPair (typeOf a) (typeOf a)) IZ a
+    a' <- ppExpr' 0 (Const name `SCons` val) a
+    b' <- ppExpr' 11 val b
+    c' <- ppExpr' 11 val c
+    let opname = "fold1i" ++ ppCommut cm
+    return $ ppParen (d > 10) $
+      ppApp (annotate AHighlight (ppString opname) <> ppX expr) [ppLam [ppString name] a', b', c']
+
+  ESum1Inner _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "sum1i" <> ppX expr <+> e'
+
+  EUnit _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "unit" <> ppX expr <+> e'
+
+  EReplicate1Inner _ a b -> do
+    a' <- ppExpr' 11 val a
+    b' <- ppExpr' 11 val b
+    return $ ppParen (d > 10) $ ppApp (ppString "replicate1i" <> ppX expr) [a', b']
+
+  EMaximum1Inner _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "maximum1i" <> ppX expr <+> e'
+
+  EMinimum1Inner _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "minimum1i" <> ppX expr <+> e'
+
+  EReshape _ n esh e -> do
+    esh' <- ppExpr' 11 val esh
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppApp (ppString ("reshape" ++ intSubscript (fromSNat n)) <> ppX expr) [esh', e']
+
+  EZip _ e1 e2 -> do
+    e1' <- ppExpr' 11 val e1
+    e2' <- ppExpr' 11 val e2
+    return $ ppParen (d > 10) $ ppApp (ppString "zip" <> ppX expr) [e1', e2']
+
+  EFold1InnerD1 _ cm a b c -> do
+    name <- genNameIfUsedIn (STPair (typeOf b) (typeOf b)) IZ a
+    a' <- ppExpr' 0 (Const name `SCons` val) a
+    b' <- ppExpr' 11 val b
+    c' <- ppExpr' 11 val c
+    let opname = "fold1iD1" ++ ppCommut cm
+    return $ ppParen (d > 10) $
+      ppApp (annotate AHighlight (ppString opname) <> ppX expr) [ppLam [ppString name] a', b', c']
+
+  EFold1InnerD2 _ cm ef ebog ed -> do
+    let STArr _ tB = typeOf ebog
+        STArr _ t2 = typeOf ed
+    namef1 <- genNameIfUsedIn tB (IS IZ) ef
+    namef2 <- genNameIfUsedIn t2 IZ ef
+    ef' <- ppExpr' 0 (Const namef2 `SCons` Const namef1 `SCons` val) ef
+    ebog' <- ppExpr' 11 val ebog
+    ed' <- ppExpr' 11 val ed
+    let opname = "fold1iD2" ++ ppCommut cm
+    return $ ppParen (d > 10) $
+      ppApp (annotate AHighlight (ppString opname) <> ppX expr)
+            [ppLam [ppString namef1, ppString namef2] ef', ebog', ed']
+
+  EConst _ ty v
+    | Dict <- scalRepIsShow ty -> return $ ppString (showsPrec d v "") <> ppX expr
+
+  EIdx0 _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "idx0" <> ppX expr <+> e'
+
+  EIdx1 _ a b -> do
+    a' <- ppExpr' 9 val a
+    b' <- ppExpr' 9 val b
+    return $ ppParen (d > 8) $ a' <+> ppString ".!" <> ppX expr <+> b'
+
+  EIdx _ a b -> do
+    a' <- ppExpr' 9 val a
+    b' <- ppExpr' 10 val b
+    return $ ppParen (d > 8) $
+      a' <+> ppString "!" <> ppX expr <+> b'
+
+  EShape _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppString "shape" <> ppX expr <+> e'
+
+  EOp _ op (EPair _ a b)
+    | (Infix, ops) <- operator op -> do
+        a' <- ppExpr' 9 val a
+        b' <- ppExpr' 9 val b
+        return $ ppParen (d > 8) $ a' <+> ppString ops <> ppX expr <+> b'
+
+  EOp _ op e -> do
+    e' <- ppExpr' 11 val e
+    let ops = case operator op of
+                (Infix, s) -> "(" ++ s ++ ")"
+                (Prefix, s) -> s
+    return $ ppParen (d > 10) $ ppString ops <> ppX expr <+> e'
+
+  ECustom _ t1 t2 t3 a b c e1 e2 -> do
+    en1 <- genNameIfUsedIn t1 (IS IZ) a
+    en2 <- genNameIfUsedIn t2 IZ a
+    pn1 <- genNameIfUsedIn (d1 t1) (IS IZ) b
+    pn2 <- genNameIfUsedIn (d1 t2) IZ b
+    dn1 <- genNameIfUsedIn' "tape" t3 (IS IZ) c
+    dn2 <- genNameIfUsedIn' "d" (d2 (typeOf a)) IZ c
+    a' <- ppExpr' 11 (Const pn2 `SCons` Const pn1 `SCons` SNil) a
+    b' <- ppExpr' 11 (Const pn2 `SCons` Const pn1 `SCons` SNil) b
+    c' <- ppExpr' 11 (Const dn2 `SCons` Const dn1 `SCons` SNil) c
+    e1' <- ppExpr' 11 val e1
+    e2' <- ppExpr' 11 val e2
+    return $ ppParen (d > 10) $
+      ppApp (ppString "custom" <> ppX expr)
+        [ppLam [ppString en1, ppString en2] a'
+        ,ppLam [ppString pn1, ppString pn2] b'
+        ,ppLam [ppString dn1, ppString dn2] c'
+        ,e1'
+        ,e2']
+
+  ERecompute _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ ppParen (d > 10) $ ppApp (ppString "recompute" <> ppX expr) [e']
+
+  EWith _ t e1 e2 -> do
+    e1' <- ppExpr' 11 val e1
+    name <- genNameIfUsedIn' "ac" (STAccum t) IZ e2
+    e2' <- ppExpr' 0 (Const name `SCons` val) e2
+    return $ ppParen (d > 0) $
+      group $ flatAlt
+        (hang 2 $
+           annotate AWith (ppString "with") <> ppX expr <+> e1'
+           <+> ppString "$" <+> ppString "\\" <> ppString name <+> ppString "->"
+           <> hardline <> e2')
+        (ppApp (annotate AWith (ppString "with") <> ppX expr) [e1', ppLam [ppString name] e2'])
+
+  EAccum _ t prj e1 sp e2 e3 -> do
+    e1' <- ppExpr' 11 val e1
+    e2' <- ppExpr' 11 val e2
+    e3' <- ppExpr' 11 val e3
+    return $ ppParen (d > 10) $
+      ppApp (annotate AMonoid (ppString "accum") <> ppX expr <+> ppString "@" <> ppSMTy' 11 (applySparse sp (acPrjTy prj t)))
+            [ppString (ppAcPrj t prj), ppString (ppSparse (acPrjTy prj t) sp), e1', e2', e3']
+
+  EZero _ t e1 -> do
+    e1' <- ppExpr' 11 val e1
+    return $ ppParen (d > 0) $
+      annotate AMonoid (ppString "zero") <> ppX expr <+> ppString "@" <> ppSMTy' 11 t <+> e1'
+
+  EDeepZero _ t e1 -> do
+    e1' <- ppExpr' 11 val e1
+    return $ ppParen (d > 0) $
+      annotate AMonoid (ppString "deepzero") <> ppX expr <+> ppString "@" <> ppSMTy' 11 t <+> e1'
+
+  EPlus _ t a b -> do
+    a' <- ppExpr' 11 val a
+    b' <- ppExpr' 11 val b
+    return $ ppParen (d > 10) $
+      ppApp (annotate AMonoid (ppString "plus") <> ppX expr <+> ppString "@" <> ppSMTy' 11 t) [a', b']
+
+  EOneHot _ t prj a b -> do
+    a' <- ppExpr' 11 val a
+    b' <- ppExpr' 11 val b
+    return $ ppParen (d > 10) $
+      ppApp (annotate AMonoid (ppString "onehot") <> ppX expr <+> ppString "@" <> ppSMTy' 11 (acPrjTy prj t)) [ppString (ppAcPrj t prj), a', b']
+
+  EError _ _ s -> return $ ppParen (d > 10) $ ppString "error" <> ppX expr <+> ppString (show s)
+
+ppExprLet :: PrettyX x => Int -> SVal env -> Expr x env t -> M ADoc
+ppExprLet d val etop = do
+  let collect :: PrettyX x => SVal env -> Expr x env t -> M ([(String, Occ, ADoc)], ADoc)
+      collect val' (ELet _ rhs body) = do
+        let occ = occCount IZ body
+        name <- genNameIfUsedIn (typeOf rhs) IZ body
+        rhs' <- ppExpr' 0 val' rhs
+        (binds, core) <- collect (Const name `SCons` val') body
+        return ((name, occ, rhs') : binds, core)
+      collect val' e = ([],) <$> ppExpr' 0 val' e
+
+  (binds, core) <- collect val etop
+
+  return $ ppParen (d > 0) $
+    align $
+      annotate AKey (ppString "let")
+      <+> align (mconcat $ intersperse hardline $
+                   map (\(name, _occ, rhs) ->
+                          ppString (name ++ {- " (" ++ show _occ ++ ")" ++ -} " = ") <> rhs)
+                       binds)
+      <> hardline <> annotate AKey (ppString "in") <+> core
+
+ppApp :: ADoc -> [ADoc] -> ADoc
+ppApp fun args = group $ fun <+> align (sep args)
+
+ppLam :: [ADoc] -> ADoc -> ADoc
+ppLam args body = ppString "(" <> hang 2 (ppString "\\" <> sep (args ++ [ppString "->"])
+                                          <> softline <> body <> ppString ")")
+
+ppAcPrj :: SMTy a -> SAcPrj p a b -> String
+ppAcPrj _ SAPHere = "."
+ppAcPrj (SMTPair t _) (SAPFst prj) = "(" ++ ppAcPrj t prj ++ ",)"
+ppAcPrj (SMTPair _ t) (SAPSnd prj) = "(," ++ ppAcPrj t prj ++ ")"
+ppAcPrj (SMTLEither t _) (SAPLeft prj) = "(" ++ ppAcPrj t prj ++ "|)"
+ppAcPrj (SMTLEither _ t) (SAPRight prj) = "(|" ++ ppAcPrj t prj ++ ")"
+ppAcPrj (SMTMaybe t) (SAPJust prj) = "J" ++ ppAcPrj t prj
+ppAcPrj (SMTArr n t) (SAPArrIdx prj) = "[" ++ ppAcPrj t prj ++ "]" ++ intSubscript (fromSNat n)
+
+ppSparse :: SMTy a -> Sparse a b -> String
+ppSparse t sp | Just Refl <- isDense t sp = "D"
+ppSparse _ SpAbsent = "A"
+ppSparse t (SpSparse s) = "S" ++ ppSparse t s
+ppSparse (SMTPair t1 t2) (SpPair s1 s2) = "(" ++ ppSparse t1 s1 ++ "," ++ ppSparse t2 s2 ++ ")"
+ppSparse (SMTLEither t1 t2) (SpLEither s1 s2) = "(" ++ ppSparse t1 s1 ++ "|" ++ ppSparse t2 s2 ++ ")"
+ppSparse (SMTMaybe t) (SpMaybe s) = "M" ++ ppSparse t s
+ppSparse (SMTArr _ t) (SpArr s) = "A" ++ ppSparse t s
+ppSparse (SMTScal _) SpScal = "."
+
+ppCommut :: Commutative -> String
+ppCommut Commut = "(C)"
+ppCommut Noncommut = ""
+
+ppX :: PrettyX x => Expr x env t -> ADoc
+ppX expr = annotate AExt $ ppString $ prettyXsuffix (extOf expr)
+
+data Fixity = Prefix | Infix
+  deriving (Show)
+
+operator :: SOp a t -> (Fixity, String)
+operator OAdd{} = (Infix, "+")
+operator OMul{} = (Infix, "*")
+operator ONeg{} = (Prefix, "negate")
+operator OLt{} = (Infix, "<")
+operator OLe{} = (Infix, "<=")
+operator OEq{} = (Infix, "==")
+operator ONot = (Prefix, "not")
+operator OAnd = (Infix, "&&")
+operator OOr = (Infix, "||")
+operator OIf = (Prefix, "ifB")
+operator ORound64 = (Prefix, "round")
+operator OToFl64 = (Prefix, "toFl64")
+operator ORecip{} = (Prefix, "recip")
+operator OExp{} = (Prefix, "exp")
+operator OLog{} = (Prefix, "log")
+operator OIDiv{} = (Infix, "`div`")
+operator OMod{} = (Infix, "`mod`")
+
+ppSTy :: Int -> STy t -> String
+ppSTy d ty = render $ ppSTy' d ty
+
+ppSTy' :: Int -> STy t -> Doc q
+ppSTy' _ STNil = ppString "1"
+ppSTy' d (STPair a b) = ppParen (d > 7) $ ppSTy' 8 a <> ppString " * " <> ppSTy' 8 b
+ppSTy' d (STEither a b) = ppParen (d > 6) $ ppSTy' 7 a <> ppString " + " <> ppSTy' 7 b
+ppSTy' d (STLEither a b) = ppParen (d > 6) $ ppSTy' 7 a <> ppString " ⊕ " <> ppSTy' 7 b
+ppSTy' d (STMaybe t) = ppParen (d > 10) $ ppString "Maybe " <> ppSTy' 11 t
+ppSTy' d (STArr n t) = ppParen (d > 10) $
+  ppString "Arr " <> ppString (show (fromSNat n)) <> ppString " " <> ppSTy' 11 t
+ppSTy' _ (STScal sty) = ppString $ case sty of
+  STI32 -> "i32"
+  STI64 -> "i64"
+  STF32 -> "f32"
+  STF64 -> "f64"
+  STBool -> "bool"
+ppSTy' d (STAccum t) = ppParen (d > 10) $ ppString "Accum " <> ppSMTy' 11 t
+
+ppSMTy :: Int -> SMTy t -> String
+ppSMTy d ty = render $ ppSMTy' d ty
+
+ppSMTy' :: Int -> SMTy t -> Doc q
+ppSMTy' _ SMTNil = ppString "1"
+ppSMTy' d (SMTPair a b) = ppParen (d > 7) $ ppSMTy' 8 a <> ppString " * " <> ppSMTy' 8 b
+ppSMTy' d (SMTLEither a b) = ppParen (d > 6) $ ppSMTy' 7 a <> ppString " ⊕ " <> ppSMTy' 7 b
+ppSMTy' d (SMTMaybe t) = ppParen (d > 10) $ ppString "Maybe " <> ppSMTy' 11 t
+ppSMTy' d (SMTArr n t) = ppParen (d > 10) $
+  ppString "Arr " <> ppString (show (fromSNat n)) <> ppString " " <> ppSMTy' 11 t
+ppSMTy' _ (SMTScal sty) = ppString $ case sty of
+  STI32 -> "i32"
+  STI64 -> "i64"
+  STF32 -> "f32"
+  STF64 -> "f64"
+
+ppString :: String -> Doc x
+ppString = fromString
+
+ppParen :: Bool -> Doc x -> Doc x
+ppParen True = parens
+ppParen False = id
+
+intSubscript :: Int -> String
+intSubscript = \case 0 -> "₀"
+                     n | n < 0 -> '₋' : go (-n) ""
+                       | otherwise -> go n ""
+  where go 0 suff = suff
+        go n suff = let (q, r) = n `quotRem` 10
+                    in go q ("₀₁₂₃₄₅₆₇₈₉" !! r : suff)
+
+data Annot = AKey | AWith | AHighlight | AMonoid | AExt
+  deriving (Show)
+
+annotToANSI :: Annot -> PT.AnsiStyle
+annotToANSI AKey = PT.bold
+annotToANSI AWith = PT.color PT.Red <> PT.underlined
+annotToANSI AHighlight = PT.color PT.Blue
+annotToANSI AMonoid = PT.color PT.Green
+annotToANSI AExt = PT.colorDull PT.White
+
+type ADoc = Doc Annot
+
+render :: Doc Annot -> String
+render =
+  (if stdoutTTY then TL.unpack . PT.renderLazy . reAnnotateS annotToANSI
+                else renderString)
+  . layoutPretty LayoutOptions { layoutPageWidth = AvailablePerLine 120 1.0 }
+  where
+    stdoutTTY = unsafePerformIO $ hSupportsANSI stdout