{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE GADTs, EmptyDataDecls, RankNTypes, TypeOperators #-}
module SystemF where

infixr 0 :->
newtype a :-> b = TypFun { typFunExtract :: b }

type Decl typ rest = (typ, rest)
type TyVarDecl var rest = var :-> rest
type NDecl = ()
type Name = String

data DeclWitness a b c d where
    Z :: DeclWitness a as (Decl a as) ts
    P :: DeclWitness a as bs ts -> DeclWitness a (Decl b as) (Decl b bs) ts
    T :: DeclWitness a as bs ts -> DeclWitness a (TyVarDecl t as) (TyVarDecl t bs) (TyVarDecl t ts)

data DeclWitnessT a b c d e where
    ZT :: DeclWitnessT a as (TyVarDecl a as) ts (TyVarDecl a ts)
    PT :: DeclWitnessT a as bs ats bts -> DeclWitnessT a (Decl b as) (Decl b bs) ats bts
    TT :: DeclWitnessT a as bs ats bts -> DeclWitnessT a (TyVarDecl t as) (TyVarDecl t bs) (TyVarDecl t ats) (TyVarDecl t bts)

data TypRep a ts where
    TPrim   :: Name -> a -> TypRep a ts -- a should be "undefined `asTypeOf` somePrimitiveType"
    TPi     :: Name -> TypRep a (TyVarDecl b ts) -> TypRep a ts
    TVar    :: TypRep a (TyVarDecl a ts)
    TArrow  :: TypRep a ts -> TypRep b ts -> TypRep (a -> b) ts
    TPush   :: TypRep a ts -> TypRep a (TyVarDecl b ts)
type NTypRep a = TypRep a ()

precTyNone, precTyArrow, precTyArrowR, precTyPi :: Int
precTyPi = 0
precTyArrowR = 1
precTyArrow = 2
precTyNone = 3

showTyp :: [Name] -> Int -> TypRep a ts -> ShowS
showTyp _  _ (TPrim n _) = showString n
showTyp ns p t@(TPi _ _) = showParen (p > precTyPi) (showString "#" . showPis ns t)
showTyp ns _ (TVar)      = showString $ safeHead "_" ns
showTyp ns p (TArrow l r) = showParen (p >= precTyArrow) (showTyp ns precTyArrow l . showString "->" . showTyp ns precTyArrowR r)
showTyp ns p (TPush t)   = showTyp (safeTail ns) p t

showPis :: [Name] -> TypRep a ts -> ShowS
showPis ns (TPi n t) = showString n' . showString " " . showPis (n':ns) t
    where n' = mkName n ns
showPis ns t = showString ". " . showTyp ns precTyPi t

data Expr a vs ts where
    EPrim   :: Name -> a -> Expr a vs ts
    ELam    :: Name -> TypRep b ts -> Expr a (Decl b vs) ts -> Expr (b -> a) vs ts
    EAp     :: Expr (b -> a) vs ts -> Expr b vs ts -> Expr a vs ts
    EVar    :: Expr a (Decl a vs) ts
    EPush   :: Expr a vs ts -> Expr a (Decl b vs) ts
    EFix    :: Expr (a -> a) vs ts -> Expr a vs ts
    EGamma  :: Name -> Expr a (TyVarDecl b vs) (TyVarDecl b ts) -> Expr (b :-> a) vs ts
    ETypAp  :: Expr (b :-> a) vs ts -> TypRep b ts -> Expr a vs ts
    ETPush  :: Expr a vs ts -> Expr a (TyVarDecl b vs) (TyVarDecl b ts)

type NExpr a = Expr a NDecl NDecl

safeTail :: [a] -> [a]
safeTail [] = []
safeTail (_:as) = as

safeHead :: a -> [a] -> a
safeHead d []    = d
safeHead _ (a:_) = a

mkName :: Name -> [Name] -> Name
mkName n ns
	| n `elem` ns = mkName0 1 n ns
        | otherwise   = n
mkName0 :: Integer -> Name -> [Name] -> Name
mkName0 x n ns
	| (name) `elem` ns = mkName0 (x+1) n ns
        | otherwise        = name
	where name = n ++ show x

precNone, precLam, precApL, precAp :: Int
precNone = 0
precLam = 1
precApL = 2
precAp  = 3

showLambdas :: [Name] -> [Name] -> Expr a vs ts -> ShowS
showLambdas ns ts (ELam n t e) = showString n' . showString ":" . showTyp ts precTyNone t . showString " " . showLambdas (n':ns) ts e
    where n' = mkName n ns
showLambdas ns ts (EGamma t e) = showString t' . showString " " . showLambdas ns (t':ts) e
    where t' = mkName t ts
showLambdas ns ts e          = showString "-> " . showExpr ns ts precLam e
{- this is generally OK but it makes name collisions look bad in, for example
   ELam "x" $ EPush $ ELam "x" EVar
   which, with the following line commented shows
      \x -> \x -> x
   but with uncommented
      \x x -> x
   which, while not technically ambiguous (see below) still looks bad.
      (ELam "x" $ ELam "x" $ EPush EVar) shows as \x x1 -> x
      (ELam "x" $ ELam "x" $ EVar)       shows as \x x1 -> x1
      (ELam "x" $ EPush $ ELam "x" EVar) shows as \x -> \x -> x
 -}
-- showLambdas ns (EPush e)  = showLambdas (safeTail ns) e

showExpr :: [Name] -> [Name] -> Int -> Expr a vs ts -> ShowS
showExpr ns ts p (EPush e)  = showExpr (safeTail ns) ts p e
showExpr _  _  _ (EPrim n _)= showString n
showExpr ns ts p e@(ELam _ _ _) = showParen (p > precLam) (showString "\\" . showLambdas ns ts e)
showExpr ns _  _ (EVar)     = showString (safeHead "_" ns)
showExpr ns ts p (EFix e)   = showParen (p >= precAp)  (showString "fix " . showExpr ns ts precAp e)
showExpr ns ts p (EAp l r)  = showParen (p >= precAp)  (showExpr ns ts precApL l . showString " " . showExpr ns ts precAp r)
showExpr ns ts p e@(EGamma _ _) = showParen (p > precLam) (showString "\\" . showLambdas ns ts e)
showExpr ns ts p (ETypAp l r) = showParen (p >= precAp) (showExpr ns ts precApL l . showString " " . showTyp ts precTyNone r)
showExpr ns ts p (ETPush e) = showExpr ns (safeTail ts) p e

showExprExplicit :: Expr a ts vs -> String
showExprExplicit (EPush e) = "(EPush " ++ showExprExplicit e ++ ")"
showExprExplicit (EPrim n _) = "(EPrim " ++ n ++ ")"
showExprExplicit (ELam n t e) = "(ELam " ++ show n ++ " " ++ showTypExplicit t ++ " " ++ showExprExplicit e ++ ")"
showExprExplicit (EVar) = "EVar"
showExprExplicit (EFix e) = "(EFix " ++ showExprExplicit e ++ ")"
showExprExplicit (EAp l r) = "(EAp " ++ showExprExplicit l ++ " " ++ showExprExplicit r ++ ")"
showExprExplicit (EGamma n e) = "(EGamma " ++ show n ++ " " ++ showExprExplicit e ++ ")"
showExprExplicit (ETypAp l r) = "(ETypAp " ++ showExprExplicit l ++ " " ++ showTypExplicit r ++ ")"
showExprExplicit (ETPush e) = "(ETPush " ++ showExprExplicit e ++ ")"

showTypExplicit :: TypRep a ts -> String
showTypExplicit (TPrim n _) = "(TPrim " ++ n ++ ")"
showTypExplicit (TPi n e) = "(TPi " ++ show n ++ " " ++ showTypExplicit e ++ ")"
showTypExplicit (TVar) = "TVar"
showTypExplicit (TArrow l r) = "(TArrow " ++ showTypExplicit l ++ " " ++ showTypExplicit r ++ ")"
showTypExplicit (TPush e) = "(TPush " ++ showTypExplicit e ++ ")"

instance Show (Expr a vs ts) where
    showsPrec = showExpr [] []

instance Show (TypRep a ts) where
    showsPrec = showTyp []

combS :: (x -> y -> z) -> (x -> y) -> (x -> z)
combS a b c = a c $ b c
combK :: (x -> y -> x)
combK = const

compile :: NExpr a -> a
compile e = compile' e ()

compile' :: Expr a decls ts -> decls -> a
compile' (EPrim _ p) = combK p
compile' (EAp a b)   = let { ac = compile' a ; bc = compile' b } in
		       combS ac bc
compile' (ELam _ _ e)  = let f = compile' e in \r x -> f (x,r)
compile' (EPush e)   = let f = compile' e in f . snd
compile' (EVar)      = fst
compile' (EFix e)    = let f = compile' e in 
	\r -> let { f' = f r ; x = f' x } in x
compile' (EGamma _ e) = let f = compile' e in TypFun . f . TypFun
compile' (ETypAp e _) = let f = compile' e in typFunExtract . f
compile' (ETPush e) = let f = compile' e in f . typFunExtract

type Prim a = forall vs ts. Expr a vs ts

mkPrim0 :: String -> a -> Expr a vs ts
mkPrim0 n p = EPrim n p

mkPrim1 :: Name -> Name -> TypRep a ts -> (a -> b) -> Expr (a -> b) vs ts
--mkPrim1 n v t p = ELam v t $ EAp (EPrim n p) EVar
mkPrim1 n _ _ p = EPrim n p

mkPrim2 :: Name -> Name -> TypRep a ts -> Name -> TypRep b ts -> (a -> b -> c) -> Expr (a -> b -> c) vs ts
--mkPrim2 n v1 t1 v2 t2 p = ELam v1 t1 $ ELam v2 t2 $ EAp (EAp (EPrim n p) (EPush EVar)) EVar
mkPrim2 n _ _ _ _ p = EPrim n p

mkPrim3 :: Name -> Name -> TypRep a ts -> Name -> TypRep b ts -> Name -> TypRep c ts -> (a -> b -> c -> d) -> Expr (a -> b -> c -> d) vs ts
--mkPrim3 n v1 t1 v2 t2 v3 t3 p = ELam v1 t1 $ ELam v2 t2 $ ELam v3 t3 $ EAp (EAp (EAp (EPrim n p) (EPush $ EPush EVar)) (EPush EVar)) EVar
mkPrim3 n _ _ _ _ _ _ p = EPrim n p

{- some useful primitives -}
pIf :: Prim (a :-> Bool -> a -> a -> a)
pIf = EGamma "A" $ mkPrim3 "if#" "b" tBool "t" TVar "f" TVar $ \b t f -> case b of True -> t ; False -> f

pPlus :: Prim (Integer -> Integer -> Integer)
pPlus = mkPrim2 "plus#" "x" tInt "y" tInt (+)

pMinus :: Prim (Integer -> Integer -> Integer)
pMinus = mkPrim2 "minus#" "x" tInt "y" tInt (-)

pTimes :: Prim (Integer -> Integer -> Integer)
pTimes = mkPrim2 "times#" "x" tInt "y" tInt (*)

pEq :: Prim (Integer -> Integer -> Bool)
pEq = mkPrim2 "eq#" "x" tInt "y" tInt (==)

pLt :: Prim (Integer -> Integer -> Bool)
pLt = mkPrim2 "lt#" "x" tInt "y" tInt (<)

pInt :: Integer -> Prim Integer
pInt x = mkPrim0 (show x) x

tUnit :: TypRep () ts
tUnit = TPrim "()" (undefined :: ())

tInt :: TypRep Integer ts
tInt = TPrim "I#" (undefined :: Integer)

tBool :: TypRep Bool ts
tBool = TPrim "B#" (undefined :: Bool)

tAny :: TypRep a ts
tAny = TPrim "ANY#" undefined

{- some tests -}
eFact :: NExpr (Integer -> Integer)
eFact = EFix $ ELam "rec" (TArrow tInt tInt) $ ELam "z" tInt $ EAp (EAp (EAp (ETypAp pIf tInt) (EAp (EAp pEq EVar) (pInt 0))) (pInt 1)) (EAp (EAp pTimes EVar) (EAp (EPush EVar) (EAp (EAp pMinus EVar) (pInt 1))))

eNFib :: NExpr (Integer -> Integer)
eNFib = EFix $ ELam "rec" (TArrow tInt tInt) $ ELam "z" tInt $ EAp (EAp (EAp (ETypAp pIf tInt) (EAp (EAp pLt EVar) (pInt 2))) (pInt 1)) (EAp (EAp pPlus (EAp (EPush EVar) (EAp (EAp pMinus EVar) (pInt 1)))) (EAp (EPush EVar) (EAp (EAp pMinus EVar) (pInt 2))))

eK :: NExpr (a :-> b :-> a -> b -> a)
eK = EGamma "A" $ EGamma "B" $ ELam "a" (TPush TVar) $ ELam "b" TVar $ EPush EVar

nfibCompile :: Integer -> Integer
nfibCompile = compile eNFib

fix :: (a -> a) -> a
fix f = let x = f x in x

nfibGHC :: Integer -> Integer
nfibGHC = fix $ \rec z -> if (z < 2) then 1 else rec (z-1) + rec (z-2)

-- the goal of "optimize" is to remove unnecessary "EPush" objects which don't contain
-- any free variables, 

interpret :: NExpr a -> a
interpret (EPrim _ p) = p
interpret (ELam x _ e)  = \v -> interpret (subst Z (EPrim x v) e)
interpret (EAp a b)   = interpret a $ interpret b
interpret (EFix e)    = x where
   f = interpret e
   x = f x
interpret (EGamma _ e) = TypFun $ interpret (substTExpr ZT (TPrim "any" undefined) e)
interpret (ETypAp e t) = typFunExtract $ interpret e
-- interpret (EVar)      = error "unreachable"
-- interpret (ETPush _) = error "unreachable"
-- interpret (EPush _)   = error "unreachable"
interpret _ = error "unreachable"

subst :: DeclWitness a vs vsa ts -> Expr a vs ts -> Expr b vsa ts -> Expr b vs ts
subst Z     v         (EVar)      = v
--    (T _) _         EVar        = error "unreachable"
subst (P _) _         (EVar)      = (EVar)
subst Z     _         (EPush e)   = e
--    (T _) _         (EPush _)   = error "unreachable"
subst (P w) (EPush v) (EPush e)   = EPush (subst w v e)
subst (P _) _         (EPush _)   = error $ unwords 
	["subst called with invalid arguments",
         "-- all EPush to match P(P(P(Z))) etc.",
	 "must be at the top level of the first expression."]
--    Z     _         (ETPush _)  = error "unreachable"
subst (T w) (ETPush v) (ETPush e) = ETPush (subst w v e)
subst (T _) _         (ETPush _)  = error $ unwords
	["subst called with invalid arguments",
         "-- all ETPush to match T(T(T(Z))) etc.",
	 "must be at the top level of the first expression."]
--    (P _) _         (ETPush _)  = error "unreachable"
subst _     _         (EPrim n x) = EPrim n x
subst w     v         (ELam n t e)  = ELam n t (subst (P w) (EPush v) e)
subst w     v         (EAp e1 e2) = EAp (subst w v e1) (subst w v e2)
subst w     v         (EFix e)    = EFix (subst w v e)
subst w     v         (EGamma n e) = EGamma n (subst (T w) (ETPush v) e)
subst w     v         (ETypAp e t) = ETypAp (subst w v e) t

substTExpr :: DeclWitnessT a vs vsa ts tsa -> TypRep a ts -> Expr b vsa tsa -> Expr b vs ts
-- substTExpr ZT _       (EVar) = EVar
substTExpr (PT _) _   (EVar) = EVar
-- substTExpr (TT _)  _   (EVar) = EVar
-- substTExpr ZT _ (EPush e) = error "unreachable"
substTExpr (PT w) v (EPush e) = EPush (substTExpr w v e)
-- substTExpr (TT _) _ (EPush e) = error "unreachable"
substTExpr ZT _ (ETPush e) = e
-- substTExpr (PT w) _ (ETPush e) = error "unreachable"
substTExpr (TT w) (TPush v) (ETPush e) = ETPush $ substTExpr w v e
substTExpr (TT w) _ (ETPush e) = error "invalid arguments to substTExpr"
substTExpr _ _ (EPrim n x) = EPrim n x
substTExpr w v (ELam n t e) = ELam n (substTTyp w v t) (substTExpr (PT w) v e)
substTExpr w v (EAp e1 e2) = EAp (substTExpr w v e1) (substTExpr w v e2)
substTExpr w v (EFix e) = EFix (substTExpr w v e)
substTExpr w v (EGamma n e) = EGamma n (substTExpr (TT w) (TPush v) e)
substTExpr w v (ETypAp e t) = ETypAp (substTExpr w v e) (substTTyp w v t)

substTTyp :: DeclWitnessT a vs vsa ts tsa -> TypRep a ts -> TypRep b tsa -> TypRep b ts
substTTyp (PT w) v         e            = substTTyp w v e
substTTyp _      _         (TPrim n a)  = TPrim n a
substTTyp w      v         (TPi n t)    = TPi n $ substTTyp (TT w) (TPush v) t
substTTyp ZT     v         TVar         = v
substTTyp (TT _) _         TVar         = TVar
substTTyp w      v         (TArrow l r) = TArrow (substTTyp w v l) (substTTyp w v r)
substTTyp ZT     _         (TPush t)    = t
substTTyp (TT w) (TPush v) (TPush t)    = TPush $ substTTyp w v t
substTTyp (TT _) _         (TPush t)    = error "invalid arguments to substTTyp"


eId :: NExpr (a :-> a -> a)
eId = EGamma "A" $ ELam "a" (TVar) EVar

eBottom :: NExpr (a :-> a)
eBottom = EGamma "A" $ EFix $ ELam "a" TVar EVar

subAp :: Expr a vs ts -> Expr a vs ts
subAp (EAp (ELam _ _ e) v) = subst Z v e
subAp (ETypAp (EGamma _ e) t) = substTExpr ZT t e
subAp e                  = e

whnf :: Expr a vs ts -> Expr a vs ts
whnf e@(EAp l r) = case whnf l of
    l'@(ELam _ _ _) -> whnf $ subAp (EAp l' r)
    (EPrim n p)     -> EPrim (n++"*") $ p $ compile' (whnf r) undefined -- should be ok
    l'              -> (EAp l' r)
whnf (EFix e)  = whnf $ EAp e $ EFix e
whnf (ETypAp l r) = let l' = whnf l in whnf $ subAp $ ETypAp l' r
whnf e = e

extractPrim :: Expr a vs ts -> Maybe a
extractPrim (EPrim _ a) = Just a
extractPrim _           = Nothing