YQL-20339 merge blocks optimizer

init
commit_hash:fee9a86a1f95e9271bf49205dad4b6eb0b934dec

5 files changed, 147 insertions, 0 deletions

diff --git a/yql/essentials/core/common_opt/ya.make b/yql/essentials/core/common_opt/ya.make
index 31ae6c7d324..e301b885994 100644
--- a/yql/essentials/core/common_opt/ya.make
+++ b/yql/essentials/core/common_opt/ya.make
@@ -2,6 +2,7 @@ LIBRARY()
 
 SRCS(
     yql_co.h
+    yql_co_blocks.cpp
     yql_co_extr_members.cpp
     yql_flatmap_over_join.cpp
     yql_co_finalizers.cpp
diff --git a/yql/essentials/core/common_opt/yql_co_blocks.cpp b/yql/essentials/core/common_opt/yql_co_blocks.cpp
new file mode 100644
index 00000000000..94feed3727d
--- /dev/null
+++ b/yql/essentials/core/common_opt/yql_co_blocks.cpp
@@ -0,0 +1,111 @@
+#include "yql_co_blocks.h"
+
+#include <yql/essentials/core/yql_expr_optimize.h>
+#include <yql/essentials/core/yql_type_annotation.h>
+#include <yql/essentials/core/expr_nodes/yql_expr_nodes.h>
+
+#include <yql/essentials/utils/log/log.h>
+
+namespace NYql {
+
+namespace {
+
+class TBlockVisitor {
+public:
+    void Visit(const TExprNode& node) {
+        auto [it, inserted] = Visited_.emplace(&node, nullptr);
+        if (!inserted) {
+            it->second = nullptr; // multiple paths
+            return;
+        }
+
+        const bool isBlock = node.IsCallable("Block");
+        if (isBlock) {
+            Blocks_.emplace(&node);
+            if (!BlockStack_.empty()) {
+                it->second = BlockStack_.back();
+                YQL_ENSURE(it->second);
+            }
+
+            BlockStack_.push_back(&node);
+            Visit(node.Head().Tail());
+            BlockStack_.pop_back();
+        }
+
+        if (node.Type() == TExprNode::Lambda) {
+            TVector<const TExprNode*> savedBlocks = std::move(BlockStack_);
+            BlockStack_.clear();
+            for (ui32 i = 1; i < node.ChildrenSize(); ++i) {
+                Visit(*node.Child(i));
+            }
+
+            BlockStack_ = std::move(savedBlocks);
+        } else {
+            for (const auto& child : node.Children()) {
+                Visit(*child);
+            }
+        }
+    }
+
+    TNodeMap<const TExprNode*> GetBlockParents() const {
+        TNodeMap<const TExprNode*> ret;
+        for (const auto b : Blocks_) {
+            auto it = Visited_.find(b);
+            YQL_ENSURE(it != Visited_.cend());
+            if (it->second != nullptr) {
+                ret.emplace(b, it->second);
+            }
+        }
+
+        for (auto& [node,parent] : ret) {
+            for (;;) {
+                auto parentIt = ret.find(parent);
+                if (parentIt == ret.cend()) {
+                    break;
+                }
+
+                if (parentIt->second == nullptr) {
+                    break;
+                }
+
+                parent = parentIt->second;
+            }
+        }
+
+        return ret;
+    }
+
+private:
+    TNodeMap<const TExprNode*> Visited_;
+    TVector<const TExprNode*> BlockStack_;
+    TNodeSet Blocks_;
+};
+
+}
+
+IGraphTransformer::TStatus OptimizeBlocks(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx,
+    TTypeAnnotationContext& typeCtx) {
+    output = input;
+    TBlockVisitor visitor;
+    visitor.Visit(*input);
+    auto blockParents = visitor.GetBlockParents();
+    if (blockParents.empty()) {
+        return IGraphTransformer::TStatus::Ok;
+    }
+
+    YQL_CLOG(DEBUG, Core) << "Found " << blockParents.size() << " nested blocks";
+    TNodeOnNodeOwnedMap toOptimize;
+    for (const auto [node, parent]: blockParents) {
+        auto lambda = NNodes::TCoLambda(node->HeadPtr());
+        auto parentLambda = NNodes::TCoLambda(parent->HeadPtr());
+        toOptimize[lambda.Args().Arg(0).Raw()] = parentLambda.Args().Arg(0).Ptr();
+        toOptimize[node] = lambda.Body().Ptr();
+    }
+
+    TOptimizeExprSettings settings(&typeCtx);
+    settings.CustomInstantTypeTransformer = typeCtx.CustomInstantTypeTransformer.Get();
+    return RemapExpr(input, output, toOptimize, ctx, settings);
+}
+
+}
+
diff --git a/yql/essentials/core/common_opt/yql_co_blocks.h b/yql/essentials/core/common_opt/yql_co_blocks.h
new file mode 100644
index 00000000000..bbc55d32c3f
--- /dev/null
+++ b/yql/essentials/core/common_opt/yql_co_blocks.h
@@ -0,0 +1,13 @@
+#pragma once
+
+#include <yql/essentials/ast/yql_expr.h>
+#include <yql/essentials/core/yql_graph_transformer.h>
+
+namespace NYql {
+
+struct TTypeAnnotationContext;
+
+IGraphTransformer::TStatus OptimizeBlocks(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx,
+    TTypeAnnotationContext& typeCtx);
+
+} // NYql
diff --git a/yql/essentials/core/common_opt/yql_co_flow1.cpp b/yql/essentials/core/common_opt/yql_co_flow1.cpp
index 4542c707072..82a1dede75d 100644
--- a/yql/essentials/core/common_opt/yql_co_flow1.cpp
+++ b/yql/essentials/core/common_opt/yql_co_flow1.cpp
@@ -2104,6 +2104,20 @@ void RegisterCoFlowCallables1(TCallableOptimizerMap& map) {
         }
         return node;
     };
+
+    map["ToMutDict"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
+        Y_UNUSED(ctx);
+        if (!optCtx.IsSingleUsage(node->Head())) {
+            return node;
+        }
+
+        if (node->Head().IsCallable("FromMutDict")) {
+            YQL_CLOG(DEBUG, Core) << "Skip " << node->Content() << " over " << node->Head().Content();
+            return node->Head().HeadPtr();
+        }
+
+        return node;
+    };
 }
 
 }
diff --git a/yql/essentials/core/common_opt/yql_co_transformer.cpp b/yql/essentials/core/common_opt/yql_co_transformer.cpp
index 1a807f76858..3287e20f146 100644
--- a/yql/essentials/core/common_opt/yql_co_transformer.cpp
+++ b/yql/essentials/core/common_opt/yql_co_transformer.cpp
@@ -1,5 +1,6 @@
 #include "yql_co_transformer.h"
 #include "yql_co.h"
+#include "yql_co_blocks.h"
 
 #include <yql/essentials/core/yql_expr_optimize.h>
 #include <yql/essentials/core/yql_expr_type_annotation.h>
@@ -93,6 +94,13 @@ IGraphTransformer::TStatus TCommonOptTransformer::DoTransform(TExprNode::TPtr in
         return status;
     }
 
+    if (TypeCtx_->LangVer >= MakeLangVersion(2025, 4)) {
+        status = OptimizeBlocks(input = std::move(output), output, ctx, *TypeCtx_);
+        if (status.Level != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+    }
+
     if (!ScanErrors(*output, ctx)) {
         return IGraphTransformer::TStatus::Error;
     }