Optimize Coalesce(Just(...) ...) on peephole final stage. (#3957)

Author: Tony-Romanov

ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp

@@ -8147,16 +8147,23 @@ TExprNode::TPtr DropAssume(const TExprNode::TPtr& node, TExprContext&) {
 }
 
 TExprNode::TPtr OptimizeCoalesce(const TExprNode::TPtr& node, TExprContext& ctx) {
-    if (const auto& input = node->Head(); input.IsCallable("If") && (input.Child(1U)->IsComplete() || input.Child(2U)->IsComplete())) {
+    if (const auto& input = node->Head(); input.IsCallable("If") && input.ChildrenSize() == 3U &&
+        (input.Child(1U)->IsComplete() || input.Child(2U)->IsComplete())) {
         YQL_CLOG(DEBUG, CorePeepHole) << "Dive " << node->Content() << " into " << input.Content();
         return ctx.ChangeChildren(input, {
             input.HeadPtr(),
             ctx.ChangeChild(*node, 0U, input.ChildPtr(1U)),
             ctx.ChangeChild(*node, 0U, input.ChildPtr(2U))
         });
-    } else if (input.IsCallable("Nothing")) {
-        YQL_CLOG(DEBUG, CorePeepHole) << "Drop " << node->Content() << " over " << input.Content();
-        return node->TailPtr();
+    } else if (node->ChildrenSize() == 2U) {
+        if (input.IsCallable("Nothing")) {
+            YQL_CLOG(DEBUG, CorePeepHole) << "Drop " << node->Content() << " over " << input.Content();
+            return node->TailPtr();
+        } else if (input.IsCallable("Just")) {
+            YQL_CLOG(DEBUG, CorePeepHole) << "Drop " << node->Content() << " over " << input.Content();
+            return IsSameAnnotation(*node->GetTypeAnn(), *input.GetTypeAnn()) ?
+                node->HeadPtr() : input.HeadPtr();
+        }
     }
     return node;
 }