// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements. See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership. The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License. You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied. See the License for the

// specific language governing permissions and limitations

// under the License.

#include <bit>

#include <bitset>

#include "common/status.h"

#include "core/assert_cast.h"

#include "core/column/column.h"

#include "core/column/column_vector.h"

#include "core/data_type/data_type_number.h"

#include "core/types.h"

#include "exprs/function/cast_type_to_either.h"

#include "exprs/function/simple_function_factory.h"

namespace doris {

class FunctionBitTest : public IFunction {

public:

static constexpr auto name = "bit_test";

static FunctionPtr create() { return std::make_shared<FunctionBitTest>(); }

String get_name() const override { return name; }

size_t get_number_of_arguments() const override { return 0; }

bool is_variadic() const override { return true; }

DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {

return std::make_shared<DataTypeInt8>();

}

Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,

uint32_t result, size_t input_rows_count) const override {

bool valid =

cast_type(block.get_by_position(arguments[0]).type.get(), [&](const auto& type) {

using DataType = std::decay_t<decltype(type)>;

if (auto col = check_and_get_column<ColumnVector<DataType::PType>>(

block.get_by_position(arguments[0]).column.get()) ||

is_column_const(*block.get_by_position(arguments[0]).column)) {

execute_inner<DataType::PType>(block, arguments, result, input_rows_count);

return true;

}

return false;

});

if (!valid) {

return Status::RuntimeError(

"{}'s argument does not match the expected data type, type: {}, column: {}",

get_name(), block.get_by_position(arguments[0]).type->get_name(),

block.get_by_position(arguments[0]).column->dump_structure());

}

return Status::OK();

}

template <typename F>

static bool cast_type(const IDataType* type, F&& f) {

return cast_type_to_either<DataTypeInt8, DataTypeInt16, DataTypeInt32, DataTypeInt64,

DataTypeInt128>(type, std::forward<F>(f));

}

template <PrimitiveType T>

void execute_inner(Block& block, const ColumnNumbers& arguments, uint32_t result,

size_t input_rows_count) const {

size_t argument_size = arguments.size();

std::vector<ColumnPtr> argument_columns(argument_size);

auto result_data_column = ColumnInt8::create(input_rows_count, 1);

auto& res_data = result_data_column->get_data();

// maybe most user is bit_test(column, const), so only handle this case

if (argument_size == 2) {

std::vector<uint8_t> is_consts(argument_size);

std::tie(argument_columns[0], is_consts[0]) =

unpack_if_const(block.get_by_position(arguments[0]).column);

std::tie(argument_columns[1], is_consts[1]) =

unpack_if_const(block.get_by_position(arguments[1]).column);

execute_for_two_argument<T>(argument_columns, is_consts, res_data, input_rows_count);

} else {

for (size_t i = 0; i < argument_size; ++i) {

argument_columns[i] = block.get_by_position(arguments[i])

.column->convert_to_full_column_if_const();

}

execute_for_others_arg<T>(argument_columns, res_data, argument_size, input_rows_count);

}

block.replace_by_position(result, std::move(result_data_column));

}

template <PrimitiveType T>

void execute_for_two_argument(std::vector<ColumnPtr>& argument_columns,

std::vector<uint8_t>& is_consts, ColumnInt8::Container& res_data,

size_t input_rows_count) const {

const auto& first_column_data =

assert_cast<const ColumnVector<T>&>(*argument_columns[0].get()).get_data();

const auto& second_column_data =

assert_cast<const ColumnVector<T>&>(*argument_columns[1].get()).get_data();

for (int i = 0; i < input_rows_count; ++i) {

auto first_value = first_column_data[index_check_const(i, is_consts[0])];

auto second_value = second_column_data[index_check_const(i, is_consts[1])];

// the pos is invalid, set result = 0

if (second_value < 0 ||

second_value >= sizeof(typename PrimitiveTypeTraits<T>::CppType) * 8) {

res_data[i] = 0;

continue;

}

res_data[i] = ((first_value >> second_value) & 1);

}

}

template <PrimitiveType T>

void execute_for_others_arg(std::vector<ColumnPtr>& argument_columns,

ColumnInt8::Container& res_data, size_t argument_size,

size_t input_rows_count) const {

const auto& first_column_data =

assert_cast<const ColumnVector<T>&>(*argument_columns[0].get()).get_data();

for (int i = 0; i < input_rows_count; ++i) {

auto first_value = first_column_data[i];

for (int col = 1; col < argument_size; ++col) {

const auto& arg_column_data =

assert_cast<const ColumnVector<T>&>(*argument_columns[col].get())

.get_data();

// the pos is invalid, set result = 0

if (arg_column_data[i] < 0 ||

arg_column_data[i] >= sizeof(typename PrimitiveTypeTraits<T>::CppType) * 8) {

res_data[i] = 0;

break;

}

// if one of pos & result is 0, could set res = 0, and return directly.

if (!((first_value >> arg_column_data[i]) & 1)) {

res_data[i] = 0;

break;

}

}

}

}

};

void register_function_bit_test(SimpleFunctionFactory& factory) {

factory.register_function<FunctionBitTest>();

factory.register_alias("bit_test", "bit_test_all");

}

} // namespace doris