remove redundant semicolons, always add semicolon after last statement in lambdas and actions

commit_hash:d2b19f44d3c6f6e89f03c2c27e897d2b55e6174e

2 files changed, 81 insertions, 14 deletions

diff --git a/yql/essentials/sql/v1/format/sql_format.cpp b/yql/essentials/sql/v1/format/sql_format.cpp
index 6a2c0e13bf5..c5052c44c88 100644
--- a/yql/essentials/sql/v1/format/sql_format.cpp
+++ b/yql/essentials/sql/v1/format/sql_format.cpp
@@ -43,6 +43,39 @@ TTokenIterator SkipWSOrComment(TTokenIterator curr, TTokenIterator end) {
     return curr;
 }
 
+TTokenIterator SkipWSOrCommentBackward(TTokenIterator curr, TTokenIterator begin) {
+    while (curr != begin && (curr->Name == "WS" || curr->Name == "COMMENT")) {
+        --curr;
+    }
+    return curr;
+}
+
+void SkipForValidate(
+    TTokenIterator& in,
+    TTokenIterator& out,
+    const TParsedTokenList& query,
+    const TParsedTokenList& formattedQuery
+) {
+    in = SkipWS(in, query.end());
+    out = SkipWS(out, formattedQuery.end());
+
+    while (
+        in != query.end() && in->Name == "SEMICOLON" &&
+        (out == formattedQuery.end() || out->Name != "SEMICOLON") &&
+        in != query.begin() && IsIn({"SEMICOLON", "LBRACE_CURLY", "AS"}, SkipWSOrCommentBackward(in - 1, query.begin())->Name)
+    ) {
+        in = SkipWS(++in, query.end());
+    }
+
+    auto inSkippedComments = SkipWSOrComment(in, query.end());
+    if (
+        out != formattedQuery.end() && out->Name == "SEMICOLON" &&
+        inSkippedComments != query.end() && IsIn({"RBRACE_CURLY", "END"}, inSkippedComments->Name)
+    ) {
+        out = SkipWS(++out, formattedQuery.end());
+    }
+}
+
 TParsedToken TransformTokenForValidate(TParsedToken token) {
     if (token.Name == "EQUALS2") {
         token.Name = "EQUALS";
@@ -61,8 +94,7 @@ bool Validate(const TParsedTokenList& query, const TParsedTokenList& formattedQu
     auto outEnd = formattedQuery.end();
 
     while (in != inEnd && out != outEnd) {
-        in = SkipWS(in, inEnd);
-        out = SkipWS(out, outEnd);
+        SkipForValidate(in, out, query, formattedQuery);
         if (in != inEnd && out != outEnd) {
             auto inToken = TransformTokenForValidate(*in);
             auto outToken = TransformTokenForValidate(*out);
@@ -82,8 +114,7 @@ bool Validate(const TParsedTokenList& query, const TParsedTokenList& formattedQu
             ++out;
         }
     }
-    in = SkipWS(in, inEnd);
-    out = SkipWS(out, outEnd);
+    SkipForValidate(in, out, query, formattedQuery);
     return in == inEnd && out == outEnd;
 }
 
@@ -777,20 +808,34 @@ private:
         }
     }
 
+    template <typename T>
+    void SkipSemicolons(const ::google::protobuf::RepeatedPtrField<T>& field, bool printOne = false) {
+        for (const auto& m : field) {
+            if (printOne) {
+                Visit(m);
+                printOne = false;
+            } else {
+                ++TokenIndex;
+            }
+        }
+        if (printOne) {
+            Out(';');
+        }
+    }
+
     void VisitDefineActionOrSubqueryBody(const TRule_define_action_or_subquery_body& msg) {
-        VisitRepeated(msg.GetBlock1());
+        SkipSemicolons(msg.GetBlock1());
         if (msg.HasBlock2()) {
             const auto& b = msg.GetBlock2();
             Visit(b.GetRule_sql_stmt_core1());
             for (auto block : b.GetBlock2()) {
-                VisitRepeated(block.GetBlock1());
+                SkipSemicolons(block.GetBlock1(), /* printOne = */ true);
                 if (!IsSimpleStatement(block.GetRule_sql_stmt_core2()).GetOrElse(false)) {
                     Out('\n');
                 }
                 Visit(block.GetRule_sql_stmt_core2());
             }
-
-            VisitRepeated(b.GetBlock3());
+            SkipSemicolons(b.GetBlock3(), /* printOne = */ true);
         }
     }
 
@@ -2344,9 +2389,10 @@ private:
     void VisitLambdaBody(const TRule_lambda_body& msg) {
         PushCurrentIndent();
         NewLine();
-        VisitRepeated(msg.GetBlock1());
+        SkipSemicolons(msg.GetBlock1());
         for (const auto& block : msg.GetBlock2()) {
-            Visit(block);
+            Visit(block.GetRule_lambda_stmt1());
+            SkipSemicolons(block.GetBlock2(), /* printOne = */ true);
             NewLine();
         }
 
@@ -2354,8 +2400,8 @@ private:
         ExprLineIndent = CurrentIndent;
 
         Visit(msg.GetRule_expr4());
-        VisitRepeated(msg.GetBlock5());
 
+        SkipSemicolons(msg.GetBlock5(), /* printOne = */ true);
         ExprLineIndent = 0;
 
         PopCurrentIndent();
@@ -3068,7 +3114,7 @@ public:
                 return false;
             }
 
-            if (!Validate(stmtFormattedTokens, stmtTokens)) {
+            if (!Validate(stmtTokens, stmtFormattedTokens)) {
                 issues.AddIssue(NYql::TIssue({}, TStringBuilder() << "Validation failed: " << currentQuery.Quote() << " != " << currentFormattedQuery.Quote()));
                 return false;
             }
diff --git a/yql/essentials/sql/v1/format/sql_format_ut.h b/yql/essentials/sql/v1/format/sql_format_ut.h
index 27a06119ccf..1503bb0414f 100644
--- a/yql/essentials/sql/v1/format/sql_format_ut.h
+++ b/yql/essentials/sql/v1/format/sql_format_ut.h
@@ -557,6 +557,25 @@ Y_UNIT_TEST(DefineActionOrSubquery) {
             "VALUES\n\t\t\t(1)\n\t\t;\n\tEND DEFINE;\n\n\t"
             "DEFINE SUBQUERY $c() AS\n\t\tSELECT\n\t\t\t1\n\t\t;\n\t"
             "END DEFINE;\n\tDO\n\t\t$b()\n\t;\n\n\tPROCESS $c();\nEND DEFINE;\n"},
+        {"define action $foo($bar) as;"
+            "$a = 10;; "
+            "$b = 20;;; "
+            "$c = $a + $b "
+            "end define",
+            "DEFINE ACTION $foo($bar) AS\n\t"
+            "$a = 10;\n\t"
+            "$b = 20;\n\t"
+            "$c = $a + $b;\n"
+            "END DEFINE;\n"},
+        {"define subquery $s() as;"
+            "select * from $t1 "
+            "union all select * from $t2 "
+            "end define",
+            "DEFINE SUBQUERY $s() AS\n\t"
+            "SELECT\n\t\t*\n\tFROM\n\t\t$t1\n\t"
+            "UNION ALL\n\t"
+            "SELECT\n\t\t*\n\tFROM\n\t\t$t2\n\t;\n"
+            "END DEFINE;\n"},
     };
 
     TSetup setup;
@@ -872,9 +891,11 @@ Y_UNIT_TEST(CompositeTypesAndQuestions) {
 Y_UNIT_TEST(Lambda) {
     TCases cases = {
         {"$f=($a,$b)->{$x=$a+$b;return $a*$x};$g=($a,$b?)->($a+$b??0);select $f(10,4),$g(1,2);",
-            "$f = ($a, $b) -> {\n\t$x = $a + $b;\n\tRETURN $a * $x\n};\n\n"
+            "$f = ($a, $b) -> {\n\t$x = $a + $b;\n\tRETURN $a * $x;\n};\n\n"
             "$g = ($a, $b?) -> ($a + $b ?? 0);\n\n"
             "SELECT\n\t$f(10, 4),\n\t$g(1, 2)\n;\n"},
+        {"$f=($arg)->{;$a=10;;$b=20;;;RETURN $a+$b}",
+            "$f = ($arg) -> {\n\t$a = 10;\n\t$b = 20;\n\tRETURN $a + $b;\n};\n"},
     };
 
     TSetup setup;
@@ -1444,7 +1465,7 @@ Y_UNIT_TEST(ExistsExpr) {
 Y_UNIT_TEST(LambdaInsideExpr) {
     TCases cases = {
         {"SELECT ListMap(AsList(1,2),($x)->{return $x+1});",
-            "SELECT\n\tListMap(\n\t\tAsList(1, 2), ($x) -> {\n\t\t\tRETURN $x + 1\n\t\t}\n\t)\n;\n"},
+            "SELECT\n\tListMap(\n\t\tAsList(1, 2), ($x) -> {\n\t\t\tRETURN $x + 1;\n\t\t}\n\t)\n;\n"},
     };
 
     TSetup setup;