Generalize DoGenGetValues of wide While-variations

Author: Vadim Averin

ydb/library/yql/minikql/comp_nodes/mkql_wide_filter.cpp

@@ -46,6 +46,16 @@ class TBaseWideFilterWrapper {
                 if (Predicate == Items[i] || Items[i]->GetDependencesCount() > 0U)
                     *out = *fields[i];
     }
+
+    bool ApplyPredicate(TComputationContext& ctx, NUdf::TUnboxedValue*const* values) const {
+        auto **fields = GetFields(ctx);
+        PrepareArguments(ctx, values);
+        for (size_t idx = 0; idx < Items.size(); idx++) {
+            *fields[idx] = *values[idx];
+        }
+        return Predicate->GetValue(ctx).Get<bool>();
+    }
+
 #ifndef MKQL_DISABLE_CODEGEN
     template<bool ReplaceOriginalGetter = true>
     Value* GenGetPredicate(const TCodegenContext& ctx,
@@ -154,12 +164,7 @@ using TBaseComputation = TSimpleStatefulWideFlowCodegeneratorNode<TWideFilterWit
             return EProcessResult::Finish;
         }
         if (fetchRes == EFetchResult::One) {
-            auto **fields = GetFields(ctx);
-            PrepareArguments(ctx, values);
-            for (size_t idx = 0; idx < Items.size(); idx++) {
-                *fields[idx] = *values[idx];
-            }
-            if (Predicate->GetValue(ctx).Get<bool>()) {
+            if (ApplyPredicate(ctx, values)) {
                 FillOutputs(ctx, values);
                 limit--;
                 return EProcessResult::One;
@@ -182,80 +187,37 @@ using TBaseComputation = TSimpleStatefulWideFlowCodegeneratorNode<TWideFilterWit
 };
 
 template<bool Inclusive>
-class TWideTakeWhileWrapper : public TStatefulWideFlowCodegeneratorNode<TWideTakeWhileWrapper<Inclusive>>,  public TBaseWideFilterWrapper {
-using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideTakeWhileWrapper<Inclusive>>;
+class TWideTakeWhileWrapper : public TSimpleStatefulWideFlowCodegeneratorNode<TWideTakeWhileWrapper<Inclusive>, bool>,  public TBaseWideFilterWrapper {
+using TBaseComputation = TSimpleStatefulWideFlowCodegeneratorNode<TWideTakeWhileWrapper<Inclusive>, bool>;
 public:
      TWideTakeWhileWrapper(TComputationMutables& mutables, IComputationWideFlowNode* flow, TComputationExternalNodePtrVector&& items,
             IComputationNode* predicate)
         : TBaseComputation(mutables, flow, EValueRepresentation::Embedded)
         , TBaseWideFilterWrapper(mutables, flow, std::move(items), predicate)
     {}
 
-    EFetchResult DoCalculate(NUdf::TUnboxedValue& state, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
-        if (!state.IsInvalid()) {
-            return EFetchResult::Finish;
-        }
-
-        PrepareArguments(ctx, output);
-
-        auto **fields = GetFields(ctx);
-
-        if (const auto result = Flow->FetchValues(ctx, fields); EFetchResult::One != result)
-            return result;
-
-        const bool predicate = Predicate->GetValue(ctx).Get<bool>();
-        if (!predicate)
-            state = NUdf::TUnboxedValuePod();
+    void InitState(bool& stop, TComputationContext& ctx) const {
+        stop = false;
+    }
 
-        if (Inclusive || predicate) {
-            FillOutputs(ctx, output);
-            return EFetchResult::One;
+    TBaseComputation::EProcessResult DoProcess(bool& stop, TComputationContext& ctx, EFetchResult fetchRes, NUdf::TUnboxedValue*const* values) const {
+        if (stop) {
+            return TBaseComputation::EProcessResult::Finish;
         }
-
-        return EFetchResult::Finish;
+        if (fetchRes == EFetchResult::One) {
+            const bool predicate = ApplyPredicate(ctx, values);
+            if (!predicate) {
+                stop = true;
+            }
+            if (Inclusive || predicate) {
+                FillOutputs(ctx, values);
+                return TBaseComputation::EProcessResult::One;
+            }
+            return TBaseComputation::EProcessResult::Finish;
+        }
+        return static_cast<TBaseComputation::EProcessResult>(fetchRes);
     }
-#ifndef MKQL_DISABLE_CODEGEN
-    ICodegeneratorInlineWideNode::TGenerateResult DoGenGetValues(const TCodegenContext& ctx, Value* statePtr, BasicBlock*& block) const {
-        auto& context = ctx.Codegen.GetContext();
-
-        const auto resultType = Type::getInt32Ty(context);
-
-        const auto work = BasicBlock::Create(context, "work", ctx.Func);
-        const auto test = BasicBlock::Create(context, "test", ctx.Func);
-        const auto stop = BasicBlock::Create(context, "stop", ctx.Func);
-        const auto done = BasicBlock::Create(context, "done", ctx.Func);
-
-        const auto result = PHINode::Create(resultType, 4U, "result", done);
-        result->addIncoming(ConstantInt::get(resultType, static_cast<i32>(EFetchResult::Finish)), block);
-
-        const auto state = new LoadInst(Type::getInt128Ty(context), statePtr, "state", block);
-        const auto finished = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, state, GetTrue(context), "finished", block);
-
-        BranchInst::Create(done, work, IsValid(statePtr, block), block);
-
-        block = work;
-        auto status = GetNodeValues(Flow, ctx, block);
-        const auto special = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_SLE, status.first, ConstantInt::get(resultType, 0), "special", block);
-        result->addIncoming(status.first, block);
-        BranchInst::Create(done, test, special, block);
-
-        block = test;
-
-        const auto predicate = GenGetPredicate(ctx, status.second, block);
-        result->addIncoming(status.first, block);
-        BranchInst::Create(done, stop, predicate, block);
-
-        block = stop;
 
-        new StoreInst(GetEmpty(context), statePtr, block);
-        result->addIncoming(ConstantInt::get(resultType, static_cast<i32>(Inclusive ? EFetchResult::One: EFetchResult::Finish)), block);
-
-        BranchInst::Create(done, block);
-
-        block = done;
-        return {result, std::move(status.second)};
-    }
-#endif
 private:
     void RegisterDependencies() const final {
         if (const auto flow = this->FlowDependsOn(Flow)) {
@@ -266,72 +228,33 @@ using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideTakeWhileWrappe
 };
 
 template<bool Inclusive>
-class TWideSkipWhileWrapper : public TStatefulWideFlowCodegeneratorNode<TWideSkipWhileWrapper<Inclusive>>,  public TBaseWideFilterWrapper {
-using TBaseComputation = TStatefulWideFlowCodegeneratorNode<TWideSkipWhileWrapper<Inclusive>>;
+class TWideSkipWhileWrapper : public TSimpleStatefulWideFlowCodegeneratorNode<TWideSkipWhileWrapper<Inclusive>, bool>,  public TBaseWideFilterWrapper {
+using TBaseComputation = TSimpleStatefulWideFlowCodegeneratorNode<TWideSkipWhileWrapper<Inclusive>, bool>;
 public:
      TWideSkipWhileWrapper(TComputationMutables& mutables, IComputationWideFlowNode* flow, TComputationExternalNodePtrVector&& items, IComputationNode* predicate)
         : TBaseComputation(mutables, flow, EValueRepresentation::Embedded)
         , TBaseWideFilterWrapper(mutables, flow, std::move(items), predicate)
     {}
 
-    EFetchResult DoCalculate(NUdf::TUnboxedValue& state, TComputationContext& ctx, NUdf::TUnboxedValue*const* output) const {
-        if (!state.IsInvalid()) {
-            return Flow->FetchValues(ctx, output);
-        }
-
-        auto **fields = GetFields(ctx);
-
-        do {
-            PrepareArguments(ctx, output);
-            if (const auto result = Flow->FetchValues(ctx, fields); EFetchResult::One != result)
-                return result;
-        } while (Predicate->GetValue(ctx).Get<bool>());
-
-        state = NUdf::TUnboxedValuePod();
+    void InitState(bool& start, TComputationContext& ctx) const {
+        start = false;
+    }
 
-        if constexpr (Inclusive)
-            return Flow->FetchValues(ctx, output);
-        else {
-            FillOutputs(ctx, output);
-            return EFetchResult::One;
+    TBaseComputation::EProcessResult DoProcess(bool& start, TComputationContext& ctx, EFetchResult fetchRes, NUdf::TUnboxedValue*const* values) const {
+        if (!start && fetchRes == EFetchResult::One) {
+            const bool predicate = ApplyPredicate(ctx, values);
+            if (!predicate) {
+                start = true;
+            }
+            if (!Inclusive && !predicate) {
+                FillOutputs(ctx, values);
+                return TBaseComputation::EProcessResult::One;
+            }
+            return TBaseComputation::EProcessResult::Fetch;
         }
+        return static_cast<TBaseComputation::EProcessResult>(fetchRes);
     }
-#ifndef MKQL_DISABLE_CODEGEN
-    ICodegeneratorInlineWideNode::TGenerateResult DoGenGetValues(const TCodegenContext& ctx, Value* statePtr, BasicBlock*& block) const {
-        auto& context = ctx.Codegen.GetContext();
-
-        const auto resultType = Type::getInt32Ty(context);
-
-        const auto work = BasicBlock::Create(context, "work", ctx.Func);
-        const auto test = BasicBlock::Create(context, "test", ctx.Func);
-        const auto stop = BasicBlock::Create(context, "stop", ctx.Func);
-        const auto done = BasicBlock::Create(context, "done", ctx.Func);
-
-        BranchInst::Create(work, block);
-
-        block = work;
-
-        const auto status = GetNodeValues(Flow, ctx, block);
-
-        const auto special = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_SLE, status.first, ConstantInt::get(resultType, 0), "special", block);
-        const auto passtrought = BinaryOperator::CreateOr(special, IsValid(statePtr, block), "passtrought", block);
-        BranchInst::Create(done, test, passtrought, block);
 
-        block = test;
-
-        const auto predicate = GenGetPredicate<false>(ctx, status.second, block);
-        BranchInst::Create(work, stop, predicate, block);
-
-        block = stop;
-
-        new StoreInst(GetEmpty(context), statePtr, block);
-
-        BranchInst::Create(Inclusive ? work : done, block);
-
-        block = done;
-        return status;
-    }
-#endif
 private:
     void RegisterDependencies() const final {
         if (const auto flow = this->FlowDependsOn(Flow)) {