iris/BidirTT/Pretty.lean

import BidirTT.Syntax

namespace BidirTT

mutual
  partial def prettyTmWith (names : List Name) (nextFresh : Nat) : Tm → String × Nat
    | .var i =>
        match names[i]? with
        | some name => (name, nextFresh)
        | none      => (s!"#{i}", nextFresh)
    | .lam t =>
        let x := s!"x{nextFresh}"
        let (body, nextFresh) := prettyTmWith (x :: names) (nextFresh + 1) t
        (s!"(fun {x} => {body})", nextFresh)
    | .app t u    =>
        let (pt, nextFresh) := prettyTmWith names nextFresh t
        let (pu, nextFresh) := prettyTmWith names nextFresh u
        (s!"({pt} {pu})", nextFresh)
    | .pi a b     =>
        let x := s!"x{nextFresh}"
        let (pa, nextFresh) := prettyTmWith names (nextFresh + 1) a
        let (pb, nextFresh) := prettyTmWith (x :: names) nextFresh b
        (s!"(Pi ({x} : {pa}) -> {pb})", nextFresh)
    | .sig a b    =>
        let x := s!"x{nextFresh}"
        let (pa, nextFresh) := prettyTmWith names (nextFresh + 1) a
        let (pb, nextFresh) := prettyTmWith (x :: names) nextFresh b
        (s!"(Sigma ({x} : {pa}) * {pb})", nextFresh)
    | .pair t u   =>
        let (pt, nextFresh) := prettyTmWith names nextFresh t
        let (pu, nextFresh) := prettyTmWith names nextFresh u
        (s!"({pt}, {pu})", nextFresh)
    | .fst t      =>
        let (pt, nextFresh) := prettyTmWith names nextFresh t
        (s!"({pt}.1)", nextFresh)
    | .snd t      =>
        let (pt, nextFresh) := prettyTmWith names nextFresh t
        (s!"({pt}.2)", nextFresh)
    | .nat        => ("Nat", nextFresh)
    | .zero       => ("zero", nextFresh)
    | .succ t     =>
        let (pt, nextFresh) := prettyTmWith names nextFresh t
        (s!"(succ {pt})", nextFresh)
    | .natElim m z s n =>
        let (pm, nextFresh) := prettyTmWith names nextFresh m
        let (pz, nextFresh) := prettyTmWith names nextFresh z
        let (ps, nextFresh) := prettyTmWith names nextFresh s
        let (pn, nextFresh) := prettyTmWith names nextFresh n
        (s!"(natElim {pm} {pz} {ps} {pn})", nextFresh)
    | .unit       => ("Unit", nextFresh)
    | .triv       => ("tt", nextFresh)
    | .unitElim m t u =>
        let (pm, nextFresh) := prettyTmWith names nextFresh m
        let (pt, nextFresh) := prettyTmWith names nextFresh t
        let (pu, nextFresh) := prettyTmWith names nextFresh u
        (s!"(unitElim {pm} {pt} {pu})", nextFresh)
    | .empty      => ("Empty", nextFresh)
    | .emptyElim m e =>
        let (pm, nextFresh) := prettyTmWith names nextFresh m
        let (pe, nextFresh) := prettyTmWith names nextFresh e
        (s!"(emptyElim {pm} {pe})", nextFresh)
    | .id a t u   =>
        let (pa, nextFresh) := prettyTmWith names nextFresh a
        let (pt, nextFresh) := prettyTmWith names nextFresh t
        let (pu, nextFresh) := prettyTmWith names nextFresh u
        (s!"(Id {pa} {pt} {pu})", nextFresh)
    | .refl       => ("refl", nextFresh)
    | .idElim m r y p =>
        let (pm, nextFresh) := prettyTmWith names nextFresh m
        let (pr, nextFresh) := prettyTmWith names nextFresh r
        let (py, nextFresh) := prettyTmWith names nextFresh y
        let (pp, nextFresh) := prettyTmWith names nextFresh p
        (s!"(idElim {pm} {pr} {py} {pp})", nextFresh)
    | .univ i     => (s!"U{i.pretty}", nextFresh)
    | .letE a t u =>
        let x := s!"x{nextFresh}"
        let (pa, nextFresh) := prettyTmWith names (nextFresh + 1) a
        let (pt, nextFresh) := prettyTmWith names nextFresh t
        let (pu, nextFresh) := prettyTmWith (x :: names) nextFresh u
        (s!"(let {x} : {pa} := {pt}; {pu})", nextFresh)
end

def prettyTm (tm : Tm) : String :=
  (prettyTmWith [] 0 tm).fst

end BidirTT