Better LLVM for WideTop[Sort].

2 files changed, 24 insertions, 8 deletions

diff --git a/ydb/library/yql/minikql/comp_nodes/mkql_wide_top_sort.cpp b/ydb/library/yql/minikql/comp_nodes/mkql_wide_top_sort.cpp
index 485b96a9aa..cdff692499 100644
--- a/ydb/library/yql/minikql/comp_nodes/mkql_wide_top_sort.cpp
+++ b/ydb/library/yql/minikql/comp_nodes/mkql_wide_top_sort.cpp
@@ -69,23 +69,24 @@ public:
         return Fields.data();
     }
 
-    void Push() {
+    NUdf::TUnboxedValue* Push() {
         if (Full.size() < Count) {
             Full.emplace_back(Free.back());
             Free.pop_back();
             ResetFields();
-            return;
+            return static_cast<NUdf::TUnboxedValue*>(Full.back());
         } else if (!Throat) {
             Throat = *std::max_element(Full.cbegin(), Full.cend(), LessFunc);
         }
 
         if (!LessFunc(Tongue, Throat)) {
             std::fill_n(static_cast<NUdf::TUnboxedValue*>(Tongue), Indexes.size(), NUdf::TUnboxedValuePod());
-            return;
+            return nullptr;
         }
 
         Full.emplace_back(Free.back());
         Free.pop_back();
+        const auto out = Full.back();
 
         if (Full.size() == GetStorageSize()) {
             std::nth_element(Full.begin(), Full.begin() + Count, Full.end(), LessFunc);
@@ -98,6 +99,7 @@ public:
         }
 
         ResetFields();
+        return static_cast<NUdf::TUnboxedValue*>(out);
     }
 
     template<bool Sort>
@@ -325,17 +327,31 @@ public:
             const auto tonguePtr = GetElementPtrInst::CreateInBounds(stateArg, { stateFields.This(), stateFields.GetTongue() }, "tongue_ptr", block);
             const auto tongue = new LoadInst(tonguePtr, "tongue", block);
 
-            for (auto i = 0U; i < Representations.size(); ++i) {
+            for (auto i = 0U; i < KeyTypes.size(); ++i) {
                 const auto item = getres.second[Indexes[i]](ctx, block);
                 ValueAddRef(Representations[i], item, ctx, block);
-                const auto ptr = GetElementPtrInst::CreateInBounds(tongue, {ConstantInt::get(Type::getInt32Ty(context), i)}, (TString("ptr_") += ToString(i)).c_str(), block);
+                const auto ptr = GetElementPtrInst::CreateInBounds(tongue, {ConstantInt::get(Type::getInt32Ty(context), i)}, (TString("key_ptr_") += ToString(i)).c_str(), block);
                 new StoreInst(item, ptr, block);
             }
 
             const auto pushFunc = ConstantInt::get(Type::getInt64Ty(context), GetMethodPtr(&TState::Push));
-            const auto pushType = FunctionType::get(Type::getVoidTy(context), {stateArg->getType()}, false);
+            const auto pushType = FunctionType::get(outputPtrType, {stateArg->getType()}, false);
             const auto pushPtr = CastInst::Create(Instruction::IntToPtr, pushFunc, PointerType::getUnqual(pushType), "function", block);
-            CallInst::Create(pushPtr, {stateArg}, "", block);
+            const auto inputPtr = CallInst::Create(pushPtr, {stateArg}, "input", block);
+            const auto accepted = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_NE, inputPtr, ConstantPointerNull::get(outputPtrType), "accepted", block);
+
+            const auto push = BasicBlock::Create(context, "push", ctx.Func);
+
+            BranchInst::Create(push, loop, accepted, block);
+
+            block = push;
+
+            for (auto i = KeyTypes.size(); i < Representations.size(); ++i) {
+                const auto item = getres.second[Indexes[i]](ctx, block);
+                ValueAddRef(Representations[i], item, ctx, block);
+                const auto ptr = GetElementPtrInst::CreateInBounds(inputPtr, {ConstantInt::get(Type::getInt32Ty(context), 0), ConstantInt::get(Type::getInt32Ty(context), i)}, (TString("pay_ptr_") += ToString(i)).c_str(), block);
+                new StoreInst(item, ptr, block);
+            }
 
             BranchInst::Create(loop, block);
         }
diff --git a/ydb/library/yql/minikql/computation/mkql_computation_node_codegen.cpp b/ydb/library/yql/minikql/computation/mkql_computation_node_codegen.cpp
index 3479db37b5..7eedc6bbe5 100644
--- a/ydb/library/yql/minikql/computation/mkql_computation_node_codegen.cpp
+++ b/ydb/library/yql/minikql/computation/mkql_computation_node_codegen.cpp
@@ -428,7 +428,7 @@ Value* GenCompareFunction<true>(NUdf::EDataSlot slot, Value* lv, Value* rv, TCod
     const auto done = BasicBlock::Create(context, "done", ctx.Func);
 
     const auto resultType = Type::getInt32Ty(context);
-    const auto res = PHINode::Create(resultType, 2U, "result", done);
+    const auto res = PHINode::Create(resultType, 3U, "result", done);
 
     const auto le = IsEmpty(lv, block);
     const auto re = IsEmpty(rv, block);