perga/lib/Expr.hs

110 lines
3.4 KiB
Haskell

module Expr where
import Data.Function (on)
import Data.Text (Text)
import qualified Data.Text as T
data Expr where
Var :: Integer -> Text -> Expr
Free :: Text -> Expr
Star :: Expr
Square :: Expr
App :: Expr -> Expr -> Expr
Abs :: Text -> Expr -> Expr -> Expr
Pi :: Text -> Expr -> Expr -> Expr
deriving (Show, Ord)
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
occursFree :: Integer -> Expr -> Bool
occursFree n (Var k _) = n == k
occursFree _ (Free _) = False
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
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 _ _ (Free s) = Free s
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
{- --------------------- PRETTY PRINTING ----------------------------- -}
parenthesize :: Text -> Text
parenthesize s = T.concat ["(", s, ")"]
collectLambdas :: Expr -> ([(Text, Expr)], Expr)
collectLambdas (Abs x ty body) = ((x, ty) : params, final)
where
(params, final) = collectLambdas body
collectLambdas e = ([], e)
collectPis :: Expr -> ([(Text, Expr)], Expr)
collectPis p@(Pi "" _ _) = ([], p)
collectPis (Pi x ty body) = ((x, ty) : params, final)
where
(params, final) = collectPis body
collectPis e = ([], e)
groupParams :: [(Text, Expr)] -> [([Text], Expr)]
groupParams = foldr addParam []
where
addParam :: (Text, Expr) -> [([Text], Expr)] -> [([Text], 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 :: ([Text], Expr) -> Text
showParamGroup (ids, ty) = parenthesize $ T.unwords ids <> " : " <> pretty ty
helper :: Integer -> Expr -> Text
helper _ (Var _ s) = s
helper _ (Free 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 = "" <> T.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 = "λ " <> T.unwords grouped <> " . " <> pretty body
pretty :: Expr -> Text
pretty = helper 0
prettyS :: Expr -> String
prettyS = T.unpack . pretty