Handle repeats from inlined ExpandApply calls

init
commit_hash:c6373ff2a9e9562c16afe8d56cbd495790545274

6 files changed, 24 insertions, 11 deletions

diff --git a/yql/essentials/core/common_opt/yql_co_pgselect.cpp b/yql/essentials/core/common_opt/yql_co_pgselect.cpp
index 26208b1c49..135eb7d46c 100644
--- a/yql/essentials/core/common_opt/yql_co_pgselect.cpp
+++ b/yql/essentials/core/common_opt/yql_co_pgselect.cpp
@@ -494,8 +494,9 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinksPartial(TPositionHand
                     });
 
                     ctx.Step.Repeat(TExprStep::ExpandApplyForLambdas);
-                    auto status = ExpandApply(traits, traits, ctx);
+                    auto status = ExpandApplyNoRepeat(traits, traits, ctx);
                     YQL_ENSURE(status != IGraphTransformer::TStatus::Error);
+
                     switch (factoryIndex) {
                     case 0:
                         countAllTraits = traits;
@@ -2791,8 +2792,9 @@ TExprNode::TPtr BuildWindows(TPositionHandle pos, const TExprNode::TPtr& list, c
 
                     auto traits = ctx.ReplaceNodes(lambda->TailPtr(), replaces);
                     ctx.Step.Repeat(TExprStep::ExpandApplyForLambdas);
-                    auto status = ExpandApply(traits, traits, ctx);
+                    auto status = ExpandApplyNoRepeat(traits, traits, ctx);
                     YQL_ENSURE(status != IGraphTransformer::TStatus::Error);
+
                     value = traits;
                 } else {
                     ythrow yexception() << "Not supported function: " << name;
diff --git a/yql/essentials/core/type_ann/type_ann_list.cpp b/yql/essentials/core/type_ann/type_ann_list.cpp
index b8c6782267..efcd41ac76 100644
--- a/yql/essentials/core/type_ann/type_ann_list.cpp
+++ b/yql/essentials/core/type_ann/type_ann_list.cpp
@@ -154,7 +154,7 @@ namespace {
 
             auto traits = ctx.ReplaceNodes(TExprNode::TPtr(traitsFactoryBody), factoryReplaces);
             ctx.Step.Repeat(TExprStep::ExpandApplyForLambdas);
-            auto status = ExpandApply(traits, traits, ctx);
+            auto status = ExpandApplyNoRepeat(traits, traits, ctx);
             if (status == IGraphTransformer::TStatus::Error) {
                 return nullptr;
             }
diff --git a/yql/essentials/core/yql_aggregate_expander.cpp b/yql/essentials/core/yql_aggregate_expander.cpp
index a9b5a16162..0ebfaa988b 100644
--- a/yql/essentials/core/yql_aggregate_expander.cpp
+++ b/yql/essentials/core/yql_aggregate_expander.cpp
@@ -145,7 +145,7 @@ TExprNode::TPtr TAggregateExpander::ExpandAggApply(const TExprNode::TPtr& node)
         });
 
     Ctx.Step.Repeat(TExprStep::ExpandApplyForLambdas);
-    auto status = ExpandApply(traits, traits, Ctx);
+    auto status = ExpandApplyNoRepeat(traits, traits, Ctx);
     YQL_ENSURE(status != IGraphTransformer::TStatus::Error);
     return traits;
 }
diff --git a/yql/essentials/core/yql_expr_optimize.cpp b/yql/essentials/core/yql_expr_optimize.cpp
index 91608eed80..24716a1668 100644
--- a/yql/essentials/core/yql_expr_optimize.cpp
+++ b/yql/essentials/core/yql_expr_optimize.cpp
@@ -862,6 +862,16 @@ IGraphTransformer::TStatus ExpandApply(const TExprNode::TPtr& input, TExprNode::
     return ret;
 }
 
+IGraphTransformer::TStatus ExpandApplyNoRepeat(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
+    output = input;
+    for (;;) {
+        auto status = ExpandApply(output, output, ctx);
+        if (status.Level != IGraphTransformer::TStatus::Repeat) {
+            return status;
+        }
+    }
+}
+
 TExprNode::TPtr ApplySyncListToWorld(const TExprNode::TPtr& main, const TSyncMap& syncList, TExprContext& ctx) {
     if (syncList.empty()) {
         return main;
diff --git a/yql/essentials/core/yql_expr_optimize.h b/yql/essentials/core/yql_expr_optimize.h
index a73548fa74..9a6e6d9f71 100644
--- a/yql/essentials/core/yql_expr_optimize.h
+++ b/yql/essentials/core/yql_expr_optimize.h
@@ -49,6 +49,7 @@ IGraphTransformer::TStatus OptimizeExprEx(const TExprNode::TPtr& input, TExprNod
     TExprContext& ctx, const TOptimizeExprSettings& settings);
 
 IGraphTransformer::TStatus ExpandApply(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx);
+IGraphTransformer::TStatus ExpandApplyNoRepeat(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx);
 TExprNode::TPtr ApplySyncListToWorld(const TExprNode::TPtr& main, const TSyncMap& syncList, TExprContext& ctx);
 
 typedef std::function<bool (const TExprNode::TPtr&)> TExprVisitPtrFunc;
diff --git a/yql/essentials/core/yql_opt_window.cpp b/yql/essentials/core/yql_opt_window.cpp
index d05f64d516..80b4bfeeac 100644
--- a/yql/essentials/core/yql_opt_window.cpp
+++ b/yql/essentials/core/yql_opt_window.cpp
@@ -250,7 +250,7 @@ TExprNode::TPtr ApplyDistinctForUpdateLambda(TExprNode::TPtr updateLambda, const
 
     auto expandedDistinctKeyType = ExpandType(updateLambda->Pos(), distinctKeyType, ctx);
     auto expandedDistinctKeyOrigType = ExpandType(updateLambda->Pos(), distinctKeyOrigType, ctx);
-    
+
     auto setAddValueUdf = ctx.Builder(updateLambda->Pos())
         .Callable("Udf")
             .Atom(0, "Set.AddValue")
@@ -480,7 +480,7 @@ TCalcOverWindowTraits ExtractCalcOverWindowTraits(const TExprNode::TPtr& frames,
                 YQL_ENSURE(rawTraits.OutputType);
 
                 auto lambdaInputType = traits->Child(0)->GetTypeAnn()->Cast<TTypeExprType>()->GetType();
-                
+
                 if (item->ChildrenSize() == 3) {
                     auto distinctKey = item->Child(2)->Content();
 
@@ -2293,7 +2293,7 @@ TExprNode::TPtr BuildFold1Lambda(TPositionHandle pos, const TExprNode::TPtr& fra
                         lambda = ReplaceLastLambdaArgWithUnsignedLiteral(*lambda, i, ctx);
                     }
                     YQL_ENSURE(lambda->Child(0)->ChildrenSize() == 2);
-                    
+
                     applied = ctx.Builder(pos)
                         .Apply(lambda)
                             .With(0, arg1)
@@ -2822,7 +2822,7 @@ bool NeedPartitionRows(const TExprNode::TPtr& frames, const TStructExprType& row
             return true;
         }
     }
-    
+
     return false;
 }
 
@@ -2838,14 +2838,14 @@ TString AllocatePartitionRowsColumn(const TStructExprType& rowType) {
     }
 }
 
-TExprNode::TPtr AddPartitionRowsColumn(TPositionHandle pos, const TExprNode::TPtr& input, const TExprNode::TPtr& keyColumns, 
+TExprNode::TPtr AddPartitionRowsColumn(TPositionHandle pos, const TExprNode::TPtr& input, const TExprNode::TPtr& keyColumns,
     const TString& columnName, TExprContext& ctx, TTypeAnnotationContext& types) {
     auto exportsPtr = types.Modules->GetModule("/lib/yql/window.yql");
     YQL_ENSURE(exportsPtr);
     const auto& exports = exportsPtr->Symbols();
     const auto ex = exports.find("count_traits_factory");
     YQL_ENSURE(exports.cend() != ex);
-    TNodeOnNodeOwnedMap deepClones;    
+    TNodeOnNodeOwnedMap deepClones;
     auto lambda = ctx.DeepCopy(*ex->second, exportsPtr->ExprCtx(), deepClones, true, false);
     auto listTypeNode = ctx.NewCallable(pos, "TypeOf", {input});
     auto extractor = ctx.Builder(pos)
@@ -2862,7 +2862,7 @@ TExprNode::TPtr AddPartitionRowsColumn(TPositionHandle pos, const TExprNode::TPt
     });
 
     ctx.Step.Repeat(TExprStep::ExpandApplyForLambdas);
-    auto status = ExpandApply(traits, traits, ctx);
+    auto status = ExpandApplyNoRepeat(traits, traits, ctx);
     YQL_ENSURE(status != IGraphTransformer::TStatus::Error);
 
     return ctx.Builder(pos)