move Meta and Assign into a separate type

2 files changed, 96 insertions, 75 deletions

diff --git a/src/LambdaCube/Compiler/Infer.hs b/src/LambdaCube/Compiler/Infer.hs
index 12b5a680..eee5f6e6 100644
--- a/src/LambdaCube/Compiler/Infer.hs
+++ b/src/LambdaCube/Compiler/Infer.hs
@@ -53,8 +53,16 @@ data Exp
     | Neut MaxDB Neutral
     | Label LabelKind Exp{-folded expression-} Exp{-unfolded expression-}
     | LabelEnd_ LEKind Exp
-    | Meta Exp Exp
-    | Assign !Int Exp Exp       -- De Bruijn index decreasing assign reservedOp, only for metavariables (non-recursive) -- TODO: remove
+  deriving (Show)
+
+-- expressions with constraints
+type MExp = CEnv Exp
+
+-- constraints env
+data CEnv a
+    = MEnd a
+    | Meta Exp (CEnv a)
+    | Assign !Int Exp (CEnv a)       -- De Bruijn index decreasing assign reservedOp, only for metavariables (non-recursive)
   deriving (Show)
 
 data Neutral
@@ -266,8 +274,6 @@ instance Eq Exp where
     ELit l == ELit l' = l == l'
     App a b == App a' b' = (a, b) == (a', b')
     Var a == Var a' = a == a'
-    -- Meta b c == Meta b' c' = (b, c) == (b', c')
-    -- Assign a b c == Assign a' b' c' = (a, b, c) == (a', b', c')
     _ == _ = False
 
 newtype MaxDB = MaxDB {getMaxDB{-, getMaxDB' -} :: Maybe Int}
@@ -308,8 +314,6 @@ maxDB_ = \case
     TType -> mempty
     ELit _ -> mempty
     LabelEnd_ _ x -> maxDB_ x
-    Meta a b -> maxDB_ a <> lowerDB (maxDB_ b)
-    Assign j x a -> lowerDB' j $ maxDB_ x <> maxDB_ a
 
 closedExp = \case
     Lam_ _ a b c -> Lam_ closed a b c
@@ -319,8 +323,8 @@ closedExp = \case
     Neut _ a -> Neut closed a
     e -> e
 
-assign :: (Int -> Exp -> Exp -> a) -> (Int -> Exp -> Exp -> a) -> Int -> Exp -> Exp -> a
-assign _ clet i (Var j) b | i > j = clet j (Var (i-1)) $ substE "assign" j (Var (i-1)) $ up1E i b
+assign :: (Int -> Exp -> MExp -> a) -> (Int -> Exp -> MExp -> a) -> Int -> Exp -> MExp -> a
+assign _ clet i (Var j) b | i > j = clet j (Var (i-1)) $ substE2 "assign" j (Var (i-1)) $ up1EM i b
 assign clet _ i a b = clet i a b
 
 handleLet i j f
@@ -342,8 +346,6 @@ foldE f i = \case
     ELit _ -> mempty
     App a b -> foldE f i a <> foldE f i b
     LabelEnd_ _ x -> foldE f i x
-    Meta a b -> foldE f i a <> foldE f (i+1) b
-    Assign j x a -> handleLet i j $ \i' j' -> foldE f i' x <> foldE f i' a
 
 freeE x | isClosed x = mempty
 freeE x = foldE (\i k -> Set.fromList [k - i | k >= i]) 0 x
@@ -367,10 +369,17 @@ up1E i e = case e of
     App a b -> App (up1E i a) (up1E i b)
     Label lk x y -> Label lk (up1E i x) $ up1E i y
     LabelEnd_ k x -> LabelEnd_ k $ up1E i x
-    Meta a b -> Meta (up1E i a) (up1E (i+1) b)
-    Assign j a b -> handleLet i j $ \i' j' -> join assign Assign j' (up1E i' a) (up1E i' b)
+
+bothUpE :: Int -> Int -> (MExp, Exp) -> (MExp, Exp)
+bothUpE i j (e, et) = (upEM i j e, upE i j et)
+
+up1EM i = \case
+    MEnd a -> MEnd $ up1E i a
+    Meta a b -> Meta (up1E i a) (up1EM (i+1) b)
+    Assign j a b -> handleLet i j $ \i' j' -> join assign Assign j' (up1E i' a) (up1EM i' b)
 
 upE i n = iterateN n (up1E i)
+upEM i n = iterateN n (up1EM i)
 
 substE err = substE_ (error $ "substE: todo: environment required in " ++ err)  -- todo: remove
 
@@ -390,13 +399,23 @@ substE_ te i x e = case e of
     ELit l -> ELit l
     App a b  -> app_ (substE_ te i x a) (substE_ te i x b)  -- todo: precise env?
     LabelEnd_ k a -> LabelEnd_ k $ substE_ te i x a
-    Meta a b  -> Meta (substE_ te i x a) (substE_ te (i+1) (up1E 0 x) b)
+
+bothSubstE err = bothSubstE_ (error $ "bothSubstE: todo: environment required in " ++ err)  -- todo: remove
+substE2 err = substE2_ (error $ "bothSubstE: todo: environment required in " ++ err)  -- todo: remove
+
+bothSubstE_ :: Env -> Int -> Exp -> (MExp, Exp) -> (MExp, Exp)
+bothSubstE_ env i x (e, et) = (substE2_ env i x e, substE_ env i x et)
+
+substE2_ :: Env -> Int -> Exp -> MExp -> MExp
+substE2_ te i x = \case
+    MEnd a -> MEnd $ substE_ te i x a
+    Meta a b  -> Meta (substE_ te i x a) (substE2_ te (i+1) (up1E 0 x) b)
     Assign j a b
-        | j > i, Just a' <- downE i a       -> join assign Assign (j-1) a' (substE "sa" i (substE "sa" (j-1) a' x) b)
-        | j > i, Just x' <- downE (j-1) x   -> join assign Assign (j-1) (substE "sa" i x' a) (substE "sa" i x' b)
-        | j < i, Just a' <- downE (i-1) a   -> join assign Assign j a' (substE "sa" (i-1) (substE "sa" j a' x) b)
-        | j < i, Just x' <- downE j x       -> join assign Assign j (substE "sa" (i-1) x' a) (substE "sa" (i-1) x' b)
-        | j == i    -> Meta (cstr x a) $ up1E 0 b
+        | j > i, Just a' <- downE i a       -> join assign Assign (j-1) a' (substE2 "sa" i (substE "sa" (j-1) a' x) b)
+        | j > i, Just x' <- downE (j-1) x   -> join assign Assign (j-1) (substE "sa" i x' a) (substE2 "sa" i x' b)
+        | j < i, Just a' <- downE (i-1) a   -> join assign Assign j a' (substE2 "sa" (i-1) (substE "sa" j a' x) b)
+        | j < i, Just x' <- downE j x       -> join assign Assign j (substE "sa" (i-1) x' a) (substE2 "sa" (i-1) x' b)
+        | j == i    -> Meta (cstr x a) $ up1EM 0 b
 
 downE t x | usedE t x = Nothing
           | otherwise = Just $ substE_ (error "impossible") t (error "impossible: downE") x
@@ -701,8 +720,6 @@ expType_ msg te = \case
     TType -> TType
     ELit l -> litType l
     LabelEnd_ k x -> expType te x
-    Meta{} -> error "meta type"
-    Assign{} -> error "let type"
   where
     expType = expType_ msg
     app (Pi _ a b) x = substE "expType_" 0 x b
@@ -728,8 +745,6 @@ take' e n xs = case splitAt n xs of
     (as, []) -> as
     (as, _) -> as ++ [e]
 -}
-both f = f *** f
-
 showSI :: Env -> SI -> String
 showSI e = showSI_ (fst $ extractEnv e)
 
@@ -805,57 +820,54 @@ inferN tracelevel lp = infer  where
 
     focus_ :: Env -> ExpType -> TCM m ExpType'
     focus_ env eet@(e, et) = (if tracelevel >= 1 then trace_ $ "focus: " ++ showEnvExp env e else id) $ (if debug then fmap ({-first $ -}recheck' "focus" env) else id) $ case env of
-        ELabelEnd te | nm e -> focus_ te (LabelEnd e, et)
-        CheckSame x te | nm e -> focus_ (EBind2_ (debugSI "focus_ CheckSame") BMeta (cstr x e) te) (up1E 0 e, up1E 0 et)
+        ELabelEnd te -> focus_ te (LabelEnd e, et)
+        CheckSame x te -> focus_ (EBind2_ (debugSI "focus_ CheckSame") BMeta (cstr x e) te) (up1E 0 e, up1E 0 et)
         CheckAppType si h t te b   -- App1 h (CheckType t te) b
-            | nm e, Pi h' x (downE 0 -> Just y) <- et, h == h' -> case t of
+            | Pi h' x (downE 0 -> Just y) <- et, h == h' -> case t of
                 Pi Hidden t1 t2 | h == Visible -> focus_ (EApp1 si h (CheckType_ (sourceInfo b) t te) b) eet  -- <<e>> b : {t1} -> {t2}
                 _ -> focus_ (EBind2_ (sourceInfo b) BMeta (cstr t y) $ EApp1 si h te b) (up1E 0 e, up1E 0 et)
-            | nm e, otherwise -> focus_ (EApp1 si h (CheckType_ (sourceInfo b) t te) b) eet
+            | otherwise -> focus_ (EApp1 si h (CheckType_ (sourceInfo b) t te) b) eet
         EApp1 si h te b
-            | nm e, Pi h' x y <- et, h == h' -> checkN (EApp2 si h e te) b x
-            | nm e, Pi Hidden x y  <- et, h == Visible -> focus_ (EApp1 mempty Hidden env $ Wildcard $ Wildcard SType) eet  --  e b --> e _ b
+            | Pi h' x y <- et, h == h' -> checkN (EApp2 si h e te) b x
+            | Pi Hidden x y  <- et, h == Visible -> focus_ (EApp1 mempty Hidden env $ Wildcard $ Wildcard SType) eet  --  e b --> e _ b
 --            | CheckType (Pi Hidden _ _) te' <- te -> error "ok"
 --            | CheckAppType Hidden _ te' _ <- te -> error "ok"
-            | nm e, otherwise -> infer (CheckType_ (sourceInfo b) (Var 2) $ cstr' h (upE 0 2 et) (Pi Visible (Var 1) (Var 1)) (upE 0 2 e) $ EBind2_ (sourceInfo b) BMeta TType $ EBind2_ (sourceInfo b) BMeta TType te) (upS__ 0 3 b)
+            | otherwise -> infer (CheckType_ (sourceInfo b) (Var 2) $ cstr' h (upE 0 2 et) (Pi Visible (Var 1) (Var 1)) (upE 0 2 e) $ EBind2_ (sourceInfo b) BMeta TType $ EBind2_ (sourceInfo b) BMeta TType te) (upS__ 0 3 b)
           where
             cstr' h x y e = EApp2 mempty h (eval (error "cstr'") $ Coe (up1E 0 x) (up1E 0 y) (Var 0) (up1E 0 e)) . EBind2_ (sourceInfo b) BMeta (cstr x y)
-        ELet2 le (x{-let-}, xt) te | nm e -> focus_ te (substE "app2" 0 (mkELet le x xt){-let-} e{-in-}, et)
-        CheckIType x te | nm e -> checkN te x e
+        ELet2 le (x{-let-}, xt) te -> focus_ te (substE "app2" 0 (mkELet le x xt){-let-} e{-in-}, et)
+        CheckIType x te -> checkN te x e
         CheckType_ si t te
-            | nm e, hArgs et > hArgs t
+            | hArgs et > hArgs t
                             -> focus_ (EApp1 mempty Hidden (CheckType_ si t te) $ Wildcard $ Wildcard SType) eet
-            | nm e, hArgs et < hArgs t, Pi Hidden t1 t2 <- t
+            | hArgs et < hArgs t, Pi Hidden t1 t2 <- t
                             -> focus_ (CheckType_ si t2 $ EBind2 (BLam Hidden) t1 te) eet
-            | nm e, otherwise    -> focus_ (EBind2_ si BMeta (cstr t et) te) (up1E 0 e, up1E 0 et)
-        EApp2 si h a te | nm e    -> focus' te si $ app_ a e        --  h??
-        EBind1 si h te b | nm e    -> infer (EBind2_ (sourceInfo b) h e te) b
-        EBind2_ si (BLam h) a te | nm e -> focus_ te (Lam h a e, Pi h a e)
-        EBind2_ si (BPi h) a te | nm e -> focus_' te si (Pi h a e, TType)
-        _ -> focus2 env eet
-      where
-        nm (Meta _ _) = False
-        nm (Assign _ _ _) = False
-        nm _ = True
-
+            | otherwise    -> focus_ (EBind2_ si BMeta (cstr t et) te) (up1E 0 e, up1E 0 et)
+        EApp2 si h a te    -> focus' te si $ app_ a e        --  h??
+        EBind1 si h te b   -> infer (EBind2_ (sourceInfo b) h e te) b
+        EBind2_ si (BLam h) a te -> focus_ te (Lam h a e, Pi h a e)
+        EBind2_ si (BPi h) a te -> focus_' te si (Pi h a e, TType)
+        _ -> focus2 env (MEnd e, et)
+
+    focus2 :: Env -> (MExp, Type) -> TCM m ExpType'
     focus2 env eet = case env of
-        ELet1 le te b{-in-} -> replaceMetas "let" (Lam Hidden) (fst eet) >>= \e' -> infer (ELet2 le (addType_ te e'{-let-}) te) b{-in-}
+        ELet1 le te b{-in-} -> infer (ELet2 le (addType_ te $ replaceMetas (Lam Hidden) (fst eet){-let-}) te) b{-in-}
         EBind2_ si BMeta tt te
-            | Unit <- tt    -> refocus te $ both (substE_ te 0 TT) eet
+            | Unit <- tt    -> refocus te $ bothSubstE_ te 0 TT eet
             | Empty msg <- tt   -> throwError $ "type error: " ++ msg ++ "\nin " ++ showSI te si ++ "\n"-- todo: better error msg
             | T2 x y <- tt, let te' = EBind2_ si BMeta (up1E 0 y) $ EBind2_ si BMeta x te
-                            -> refocus te' $ both (substE_ te' 2 (t2C (Var 1) (Var 0)) . upE 0 2) eet
-            | CstrT t a b <- tt, a == b  -> refocus te $ both (substE "inferN2" 0 TT) eet
+                            -> refocus te' $ bothSubstE_ te' 2 (t2C (Var 1) (Var 0)) $ bothUpE 0 2 eet
+            | CstrT t a b <- tt, a == b  -> refocus te $ bothSubstE "inferN2" 0 TT eet
             | CstrT t a b <- tt, Just r <- cst a b -> r
             | CstrT t a b <- tt, Just r <- cst b a -> r
             | isCstr tt, EBind2 h x te' <- te{-, h /= BMeta todo: remove-}, Just x' <- downE 0 tt, x == x'
-                            -> refocus te $ both (substE "inferN3" 1 (Var 0)) eet
+                            -> refocus te $ bothSubstE "inferN3" 1 (Var 0) eet
             | EBind2 h x te' <- te, h /= BMeta, Just b' <- downE 0 tt
-                            -> refocus (EBind2_ si h (up1E 0 x) $ EBind2_ si BMeta b' te') $ both (substE "inferN3" 2 (Var 0) . up1E 0) eet
+                            -> refocus (EBind2_ si h (up1E 0 x) $ EBind2_ si BMeta b' te') $ bothSubstE "inferN3" 2 (Var 0) $ bothUpE 0 1 eet
             | ELet2 le (x, xt) te' <- te, Just b' <- downE 0 tt
-                            -> refocus (ELet2 le (up1E 0 x, up1E 0 xt) $ EBind2_ si BMeta b' te') $ both (substE "inferN32" 2 (Var 0) . up1E 0) eet
+                            -> refocus (ELet2 le (up1E 0 x, up1E 0 xt) $ EBind2_ si BMeta b' te') $ bothSubstE "inferN32" 2 (Var 0) $ bothUpE 0 1 eet
             | EBind1 si h te' x <- te -> refocus (EBind1 si h (EBind2_ si BMeta tt te') $ upS__ 1 1 x) eet
-            | ELet1 le te' x     <- te, floatLetMeta $ expType_ "3" env $ Lam Hidden tt $ replaceMetas' (Lam Hidden) (fst eet)
+            | ELet1 le te' x     <- te, floatLetMeta $ expType_ "3" env $ Lam Hidden tt $ replaceMetas (Lam Hidden) (fst eet)
                                     -> refocus (ELet1 le (EBind2_ si BMeta tt te') $ upS__ 1 1 x) eet
             | CheckAppType si h t te' x <- te -> refocus (CheckAppType si h (up1E 0 t) (EBind2_ si BMeta tt te') $ upS x) eet
             | EApp1 si h te' x <- te -> refocus (EApp1 si h (EBind2_ si BMeta tt te') $ upS x) eet
@@ -865,13 +877,14 @@ inferN tracelevel lp = infer  where
             | ELabelEnd te'   <- te -> refocus (ELabelEnd $ EBind2_ si BMeta tt te') eet
             | otherwise             -> focus2 te (Meta tt $ fst eet, snd eet {-???-})
           where
-            refocus = refocus_ focus_
-            refocus'' te = refocus te . addType_ te
+            refocus = refocus_ focus2
+            refocus'' te = refocus te . addType'' te
             cst x = \case
                 Var i | fst (varType "X" i te) == BMeta
                       , Just y <- downE i x
-                      -> Just $ join assign (\i x -> refocus'' $ EAssign i x te) i y $ fst $ both (substE_ te 0 {-ReflCstr y-}TT . substE_ te (i+1) (up1E 0 y)) eet
+                      -> Just $ join assign (\i x -> refocus'' $ EAssign i x te) i y $ fst $ bothSubstE_ te 0 {-ReflCstr y-}TT $ bothSubstE_ te (i+1) (up1E 0 y) eet
                 _ -> Nothing
+
         EAssign i b te -> case te of
             EBind2_ si h x te' | i > 0, Just b' <- downE 0 b
                               -> refocus' (EBind2_ si h (substE "inferN5" (i-1) b' x) (EAssign (i-1) b' te')) eet
@@ -890,32 +903,31 @@ inferN tracelevel lp = infer  where
             -- todo: CheckSame Exp Env
           where
             refocus' = refocus_ refocus'
-            focus2' te = focus2 te . addType_ te
             assign'  te i x (e, t) = assign (\i x e -> focus2' te (Assign i x e)) (foc te) i x e
-            foc te i x = refocus' te' . addType_ te' where te' = EAssign i x te
+            foc te i x = refocus' te' . addType'' te' where te' = EAssign i x te
             pull i = \case
                 EBind2 BMeta _ te | i == 0 -> Just te
                 EBind2_ si h x te   -> EBind2_ si h <$> downE (i-1) x <*> pull (i-1) te
                 EAssign j b te  -> EAssign (if j <= i then j else j-1) <$> downE i b <*> pull (if j <= i then i+1 else i) te
                 _               -> Nothing
-        EGlobal{} -> replaceMetas "inferN" lp $ fst eet --return (e, et)
-        _ | nm (fst eet) -> throwError_ $ "focus impossible: " ++ show env
-        _ -> throwError_ $ "focus checkMetas: " ++ ppShow (fst eet)
-      where
-        refocus_ f e (Meta x a, t) = f (EBind2 BMeta x e) (a, t)
-        refocus_ _ e (Assign i x a, t) = focus (EAssign i x e) a
-        refocus_ _ e (a, t) = focus e a
 
-        replaceMetas err bind = return . replaceMetas' bind
+        EGlobal{} -> return $ replaceMetas lp $ fst eet --return (e, et)
+        _ -> case fst eet of
+            MEnd x -> throwError_ $ "focus todo: " ++ ppShow x
+            _ -> throwError_ $ "focus checkMetas: " ++ ppShow (fst eet)
+      where
+        addType'' te x = (x, error "addType''")
+        focus2' te = focus2 te . addType'' te
 
-        replaceMetas' bind = \case
-            Meta a t -> bind a $ replaceMetas' bind t
-            Assign i x t -> bind (cstr (Var i) $ upE i 1 x) . up1E 0 . upE i 1 $ replaceMetas' bind t
-            t ->  t
+        refocus_ :: (Env -> (MExp, Type) -> TCM m ExpType') -> Env -> (MExp, Type) -> TCM m ExpType'
+        refocus_ f e (Meta x a, t) = f (EBind2 BMeta x e) (a, t)
+        refocus_ _ e (Assign i x a, t) = focus2' (EAssign i x e) a
+        refocus_ _ e (MEnd a, t) = focus e a
 
-        nm (Meta _ _) = False
-        nm (Assign _ _ _) = False
-        nm _ = True
+        replaceMetas bind = \case
+            Meta a t -> bind a $ replaceMetas bind t
+            Assign i x t -> bind (cstr (Var i) $ upE i 1 x) . up1E 0 . upE i 1 $ replaceMetas bind t
+            MEnd t ->  t
 
 
 isCstr CstrT{} = True
@@ -1249,6 +1261,9 @@ joinPolyEnvs _ = return . foldr mappend' mempty'           -- todo
 instance PShow Exp where
     pShowPrec _ = showDoc_ . expDoc
 
+instance PShow MExp where
+    pShowPrec _ = showDoc_ . mExpDoc
+
 instance PShow Env where
     pShowPrec _ e = showDoc_ $ envDoc e $ pure $ shAtom $ underlined "<<HERE>>"
 
@@ -1333,11 +1348,17 @@ expDoc_ ts e = fmap inGreen <$> f e
         TType           -> pure $ shAtom "Type"
         ELit l          -> pure $ shAtom $ show l
         LabelEnd_ k x   -> shApp Visible (shAtom $ "labend" ++ show k) <$> f x
-        Meta a b        -> join $ shLam (usedE 0 b) BMeta <$> f a <*> pure (f b)
-        Assign i x e    -> shLet i (f x) (f e)
 
     shAtom_ = shAtom . if ts then switchTick else id
 
+mExpDoc e = fmap inGreen <$> f e
+  where
+    f :: MExp -> Doc
+    f = \case
+        MEnd a          -> expDoc a
+        Meta a b        -> join $ shLam True BMeta <$> expDoc a <*> pure (f b)
+        Assign i x e    -> shLet i (expDoc x) (f e)
+
 -------------------------------------------------------------------------------- main
 
 smartTrace :: MonadError String m => Extensions -> (Bool -> m a) -> m a
diff --git a/testdata/typesig.reject.out b/testdata/typesig.reject.out
index da3c6644..58881830 100644
--- a/testdata/typesig.reject.out
+++ b/testdata/typesig.reject.out
@@ -1 +1 @@
-focus checkMetas: \(a : V0~'X)->X
\ No newline at end of file
+focus checkMetas: \(a : V0~'X)->X
\ No newline at end of file