{-# LANGUAGE GADTs #-}

module Expr where

import Data.Function (on)

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

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

mapIndices :: (Integer -> Expr) -> Expr -> Expr
mapIndices f (Var n) = f n
mapIndices _ Star = Star
mapIndices _ Square = Square
mapIndices f (App m n) = App (mapIndices f m) (mapIndices f n)
mapIndices f (Abs m n) = Abs (mapIndices f m) (mapIndices f n)
mapIndices f (Pi m n) = Pi (mapIndices f m) (mapIndices f n)

incIndices :: Expr -> Expr
incIndices = mapIndices (Var . (+ 1))

decIndices :: Expr -> Expr
decIndices = mapIndices (Var . subtract 1)

subst :: Integer -> Expr -> Expr -> Expr
subst k s (Var n) = if k == n then s else Var n
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 m n) = Abs (subst k s m) (subst k (incIndices s) n)
subst k s (Pi m n) = Pi (subst k s m) (subst k (incIndices s) n)

betaReduce :: Expr -> Expr
betaReduce (Var k) = Var k
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 t v) = Abs (betaReduce t) (betaReduce v)
betaReduce (Pi t v) = Pi (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