perga/app/Expr.hs

{-# LANGUAGE GADTs #-}

module Expr where

import Data.Function (on)

data Expr where
    Var :: Integer -> String -> Expr
    Star :: Expr
    Square :: Expr
    App :: Expr -> Expr -> Expr
    Abs :: String -> Expr -> Expr -> Expr
    Pi :: String -> Expr -> Expr -> Expr
    deriving (Show)

instance Eq Expr where
    (Var n _) == (Var m _) = n == m
    Star == Star = True
    Square == Square = True
    (App e1 e2) == (App f1 f2) = e1 == f1 && e2 == f2
    (Abs _ t1 b1) == (Abs _ t2 b2) = t1 == t2 && b1 == b2
    (Pi _ t1 b1) == (Pi _ t2 b2) = t1 == t2 && b1 == b2
    _ == _ = False

infixl 4 <.>

(<.>) :: Expr -> Expr -> Expr
(<.>) = App

infixr 2 .->

(.->) :: Expr -> Expr -> Expr
a .-> b = Pi "" a (incIndices b)

occursFree :: Integer -> Expr -> Bool
occursFree n (Var k _) = n == k
occursFree _ Star = False
occursFree _ Square = False
occursFree n (App a b) = on (||) (occursFree n) a b
occursFree n (Abs _ a b) = occursFree n a || occursFree (n + 1) b
occursFree n (Pi _ a b) = occursFree n a || occursFree (n + 1) b

{- --------------------- PRETTY PRINTING ----------------------------- -}

parenthesize :: String -> String
parenthesize s = "(" ++ s ++ ")"

collectLambdas :: Expr -> ([(String, Expr)], Expr)
collectLambdas (Abs x ty body) = ((x, ty) : params, final)
  where
    (params, final) = collectLambdas body
collectLambdas e = ([], e)

collectPis :: Expr -> ([(String, Expr)], Expr)
collectPis p@(Pi "" _ _) = ([], p)
collectPis (Pi x ty body) = ((x, ty) : params, final)
  where
    (params, final) = collectPis body
collectPis e = ([], e)

groupParams :: [(String, Expr)] -> [([String], Expr)]
groupParams = foldr addParam []
  where
    addParam :: (String, Expr) -> [([String], Expr)] -> [([String], Expr)]
    addParam (x, t) [] = [([x], t)]
    addParam (x, t) l@((xs, s) : rest)
        | incIndices t == s = (x : xs, t) : rest
        | otherwise = ([x], t) : l

showParamGroup :: ([String], Expr) -> String
showParamGroup (ids, ty) = parenthesize $ unwords ids ++ " : " ++ pretty ty

helper :: Integer -> Expr -> String
helper _ (Var _ s) = s
helper _ Star = "*"
helper _ Square = "□"
helper k (App e1 e2) = if k > 3 then parenthesize res else res
  where
    res = helper 3 e1 ++ " " ++ helper 4 e2
helper k (Pi "" t1 t2) = if k > 2 then parenthesize res else res
  where
    res = helper 3 t1 ++ " -> " ++ helper 2 t2
helper k e@(Pi{}) = if k > 2 then parenthesize res else res
  where
    (params, body) = collectPis e
    grouped = showParamGroup <$> groupParams params
    res = "∏ " ++ unwords grouped ++ " . " ++ pretty body
helper k e@(Abs{}) = if k >= 1 then parenthesize res else res
  where
    (params, body) = collectLambdas e
    grouped = showParamGroup <$> groupParams params
    res = "λ " ++ unwords grouped ++ " . " ++ pretty body

pretty :: Expr -> String
pretty = helper 0

{- --------------- ACTUAL MATH STUFF ---------------- -}

isSort :: Expr -> Bool
isSort Star = True
isSort Square = True
isSort _ = False

shiftIndices :: Integer -> Integer -> Expr -> Expr
shiftIndices d c (Var k x)
    | k >= c = Var (k + d) x
    | otherwise = Var k x
shiftIndices _ _ Star = Star
shiftIndices _ _ Square = Square
shiftIndices d c (App m n) = App (shiftIndices d c m) (shiftIndices d c n)
shiftIndices d c (Abs x m n) = Abs x (shiftIndices d c m) (shiftIndices d (c + 1) n)
shiftIndices d c (Pi x m n) = Pi x (shiftIndices d c m) (shiftIndices d (c + 1) n)

incIndices :: Expr -> Expr
incIndices = shiftIndices 1 0

-- substitute s for k *AND* decrement indices; only use after reducing a redex.
subst :: Integer -> Expr -> Expr -> Expr
subst k s (Var n x)
    | k == n = s
    | otherwise = Var (n - 1) x
subst _ _ Star = Star
subst _ _ Square = Square
subst k s (App m n) = App (subst k s m) (subst k s n)
subst k s (Abs x m n) = Abs x (subst k s m) (subst (k + 1) (incIndices s) n)
subst k s (Pi x m n) = Pi x (subst k s m) (subst (k + 1) (incIndices s) n)

betaReduce :: Expr -> Expr
betaReduce (Var k s) = Var k s
betaReduce Star = Star
betaReduce Square = Square
betaReduce (App (Abs _ _ v) n) = subst 0 n v
betaReduce (App m n) = App (betaReduce m) (betaReduce n)
betaReduce (Abs x t v) = Abs x (betaReduce t) (betaReduce v)
betaReduce (Pi x t v) = Pi x (betaReduce t) (betaReduce v)

betaNF :: Expr -> Expr
betaNF e = if e == e' then e else betaNF e'
  where
    e' = betaReduce e

betaEquiv :: Expr -> Expr -> Bool
betaEquiv = on (==) betaNF