Merge pull request #160 from kmyk/assert

Use assertions in core and C++

examples/dp_z-kubaru.py

@@ -4,7 +4,7 @@
 INF = 10 ** 18
 
 def solve(n: int, c: int, h: List[int]) -> int:
-    assert 2 <= n <= 10 ** 5
+    assert 2 <= n <= 2 * 10 ** 5
     assert 1 <= c <= 10 ** 12
     assert len(h) == n
     assert all(1 <= h_i <= 10 ** 6 for h_i in h)

examples/dp_z-morau.py

@@ -4,7 +4,7 @@
 INF = 10 ** 18
 
 def solve(n: int, c: int, h: List[int]) -> int:
-    assert 2 <= n <= 10 ** 5
+    assert 2 <= n <= 2 * 10 ** 5
     assert 1 <= c <= 10 ** 12
     assert len(h) == n
     assert all(1 <= h_i <= 10 ** 6 for h_i in h)

src/Jikka/CPlusPlus/Convert/FromCore.hs

@@ -481,14 +481,14 @@ runAppBuiltin env f ts args = wrapError' ("converting builtin " ++ X.formatBuilt
     X.All -> go01' $ \xs -> do
       y <- Y.newFreshName Y.LocalNameKind
       return
-        ( [ Y.Declare Y.TyBool y (Y.DeclareCopy (Y.BinOp Y.Equal (Y.callFunction "std::find" [] [Y.begin xs, Y.end xs, Y.Lit (Y.LitBool True)]) (Y.end xs)))
+        ( [ Y.Declare Y.TyBool y (Y.DeclareCopy (Y.BinOp Y.Equal (Y.callFunction "std::find" [] [Y.begin xs, Y.end xs, Y.Lit (Y.LitBool False)]) (Y.end xs)))
           ],
           Y.Var y
         )
     X.Any -> go01' $ \xs -> do
       y <- Y.newFreshName Y.LocalNameKind
       return
-        ( [ Y.Declare Y.TyBool y (Y.DeclareCopy (Y.BinOp Y.NotEqual (Y.callFunction "std::find" [] [Y.begin xs, Y.end xs, Y.Lit (Y.LitBool False)]) (Y.end xs)))
+        ( [ Y.Declare Y.TyBool y (Y.DeclareCopy (Y.BinOp Y.NotEqual (Y.callFunction "std::find" [] [Y.begin xs, Y.end xs, Y.Lit (Y.LitBool True)]) (Y.end xs)))
           ],
           Y.Var y
         )
@@ -637,6 +637,10 @@ runExpr env = \case
     (stmts1, e1) <- runExpr env e1
     (stmts2, e2) <- runExpr ((x, t, y) : env) e2
     return (stmts1 ++ Y.Declare t' y (Y.DeclareCopy e1) : stmts2, e2)
+  X.Assert e1 e2 -> do
+    (stmts1, e1) <- runExpr env e1
+    (stmts2, e2) <- runExpr env e2
+    return (stmts1 ++ Y.Assert e1 : stmts2, e2)
 
 runToplevelFunDef :: (MonadAlpha m, MonadError Error m) => Env -> Y.VarName -> [(X.VarName, X.Type)] -> X.Type -> X.Expr -> m [Y.ToplevelStatement]
 runToplevelFunDef env f args ret body = do
@@ -713,6 +717,11 @@ runToplevelExpr env = \case
     stmt <- runToplevelFunDef ((f, t, g) : env) g args ret body
     cont <- runToplevelExpr ((f, t, g) : env) cont
     return $ stmt ++ cont
+  X.ToplevelAssert e cont -> do
+    (stmts, e) <- runExpr env e
+    let stmt = Y.StaticAssert (Y.CallExpr (Y.Lam [] Y.TyBool (stmts ++ [Y.Return e])) []) ""
+    cont <- runToplevelExpr env cont
+    return $ stmt : cont
 
 runProgram :: (MonadAlpha m, MonadError Error m) => X.Program -> m Y.Program
 runProgram prog = Y.Program <$> runToplevelExpr [] prog

src/Jikka/CPlusPlus/Convert/MoveSemantics.hs

@@ -152,6 +152,7 @@ runToplevelStatement :: MonadState (M.Map VarName VarName) m => ToplevelStatemen
 runToplevelStatement = \case
   VarDef t x e -> VarDef t x <$> runExpr e
   FunDef ret f args body -> FunDef ret f args <$> runStatements body []
+  StaticAssert e msg -> StaticAssert <$> runExpr e <*> pure msg
 
 runProgram :: Monad m => Program -> m Program
 runProgram (Program decls) = (`evalStateT` M.empty) $ do

src/Jikka/CPlusPlus/Convert/UnpackTuples.hs

@@ -203,6 +203,7 @@ runToplevelStatement :: (MonadAlpha m, MonadError Error m, MonadState (M.Map Var
 runToplevelStatement = \case
   VarDef t x e -> VarDef t x <$> runExpr e
   FunDef ret f args body -> FunDef ret f args <$> runStatements body []
+  StaticAssert e msg -> StaticAssert <$> runExpr e <*> pure msg
 
 runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program
 runProgram (Program decls) = (`evalStateT` M.empty) $ do

src/Jikka/CPlusPlus/Format.hs

@@ -320,13 +320,16 @@ formatToplevelStatement = \case
         args' = intercalate ", " $ map (\(t, x) -> formatType t ++ " " ++ unVarName x) args
         body' = concatMap formatStatement body
      in [ret' ++ " " ++ unVarName f ++ "(" ++ args' ++ ") {"] ++ body' ++ ["}"]
+  StaticAssert e msg ->
+    ["static_assert (" ++ resolvePrec CommaPrec (formatExpr e) ++ ", " ++ formatLiteral (LitString msg) ++ ");"]
 
 formatProgram :: Program -> [Code]
 formatProgram prog =
   let body = concatMap formatToplevelStatement (decls prog)
       standardHeaders =
         [ "#include <algorithm>",
           "#include <array>",
+          "#include <cassert>",
           "#include <cstdint>",
           "#include <functional>",
           "#include <iostream>",

src/Jikka/CPlusPlus/Language/Expr.hs

@@ -249,6 +249,8 @@ data ToplevelStatement
     VarDef Type VarName Expr
   | -- | @T f(T1 x1, T2 x2, ...) { stmt1; stmt2; ... }@
     FunDef Type VarName [(Type, VarName)] [Statement]
+  | -- | @static_assert(e, msg);@
+    StaticAssert Expr String
   deriving (Eq, Ord, Show, Read)
 
 newtype Program = Program

src/Jikka/CPlusPlus/Language/Util.hs

@@ -198,6 +198,7 @@ mapExprStatementToplevelStatementM :: Monad m => (Expr -> m Expr) -> (Statement
 mapExprStatementToplevelStatementM f g = \case
   VarDef t x e -> VarDef t x <$> mapExprStatementExprM f g e
   FunDef ret h args body -> FunDef ret h args <$> mapM (mapExprStatementStatementM f g) body
+  StaticAssert e msg -> StaticAssert <$> mapExprStatementExprM f g e <*> pure msg
 
 mapExprStatementProgramM :: Monad m => (Expr -> m Expr) -> (Statement -> m Statement) -> Program -> m Program
 mapExprStatementProgramM f g (Program decls) = Program <$> mapM (mapExprStatementToplevelStatementM f g) decls

src/Jikka/Core/Convert/ANormal.hs

@@ -22,7 +22,7 @@ import Jikka.Core.Language.Lint
 import Jikka.Core.Language.TypeCheck
 import Jikka.Core.Language.Util
 
-destruct :: (MonadAlpha m, MonadError Error m) => TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr)
+destruct :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> m ([(VarName, Type)], Expr -> Expr, Expr)
 destruct env = \case
   e@Var {} -> return (env, id, e)
   e@Lit {} -> return (env, id, e)
@@ -38,8 +38,12 @@ destruct env = \case
     (env, ctx, e1) <- destruct env e1
     (env, ctx', e2) <- destruct ((x, t) : env) e2
     return (env, ctx . Let x t e1 . ctx', e2)
+  Assert e1 e2 -> do
+    (env, ctx, e1) <- destruct env e1
+    (env, ctx', e2) <- destruct env e2
+    return (env, ctx . Assert e1 . ctx', e2)
 
-runApp :: (MonadAlpha m, MonadError Error m) => TypeEnv -> Expr -> [Expr] -> m Expr
+runApp :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> [Expr] -> m Expr
 runApp env f args = go env id args
   where
     go :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> ([Expr] -> [Expr]) -> [Expr] -> m Expr
@@ -51,7 +55,7 @@ runApp env f args = go env id args
       e <- go env (acc . (arg :)) args
       return $ ctx e
 
-runExpr :: (MonadAlpha m, MonadError Error m) => TypeEnv -> Expr -> m Expr
+runExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> m Expr
 runExpr env = \case
   Var x -> return $ Var x
   Lit lit -> return $ Lit lit
@@ -70,19 +74,24 @@ runExpr env = \case
     (env, ctx, e1) <- destruct env e1
     e2 <- runExpr ((x, t) : env) e2
     return $ ctx (Let x t e1 e2)
+  Assert e1 e2 -> do
+    e1 <- runExpr env e1
+    (env, ctx, e1) <- destruct env e1
+    e2 <- runExpr env e2
+    return $ ctx (Assert e1 e2)
 
-runToplevelExpr :: (MonadAlpha m, MonadError Error m) => TypeEnv -> ToplevelExpr -> m ToplevelExpr
-runToplevelExpr env = \case
-  ResultExpr e -> ResultExpr <$> runExpr env e
-  ToplevelLet x t e cont -> do
-    e <- runExpr env e
-    cont <- runToplevelExpr ((x, t) : env) cont
-    return $ ToplevelLet x t e cont
-  ToplevelLetRec f args ret body cont -> do
-    let t = curryFunTy (map snd args) ret
-    body <- runExpr (reverse args ++ (f, t) : env) body
-    cont <- runToplevelExpr ((f, t) : env) cont
-    return $ ToplevelLetRec f args ret body cont
+-- | TODO: convert `ToplevelExpr` too
+runProgram :: (MonadAlpha m, MonadError Error m) => ToplevelExpr -> m ToplevelExpr
+runProgram = mapToplevelExprProgramM go
+  where
+    go env = \case
+      ResultExpr e -> ResultExpr <$> runExpr env e
+      ToplevelLet x t e cont -> ToplevelLet x t <$> runExpr env e <*> pure cont
+      ToplevelLetRec f args ret body cont -> do
+        let t = curryFunTy (map snd args) ret
+        let env' = reverse args ++ (f, t) : env
+        ToplevelLetRec f args ret <$> runExpr env' body <*> pure cont
+      ToplevelAssert e cont -> ToplevelAssert <$> runExpr env e <*> pure cont
 
 -- | `run` makes a given program A-normal form.
 -- A program is an A-normal form iff assigned exprs of all let-statements are values or function applications.
@@ -99,6 +108,6 @@ runToplevelExpr env = \case
 run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
 run prog = wrapError' "Jikka.Core.Convert.ANormal" $ do
   prog <- Alpha.runProgram prog
-  prog <- runToplevelExpr [] prog
+  prog <- runProgram prog
   ensureWellTyped prog
   return prog

src/Jikka/Core/Convert/Alpha.hs

@@ -38,6 +38,7 @@ runExpr env = \case
     y <- rename x
     e2 <- runExpr ((x, y) : env) e2
     return $ Let y t e1 e2
+  Assert e1 e2 -> Assert <$> runExpr env e1 <*> runExpr env e2
 
 runToplevelExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, VarName)] -> ToplevelExpr -> m ToplevelExpr
 runToplevelExpr env = \case
@@ -57,6 +58,7 @@ runToplevelExpr env = \case
     body <- runExpr (args1 ++ (f, g) : env) body
     cont <- runToplevelExpr ((f, g) : env) cont
     return $ ToplevelLetRec g args2 ret body cont
+  ToplevelAssert e1 e2 -> ToplevelAssert <$> runExpr env e1 <*> runToplevelExpr env e2
 
 runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program
 runProgram = runToplevelExpr []

src/Jikka/Core/Convert/ConstantPropagation.hs

@@ -35,6 +35,7 @@ runExpr env = \case
      in if isConstantTimeExpr e1'
           then runExpr (M.insert x e1' env) e2
           else Let x t e1' (runExpr env e2)
+  Assert e1 e2 -> Assert (runExpr env e1) (runExpr env e2)
 
 runToplevelExpr :: Env -> ToplevelExpr -> ToplevelExpr
 runToplevelExpr env = \case
@@ -46,6 +47,7 @@ runToplevelExpr env = \case
           else ToplevelLet x t e' (runToplevelExpr env cont)
   ToplevelLetRec f args ret body cont ->
     ToplevelLetRec f args ret (runExpr env body) (runToplevelExpr env cont)
+  ToplevelAssert e1 e2 -> ToplevelAssert (runExpr env e1) (runToplevelExpr env e2)
 
 run' :: Program -> Program
 run' = runToplevelExpr M.empty

src/Jikka/Core/Convert/KubaruToMorau.hs

@@ -50,6 +50,7 @@ runFunctionBody c i j step y x k = do
         Lam x t e
           | x == c || x == i || x == j -> throwRuntimeError "name confliction found"
           | otherwise -> Lam x t <$> go e
+        Assert e1 e2 -> Assert <$> go e1 <*> go e2
   go step
 
 -- | TODO: remove the assumption that the length of @a@ is equals to @n@

src/Jikka/Core/Convert/MakeScanl.hs

@@ -74,6 +74,7 @@ getRecurrenceFormulaStep1 shift t a i body = do
         App f e -> App <$> go f <*> go e
         Lam x t e -> Lam x t <$> if x == a then Just e else go e
         Let x t e1 e2 -> Let x t <$> go e1 <*> if x == a then Just e2 else go e2
+        Assert f e -> Assert <$> go f <*> go e
   return $ case go body of
     Just body -> Just $ Lam2 x t i IntTy body
     Nothing -> Nothing
@@ -97,6 +98,7 @@ getRecurrenceFormulaStep shift size t a i body = do
         App f e -> App <$> go f <*> go e
         Lam x t e -> Lam x t <$> if x == a then Just e else go e
         Let x t e1 e2 -> Let x t <$> go e1 <*> if x == a then Just e2 else go e2
+        Assert f e -> Assert <$> go f <*> go e
   return $ case go body of
     Just body -> Just $ Lam2 x (TupleTy ts) i IntTy (uncurryApp (Tuple' ts) (map (\i -> Proj' ts i (Var x)) [1 .. size - 1] ++ [body]))
     Nothing -> Nothing
@@ -148,6 +150,7 @@ checkAccumulationFormulaStep a i = go
       App f e -> go f && go e
       Lam x _ e -> x == a || go e
       Let x _ e1 e2 -> go e1 && (x == a || go e2)
+      Assert e1 e2 -> go e1 && go e2
 
 -- |
 -- * This assumes that `Range2` and `Range3` are already converted to `Range1` (`Jikka.Core.Convert.ShortCutFusion`).
@@ -196,6 +199,7 @@ checkGenericRecurrenceFormulaStep a = \i k -> go (M.fromList [(i, k - 1)])
       App f e -> go env f && go env e
       Lam x _ e -> x == a || go env e
       Let x _ e1 e2 -> go env e1 && (x == a || go env e2)
+      Assert e1 e2 -> go env e1 && go env e2
 
 reduceFoldlSetAtGeneric :: MonadAlpha m => RewriteRule m
 reduceFoldlSetAtGeneric = RewriteRule $ \_ -> \case

src/Jikka/Core/Convert/RemoveUnusedVars.hs

@@ -23,32 +23,24 @@ import Jikka.Core.Language.FreeVars (isUnusedVar)
 import Jikka.Core.Language.Lint
 import Jikka.Core.Language.Util
 
-runLet :: VarName -> Type -> Expr -> Expr -> Expr
-runLet x t e1 e2
-  | isUnusedVar x e2 = e2
-  | otherwise = Let x t e1 e2
+runExpr :: [(VarName, Type)] -> Expr -> Expr
+runExpr _ = mapExpr go []
+  where
+    go _ = \case
+      Let x _ _ e2 | x `isUnusedVar` e2 -> e2
+      e -> e
 
-runExpr :: Expr -> Expr
-runExpr = \case
-  Var x -> Var x
-  Lit lit -> Lit lit
-  App f e -> App (runExpr f) (runExpr e)
-  Lam x t e -> Lam x t (runExpr e)
-  Let x t e1 e2 -> runLet x t (runExpr e1) (runExpr e2)
-
-runToplevelExpr :: ToplevelExpr -> ToplevelExpr
-runToplevelExpr = \case
-  ResultExpr e -> ResultExpr $ runExpr e
-  ToplevelLet x t e cont -> ToplevelLet x t (runExpr e) (runToplevelExpr cont)
+-- | TODO: Remove `ToplevelLet` if its variable is not used.
+runToplevelExpr :: [(VarName, Type)] -> ToplevelExpr -> ToplevelExpr
+runToplevelExpr _ = \case
   ToplevelLetRec f args ret body cont ->
-    let body' = runExpr body
-        cont' = runToplevelExpr cont
-     in if isUnusedVar f body'
-          then ToplevelLet f (curryFunTy (map snd args) ret) (curryLam args body') cont'
-          else ToplevelLetRec f args ret body' cont'
+    if isUnusedVar f body
+      then ToplevelLet f (curryFunTy (map snd args) ret) (curryLam args body) cont
+      else ToplevelLetRec f args ret body cont
+  e -> e
 
 run' :: Program -> Program
-run' = runToplevelExpr
+run' = mapToplevelExprProgram runToplevelExpr . mapExprProgram runExpr
 
 -- | `run` removes unused variables in given programs.
 --

src/Jikka/Core/Convert/SegmentTree.hs

@@ -113,6 +113,7 @@ replaceWithSegtrees a segtrees = go M.empty
          in case check env e1' of
               Just (e1', b, semigrp) -> go (M.insert x (e1', b, semigrp) env) e2
               Nothing -> Let x t (go env e1) (go env e2)
+      Assert e1 e2 -> Assert (go env e1) (go env e2)
     check :: M.Map VarName (Expr, Expr, Semigroup') -> Expr -> Maybe (Expr, Expr, Semigroup')
     check env = \case
       Var x -> M.lookup x env

src/Jikka/Core/Convert/TrivialLetElimination.hs

@@ -35,12 +35,14 @@ isEliminatable x = \case
   App f e -> isEliminatable x f `plus` isEliminatable x e
   Lam y _ e -> if x == y then Nothing else isEliminatable x e $> False -- moving an expr into a lambda may increase the time complexity
   Let y _ e1 e2 -> isEliminatable x e1 `plus` (if x == y then Nothing else isEliminatable x e2)
+  Assert e1 e2 -> isEliminatable x e1 `plus` isEliminatable x e2
 
 isEliminatableToplevelExpr :: VarName -> ToplevelExpr -> Maybe Bool
 isEliminatableToplevelExpr x = \case
   ResultExpr e -> isEliminatable x e
   ToplevelLet y _ e cont -> isEliminatable x e `plus` (if x == y then Nothing else isEliminatableToplevelExpr x cont)
   ToplevelLetRec f args _ body cont -> if x == f then Nothing else isEliminatableToplevelExpr x cont `plus` (if x `elem` map fst args then Nothing else isEliminatable x body)
+  ToplevelAssert e cont -> isEliminatable x e `plus` isEliminatableToplevelExpr x cont
 
 runExpr :: M.Map VarName Expr -> Expr -> Expr
 runExpr env = \case
@@ -53,6 +55,7 @@ runExpr env = \case
      in if isEliminatable x e2 /= Just False
           then runExpr (M.insert x e1' env) e2
           else Let x t e1' (runExpr env e2)
+  Assert e1 e2 -> Assert (runExpr env e1) (runExpr env e2)
 
 runToplevelExpr :: M.Map VarName Expr -> ToplevelExpr -> ToplevelExpr
 runToplevelExpr env = \case
@@ -64,6 +67,8 @@ runToplevelExpr env = \case
           else ToplevelLet x t e' (runToplevelExpr env cont)
   ToplevelLetRec f args ret body cont ->
     ToplevelLetRec f args ret (runExpr env body) (runToplevelExpr env cont)
+  ToplevelAssert e cont ->
+    ToplevelAssert (runExpr env e) (runToplevelExpr env cont)
 
 run' :: Program -> Program
 run' = runToplevelExpr M.empty

src/Jikka/Core/Convert/TypeInfer.hs

@@ -73,6 +73,9 @@ formularizeExpr = \case
     formularizeVarName x t
     formularizeExpr' e1 t
     formularizeExpr e2
+  Assert e1 e2 -> do
+    formularizeExpr' e1 BoolTy
+    formularizeExpr e2
 
 formularizeExpr' :: (MonadWriter Eqns m, MonadAlpha m, MonadError Error m) => Expr -> Type -> m ()
 formularizeExpr' e t = do
@@ -91,6 +94,9 @@ formularizeToplevelExpr = \case
     mapM_ (uncurry formularizeVarName) args
     formularizeExpr' body ret
     formularizeToplevelExpr cont
+  ToplevelAssert e cont -> do
+    formularizeExpr' e BoolTy
+    formularizeToplevelExpr cont
 
 formularizeProgram :: (MonadAlpha m, MonadError Error m) => Program -> m [Equation]
 formularizeProgram prog = getDual <$> execWriterT (formularizeToplevelExpr prog)

src/Jikka/Core/Evaluate.hs

@@ -29,7 +29,7 @@ import qualified Data.Vector as V
 import Jikka.Common.Alpha
 import Jikka.Common.Error
 import Jikka.Common.Matrix
-import Jikka.Core.Format (formatBuiltinIsolated)
+import Jikka.Core.Format (formatBuiltinIsolated, formatExpr)
 import Jikka.Core.Language.BuiltinPatterns
 import Jikka.Core.Language.Expr
 import Jikka.Core.Language.Lint
@@ -292,6 +292,11 @@ evaluateExpr env = \case
   Let x _ e1 e2 -> do
     v1 <- evaluateExpr env e1
     evaluateExpr ((x, v1) : env) e2
+  Assert e1 e2 -> do
+    p <- valueToBool =<< evaluateExpr env e1
+    if p
+      then evaluateExpr env e2
+      else throwRuntimeError $ "assertion failed: " ++ formatExpr e1
 
 callToplevelExpr :: (MonadFix m, MonadError Error m) => Env -> ToplevelExpr -> [Value] -> m Value
 callToplevelExpr env e args = case e of
@@ -301,6 +306,11 @@ callToplevelExpr env e args = case e of
   ToplevelLetRec f args' _ body cont -> do
     val <- mfix $ \val -> evaluateExpr ((f, val) : env) (curryLam args' body)
     callToplevelExpr ((f, val) : env) cont args
+  ToplevelAssert e cont -> do
+    p <- valueToBool =<< evaluateExpr env e
+    if p
+      then callToplevelExpr env cont args
+      else throwRuntimeError $ "toplevel assertion failed: " ++ formatExpr e
   ResultExpr e -> do
     val <- evaluateExpr env e
     callValue val args