123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540 |
- #include "mkql_chain1_map.h"
- #include <yql/essentials/minikql/computation/mkql_computation_node_holders.h>
- #include <yql/essentials/minikql/computation/mkql_computation_node_codegen.h> // Y_IGNORE
- #include <yql/essentials/minikql/computation/mkql_custom_list.h>
- #include <yql/essentials/minikql/mkql_node_cast.h>
- namespace NKikimr {
- namespace NMiniKQL {
- namespace {
- struct TComputationNodes {
- IComputationExternalNode* const ItemArg;
- IComputationExternalNode* const StateArg;
- IComputationNode* const InitItem;
- IComputationNode* const InitState;
- IComputationNode* const UpdateItem;
- IComputationNode* const UpdateState;
- };
- class TFold1MapFlowWrapper : public TStatefulFlowCodegeneratorNode<TFold1MapFlowWrapper> {
- typedef TStatefulFlowCodegeneratorNode<TFold1MapFlowWrapper> TBaseComputation;
- public:
- TFold1MapFlowWrapper(TComputationMutables& mutables, EValueRepresentation kind, IComputationNode* flow,
- IComputationExternalNode* itemArg, IComputationExternalNode* stateArg,
- IComputationNode* initItem, IComputationNode* initState,
- IComputationNode* updateItem, IComputationNode* updateState)
- : TBaseComputation(mutables, flow, kind, EValueRepresentation::Embedded),
- Flow(flow), ComputationNodes({itemArg, stateArg, initItem, initState, updateItem, updateState})
- {}
- NUdf::TUnboxedValue DoCalculate(NUdf::TUnboxedValue& state, TComputationContext& ctx) const {
- auto item = Flow->GetValue(ctx);
- if (item.IsSpecial()) {
- return item;
- }
- ComputationNodes.ItemArg->SetValue(ctx, std::move(item));
- const bool init = state.IsInvalid();
- const auto value = (init ? ComputationNodes.InitItem : ComputationNodes.UpdateItem)->GetValue(ctx);
- ComputationNodes.StateArg->SetValue(ctx, (init ? ComputationNodes.InitState : ComputationNodes.UpdateState)->GetValue(ctx));
- if (init) {
- state = NUdf::TUnboxedValuePod(true);
- }
- return value;
- }
- #ifndef MKQL_DISABLE_CODEGEN
- Value* DoGenerateGetValue(const TCodegenContext& ctx, Value* statePtr, BasicBlock*& block) const {
- auto& context = ctx.Codegen.GetContext();
- const auto codegenItemArg = dynamic_cast<ICodegeneratorExternalNode*>(ComputationNodes.ItemArg);
- const auto codegenStateArg = dynamic_cast<ICodegeneratorExternalNode*>(ComputationNodes.StateArg);
- MKQL_ENSURE(codegenItemArg, "Item arg must be codegenerator node.");
- MKQL_ENSURE(codegenStateArg, "State arg must be codegenerator node.");
- const auto valueType = Type::getInt128Ty(context);
- const auto good = BasicBlock::Create(context, "good", ctx.Func);
- const auto done = BasicBlock::Create(context, "done", ctx.Func);
- const auto result = PHINode::Create(valueType, 3U, "result", done);
- const auto item = GetNodeValue(Flow, ctx, block);
- result->addIncoming(item, block);
- BranchInst::Create(done, good, IsSpecial(item, block, context), block);
- block = good;
- codegenItemArg->CreateSetValue(ctx, block, item);
- const auto init = BasicBlock::Create(context, "init", ctx.Func);
- const auto next = BasicBlock::Create(context, "next", ctx.Func);
- const auto state = new LoadInst(valueType, statePtr, "load", block);
- BranchInst::Create(init, next, IsInvalid(state, block, context), block);
- block = init;
- const auto one = GetNodeValue(ComputationNodes.InitItem, ctx, block);
- codegenStateArg->CreateSetValue(ctx, block, GetNodeValue(ComputationNodes.InitState, ctx, block));
- result->addIncoming(one, block);
- new StoreInst(GetTrue(context), statePtr, block);
- BranchInst::Create(done, block);
- block = next;
- const auto two = GetNodeValue(ComputationNodes.UpdateItem, ctx, block);
- codegenStateArg->CreateSetValue(ctx, block, GetNodeValue(ComputationNodes.UpdateState, ctx, block));
- result->addIncoming(two, block);
- BranchInst::Create(done, block);
- block = done;
- return result;
- }
- #endif
- private:
- void RegisterDependencies() const final {
- if (const auto flow = FlowDependsOn(Flow)) {
- DependsOn(flow, ComputationNodes.InitItem);
- DependsOn(flow, ComputationNodes.InitState);
- DependsOn(flow, ComputationNodes.UpdateItem);
- DependsOn(flow, ComputationNodes.UpdateState);
- Own(flow, ComputationNodes.ItemArg);
- Own(flow, ComputationNodes.StateArg);
- }
- }
- IComputationNode* const Flow;
- const TComputationNodes ComputationNodes;
- };
- template <bool IsStream>
- class TBaseChain1MapWrapper {
- public:
- class TListValue : public TCustomListValue {
- public:
- class TIterator : public TComputationValue<TIterator> {
- public:
- TIterator(TMemoryUsageInfo* memInfo, TComputationContext& compCtx, NUdf::TUnboxedValue&& iter, const TComputationNodes& computationNodes)
- : TComputationValue<TIterator>(memInfo)
- , CompCtx(compCtx)
- , Iter(std::move(iter))
- , ComputationNodes(computationNodes)
- {}
- private:
- bool Next(NUdf::TUnboxedValue& value) final {
- if (!Iter.Next(ComputationNodes.ItemArg->RefValue(CompCtx))) {
- return false;
- }
- ++Length;
- auto itemNode = Length == 1 ? ComputationNodes.InitItem : ComputationNodes.UpdateItem;
- auto stateNode = Length == 1 ? ComputationNodes.InitState : ComputationNodes.UpdateState;
- value = itemNode->GetValue(CompCtx);
- ComputationNodes.StateArg->SetValue(CompCtx, stateNode->GetValue(CompCtx));
- return true;
- }
- TComputationContext& CompCtx;
- const NUdf::TUnboxedValue Iter;
- const TComputationNodes& ComputationNodes;
- ui64 Length = 0;
- };
- TListValue(TMemoryUsageInfo* memInfo, TComputationContext& compCtx, NUdf::TUnboxedValue&& list, const TComputationNodes& computationNodes)
- : TCustomListValue(memInfo)
- , CompCtx(compCtx)
- , List(std::move(list))
- , ComputationNodes(computationNodes)
- {}
- private:
- NUdf::TUnboxedValue GetListIterator() const final {
- return CompCtx.HolderFactory.Create<TIterator>(CompCtx, List.GetListIterator(), ComputationNodes);
- }
- ui64 GetListLength() const final {
- if (!Length) {
- Length = List.GetListLength();
- }
- return *Length;
- }
- bool HasListItems() const final {
- if (!HasItems) {
- HasItems = List.HasListItems();
- }
- return *HasItems;
- }
- TComputationContext& CompCtx;
- const NUdf::TUnboxedValue List;
- const TComputationNodes& ComputationNodes;
- };
- class TStreamValue : public TComputationValue<TStreamValue> {
- public:
- using TBase = TComputationValue<TStreamValue>;
- TStreamValue(TMemoryUsageInfo* memInfo, TComputationContext& compCtx, NUdf::TUnboxedValue&& list, const TComputationNodes& computationNodes)
- : TBase(memInfo)
- , CompCtx(compCtx)
- , List(std::move(list))
- , ComputationNodes(computationNodes)
- {}
- private:
- NUdf::EFetchStatus Fetch(NUdf::TUnboxedValue& value) final {
- const auto status = List.Fetch(ComputationNodes.ItemArg->RefValue(CompCtx));
- if (status != NUdf::EFetchStatus::Ok) {
- return status;
- }
- ++Length;
- auto itemNode = Length == 1 ? ComputationNodes.InitItem : ComputationNodes.UpdateItem;
- auto stateNode = Length == 1 ? ComputationNodes.InitState : ComputationNodes.UpdateState;
- value = itemNode->GetValue(CompCtx);
- ComputationNodes.StateArg->SetValue(CompCtx, stateNode->GetValue(CompCtx));
- return NUdf::EFetchStatus::Ok;
- }
- TComputationContext& CompCtx;
- const NUdf::TUnboxedValue List;
- const TComputationNodes& ComputationNodes;
- ui64 Length = 0;
- };
- TBaseChain1MapWrapper(IComputationNode* list, IComputationExternalNode* itemArg, IComputationExternalNode* stateArg,
- IComputationNode* initItem, IComputationNode* initState,
- IComputationNode* updateItem, IComputationNode* updateState)
- : List(list), ComputationNodes({itemArg, stateArg, initItem, initState, updateItem, updateState})
- {}
- #ifndef MKQL_DISABLE_CODEGEN
- template<bool IsFirst>
- Function* GenerateMapper(NYql::NCodegen::ICodegen& codegen, const TString& name) const {
- auto& module = codegen.GetModule();
- auto& context = codegen.GetContext();
- const auto newItem = IsFirst ? ComputationNodes.InitItem : ComputationNodes.UpdateItem;
- const auto newState = IsFirst ? ComputationNodes.InitState : ComputationNodes.UpdateState;
- const auto codegenItemArg = dynamic_cast<ICodegeneratorExternalNode*>(ComputationNodes.ItemArg);
- const auto codegenStateArg = dynamic_cast<ICodegeneratorExternalNode*>(ComputationNodes.StateArg);
- MKQL_ENSURE(codegenItemArg, "Item arg must be codegenerator node.");
- MKQL_ENSURE(codegenStateArg, "State arg must be codegenerator node.");
- if (const auto f = module.getFunction(name.c_str()))
- return f;
- const auto valueType = Type::getInt128Ty(context);
- const auto containerType = static_cast<Type*>(valueType);
- const auto contextType = GetCompContextType(context);
- const auto statusType = IsStream ? Type::getInt32Ty(context) : Type::getInt1Ty(context);
- const auto funcType = FunctionType::get(statusType, {PointerType::getUnqual(contextType), containerType, PointerType::getUnqual(valueType)}, false);
- TCodegenContext ctx(codegen);
- ctx.Func = cast<Function>(module.getOrInsertFunction(name.c_str(), funcType).getCallee());
- DISubprogramAnnotator annotator(ctx, ctx.Func);
- auto args = ctx.Func->arg_begin();
- ctx.Ctx = &*args;
- const auto containerArg = &*++args;
- const auto valuePtr = &*++args;
- const auto main = BasicBlock::Create(context, "main", ctx.Func);
- auto block = main;
- const auto container = static_cast<Value*>(containerArg);
- const auto good = BasicBlock::Create(context, "good", ctx.Func);
- const auto done = BasicBlock::Create(context, "done", ctx.Func);
- const auto itemPtr = codegenItemArg->CreateRefValue(ctx, block);
- const auto status = IsStream ?
- CallBoxedValueVirtualMethod<NUdf::TBoxedValueAccessor::EMethod::Fetch>(statusType, container, codegen, block, itemPtr):
- CallBoxedValueVirtualMethod<NUdf::TBoxedValueAccessor::EMethod::Next>(statusType, container, codegen, block, itemPtr);
- const auto icmp = IsStream ?
- CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_NE, status, ConstantInt::get(statusType, static_cast<ui32>(NUdf::EFetchStatus::Ok)), "cond", block):
- CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, status, ConstantInt::getFalse(context), "cond", block);
- BranchInst::Create(done, good, icmp, block);
- block = good;
- SafeUnRefUnboxedOne(valuePtr, ctx, block);
- GetNodeValue(valuePtr, newItem, ctx, block);
- const auto nextState = GetNodeValue(newState, ctx, block);
- codegenStateArg->CreateSetValue(ctx, block, nextState);
- BranchInst::Create(done, block);
- block = done;
- ReturnInst::Create(context, status, block);
- return ctx.Func;
- }
- using TChainMapPtr = std::conditional_t<IsStream, TStreamCodegenValueOne::TFetchPtr, TListCodegenValueOne::TNextPtr>;
- Function* MapFuncOne = nullptr;
- Function* MapFuncTwo = nullptr;
- TChainMapPtr MapOne = nullptr;
- TChainMapPtr MapTwo = nullptr;
- #endif
- IComputationNode* const List;
- const TComputationNodes ComputationNodes;
- };
- class TStreamChain1MapWrapper : public TCustomValueCodegeneratorNode<TStreamChain1MapWrapper>, private TBaseChain1MapWrapper<true> {
- typedef TCustomValueCodegeneratorNode<TStreamChain1MapWrapper> TBaseComputation;
- typedef TBaseChain1MapWrapper<true> TBaseWrapper;
- public:
- TStreamChain1MapWrapper(TComputationMutables& mutables, IComputationNode* list,
- IComputationExternalNode* itemArg, IComputationExternalNode* stateArg,
- IComputationNode* initItem, IComputationNode* initState,
- IComputationNode* updateItem, IComputationNode* updateState
- ) : TBaseComputation(mutables), TBaseWrapper(list, itemArg, stateArg, initItem, initState, updateItem, updateState)
- {}
- NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
- #ifndef MKQL_DISABLE_CODEGEN
- if (ctx.ExecuteLLVM && MapOne && MapTwo)
- return ctx.HolderFactory.Create<TStreamCodegenValueOne>(MapOne, MapTwo, &ctx, List->GetValue(ctx));
- #endif
- return ctx.HolderFactory.Create<TStreamValue>(ctx, List->GetValue(ctx), ComputationNodes);
- }
- private:
- void RegisterDependencies() const final {
- DependsOn(List);
- DependsOn(ComputationNodes.InitItem);
- DependsOn(ComputationNodes.InitState);
- DependsOn(ComputationNodes.UpdateItem);
- DependsOn(ComputationNodes.UpdateState);
- Own(ComputationNodes.ItemArg);
- Own(ComputationNodes.StateArg);
- }
- #ifndef MKQL_DISABLE_CODEGEN
- void GenerateFunctions(NYql::NCodegen::ICodegen& codegen) final {
- MapFuncOne = GenerateMapper<true>(codegen, TBaseComputation::MakeName("Fetch_One"));
- MapFuncTwo = GenerateMapper<false>(codegen, TBaseComputation::MakeName("Fetch_Two"));
- codegen.ExportSymbol(MapFuncOne);
- codegen.ExportSymbol(MapFuncTwo);
- }
- void FinalizeFunctions(NYql::NCodegen::ICodegen& codegen) final {
- if (MapFuncOne)
- MapOne = reinterpret_cast<TChainMapPtr>(codegen.GetPointerToFunction(MapFuncOne));
- if (MapFuncTwo)
- MapTwo = reinterpret_cast<TChainMapPtr>(codegen.GetPointerToFunction(MapFuncTwo));
- }
- #endif
- };
- class TListChain1MapWrapper : public TBothWaysCodegeneratorNode<TListChain1MapWrapper>, private TBaseChain1MapWrapper<false> {
- typedef TBothWaysCodegeneratorNode<TListChain1MapWrapper> TBaseComputation;
- typedef TBaseChain1MapWrapper<false> TBaseWrapper;
- public:
- TListChain1MapWrapper(TComputationMutables& mutables, IComputationNode* list,
- IComputationExternalNode* itemArg, IComputationExternalNode* stateArg,
- IComputationNode* initItem, IComputationNode* initState,
- IComputationNode* updateItem, IComputationNode* updateState
- ) : TBaseComputation(mutables), TBaseWrapper(list, itemArg, stateArg, initItem, initState, updateItem, updateState)
- {}
- NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
- auto list = List->GetValue(ctx);
- if (auto elements = list.GetElements()) {
- auto size = list.GetListLength();
- NUdf::TUnboxedValue* items = nullptr;
- const auto result = ctx.HolderFactory.CreateDirectArrayHolder(size, items);
- if (size) {
- ComputationNodes.ItemArg->SetValue(ctx, NUdf::TUnboxedValue(*elements++));
- *items++ = ComputationNodes.InitItem->GetValue(ctx);
- ComputationNodes.StateArg->SetValue(ctx, ComputationNodes.InitState->GetValue(ctx));
- while (--size) {
- ComputationNodes.ItemArg->SetValue(ctx, NUdf::TUnboxedValue(*elements++));
- *items++ = ComputationNodes.UpdateItem->GetValue(ctx);
- ComputationNodes.StateArg->SetValue(ctx, ComputationNodes.UpdateState->GetValue(ctx));
- }
- }
- return result;
- }
- return ctx.HolderFactory.Create<TListValue>(ctx, std::move(list), ComputationNodes);
- }
- #ifndef MKQL_DISABLE_CODEGEN
- NUdf::TUnboxedValuePod MakeLazyList(TComputationContext& ctx, const NUdf::TUnboxedValuePod value) const {
- return ctx.HolderFactory.Create<TListCodegenValueOne>(MapOne, MapTwo, &ctx, value);
- }
- Value* DoGenerateGetValue(const TCodegenContext& ctx, BasicBlock*& block) const {
- auto& context = ctx.Codegen.GetContext();
- const auto codegenItemArg = dynamic_cast<ICodegeneratorExternalNode*>(ComputationNodes.ItemArg);
- const auto codegenStateArg = dynamic_cast<ICodegeneratorExternalNode*>(ComputationNodes.StateArg);
- MKQL_ENSURE(codegenItemArg, "Item arg must be codegenerator node.");
- MKQL_ENSURE(codegenStateArg, "State arg must be codegenerator node.");
- const auto list = GetNodeValue(List, ctx, block);
- const auto lazy = BasicBlock::Create(context, "lazy", ctx.Func);
- const auto hard = BasicBlock::Create(context, "hard", ctx.Func);
- const auto done = BasicBlock::Create(context, "done", ctx.Func);
- const auto map = PHINode::Create(list->getType(), 3U, "map", done);
- const auto elementsType = PointerType::getUnqual(list->getType());
- const auto elements = CallBoxedValueVirtualMethod<NUdf::TBoxedValueAccessor::EMethod::GetElements>(elementsType, list, ctx.Codegen, block);
- const auto fill = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_NE, elements, ConstantPointerNull::get(elementsType), "fill", block);
- BranchInst::Create(hard, lazy, fill, block);
- {
- block = hard;
- const auto size = CallBoxedValueVirtualMethod<NUdf::TBoxedValueAccessor::EMethod::GetListLength>(Type::getInt64Ty(context), list, ctx.Codegen, block);
- const auto itemsPtr = *Stateless || ctx.AlwaysInline ?
- new AllocaInst(elementsType, 0U, "items_ptr", &ctx.Func->getEntryBlock().back()):
- new AllocaInst(elementsType, 0U, "items_ptr", block);
- const auto array = GenNewArray(ctx, size, itemsPtr, block);
- const auto items = new LoadInst(elementsType, itemsPtr, "items", block);
- const auto init = BasicBlock::Create(context, "init", ctx.Func);
- const auto loop = BasicBlock::Create(context, "loop", ctx.Func);
- const auto next = BasicBlock::Create(context, "next", ctx.Func);
- const auto stop = BasicBlock::Create(context, "stop", ctx.Func);
- const auto good = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_UGT, size, ConstantInt::get(size->getType(), 0), "good", block);
- map->addIncoming(array, block);
- BranchInst::Create(init, done, good, block);
- block = init;
- const auto head = new LoadInst(list->getType(), elements, "head", block);
- codegenItemArg->CreateSetValue(ctx, block, head);
- GetNodeValue(items, ComputationNodes.InitItem, ctx, block);
- const auto state = GetNodeValue(ComputationNodes.InitState, ctx, block);
- codegenStateArg->CreateSetValue(ctx, block, state);
- const auto index = PHINode::Create(size->getType(), 2U, "index", loop);
- index->addIncoming(ConstantInt::get(size->getType(), 1), block);
- BranchInst::Create(loop, block);
- block = loop;
- const auto more = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_UGT, size, index, "more", block);
- BranchInst::Create(next, stop, more, block);
- block = next;
- const auto src = GetElementPtrInst::CreateInBounds(list->getType(), elements, {index}, "src", block);
- const auto item = new LoadInst(list->getType(), src, "item", block);
- codegenItemArg->CreateSetValue(ctx, block, item);
- const auto dst = GetElementPtrInst::CreateInBounds(list->getType(), items, {index}, "dst", block);
- GetNodeValue(dst, ComputationNodes.UpdateItem, ctx, block);
- const auto newState = GetNodeValue(ComputationNodes.UpdateState, ctx, block);
- codegenStateArg->CreateSetValue(ctx, block, newState);
- const auto plus = BinaryOperator::CreateAdd(index, ConstantInt::get(size->getType(), 1), "plus", block);
- index->addIncoming(plus, block);
- BranchInst::Create(loop, block);
- block = stop;
- if (List->IsTemporaryValue()) {
- CleanupBoxed(list, ctx, block);
- }
- map->addIncoming(array, block);
- BranchInst::Create(done, block);
- }
- {
- block = lazy;
- const auto doFunc = ConstantInt::get(Type::getInt64Ty(context), GetMethodPtr(&TListChain1MapWrapper::MakeLazyList));
- const auto ptrType = PointerType::getUnqual(StructType::get(context));
- const auto self = CastInst::Create(Instruction::IntToPtr, ConstantInt::get(Type::getInt64Ty(context), uintptr_t(this)), ptrType, "self", block);
- const auto funType = FunctionType::get(list->getType() , {self->getType(), ctx.Ctx->getType(), list->getType()}, false);
- const auto doFuncPtr = CastInst::Create(Instruction::IntToPtr, doFunc, PointerType::getUnqual(funType), "function", block);
- const auto value = CallInst::Create(funType, doFuncPtr, {self, ctx.Ctx, list}, "value", block);
- map->addIncoming(value, block);
- BranchInst::Create(done, block);
- }
- block = done;
- return map;
- }
- #endif
- private:
- void RegisterDependencies() const final {
- DependsOn(List);
- DependsOn(ComputationNodes.InitItem);
- DependsOn(ComputationNodes.InitState);
- DependsOn(ComputationNodes.UpdateItem);
- DependsOn(ComputationNodes.UpdateState);
- Own(ComputationNodes.ItemArg);
- Own(ComputationNodes.StateArg);
- }
- #ifndef MKQL_DISABLE_CODEGEN
- void GenerateFunctions(NYql::NCodegen::ICodegen& codegen) final {
- TMutableCodegeneratorRootNode<TListChain1MapWrapper>::GenerateFunctions(codegen);
- MapFuncOne = GenerateMapper<true>(codegen, TBaseComputation::MakeName("Next_One"));
- MapFuncTwo = GenerateMapper<false>(codegen, TBaseComputation::MakeName("Next_Two"));
- codegen.ExportSymbol(MapFuncOne);
- codegen.ExportSymbol(MapFuncTwo);
- }
- void FinalizeFunctions(NYql::NCodegen::ICodegen& codegen) final {
- TMutableCodegeneratorRootNode<TListChain1MapWrapper>::FinalizeFunctions(codegen);
- if (MapFuncOne)
- MapOne = reinterpret_cast<TChainMapPtr>(codegen.GetPointerToFunction(MapFuncOne));
- if (MapFuncTwo)
- MapTwo = reinterpret_cast<TChainMapPtr>(codegen.GetPointerToFunction(MapFuncTwo));
- }
- #endif
- };
- }
- IComputationNode* WrapChain1Map(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
- MKQL_ENSURE(callable.GetInputsCount() == 7, "Expected 7 args");
- const auto type = callable.GetType()->GetReturnType();
- const auto flow = LocateNode(ctx.NodeLocator, callable, 0);
- const auto initItem = LocateNode(ctx.NodeLocator, callable, 2);
- const auto initState = LocateNode(ctx.NodeLocator, callable, 3);
- const auto updateItem = LocateNode(ctx.NodeLocator, callable, 5);
- const auto updateState = LocateNode(ctx.NodeLocator, callable, 6);
- const auto itemArg = LocateExternalNode(ctx.NodeLocator, callable, 1);
- const auto stateArg = LocateExternalNode(ctx.NodeLocator, callable, 4);
- if (type->IsFlow()) {
- return new TFold1MapFlowWrapper(ctx.Mutables, GetValueRepresentation(type), flow, itemArg, stateArg, initItem, initState, updateItem, updateState);
- } else if (type->IsStream()) {
- return new TStreamChain1MapWrapper(ctx.Mutables, flow, itemArg, stateArg, initItem, initState, updateItem, updateState);
- } else if (type->IsList()) {
- return new TListChain1MapWrapper(ctx.Mutables, flow, itemArg, stateArg, initItem, initState, updateItem, updateState);
- }
- THROW yexception() << "Expected flow, list or stream.";
- }
- }
- }
|