aboutsummaryrefslogtreecommitdiffstats
path: root/yql/essentials/minikql/comp_nodes/mkql_size.cpp
blob: 6ea1f40f423a79c40d9e3302f887a3a8be39dc3a (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
#include "mkql_size.h"
#include <yql/essentials/minikql/computation/mkql_computation_node_codegen.h>  // Y_IGNORE
#include <yql/essentials/minikql/computation/mkql_computation_node_holders.h>
#include <yql/essentials/minikql/computation/mkql_computation_node_holders_codegen.h>
#include <yql/essentials/minikql/mkql_node_cast.h>
#include <yql/essentials/minikql/mkql_node_builder.h>

namespace NKikimr {
namespace NMiniKQL {

namespace {

extern "C" void DeleteString(void* strData);

template<size_t Size>
class TSizePrimitiveTypeWrapper : public TDecoratorCodegeneratorNode<TSizePrimitiveTypeWrapper<Size>> {
    typedef TDecoratorCodegeneratorNode<TSizePrimitiveTypeWrapper<Size>> TBaseComputation;
public:
    TSizePrimitiveTypeWrapper(IComputationNode* data)
        : TBaseComputation(data)
    {}

    NUdf::TUnboxedValuePod DoCalculate(TComputationContext&, const NUdf::TUnboxedValuePod& value) const {
        return value ? NUdf::TUnboxedValuePod(ui32(Size)) : NUdf::TUnboxedValuePod();
    }

#ifndef MKQL_DISABLE_CODEGEN
    Value* DoGenerateGetValue(const TCodegenContext& ctx, Value* value, BasicBlock*& block) const {
        Y_UNUSED(ctx);
        const uint64_t init[] = {Size, 0x100000000000000ULL};
        const auto size = ConstantInt::get(value->getType(), APInt(128, 2, init));
        return SelectInst::Create(IsEmpty(value, block), value, size, "size", block);
    }
#endif
};

template<bool IsOptional>
class TSizeWrapper : public TMutableCodegeneratorNode<TSizeWrapper<IsOptional>> {
    typedef TMutableCodegeneratorNode<TSizeWrapper<IsOptional>> TBaseComputation;
public:
    TSizeWrapper(TComputationMutables& mutables, IComputationNode* data)
        : TBaseComputation(mutables, EValueRepresentation::Embedded)
        , Data(data)
    {}

    NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
        const auto& data = Data->GetValue(ctx);
        if (IsOptional && !data) {
            return NUdf::TUnboxedValuePod();
        }

        const ui32 size = data.AsStringRef().Size();
        return NUdf::TUnboxedValuePod(size);
    }

#ifndef MKQL_DISABLE_CODEGEN
    Value* DoGenerateGetValue(const TCodegenContext& ctx, BasicBlock*& block) const {
        auto& context = ctx.Codegen.GetContext();

        const auto data = GetNodeValue(this->Data, ctx, block);

        const auto type = Type::getInt8Ty(context);
        const auto embType = FixedVectorType::get(type, 16);
        const auto cast = CastInst::Create(Instruction::BitCast, data, embType, "cast", block);

        const auto mark = ExtractElementInst::Create(cast, {ConstantInt::get(type, 15)}, "mark", block);
        const auto bsize = ExtractElementInst::Create(cast, {ConstantInt::get(type, 14)}, "bsize", block);

        const auto emb = BasicBlock::Create(context, "emb", ctx.Func);
        const auto str = BasicBlock::Create(context, "str", ctx.Func);
        const auto done = BasicBlock::Create(context, "done", ctx.Func);

        const auto result = PHINode::Create(data->getType(), 4, "result", done);
        result->addIncoming(ConstantInt::get(data->getType(), 0), block);

        const auto choise = SwitchInst::Create(mark, done, 2, block);
        choise->addCase(ConstantInt::get(type, 1), emb);
        choise->addCase(ConstantInt::get(type, 2), str);

        {
            block = emb;

            const auto full = SetterFor<ui32>(bsize, context, block);
            result->addIncoming(full, block);
            BranchInst::Create(done, block);
        }

        {
            block = str;

            const auto type32 = Type::getInt32Ty(context);
            const auto strType = FixedVectorType::get(type32, 4);
            const auto four = CastInst::Create(Instruction::BitCast, data, strType, "four", block);
            const auto ssize = ExtractElementInst::Create(four, {ConstantInt::get(type, 2)}, "ssize", block);
            const auto full = SetterFor<ui32>(ssize, context, block);

            const auto half = CastInst::Create(Instruction::Trunc, data, Type::getInt64Ty(context), "half", block);
            const auto strptr = CastInst::Create(Instruction::IntToPtr, half, PointerType::getUnqual(strType), "str_ptr", block);
            const auto refptr = GetElementPtrInst::CreateInBounds(strType, strptr, {ConstantInt::get(type32, 0), ConstantInt::get(type32, 1)}, "refptr", block);
            const auto refs = new LoadInst(type32, refptr, "refs", block);
            const auto test = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_UGT, refs, ConstantInt::get(refs->getType(), 0), "test", block);

            const auto free = BasicBlock::Create(context, "free", ctx.Func);

            result->addIncoming(full, block);
            BranchInst::Create(done, free, test, block);

            block = free;

            const auto fnType = FunctionType::get(Type::getVoidTy(context), {strptr->getType()}, false);
            const auto name = "DeleteString";
            ctx.Codegen.AddGlobalMapping(name, reinterpret_cast<const void*>(&DeleteString));
            const auto func = ctx.Codegen.GetModule().getOrInsertFunction(name, fnType);
            CallInst::Create(func, {strptr}, "", block);
            result->addIncoming(full, block);
            BranchInst::Create(done, block);
        }

        block = done;
        return result;
    }
#endif
private:
    void RegisterDependencies() const final {
        this->DependsOn(Data);
    }

    IComputationNode* const Data;
};

}

IComputationNode* WrapSize(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
    MKQL_ENSURE(callable.GetInputsCount() == 1, "Expected 1 arg");
    bool isOptional;
    const auto dataType = UnpackOptionalData(callable.GetInput(0), isOptional);

    const auto data = LocateNode(ctx.NodeLocator, callable, 0);

    switch(dataType->GetSchemeType()) {
#define MAKE_PRIMITIVE_TYPE_SIZE(type, layout) \
        case NUdf::TDataType<type>::Id: \
            if (isOptional) \
                return new TSizePrimitiveTypeWrapper<sizeof(layout)>(data); \
            else \
                return ctx.NodeFactory.CreateImmutableNode(NUdf::TUnboxedValuePod(ui32(sizeof(layout))));
        KNOWN_FIXED_VALUE_TYPES(MAKE_PRIMITIVE_TYPE_SIZE)
#undef MAKE_PRIMITIVE_TYPE_SIZE
    }

    if (isOptional)
        return new TSizeWrapper<true>(ctx.Mutables, data);
    else
        return new TSizeWrapper<false>(ctx.Mutables, data);
}

}
}