speedup: tweak max de bruijn handling

3 files changed, 45 insertions, 30 deletions

diff --git a/src/LambdaCube/Compiler/Infer.hs b/src/LambdaCube/Compiler/Infer.hs
index 9afa8a36..13e18411 100644
--- a/src/LambdaCube/Compiler/Infer.hs
+++ b/src/LambdaCube/Compiler/Infer.hs
@@ -107,7 +107,12 @@ pattern NoLE <- (isNoLabelEnd -> True)
 isNoLabelEnd (LabelEnd_ _) = False
 isNoLabelEnd _ = True
 
-pattern Fun' f vs i xs n <- Fun_ _ f vs i xs n where Fun' f vs i xs n = Fun_ (foldMap maxDB_ vs <> foldMap maxDB_ xs {- <> iterateN i lowerDB (maxDB_ n)-}) f vs i xs n
+mconcat1 [] = mempty
+mconcat1 [x] = x
+mconcat1 xs = notClosed --foldl1 (<>) xs
+mconcat1' = mconcat
+
+pattern Fun' f vs i xs n <- Fun_ _ f vs i xs n where Fun' f vs i xs n = Fun_ (mconcat1 $ map maxDB_ vs ++ map maxDB_ xs {- <> iterateN i lowerDB (maxDB_ n)-}) f vs i xs n
 pattern Fun f i xs n = Fun' f [] i xs n
 pattern UTFun a t b <- (unfixlabel -> Neut (Fun (FunName a _ t) _ (reverse -> b) NoLE))
 pattern UFunN a b <- UTFun a _ b
@@ -117,7 +122,7 @@ pattern DFun_ fn xs <- Fun fn 0 (reverse -> xs) (Delta _) where
 pattern TFun' a t b = DFun_ (FunName a Nothing t) b
 pattern TFun a t b = Neut (TFun' a t b)
 
-pattern CaseFun_ a b c <- CaseFun__ _ a b c where CaseFun_ a b c = CaseFun__ (foldMap maxDB_ b <> maxDB_ c) a b c
+pattern CaseFun_ a b c <- CaseFun__ _ a b c where CaseFun_ a b c = CaseFun__ (maxDB_ c <> foldMap maxDB_ b) a b c
 pattern TyCaseFun_ a b c <- TyCaseFun__ _ a b c where TyCaseFun_ a b c = TyCaseFun__ (foldMap maxDB_ b <> maxDB_ c) a b c
 pattern App_ a b <- App__ _ a b where App_ a b = App__ (maxDB_ a <> maxDB_ b) a b
 pattern CaseFun a b c = Neut (CaseFun_ a b c)
@@ -141,7 +146,7 @@ makeCaseFunPars' te n = case neutType' te n of
 pattern Closed :: () => Up a => a -> a
 pattern Closed a <- a where Closed a = closedExp a
 
-pattern Con x n y <- Con_ _ x n y where Con x n y = Con_ (foldMap maxDB_ y) x n y
+pattern Con x n y <- Con_ _ x n y where Con x n y = Con_ (mconcat1 $ map maxDB_ y) x n y
 pattern ConN s a  <- Con (ConName s _ _) _ a
 pattern ConN' s a  <- Con (ConName _ s _) _ a
 tCon s i t a = Con (ConName s i t) 0 a
@@ -371,7 +376,7 @@ instance Up Neutral where
             CaseFun__ md s as ne -> CaseFun__ (upDB n md) s (up_ n i <$> as) (up_ n i ne)
             TyCaseFun__ md s as ne -> TyCaseFun__ (upDB n md) s (up_ n i <$> as) (up_ n i ne)
             App__ md a b -> App__ (upDB n md) (up_ n i a) (up_ n i b)
-            Fun' fn vs c x y -> Fun' fn (up_ n i <$> vs) c (up_ n i <$> x) $ up_ n (i + c) y
+            Fun_ md fn vs c x y -> Fun_ (upDB n md) fn (up_ n i <$> vs) c (up_ n i <$> x) $ up_ n (i + c) y
             LabelEnd_ x -> LabelEnd_ $ up_ n i x
             d@Delta{} -> d
 
@@ -755,6 +760,8 @@ inferN e s = do
     f ITrace{} = False
     f _ = True
 
+substTo i x = subst i x . up1_ (i+1)
+
 inferN_ :: forall m . Monad m => (forall a . String -> String -> IM m a -> IM m a) -> Env -> SExp2 -> IM m ExpType'
 inferN_ tellTrace = infer  where
 
@@ -779,12 +786,12 @@ inferN_ tellTrace = infer  where
         | x@(SGlobal (si, MatchName n)) `SAppV` SLamV (Wildcard _) `SAppV` a `SAppV` SVar siv v `SAppV` b <- e
             = infer te $ x `SAppV` SLam Visible SType (STyped mempty (subst (v+1) (Var 0) $ up 1 t, TType)) `SAppV` a `SAppV` SVar siv v `SAppV` b
             -- temporal hack
-        | x@(SGlobal (si, "'NatCase")) `SAppV` SLamV (Wildcard _) `SAppV` a `SAppV` b `SAppV` SVar siv v <- e
-            = infer te $ x `SAppV` SLamV (STyped mempty (subst (v+1) (Var 0) $ up1_ (v+2) $ up 1 t, TType)) `SAppV` a `SAppV` b `SAppV` SVar siv v
+        | x@(SGlobal (si, CaseName "'Nat")) `SAppV` SLamV (Wildcard _) `SAppV` a `SAppV` b `SAppV` SVar siv v <- e
+            = infer te $ x `SAppV` SLamV (STyped mempty (substTo (v+1) (Var 0) $ up 1 t, TType)) `SAppV` a `SAppV` b `SAppV` SVar siv v
             -- temporal hack
-        | x@(SGlobal (si, "'VecSCase")) `SAppV` SLamV (SLamV (Wildcard _)) `SAppV` a `SAppV` b `SAppV` c `SAppV` SVar siv v <- e
-        , TVec (Var n') _ <- snd $ varType "xx" v te
-            = infer te $ x `SAppV` SLamV (SLamV (STyped mempty (subst (n'+2) (Var 1) $ up1_ (n'+3) $ up 2 t, TType))) `SAppV` a `SAppV` b `SAppV` c `SAppV` SVar siv v
+        | x@(SGlobal (si, CaseName "'VecS")) `SAppV` SLamV (SLamV (Wildcard _)) `SAppV` a `SAppV` b `SAppV` c `SAppV` SVar siv v <- e
+        , TyConN "'VecS" [_, Var n'] <- snd $ varType "xx" v te
+            = infer te $ x `SAppV` SLamV (SLamV (STyped mempty (substTo (n'+2) (Var 1) $ up 2 t, TType))) `SAppV` a `SAppV` b `SAppV` c `SAppV` SVar siv v
 
 {-
             -- temporal hack
@@ -986,18 +993,18 @@ recheck' msg' e (x, xt) = (recheck_ "main" (checkEnv e) (x, xt), xt)
 
     recheck_ msg te = \case
         (Var k, zt) -> Var k    -- todo: check var type
-        (Lam b, Pi h a bt) -> Lam $ recheck_ "9" (EBind2 (BLam h) a te) (b, bt)
-        (Pi h a b, TType) -> Pi h (checkType te a) $ checkType (EBind2 (BPi h) a te) b
+        (Lam_ md b, Pi h a bt) -> Lam_ md $ recheck_ "9" (EBind2 (BLam h) a te) (b, bt)
+        (Pi_ md h a b, TType) -> Pi_ md h (checkType te a) $ checkType (EBind2 (BPi h) a te) b
         (ELit l, zt) -> ELit l  -- todo: check literal type
         (TType, TType) -> TType
-        (Neut (App_ a b), zt)
+        (Neut (App__ md a b), zt)
             | (Neut a', at) <- recheck'' "app1" te (Neut a, neutType te a)
-            -> checkApps "a" [] zt (Neut . App_ a' . head) te at [b]
-        (Con s n as, zt)      -> checkApps (show s) [] zt (Con s n . drop (conParams s)) te (conType zt s) $ mkConPars n zt ++ as
-        (TyCon s as, zt)      -> checkApps (show s) [] zt (TyCon s) te (nType s) as
+            -> checkApps "a" [] zt (Neut . App__ md a' . head) te at [b]
+        (Con_ md s n as, zt)      -> checkApps (show s) [] zt (Con_ md s n . drop (conParams s)) te (conType zt s) $ mkConPars n zt ++ as
+        (TyCon_ md s as, zt)      -> checkApps (show s) [] zt (TyCon_ md s) te (nType s) as
         (CaseFun s@(CaseFunName _ t pars) as n, zt) -> checkApps (show s) [] zt (\xs -> evalCaseFun s (init $ drop pars xs) (last xs)) te (nType s) (makeCaseFunPars te n ++ as ++ [Neut n])
         (TyCaseFun s [m, t, f] n, zt)  -> checkApps (show s) [] zt (\[m, t, n, f] -> evalTyCaseFun s [m, t, f] n) te (nType s) [m, t, Neut n, f]
-        (Neut (Fun' f vs@[] i a x), zt) -> checkApps "lab" [] zt (\xs -> Neut $ Fun' f vs i (reverse xs) x) te (nType f) $ reverse a   -- TODO: recheck x
+        (Neut (Fun_ md f vs@[] i a x), zt) -> checkApps "lab" [] zt (\xs -> Neut $ Fun_ md f vs i (reverse xs) x) te (nType f) $ reverse a   -- TODO: recheck x
         -- TODO
         (r@(Neut (Fun' f vs i a x)), zt) -> r
         (LabelEnd x, zt) -> LabelEnd $ recheck_ msg te (x, zt)
diff --git a/src/LambdaCube/Compiler/Lexer.hs b/src/LambdaCube/Compiler/Lexer.hs
index 631bbef9..83446799 100644
--- a/src/LambdaCube/Compiler/Lexer.hs
+++ b/src/LambdaCube/Compiler/Lexer.hs
@@ -87,7 +87,13 @@ instance Show Lit where
 
 type SName = String
 
+pattern CaseName :: String -> String
+pattern CaseName cs <- (getCaseName -> Just cs) where CaseName = caseName
+
 caseName (c:cs) = toLower c: cs ++ "Case"
+getCaseName cs = case splitAt 4 $ reverse cs of
+    (reverse -> "Case", xs) -> Just $ reverse xs
+    _ -> Nothing
 
 pattern MatchName cs <- (getMatchName -> Just cs) where MatchName = matchName
 
diff --git a/src/LambdaCube/Compiler/Parser.hs b/src/LambdaCube/Compiler/Parser.hs
index e4d482a8..e64327de 100644
--- a/src/LambdaCube/Compiler/Parser.hs
+++ b/src/LambdaCube/Compiler/Parser.hs
@@ -15,7 +15,7 @@ module LambdaCube.Compiler.Parser
     , pattern SVar, pattern SType, pattern Wildcard, pattern SAppV, pattern SLamV, pattern SAnn
     , pattern SBuiltin, pattern SPi, pattern Primitive, pattern SLabelEnd, pattern SLam, pattern Parens
     , pattern TyType, pattern Wildcard_
-    , debug, isPi, varDB, lowerDB, justDB, upDB, cmpDB, MaxDB (..), iterateN, traceD
+    , debug, isPi, varDB, lowerDB, upDB, notClosed, cmpDB, MaxDB(..), iterateN, traceD
     , parseLC, runDefParser
     , getParamsS, addParamsS, getApps, apps', downToS, addForalls
     , mkDesugarInfo, joinDesugarInfo
@@ -175,28 +175,30 @@ instance SetSourceInfo (SExp' a) where
         SGlobal (_, n)  -> SGlobal (si, n)
         SLit _ l        -> SLit si l
 
--------------------------------------------------------------------------------- low-level toolbox
+-------------------------------------------------------------------------------- De-Bruijn limit
 
-newtype MaxDB = MaxDB {getMaxDB :: Maybe Int}
+newtype MaxDB = MaxDB {getMaxDB :: Bool} -- True: closed
 
 instance Monoid MaxDB where
-    mempty = MaxDB Nothing
-    MaxDB a  `mappend` MaxDB a'  = MaxDB $ Just $ max (fromMaybe 0 a) (fromMaybe 0 a')
+    mempty = MaxDB True
+    MaxDB a  `mappend` MaxDB a'  = MaxDB $ a && a' -- $ Just $ max (fromMaybe 0 a) (fromMaybe 0 a')
 
 instance Show MaxDB where show _ = "MaxDB"
 
-varDB i = MaxDB $ Just $ i + 1
+varDB i = MaxDB False
 
-lowerDB (MaxDB i) = MaxDB $ (+ (- 1)) <$> i
-justDB (MaxDB i) = MaxDB $ Just $ fromMaybe 0 i
+lowerDB = id
 
 -- 0 means that no free variable is used
 -- 1 means that only var 0 is used
 --cmpDB i e = i >= maxDB e
-cmpDB _ (maxDB_ -> MaxDB x) = isNothing x
+cmpDB _ (maxDB_ -> MaxDB x) = x --isNothing x
+
+upDB n = id --(MaxDB i) = MaxDB $ (\x -> if x == 0 then x else x+n) $ i
 
-maxDB = max 0 . fromMaybe 0 . getMaxDB . maxDB_
-upDB n (MaxDB i) = MaxDB $ (\x -> if x == 0 then x else x+n) <$> i
+notClosed = MaxDB False
+
+-------------------------------------------------------------------------------- low-level toolbox
 
 class Up a where
     up_ :: Int -> Int -> a -> a
@@ -681,7 +683,7 @@ compileGuardTree ulend lend adts t = (\x -> traceD ("  !  :" ++ ppShow x) x) $ g
         ts@(GuardNode f s _ _: _) -> case Map.lookup s (snd adts) of
             Nothing -> error $ "Constructor is not defined: " ++ s
             Just (Left ((casename, inum), cns)) ->
-                foldl SAppV (SGlobal (debugSI "compileGuardTree2", casename) `SAppV` iterateN (1 + inum) SLamV (Wildcard SType))
+                foldl SAppV (SGlobal (debugSI "compileGuardTree2", casename) `SAppV` iterateN (1 + inum) SLamV (Wildcard (Wildcard SType)))
                     [ iterateN n SLamV $ guardTreeToCases $ Alts $ map (filterGuardTree (up n f) cn 0 n . upGT 0 n) ts
                     | (cn, n) <- cns
                     ]
@@ -1002,8 +1004,8 @@ mkDesugarInfo ss =
   where
     pars ty = length $ filter ((== Visible) . fst) $ fst $ getParamsS ty -- todo
 
-    hackHList ("HCons", _) = ("HCons", Left (("hlistConsCase", 0), [("HCons", 2)]))
-    hackHList ("HNil", _) = ("HNil", Left (("hlistNilCase", 0), [("HNil", 0)]))
+    hackHList ("HCons", _) = ("HCons", Left (("hlistConsCase", -1), [("HCons", 2)]))
+    hackHList ("HNil", _) = ("HNil", Left (("hlistNilCase", -1), [("HNil", 0)]))
     hackHList x = x
 
 joinDesugarInfo (fm, cm) (fm', cm') = (Map.union fm fm', Map.union cm cm')