SimplyEasy.hs

module SimplyEasy where
-- From the paper: http://strictlypositive.org/Easy.pdf
-- An attempt to add case -statement.
-- Fails on typechecking a case.

data Name
    = Const String
    | Bound Int
    | Unquoted Int
    deriving (Show, Eq)

data TermNe
    = Ann TermN TermN
    | Star
    | Pi TermN TermN
    | Var Int
    | Par Name
    | TermNe :@: TermN
    -- | Nat
    -- | NatElim TermN TermN TermN TermN
    -- | Vec TermN TermN
    -- | VecElim TermN TermN TermN TermN TermN TermN

    | Spec
    | Data TermN [TermN]
    | Case TermN TermNe [TermN]
    -- | Des TermN String [TermN]
    -- | Field TermNe Int
    deriving (Show, Eq)

data TermN
    = Inf TermNe
    | Lam TermN
    -- | Zero
    -- | Succ TermN
    -- | Nil TermN
    -- | Cons TermN TermN TermN TermN
    --
    -- | Record [TermN]
    -- | Cocase [(String, TermN)]
    | Con Int [TermN]
    -- | CSpec [TermN] [(String, ([TermN], [TermN], TermN))]
    | DSpec [TermN] [[TermN]]
    -- | RSpec [TermN] [TermN]
    deriving (Show, Eq)

------------------------------------------------------------
-- Evaluator
------------------------------------------------------------
type Env = [Value]

evalNe :: TermNe -> Env -> Value
evalNe (Ann e _) d = evalN e d
evalNe Star      d = VStar
evalNe (Pi t t') d = VPi (evalN t d) (\x -> evalN t' (x : d))
evalNe (Par x)   d = vpar x
evalNe (Var i)   d = d !! i
evalNe (e1 :@: e2) d = vapp (evalNe e1 d) (evalN e2 d)
evalNe Spec d = VSpec
evalNe (Case m n cs) d
    = let nVal = evalNe n d in
          case nVal of
          VCon i args -> foldl vapp (evalN (cs !! i) d) args
          VNeutral n -> VNeutral
                        (NCase (evalN m d) n (map (\c -> evalN c d) cs))
          _ -> error "runtime type error"
evalNe (Data s h) d = VData (evalN s d) (map (\x -> evalN x d) h)

-- evalNe Nat     d = VNat
-- evalNe (NatElim m mz ms n) d
--     = let mzVal = evalN mz d
--           msVal = evalN ms d
--           rec nVal =
--             case nVal of
--                 VZero   -> mzVal
--                 VSucc k -> msVal `vapp` k `vapp` rec k
--                 VNeutral n -> VNeutral
--                               (NNatElim (evalN m d) mzVal msVal n)
--                 _ -> error "runtime type error"
--       in rec (evalN n d)
-- evalNe (VecElim a m mn mc n xs) d
--     = let mnVal = evalN mn d
--           mcVal = evalN mc d
--           rec nVal xsVal = case xsVal of
--             VNil _         -> mnVal
--             VCons _ k x xs -> foldl vapp mcVal [k, x, xs, rec k xs]
--             VNeutral n     -> VNeutral
--                               (NVecElim (evalN a d) (evalN m d)
--                                         mnVal mcVal nVal n)
--             _              -> error "runtime type error"
--       in rec (evalN n d) (evalN xs d)
-- evalNe (Vec a n) d = VVec (evalN a d) (evalN n d)

-- evalNe (Data args) d = VNeutral (NData (evalN args d))

vapp :: Value -> Value -> Value
vapp (VLam f)     v = f v
vapp (VNeutral n) v = VNeutral (NApp n v)
vapp _            v = error "runtime type error"

evalN :: TermN -> Env -> Value
evalN (Inf i) d = evalNe i d
evalN (Lam e) d = VLam (\x -> evalN e (x : d))
evalN (Con k xs) d = VCon k (map (\x -> evalN x d) xs)
evalN (DSpec h xxs) d = VDSpec (map (\x -> evalN x d) h)
                               (map (map (\x -> evalN x d)) xxs)
-- evalN Zero     d = VZero
-- evalN (Succ k) d = VSucc (evalN k d)
-- evalN (Nil a) d = VNil (evalN a d)
-- evalN (Cons a n x xs) d = VCons (evalN a d) (evalN n d) (evalN x d) (evalN xs d)

data Value
    = VLam  (Value -> Value)
    | VStar
    | VPi Value (Value -> Value)
    | VNeutral Neutral
    | VCon Int [Value]
    | VDSpec [Value] [[Value]]
    | VData Value [Value]
    -- | VNat
    -- | VZero
    -- | VSucc Value
    -- | VNil Value
    -- | VCons Value Value Value Value
    -- | VVec Value Value
    --
    | VSpec
    -- | VData Value [Value]
    -- | VDSpec [Value] [(String, ([Value], [Value]))]

data Neutral
    = NPar Name
    | NApp Neutral Value
    | NCase Value Neutral [Value]
    -- | NNatElim Value Value Value Neutral
    -- | NVecElim Value Value Value Value Value Neutral

vpar :: Name -> Value
vpar n = VNeutral (NPar n)

------------------------------------------------------------
-- Contexts
------------------------------------------------------------

type Type = Value
type Context = [(Name, Type)]

infer0 :: Context -> TermNe -> Result Type
infer0 = infer 0

infer :: Int -> Context -> TermNe -> Result Type
infer i g (Ann e t)
    = do check i g t VStar
         let v = evalN t []
         check i g e v
         return v
infer i g Star
    = return VStar
infer i g (Pi t t')
    = do check i g t VStar
         let v = evalN t' []
         check (i + 1) ((Bound i, v) : g)
            (substN 0 (Par (Bound i)) t') VStar
         return VStar
infer i g (Var x) = error "impossible"
infer i g (Par x)
    = case lookup x g of
        Just v  -> return v
        Nothing -> throwError "unknown identifier"
infer i g (e1 :@: e2)
    = do o <- infer i g e1
         case o of
           VPi t f -> do check i g e2 t
                         return (f (evalN e2 []))
           _        -> throwError "illegal application"
infer i g Spec = return VStar
infer i g (Data spec params) = do
    check i g spec VSpec
    head <- getSpecHead spec
    checkParams i g params (map (\e -> evalN e []) head) (buildTail head id)
    pure VStar
-- Now the case statement is not going through.
-- infer i g (Case rtype arg cases) = do
--     check i g restype VStar
--     let rVal = evalN restype []
--     t <- infer i g arg
--     (VData (VDSpec _ ss) args) -> do
--         mapM (checkCase ss args rVal) cases
--     _ -> throwError "type mismatch"


getSpecHead :: TermN -> Result [TermN]
-- getSpecHead (CSpec a _) = pure a
getSpecHead (DSpec a _) = pure a
-- getSpecHead (RSpec a _) = pure a
getSpecHead _ = throwError "not a spec"
-- 
-- checkCase :: [(String, ([Value], [Value]))]
--           -> [Value]
--           -> Value
--           -> (String, TermN)
--           -> Either String ()
-- checkCase ss args restype (s, term) = throwError "TODO"


-- infer i g Nat = return VStar
-- infer i g (NatElim m mz ms n) =
--     do check i g m (VPi VNat (const VStar))
--        let mVal = evalN m []
--        check i g mz (mVal `vapp` VZero)
--        check i g ms
--          (VPi VNat (\n ->
--             VPi (mVal `vapp` n) (\_ -> mVal `vapp` VSucc n)))
--        check i g n VNat
--        let nVal = evalN n []
--        return (mVal `vapp` nVal)
-- infer i g (Vec a n)
--     = do check i g a VStar
--          check i g n VNat
--          return VStar
-- infer i g (VecElim a m mn mc n vs)
--     = do check i g a VStar
--          let aVal = evalN a []
--          check i g m
--             (VPi VNat (\n -> (VPi (VVec aVal n) (\_ -> VStar))))
--          let mVal = evalN m []
--          check i g mn (foldl vapp mVal [VZero, VNil aVal])
--          check i g mc
--             (VPi VNat (\n ->
--              (VPi aVal (\y ->
--               (VPi (VVec aVal n) (\ys ->
--                (VPi (foldl vapp mVal [n, ys]) (\_ ->
--                 (foldl vapp mVal [VSucc n, VCons aVal n y ys])))))))))
--          check i g n VNat
--          let nVal = evalN n []
--          check i g vs (VVec aVal nVal)
--          let vsVal = evalN vs []
--          return (foldl vapp mVal [nVal, vsVal])

check :: Int -> Context -> TermN -> Type -> Result ()
check i g (Inf e) t
    = do t' <- infer i g e
         checkEq t t' "type mismatch"
check i g (Lam e) (VPi v f)
    = check (i + 1) ((Bound i, v) : g)
        (substN 0 (Par (Bound i)) e) (f (vpar (Bound i)))
check i g (Con k xs) (VData (VDSpec _ xxs) zs) = do
    let f = buildAp zs id
    checkParams i g xs (xxs !! k) f
check i g (DSpec h xs) VSpec = do 
    tys' <- checkXs i g h []
    mapM_ (\ys -> checkXs i g ys tys') xs
check i g _ _
    = throwError "type mismatch"

checkXs i g [] tys = pure tys
checkXs i g (x:xs) tys = do
    check i g x (buildTy tys VStar)
    let tys' = (evalN x []) : tys
    checkXs i g xs tys'

buildTy [] ty = ty
buildTy (x:xs) ty = buildTy xs (VPi x (const ty))

buildAp :: [Value] -> (Value -> Value) -> (Value -> Value)
buildAp [] f = f
buildAp (h:hs) f = buildAp hs (\a -> f a `vapp` h)


-- check i g Zero VNat = return ()
-- check i g (Succ k) VNat = check i g k VNat
-- check i g (Nil a) (VVec bVal VZero)
--     = do check i g a VStar
--          let aVal = evalN a []
--          checkEq aVal bVal "type mismatch"
-- check i g (Cons a n x xs) (VVec bVal (VSucc k))
--     = do check i g a VStar
--          let aVal = evalN a []
--          checkEq aVal bVal "type mismatch"
--          check i g n VNat
--          let nVal = evalN n []
--          checkEq nVal k "number mismatch"
--          check i g x aVal
--          check i g xs (VVec bVal k)

checkEq :: Value -> Value -> String -> Result ()
checkEq t t' s = if (quote0 t == quote0 t')
               then return ()
               else (throwError s)

buildTail :: [TermN] -> (Value -> Value) -> Value -> Value
buildTail [] f = f
buildTail (x:xs) f = buildTail xs (\a -> vapp (f a) (evalN x []))

checkParams :: Int -> Context -> [TermN] -> [Type] -> (Type -> Type) -> Either String ()
checkParams i g [] [] f = pure ()
checkParams i g (x:xs) (t:ts) f = do
    check i g x (f t)
    checkParams i g xs ts (\a -> vapp (f a) (evalN x []))
checkParams i g _ _ f = throwError "mismatching argc"


type Result a = Either String a

throwError :: String -> Result a
throwError = Left

substNe :: Int -> TermNe -> TermNe -> TermNe
substNe i r (Ann e t) = Ann (substN i r e) t
substNe i r Star      = Star
substNe i r (Pi t t') = Pi (substN i r t) (substN (i + 1) r t')
substNe i r (Var j)   = if i == j then r else Var j
substNe i r (Par y)   = Par y
substNe i r (e1 :@: e2) = substNe i r e1 :@: substN i r e2
substNe i r Spec = Spec
substNe i r (Case a x xs) = Case (substN i r a) (substNe i r x) (substMany i r xs)
substNe i r (Data h xs) = Data (substN i r h) (map (substN i r) xs)
-- substNe i r (Data spec parms) = Data (substN i r spec) (map (substN i r) parms)
-- substNe i r (Case rv arg fd) = Case (substN i r rv) (substNe i r arg)
--                                     (map (\(s,x) -> (s, substN i r x)) fd)
-- substNe i r (Des spec s a) = Des (substN i r spec) s (map (substN i r) a)
-- substNe i r (Field a k) = Field (substNe i r a) k
-- substNe i r Nat         = Nat
-- substNe i r (NatElim m mz ms n) = NatElim
--     (substN i r m)
--     (substN i r mz)
--     (substN i r ms)
--     (substN i r n)
-- substNe i r (Vec a n) = Vec
--     (substN i r a)
--     (substN i r n)
-- substNe i r (VecElim a m mn mc n xs) = VecElim
--     (substN i r a)
--     (substN i r m)
--     (substN i r mn)
--     (substN i r mc)
--     (substN i r n)
--     (substN i r xs)

substN :: Int -> TermNe -> TermN -> TermN
substN i r (Inf e)  = Inf (substNe i r e)
substN i r (Lam e)  = Lam (substN (i + 1) r e)
-- substN i r (Record p) = Record (map (substN i r) p)
-- substN i r (Cocase c) = Cocase (map (\(s,x) -> (s, substN i r x)) c)
-- substN i r (Con spec n f) = Con (substN i r spec) n (map (substN i r) f)
-- substN i r (CSpec p dt) = CSpec (map (substN i r) p) (map (substCRow i r) dt)
substN i r (Con k x) = Con k (map (substN i r) x)
substN i r (DSpec p dt) = DSpec (map (substN i r) p) (map (substMany i r) dt)
-- substN i r (RSpec p ds) = RSpec (map (substN i r) p) (map (substN i r) ds)

-- substCRow :: Int
--           -> TermNe
--           -> (String, ([TermN], [TermN], TermN))
--           -> (String, ([TermN], [TermN], TermN))
-- substCRow i r (s, (a, p, q)) = (s,
--                                 (map (substN i r) a,
--                                  map (substN i r) p,
--                                  substN i r q))

substMany :: Int -> TermNe -> [TermN] -> [TermN]
substMany i r xs = map (substN i r) xs
--
-- substDRow :: Int
--           -> TermNe
--           -> (String, ([TermN], [TermN]))
--           -> (String, ([TermN], [TermN]))
-- substDRow i r (s, (a, p)) = (s, (map (substN i r) a,
--                                  map (substN i r) p))


-- substN i r Zero     = Zero
-- substN i r (Succ k) = Succ (substN i r k)
-- substN i r (Nil a)  = Nil (substN i r a)
-- substN i r (Cons a n x xs) = Cons (substN i r a)
--                                   (substN i r n)
--                                   (substN i r x)
--                                   (substN i r xs)

------------------------------------------------------------
-- Quotation
------------------------------------------------------------

quote0 :: Value -> TermN
quote0 = quote 0

quote :: Int -> Value -> TermN
quote i (VLam f) = Lam (quote (i + 1) (f (vpar (Unquoted i))))
quote i (VNeutral n) = Inf (neutralQuote i n)
quote i VStar = Inf Star
quote i (VPi v f)
    = Inf (Pi (quote i v) (quote (i + 1) (f (vpar (Unquoted i)))))
quote i VSpec = Inf Spec
quote i (VCon k xs) = Con k (map (quote i) xs)
quote i (VDSpec h xs) = DSpec (map (quote i) h) (map (map (quote i)) xs)
quote i (VData h xs) = Inf (Data (quote i h) (map (quote i) xs))
-- quote i VNat = Inf Nat
-- quote i VZero = Zero
-- quote i (VSucc k) = Succ (quote i k)
-- quote i (VNil n) = Nil (quote i n)
-- quote i (VCons a n x xs) = Cons (quote i a) (quote i n) (quote i x) (quote i xs)
-- quote i (VVec a n) = Inf (Vec (quote i a) (quote i n))

neutralQuote :: Int -> Neutral -> TermNe
neutralQuote i (NPar x) = varpar i x
neutralQuote i (NApp n v) = neutralQuote i n :@: quote i v
neutralQuote i (NCase a x xs) = Case (quote i a) (neutralQuote i x)
                                     (map (quote i) xs)
-- neutralQuote i (NNatElim m mz ms n) = NatElim
--     (quote i m)
--     (quote i mz)
--     (quote i ms)
--     (Inf (neutralQuote i n))
-- neutralQuote i (NVecElim a m mn mc n xs) = VecElim
--     (quote i a)
--     (quote i m)
--     (quote i mn)
--     (quote i mc)
--     (quote i n)
--     (Inf (neutralQuote i xs))

varpar :: Int -> Name -> TermNe
varpar i (Unquoted k) = Var (i - k - 1)
varpar i x            = Par x

-- id' = Lam (Inf (Var 0))
-- const' = Lam (Lam (Inf (Var 1)))
par a = Par (Const a)
-- par x = Inf (Par (Const x))
-- term1 = Ann id' (Fun (tpar "a") (tpar "a")) :@: par "y"
-- term2 = Ann const' (Fun (Fun (tpar "b") (tpar "b"))
--                         (Fun (tpar "a")
--                              (Fun (tpar "b") (tpar "b"))))
--         :@: id' :@: par "y"
-- env1 = [(Const "y", HasType (tpar "a")),
--         (Const "a", HasKind Star)]
-- env2 = [(Const "b", HasKind Star)] ++ env1

test = do
    let id' = (Lam (Inf (Var 0))) `Ann` Inf (Pi (Inf (par "a")) (Inf (par "a")))
    print id'
    print "hello"
    let result = (infer0 [(Const "a", VStar)] id') >>= pure . quote0
    print result
    --print (check 0 [] id' (Pi (Par (Const "a")) (Par (Const "a"))))