checked integral arithmeric

5 files changed, 220 insertions, 23 deletions

diff --git a/ydb/docs/ru/core/yql/reference/yql-core/syntax/_includes/pragma/global.md b/ydb/docs/ru/core/yql/reference/yql-core/syntax/_includes/pragma/global.md
index e5e0a0079e8..3ff58a43836 100644
--- a/ydb/docs/ru/core/yql/reference/yql-core/syntax/_includes/pragma/global.md
+++ b/ydb/docs/ru/core/yql/reference/yql-core/syntax/_includes/pragma/global.md
@@ -185,6 +185,17 @@ StrictJoinKeyTypes является [scoped](#pragmascope) настройкой.
 Если отключить — результат всегда становится Double.
 ClassicDivision является [scoped](#pragmascope) настройкой.
 
+### CheckedOps
+
+| Тип значения | По умолчанию |
+| --- | --- |
+| Флаг | false |
+
+При включенном режиме если в результате выполнения бинарных операций `+`,`-`,`*`,`/`,`%` или унарной операции `-` над целыми числами происходит выход за границы целевого типа аргументов или результата, то возвращается `NULL`.
+Если отключить - переполнение не проверяется. 
+Не влияет на операции с числами с  плавающей точкой или `Decimal`.
+CheckedOps является [scoped](#pragmascope) настройкой.
+
 ### AllowDotInAlias
 
 | Тип значения | По умолчанию |
diff --git a/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp b/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
index aadfc855c62..30cbbc2a3bd 100644
--- a/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
+++ b/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
@@ -6111,6 +6111,95 @@ TExprNode::TPtr DropToFlowDeps(const TExprNode::TPtr& node, TExprContext& ctx) {
     return ctx.ChangeChildren(*node, std::move(children));
 }
 
+TExprNode::TPtr BuildCheckedBinaryOpOverDecimal(TPositionHandle pos, TStringBuf op, const TExprNode::TPtr& lhs, const TExprNode::TPtr& rhs, const TTypeAnnotationNode& resultType, TExprContext& ctx) {

+    auto typeNode = ExpandType(pos, resultType, ctx);

+    return ctx.Builder(pos)

+        .Callable("SafeCast")

+            .Callable(0, op)

+                .Callable(0, "SafeCast")

+                    .Callable(0, "SafeCast")

+                        .Add(0, lhs)

+                        .Add(1, typeNode)

+                    .Seal()

+                    .Callable(1, "DataType")

+                        .Atom(0, "Decimal")

+                        .Atom(1, "20")

+                        .Atom(2, "0")

+                    .Seal()

+                .Seal()

+                .Callable(1, "SafeCast")

+                    .Callable(0, "SafeCast")

+                        .Add(0, rhs)

+                        .Add(1, typeNode)

+                    .Seal()

+                    .Callable(1, "DataType")

+                        .Atom(0, "Decimal")

+                        .Atom(1, "20")

+                        .Atom(2, "0")

+                    .Seal()

+                .Seal()

+            .Seal()

+            .Add(1, typeNode)

+        .Seal()

+        .Build();

+}

+

+TExprNode::TPtr BuildCheckedBinaryOpOverSafeCast(TPositionHandle pos, TStringBuf op, const TExprNode::TPtr& lhs, const TExprNode::TPtr& rhs, const TTypeAnnotationNode& resultType, TExprContext& ctx) {

+    auto typeNode = ExpandType(pos, resultType, ctx);

+    return ctx.Builder(pos)

+        .Callable(op)

+            .Callable(0, "SafeCast")

+                .Add(0, lhs)

+                .Add(1, typeNode)

+            .Seal()

+            .Callable(1, "SafeCast")

+                .Add(0, rhs)

+                .Add(1, typeNode)

+            .Seal()

+        .Seal()

+        .Build();

+}

+
+TExprNode::TPtr ExpandCheckedAdd(const TExprNode::TPtr& node, TExprContext& ctx) {
+    return BuildCheckedBinaryOpOverDecimal(node->Pos(), "+", node->ChildPtr(0), node->ChildPtr(1), *node->GetTypeAnn(), ctx);
+}
+
+TExprNode::TPtr ExpandCheckedSub(const TExprNode::TPtr& node, TExprContext& ctx) {
+    return BuildCheckedBinaryOpOverDecimal(node->Pos(), "-", node->ChildPtr(0), node->ChildPtr(1), *node->GetTypeAnn(), ctx);
+}
+
+TExprNode::TPtr ExpandCheckedMul(const TExprNode::TPtr& node, TExprContext& ctx) {
+    return BuildCheckedBinaryOpOverDecimal(node->Pos(), "*", node->ChildPtr(0), node->ChildPtr(1), *node->GetTypeAnn(), ctx);
+}
+
+TExprNode::TPtr ExpandCheckedDiv(const TExprNode::TPtr& node, TExprContext& ctx) {
+    return BuildCheckedBinaryOpOverSafeCast(node->Pos(), "/", node->ChildPtr(0), node->ChildPtr(1), *node->GetTypeAnn(), ctx);
+}
+
+TExprNode::TPtr ExpandCheckedMod(const TExprNode::TPtr& node, TExprContext& ctx) {
+    return BuildCheckedBinaryOpOverSafeCast(node->Pos(), "%", node->ChildPtr(0), node->ChildPtr(1), *node->GetTypeAnn(), ctx);
+}
+
+TExprNode::TPtr ExpandCheckedMinus(const TExprNode::TPtr& node, TExprContext& ctx) {
+    return ctx.Builder(node->Pos())

+        .Callable("SafeCast")

+            .Callable(0, "Minus")

+                .Callable(0, "SafeCast")

+                    .Add(0, node->HeadPtr())

+                    .Callable(1, "DataType")

+                        .Atom(0, "Decimal")

+                        .Atom(1, "20")

+                        .Atom(2, "0")

+                    .Seal()

+                .Seal()

+            .Seal()

+            .Callable(1, "TypeOf")

+                .Add(0, node->HeadPtr())

+            .Seal()

+        .Seal()

+        .Build();
+}
+
 ui64 ToDate(ui64 now)      { return std::min<ui64>(NUdf::MAX_DATE - 1U, now / 86400000000ull); }
 ui64 ToDatetime(ui64 now)  { return std::min<ui64>(NUdf::MAX_DATETIME - 1U, now / 1000000ull); }
 ui64 ToTimestamp(ui64 now) { return std::min<ui64>(NUdf::MAX_TIMESTAMP - 1ULL, now); }
@@ -6181,6 +6270,12 @@ struct TPeepHoleRules {
         {"AsRange", &ExpandAsRange},
         {"RangeFor", &ExpandRangeFor},
         {"ToFlow", &DropToFlowDeps},
+        {"CheckedAdd", &ExpandCheckedAdd},
+        {"CheckedSub", &ExpandCheckedSub},
+        {"CheckedMul", &ExpandCheckedMul},
+        {"CheckedDiv", &ExpandCheckedDiv},
+        {"CheckedMod", &ExpandCheckedMod},
+        {"CheckedMinus", &ExpandCheckedMinus},
     };
 
     static constexpr std::initializer_list<TExtPeepHoleOptimizerMap::value_type> CommonStageExtRulesInit = {
diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
index 4ec696f19c2..2ae4f45ece3 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -2250,9 +2250,10 @@ namespace NTypeAnnImpl {
             }
             return IGraphTransformer::TStatus::Ok;
         }
-    };
+    }
 
     IGraphTransformer::TStatus AddWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        const bool checked = input->Content().StartsWith("Checked");
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
@@ -2274,9 +2275,11 @@ namespace NTypeAnnImpl {
             haveOptional |= isOptional[i];
         }
 
-        auto check_result = CheckIntegralsWidth(input, ctx, dataType[0]->GetSlot(), dataType[1]->GetSlot());
-        if (check_result != IGraphTransformer::TStatus::Ok) {
-            return check_result;
+        if (!checked) {
+            auto check_result = CheckIntegralsWidth(input, ctx, dataType[0]->GetSlot(), dataType[1]->GetSlot());
+            if (check_result != IGraphTransformer::TStatus::Ok) {
+                return check_result;
+            }
         }
 
         const bool isLeftNumeric = IsDataTypeNumeric(dataType[0]->GetSlot());
@@ -2320,6 +2323,15 @@ namespace NTypeAnnImpl {
         }
 
         const TTypeAnnotationNode* resultType = commonType;
+        if (checked) {
+            if (IsDataTypeIntegral(dataType[0]->GetSlot()) && IsDataTypeIntegral(dataType[1]->GetSlot())) {
+                haveOptional = true;
+            } else {
+                output = ctx.Expr.RenameNode(*input, "+");
+                return IGraphTransformer::TStatus::Repeat;
+            }
+        }
+
         if (haveOptional) {
             resultType = ctx.Expr.MakeType<TOptionalExprType>(resultType);
         }
@@ -2329,6 +2341,8 @@ namespace NTypeAnnImpl {
     }
 
     IGraphTransformer::TStatus SubWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        const bool checked = input->Content().StartsWith("Checked");
+
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
@@ -2350,9 +2364,11 @@ namespace NTypeAnnImpl {
             haveOptional |= isOptional[i];
         }
 
-        auto check_result = CheckIntegralsWidth(input, ctx, dataType[0]->GetSlot(), dataType[1]->GetSlot());
-        if (check_result != IGraphTransformer::TStatus::Ok) {
-            return check_result;
+        if (!checked) {
+            auto check_result = CheckIntegralsWidth(input, ctx, dataType[0]->GetSlot(), dataType[1]->GetSlot());
+            if (check_result != IGraphTransformer::TStatus::Ok) {
+                return check_result;
+            }
         }
 
         const bool isLeftNumeric = IsDataTypeNumeric(dataType[0]->GetSlot());
@@ -2394,6 +2410,15 @@ namespace NTypeAnnImpl {
         }
 
         const TTypeAnnotationNode* resultType = commonType;
+        if (checked) {
+            if (IsDataTypeIntegral(dataType[0]->GetSlot()) && IsDataTypeIntegral(dataType[1]->GetSlot())) {
+                haveOptional = true;
+            } else {
+                output = ctx.Expr.RenameNode(*input, "-");
+                return IGraphTransformer::TStatus::Repeat;
+            }
+        }
+
         if (haveOptional) {
             resultType = ctx.Expr.MakeType<TOptionalExprType>(resultType);
         }
@@ -2403,6 +2428,8 @@ namespace NTypeAnnImpl {
     }
 
     IGraphTransformer::TStatus MulWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        const bool checked = input->Content().StartsWith("Checked");
+
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
@@ -2424,9 +2451,11 @@ namespace NTypeAnnImpl {
             haveOptional |= isOptional[i];
         }
 
-        auto check_result = CheckIntegralsWidth(input, ctx, dataType[0]->GetSlot(), dataType[1]->GetSlot());
-        if (check_result != IGraphTransformer::TStatus::Ok) {
-            return check_result;
+        if (!checked) {
+            auto check_result = CheckIntegralsWidth(input, ctx, dataType[0]->GetSlot(), dataType[1]->GetSlot());
+            if (check_result != IGraphTransformer::TStatus::Ok) {
+                return check_result;
+            }
         }
 
         if (IsDataTypeNumeric(dataType[0]->GetSlot()) && IsDataTypeNumeric(dataType[1]->GetSlot())) {
@@ -2455,6 +2484,15 @@ namespace NTypeAnnImpl {
         }
 
         const TTypeAnnotationNode* resultType = commonType;
+        if (checked) {
+            if (IsDataTypeIntegral(dataType[0]->GetSlot()) && IsDataTypeIntegral(dataType[1]->GetSlot())) {
+                haveOptional = true;
+            } else {
+                output = ctx.Expr.RenameNode(*input, "*");
+                return IGraphTransformer::TStatus::Repeat;
+            }
+        }
+
         if (haveOptional) {
             resultType = ctx.Expr.MakeType<TOptionalExprType>(resultType);
         }
@@ -2464,6 +2502,8 @@ namespace NTypeAnnImpl {
     }
 
     IGraphTransformer::TStatus DivWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        const bool checked = input->Content().StartsWith("Checked");
+
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
@@ -2485,9 +2525,11 @@ namespace NTypeAnnImpl {
             haveOptional |= isOptional[i];
         }
 
-        auto check_result = CheckIntegralsWidth(input, ctx, dataType[0]->GetSlot(), dataType[1]->GetSlot());
-        if (check_result != IGraphTransformer::TStatus::Ok) {
-            return check_result;
+        if (!checked) {
+            auto check_result = CheckIntegralsWidth(input, ctx, dataType[0]->GetSlot(), dataType[1]->GetSlot());
+            if (check_result != IGraphTransformer::TStatus::Ok) {
+                return check_result;
+            }
         }
 
         if (IsDataTypeNumeric(dataType[0]->GetSlot()) && IsDataTypeNumeric(dataType[1]->GetSlot())) {
@@ -2517,11 +2559,20 @@ namespace NTypeAnnImpl {
             resultType = ctx.Expr.MakeType<TOptionalExprType>(resultType);
         }
 
+        if (checked) {
+            if (!(IsDataTypeIntegral(dataType[0]->GetSlot()) && IsDataTypeIntegral(dataType[1]->GetSlot()))) {
+                output = ctx.Expr.RenameNode(*input, "/");
+                return IGraphTransformer::TStatus::Repeat;
+            }
+        }
+
         input->SetTypeAnn(resultType);
         return IGraphTransformer::TStatus::Ok;
     }
 
     IGraphTransformer::TStatus ModWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        const bool checked = input->Content().StartsWith("Checked");
+
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
@@ -2543,9 +2594,11 @@ namespace NTypeAnnImpl {
             haveOptional |= isOptional[i];
         }
 
-        auto check_result = CheckIntegralsWidth(input, ctx, dataType[0]->GetSlot(), dataType[1]->GetSlot());
-        if (check_result != IGraphTransformer::TStatus::Ok) {
-            return check_result;
+        if (!checked) {
+            auto check_result = CheckIntegralsWidth(input, ctx, dataType[0]->GetSlot(), dataType[1]->GetSlot());
+            if (check_result != IGraphTransformer::TStatus::Ok) {
+                return check_result;
+            }
         }
 
         if (IsDataTypeNumeric(dataType[0]->GetSlot()) && IsDataTypeNumeric(dataType[1]->GetSlot())) {
@@ -2572,6 +2625,13 @@ namespace NTypeAnnImpl {
             resultType = ctx.Expr.MakeType<TOptionalExprType>(resultType);
         }
 
+        if (checked) {
+            if (!(IsDataTypeIntegral(dataType[0]->GetSlot()) && IsDataTypeIntegral(dataType[1]->GetSlot()))) {
+                output = ctx.Expr.RenameNode(*input, "%");
+                return IGraphTransformer::TStatus::Repeat;
+            }
+        }
+
         input->SetTypeAnn(resultType);
         return IGraphTransformer::TStatus::Ok;
     }
@@ -2744,7 +2804,22 @@ namespace NTypeAnnImpl {
             return IGraphTransformer::TStatus::Error;
         }
 
-        input->SetTypeAnn(input->Head().GetTypeAnn());
+        bool haveOptional = false;
+        if (input->Content().StartsWith("Checked")) {
+            if (!IsDataTypeIntegral(dataSlot)) {
+                output = ctx.Expr.RenameNode(*input, "Minus");
+                return IGraphTransformer::TStatus::Repeat;
+            }
+
+            haveOptional = true;
+        }
+
+        if (!isOptional && haveOptional) {
+            input->SetTypeAnn(ctx.Expr.MakeType<TOptionalExprType>(input->Head().GetTypeAnn()));
+        } else {
+            input->SetTypeAnn(input->Head().GetTypeAnn());
+        }
+
         return IGraphTransformer::TStatus::Ok;
     }
 
@@ -11213,17 +11288,23 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["CountBits"] = &CountBitsWrapper;
         Functions["Plus"] = &PlusMinusWrapper;
         Functions["Minus"] = &PlusMinusWrapper;
+        Functions["CheckedMinus"] = &PlusMinusWrapper;
         Functions["+"] = &AddWrapper;
         Functions["Add"] = &AddWrapper;
+        Functions["CheckedAdd"] = &AddWrapper;
         Functions["AggrAdd"] = &AggrAddWrapper;
         Functions["-"] = &SubWrapper;
         Functions["Sub"] = &SubWrapper;
+        Functions["CheckedSub"] = &SubWrapper;
         Functions["*"] = &MulWrapper;
         Functions["Mul"] = &MulWrapper;
+        Functions["CheckedMul"] = &MulWrapper;
         Functions["/"] = &DivWrapper;
         Functions["Div"] = &DivWrapper;
+        Functions["CheckedDiv"] = &DivWrapper;
         Functions["%"] = &ModWrapper;
         Functions["Mod"] = &ModWrapper;
+        Functions["CheckedMod"] = &ModWrapper;
         Functions["BitAnd"] = &BitOpsWrapper<2>;
         Functions["BitOr"] = &BitOpsWrapper<2>;
         Functions["BitXor"] = &BitOpsWrapper<2>;
diff --git a/ydb/library/yql/sql/v1/context.h b/ydb/library/yql/sql/v1/context.h
index 7713a89a7b3..0a6d2c31f6c 100644
--- a/ydb/library/yql/sql/v1/context.h
+++ b/ydb/library/yql/sql/v1/context.h
@@ -48,6 +48,7 @@ namespace NSQLTranslationV1 {
         TString CurrService;
         TDeferredAtom CurrCluster;
         bool PragmaClassicDivision = true;
+        bool PragmaCheckedOps = false;
         bool StrictJoinKeyTypes = false;
         TNamedNodesMap NamedNodes;
 
diff --git a/ydb/library/yql/sql/v1/sql.cpp b/ydb/library/yql/sql/v1/sql.cpp
index e4d0ea0ea3f..b94f88e2a94 100644
--- a/ydb/library/yql/sql/v1/sql.cpp
+++ b/ydb/library/yql/sql/v1/sql.cpp
@@ -4973,7 +4973,7 @@ TNodePtr TSqlExpression::SubExpr(const TRule_con_subexpr& node, const TTrailingQ
             switch (token.GetId()) {
                 case SQLv1LexerTokens::TOKEN_NOT: opName = "Not"; break;
                 case SQLv1LexerTokens::TOKEN_PLUS: opName = "Plus"; break;
-                case SQLv1LexerTokens::TOKEN_MINUS: opName = "Minus"; break;
+                case SQLv1LexerTokens::TOKEN_MINUS: opName = Ctx.Scoped->PragmaCheckedOps ? "CheckedMinus" : "Minus"; break;
                 case SQLv1LexerTokens::TOKEN_TILDA: opName = "BitNot"; break;
                 default:
                     Ctx.IncrementMonCounter("sql_errors", "UnsupportedUnaryOperation");
@@ -5326,15 +5326,15 @@ TNodePtr TSqlExpression::BinOpList(const TNode& node, TGetNode getNode, TIter be
                 Ctx.IncrementMonCounter("sql_binary_operations", "GreaterOrEq");
                 break;
             case SQLv1LexerTokens::TOKEN_PLUS:
-                opName = "+";
+                opName = Ctx.Scoped->PragmaCheckedOps ? "CheckedAdd" : "+";
                 Ctx.IncrementMonCounter("sql_binary_operations", "Plus");
                 break;
             case SQLv1LexerTokens::TOKEN_MINUS:
-                opName = "-";
+                opName = Ctx.Scoped->PragmaCheckedOps ? "CheckedSub" : "-";
                 Ctx.IncrementMonCounter("sql_binary_operations", "Minus");
                 break;
             case SQLv1LexerTokens::TOKEN_ASTERISK:
-                opName = "*";
+                opName = Ctx.Scoped->PragmaCheckedOps ? "CheckedMul" : "*";
                 Ctx.IncrementMonCounter("sql_binary_operations", "Multiply");
                 break;
             case SQLv1LexerTokens::TOKEN_SLASH:
@@ -5342,10 +5342,12 @@ TNodePtr TSqlExpression::BinOpList(const TNode& node, TGetNode getNode, TIter be
                 Ctx.IncrementMonCounter("sql_binary_operations", "Divide");
                 if (!Ctx.Scoped->PragmaClassicDivision && partialResult) {
                     partialResult = new TCallNodeImpl(pos, "SafeCast", {std::move(partialResult), BuildDataType(pos, "Double")});
+                } else if (Ctx.Scoped->PragmaCheckedOps) {
+                    opName = "CheckedDiv";
                 }
                 break;
             case SQLv1LexerTokens::TOKEN_PERCENT:
-                opName = "%";
+                opName = Ctx.Scoped->PragmaCheckedOps ? "CheckedMod" : "%";
                 Ctx.IncrementMonCounter("sql_binary_operations", "Mod");
                 break;
             default:
@@ -9434,7 +9436,7 @@ TNodePtr TSqlQuery::PragmaStatement(const TRule_pragma_stmt& stmt, bool& success
     }
 
     const bool withConfigure = prefix || normalizedPragma == "file" || normalizedPragma == "folder" || normalizedPragma == "udf";
-    static const THashSet<TStringBuf> lexicalScopePragmas = {"classicdivision", "strictjoinkeytypes", "disablestrictjoinkeytypes"};
+    static const THashSet<TStringBuf> lexicalScopePragmas = {"classicdivision", "strictjoinkeytypes", "disablestrictjoinkeytypes", "checkedops"};
     const bool hasLexicalScope = withConfigure || lexicalScopePragmas.contains(normalizedPragma);
     for (auto pragmaValue : pragmaValues) {
         if (pragmaValue->HasAlt_pragma_value3()) {
@@ -9741,6 +9743,13 @@ TNodePtr TSqlQuery::PragmaStatement(const TRule_pragma_stmt& stmt, bool& success
                 return {};
             }
             Ctx.IncrementMonCounter("sql_pragma", "ClassicDivision");
+        } else if (normalizedPragma == "checkedops") {
+            if (values.size() != 1 || !values[0].GetLiteral() || !TryFromString(*values[0].GetLiteral(), Ctx.Scoped->PragmaCheckedOps)) {
+                Error() << "Expected boolean literal as a single argument for: " << pragma;
+                Ctx.IncrementMonCounter("sql_errors", "BadPragmaValue");
+                return {};
+            }
+            Ctx.IncrementMonCounter("sql_pragma", "CheckedOps");
         } else if (normalizedPragma == "disableunordered") {
             Ctx.Warning(Ctx.Pos(), TIssuesIds::YQL_DEPRECATED_PRAGMA)
                 << "Use of deprecated DisableUnordered pragma. It will be dropped soon";