1  module Core where
  2  
  3  import Control.Exception (assert)
  4  import Control.Monad (unless)
  5  
  6  type Idx = Int
  7  
  8  type Lvl = Int
  9  
 10  type Stage = Int
 11  
 12  type Name = String
 13  
 14  data Raw
 15    = RVar Name
 16    | RLam Name Raw
 17    | RLet Stage Name Raw Raw Raw
 18    | RApp Raw Raw
 19    | RPi Name Raw Raw
 20    | RU Stage
 21    | RQuote Raw
 22    | RSplice Raw
 23    | RLift Raw
 24  
 25  data Term
 26    = Var Idx
 27    | Lam Term -- \idx. t
 28    | Let Stage Term Term Term -- let s idx: A = t in y
 29    | App Term Term
 30    | Pi Term Term -- (idx: A) -> B
 31    | U Stage
 32    | Quote Term
 33    | Splice Term
 34    | Lift Term
 35    deriving (Eq, Show)
 36  
 37  data Value
 38    = VRigid Lvl FrameStack
 39    | VLam Closure
 40    | VPi Value Closure
 41    | VU Stage
 42    | VQuote Value
 43    | VLift Value
 44    deriving Show
 45  
 46  data FrameStack
 47    = SNil
 48    | SApp FrameStack Value
 49    | SSplice FrameStack
 50    deriving Show
 51  
 52  type Env = [Value]
 53  data Closure = Closure Env Term deriving Show
 54  
 55  pattern VVar :: Lvl -> Value
 56  pattern VVar x = VRigid x SNil
 57  
 58  ($$) :: Closure -> Value -> Value
 59  (Closure env t) $$ b = eval (b : env) t
 60  
 61  -- Evaluation
 62  
 63  app :: Value -> Value -> Value
 64  app (VLam f) u = f $$ u
 65  app (VRigid l stack) u = VRigid l (SApp stack u)
 66  app _ _ = error "bad application"
 67  
 68  eval :: Env -> Term -> Value
 69  eval env (Var x) = env !! x
 70  eval env (Lam t) = VLam $ Closure env t
 71  eval env (Let _a_stage _a_ty a b) = eval (eval env a : env) b
 72  eval env (App t u) = app (eval env t) (eval env u)
 73  eval env (Pi a b) = VPi (eval env a) (Closure env b)
 74  eval _env (U s) = VU s
 75  eval env (Quote t) = case eval env t of
 76    VRigid x (SSplice stack) -> VRigid x stack
 77    t' -> VQuote t'
 78  eval env (Splice t) = case eval env t of
 79    VQuote t' -> t'
 80    VRigid x stack -> VRigid x (SSplice stack)
 81    _ -> error "bad splice"
 82  eval env (Lift t) = VLift (eval env t)
 83  
 84  -- | idx + lvl + 1 = lvl_max <==> lvl_max - lvl - 1 = idx
 85  lvl2idx :: Lvl -> Lvl -> Idx
 86  lvl2idx lvl_max lvl = lvl_max - lvl - 1
 87  
 88  reify :: Lvl -> Value -> Term
 89  reify lvl (VRigid l sp) = reify_stack (Var $ lvl2idx lvl l) sp
 90    where
 91      reify_stack :: Term -> FrameStack -> Term
 92      reify_stack t SNil = t
 93      reify_stack t (SApp stack u) = App (reify_stack t stack) (reify lvl u)
 94      reify_stack t (SSplice stack) = Splice (reify_stack t stack)
 95  reify lvl (VLam f) = Lam (reify (lvl + 1) (f $$ VVar lvl))
 96  reify lvl (VPi a b) = Pi (reify lvl a) (reify (lvl + 1) (b $$ VVar lvl))
 97  reify _lvl (VU s) = U s
 98  reify lvl (VLift t) = Lift (reify lvl t)
 99  reify lvl (VQuote t) = Quote (reify lvl t)
100  
101  data Ctx = Ctx
102    { env :: [Value]
103    , tyenv :: [(Value, Stage)]
104    }
105  
106  nf :: Env -> Term -> Term
107  nf env = reify (length env) . eval env
108  
109  conv :: Lvl -> Value -> Value -> Bool
110  conv lvl t u = reify lvl t == reify lvl u
111  
112  -- Typechecking
113  
114  check :: Ctx -> Term -> Value -> Stage -> Either String ()
115  check ctx (Lam t) (VPi a b) stage = do
116    let ctx' = ctx{tyenv = (a, stage) : ctx.tyenv, env = VVar (length ctx.env) : ctx.env}
117    check ctx' t (b $$ VVar (length ctx.env)) stage
118  check ctx (Lift ty) (VU 1) _stage = check ctx ty (VU 0) 0
119  check ctx (Quote t) (VLift ty) _stage = check ctx t ty 0
120  check ctx (Let a_stage a_ty a b) ty stage = do
121    check ctx a_ty (VU a_stage) stage
122    let a_ty' = eval ctx.env a_ty
123    check ctx a a_ty' a_stage
124    let ctx' = ctx{env = eval ctx.env a : ctx.env, tyenv = (a_ty', a_stage) : ctx.tyenv}
125    check ctx' b ty stage
126  check ctx t ty stage = do
127    (ty', stage') <- infer ctx t
128    let lvl = length ctx.env
129    unless (conv lvl ty ty' && stage == stage') $
130      Left "Inferred type different from expected type"
131  
132  infer :: Ctx -> Term -> Either String (Value, Stage)
133  infer ctx (Var x) = pure $ ctx.tyenv !! x
134  infer ctx (App t u) = do
135    (t_ty, stage) <- infer ctx t
136    case t_ty of
137      VPi a b -> do
138        check ctx u a stage -- check function input
139        pure (b $$ eval ctx.env u, stage)
140      _ -> Left "Expected function type in application head"
141  infer _ctx (U s) = pure (VU s, s)
142  infer ctx (Pi a b) = do
143    (a_ty, stage) <- infer ctx a
144    let ctx' = ctx{tyenv = (a_ty, stage) : ctx.tyenv, env = eval ctx.env a : ctx.env}
145    check ctx' b (VU stage) stage
146    pure (VU stage, stage)
147  infer ctx (Lift t) = do check ctx t (VU 0) 0; pure (VU 1, 1)
148  infer ctx (Splice t) = do
149    (t_ty, stage) <- infer ctx t
150    unless (stage == 1) $ Left "Not meta language term in splice."
151    case t_ty of
152      VLift t_ty' -> pure (t_ty', 0)
153      _ -> Left "Bad Splice"
154  infer ctx (Quote t) = do
155    (t_ty, stage) <- infer ctx t
156    unless (stage == 0) $ Left "Not object language term in quote "
157    pure (VLift t_ty, 1)
158  infer ctx (Let a_stage a_ty a b) = do
159    check ctx a_ty (VU a_stage) a_stage
160    let a_ty' = eval ctx.env a_ty
161    check ctx a a_ty' a_stage
162    let ctx' = ctx{env = eval ctx.env a : ctx.env, tyenv = (a_ty', a_stage) : ctx.tyenv}
163    infer ctx' b
164  infer _ _ = Left "Couldn't infer type"
165  
166  -- Staging
167  
168  stage_term :: Term -> Term
169  stage_term = reify0 0 . eval0 []
170  
171  type StagingEnv = [Either Value0 Value1]
172  
173  data Value0
174    = VVar0 Idx
175    | VLam0 (Value0 -> Value0)
176    | VPi0 Value0 (Value0 -> Value0)
177    | VU0
178    | VApp0 Value0 Value0
179    | VLet0 Value0 Value0 (Value0 -> Value0)
180  
181  data Value1
182    = VLam1 (Value1 -> Value1)
183    | VQuote1 Value0
184    | -- vvv these apparently can be ignored
185      VPi1 Value1 (Value1 -> Value1)
186    | VU1
187    | VLift1 Value0
188  
189  eval1 :: StagingEnv -> Term -> Value1
190  eval1 env (Var x) = case env !! x of
191    Left _ -> error "Value0 in Value1 context"
192    Right v -> v
193  eval1 env (Lam t) = VLam1 (\u -> eval1 (Right u : env) t)
194  eval1 env (App t u) = case eval1 env t of
195    VLam1 t' -> t' $ eval1 env u
196    _ -> error "bad application during staging"
197  eval1 env (Pi a b) = VPi1 (eval1 env a) (\a' -> eval1 (Right a' : env) b)
198  eval1 _env (U s)
199    | s == 1 = VU1
200    | s == 0 = error "U 0 in Value1 context"
201    | otherwise = error "unreachable"
202  eval1 env (Quote t) = VQuote1 $ eval0 env t
203  eval1 env (Let a_stage _a_ty a b) = assert (a_stage == 1) $ eval1 (Right (eval1 env a) : env) b
204  eval1 _env (Splice _t) = error "splice in Value1 context"
205  eval1 env (Lift t) = VLift1 $ eval0 env t
206  
207  eval0 :: StagingEnv -> Term -> Value0
208  -- why does this actually do a lookup instead of constructing a VVar0?
209  eval0 env (Var x) = case env !! x of
210    Left v -> v
211    Right _ -> error "Value1 in Value0 context"
212  eval0 env (Lam t) = VLam0 $ \u -> eval0 (Left u : env) t
213  eval0 env (App t u) = VApp0 (eval0 env t) (eval0 env u)
214  eval0 env (Pi a b) = VPi0 (eval0 env a) (\a' -> eval0 (Left a' : env) b)
215  eval0 _env (U s)
216    | s == 1 = error "U 1 in Value0 context"
217    | s == 0 = VU0
218    | otherwise = error "unreachable"
219  eval0 _env (Quote _) = error "Quote in Value0 context"
220  eval0 env (Let a_stage a_ty a b) =
221    assert (a_stage == 0) $
222      VLet0 (eval0 env a_ty) (eval0 env a) (\a' -> eval0 (Left a' : env) b)
223  eval0 env (Splice t) = case eval1 env t of
224    VQuote1 v -> v
225    _ -> error "Splice didn't return quote"
226  eval0 _env (Lift _) = error "Lifted type in Value0 context"
227  
228  reify0 :: Lvl -> Value0 -> Term
229  reify0 lvl (VVar0 x) = Var $ lvl2idx lvl x
230  reify0 lvl (VLam0 t) = Lam $ reify0 (lvl + 1) (t $ VVar0 lvl)
231  reify0 lvl (VPi0 a b) = Pi (reify0 lvl a) $ reify0 (lvl + 1) (b $ VVar0 lvl)
232  reify0 _lvl VU0 = U 0
233  reify0 lvl (VApp0 t u) = App (reify0 lvl t) (reify0 lvl u)
234  reify0 lvl (VLet0 a_ty a b) = Let 0 (reify0 lvl a_ty) (reify0 lvl a) (reify0 (lvl + 1) (b $ VVar0 lvl))