Normalization of polarized lambda calculus to A-normal form

anormalform.hs

-- Polarized lambda calculus to A-normal form
-- With eta-expansion and call-saturation

type Ix = Int
type Lvl = Int

data Ty = TInt
  deriving (Show)
data TFun = TFun [Ty] Ty
  deriving (Show)

data Tm
  = Var Ix
  | App Tm Tm
  | Lam Tm
  | Let TFun Tm Tm 
  
data Let = LetApp Ix [Ix] | LetLam TFun ATm
  deriving (Show)
data ATm = ATm [Let] Ix
  deriving (Show)

type Env = [Lvl]

toIx :: Lvl -> Lvl -> Lvl
toIx dom k = dom - k - 1

impossible = undefined

norm :: Tm -> TFun -> ATm
norm tm (TFun ps _) =
  let dom = length ps in
  let args = [0..(dom - 1)] in
  go tm args dom []
  where
    go :: Tm -> [Lvl] -> Lvl -> Env -> ATm
    go (Var ix) [] dom env = ATm [] (toIx dom (env !! ix))
    go (Var ix) as dom env = ATm [LetApp (toIx dom (env !! ix)) (map (toIx dom) as)] 0
    go (Lam b) (a : as) dom env = go b as dom (a : env)
    go (Lam _) [] _ _ = impossible
    go app@(App _ _) as dom env = goApp app as [] dom env
    go (Let ty v b) as dom env =
      let lets = goLetVal v ty [] dom env in
      let innerdom = dom + (length lets) in
      let (ATm lets2 ret) = go b as innerdom ((innerdom - 1) : env) in
      ATm (lets ++ lets2) ret
    
    goApp :: Tm -> [Lvl] -> [Let] -> Lvl -> Env -> ATm
    goApp (App f a) as lets dom env =
      let (ATm lets2 ix) = go a [] dom env in
      let innerdom = dom + length lets in
      let extraArg = innerdom - ix - 1 in
      goApp f (extraArg : as) (lets ++ lets2) innerdom env
    goApp (Var ix) [] lets dom env = ATm lets (toIx dom (env !! ix))
    goApp (Var ix) as lets dom env =
      let app = LetApp (toIx dom (env !! ix)) (map (toIx dom) as) in
      ATm (lets ++ [app]) 0
    goApp (Lam b) (a : as) lets dom env =
      goApp b as lets dom (a : env)
    goApp (Lam _) [] _ _ _ = impossible
    goApp (Let ty v b) as lets dom env =
      let lets2 = goLetVal v ty [] dom env in
      let innerdom = dom + (length lets2) in
      goApp b as (lets ++ lets2) innerdom ((innerdom - 1) : env)
    
    goLetVal :: Tm -> TFun -> [Let] -> Lvl -> Env -> [Let]
    goLetVal (Let ty2 v2 b) ty lets dom env =
      let lets2 = goLetVal v2 ty2 [] dom env in
      let innerdom = dom + (length lets2) in
      goLetVal b ty (lets ++ lets2) innerdom ((innerdom - 1) : env)
    goLetVal v (TFun [] _) lets dom env =
      let (ATm lets2 _) = go v [] dom env in
      lets ++ lets2
    goLetVal v ty@(TFun ps _) lets dom env =
      let innerdom = dom + (length ps) in
      let args = [dom..(innerdom - 1)] in
      let body = go v args innerdom env in
      let lam = LetLam ty body in
      lets ++ [lam]

{-
example : Int -> Int =
  \x.
    let id : Int -> Int = (let id2 : Int -> Int = \y. y; \z. id2 z);
    id x
~>
example x =
  let id2 y = y;
  let id z = id2 z;
  let a = id x;
  a
-}
exampleTy = TFun [TInt] TInt
example = Lam $ Let (TFun [TInt] TInt) (Let (TFun [TInt] TInt) (Lam $ Var 0) $ Lam $ App (Var 1) (Var 0)) $ App (Var 0) (Var 1)

main :: IO ()
main = putStrLn (show (norm example exampleTy))