// Copyright 2022 The Tint Authors.
//
// Licensed 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 "src/tint/fuzzers/tint_ast_fuzzer/mutations/change_binary_operator.h"

#include <utility>

#include "src/tint/sem/reference.h"

namespace tint::fuzzers::ast_fuzzer {

namespace {

bool IsSuitableForShift(const sem::Type* lhs_type, const sem::Type* rhs_type) {
    // `a << b` requires b to be an unsigned scalar or vector, and `a` to be an
    // integer scalar or vector with the same width as `b`. Similar for `a >> b`.

    if (rhs_type->is_unsigned_integer_scalar()) {
        return lhs_type->is_integer_scalar();
    }
    if (rhs_type->is_unsigned_integer_vector()) {
        return lhs_type->is_unsigned_integer_vector();
    }
    return false;
}

bool CanReplaceAddSubtractWith(const sem::Type* lhs_type,
                               const sem::Type* rhs_type,
                               ast::BinaryOp new_operator) {
    // The program is assumed to be well-typed, so this method determines when
    // 'new_operator' can be used as a type-preserving replacement in an '+' or
    // '-' expression.
    switch (new_operator) {
        case ast::BinaryOp::kAdd:
        case ast::BinaryOp::kSubtract:
            // '+' and '-' are fully type compatible.
            return true;
        case ast::BinaryOp::kAnd:
        case ast::BinaryOp::kOr:
        case ast::BinaryOp::kXor:
            // These operators do not have a mixed vector-scalar form, and only work
            // on integer types.
            return lhs_type == rhs_type && lhs_type->is_integer_scalar_or_vector();
        case ast::BinaryOp::kMultiply:
            // '+' and '*' are largely type-compatible, but for matrices they are only
            // type-compatible if the matrices are square.
            return !lhs_type->is_float_matrix() || lhs_type->is_square_float_matrix();
        case ast::BinaryOp::kDivide:
        case ast::BinaryOp::kModulo:
            // '/' is not defined for matrices.
            return lhs_type->is_numeric_scalar_or_vector() &&
                   rhs_type->is_numeric_scalar_or_vector();
        case ast::BinaryOp::kShiftLeft:
        case ast::BinaryOp::kShiftRight:
            return IsSuitableForShift(lhs_type, rhs_type);
        default:
            return false;
    }
}

bool CanReplaceMultiplyWith(const sem::Type* lhs_type,
                            const sem::Type* rhs_type,
                            ast::BinaryOp new_operator) {
    // The program is assumed to be well-typed, so this method determines when
    // 'new_operator' can be used as a type-preserving replacement in a '*'
    // expression.
    switch (new_operator) {
        case ast::BinaryOp::kMultiply:
            return true;
        case ast::BinaryOp::kAdd:
        case ast::BinaryOp::kSubtract:
            // '*' is type-compatible with '+' and '-' for square matrices, and for
            // numeric scalars/vectors.
            if (lhs_type->is_square_float_matrix() && rhs_type->is_square_float_matrix()) {
                return true;
            }
            return lhs_type->is_numeric_scalar_or_vector() &&
                   rhs_type->is_numeric_scalar_or_vector();
        case ast::BinaryOp::kAnd:
        case ast::BinaryOp::kOr:
        case ast::BinaryOp::kXor:
            // These operators require homogeneous integer types.
            return lhs_type == rhs_type && lhs_type->is_integer_scalar_or_vector();
        case ast::BinaryOp::kDivide:
        case ast::BinaryOp::kModulo:
            // '/' is not defined for matrices.
            return lhs_type->is_numeric_scalar_or_vector() &&
                   rhs_type->is_numeric_scalar_or_vector();
        case ast::BinaryOp::kShiftLeft:
        case ast::BinaryOp::kShiftRight:
            return IsSuitableForShift(lhs_type, rhs_type);
        default:
            return false;
    }
}

bool CanReplaceDivideOrModuloWith(const sem::Type* lhs_type,
                                  const sem::Type* rhs_type,
                                  ast::BinaryOp new_operator) {
    // The program is assumed to be well-typed, so this method determines when
    // 'new_operator' can be used as a type-preserving replacement in a '/'
    // expression.
    switch (new_operator) {
        case ast::BinaryOp::kAdd:
        case ast::BinaryOp::kSubtract:
        case ast::BinaryOp::kMultiply:
        case ast::BinaryOp::kDivide:
        case ast::BinaryOp::kModulo:
            // These operators work in all contexts where '/' works.
            return true;
        case ast::BinaryOp::kAnd:
        case ast::BinaryOp::kOr:
        case ast::BinaryOp::kXor:
            // These operators require homogeneous integer types.
            return lhs_type == rhs_type && lhs_type->is_integer_scalar_or_vector();
        case ast::BinaryOp::kShiftLeft:
        case ast::BinaryOp::kShiftRight:
            return IsSuitableForShift(lhs_type, rhs_type);
        default:
            return false;
    }
}

bool CanReplaceLogicalAndLogicalOrWith(ast::BinaryOp new_operator) {
    switch (new_operator) {
        case ast::BinaryOp::kLogicalAnd:
        case ast::BinaryOp::kLogicalOr:
        case ast::BinaryOp::kAnd:
        case ast::BinaryOp::kOr:
        case ast::BinaryOp::kEqual:
        case ast::BinaryOp::kNotEqual:
            // These operators all work whenever '&&' and '||' work.
            return true;
        default:
            return false;
    }
}

bool CanReplaceAndOrWith(const sem::Type* lhs_type,
                         const sem::Type* rhs_type,
                         ast::BinaryOp new_operator) {
    switch (new_operator) {
        case ast::BinaryOp::kAnd:
        case ast::BinaryOp::kOr:
            // '&' and '|' work in all the same contexts.
            return true;
        case ast::BinaryOp::kAdd:
        case ast::BinaryOp::kSubtract:
        case ast::BinaryOp::kMultiply:
        case ast::BinaryOp::kDivide:
        case ast::BinaryOp::kModulo:
        case ast::BinaryOp::kXor:
            // '&' and '|' can be applied to booleans. In all other contexts,
            // integer numeric operators work.
            return !lhs_type->is_bool_scalar_or_vector();
        case ast::BinaryOp::kShiftLeft:
        case ast::BinaryOp::kShiftRight:
            return IsSuitableForShift(lhs_type, rhs_type);
        case ast::BinaryOp::kLogicalAnd:
        case ast::BinaryOp::kLogicalOr:
            // '&' and '|' can be applied to booleans, and for boolean scalar
            // scalar contexts, their logical counterparts work.
            return lhs_type->Is<sem::Bool>();
        case ast::BinaryOp::kEqual:
        case ast::BinaryOp::kNotEqual:
            // '&' and '|' can be applied to booleans, and in these contexts equality
            // comparison operators also work.
            return lhs_type->is_bool_scalar_or_vector();
        default:
            return false;
    }
}

bool CanReplaceXorWith(const sem::Type* lhs_type,
                       const sem::Type* rhs_type,
                       ast::BinaryOp new_operator) {
    switch (new_operator) {
        case ast::BinaryOp::kAdd:
        case ast::BinaryOp::kSubtract:
        case ast::BinaryOp::kMultiply:
        case ast::BinaryOp::kDivide:
        case ast::BinaryOp::kModulo:
        case ast::BinaryOp::kAnd:
        case ast::BinaryOp::kOr:
        case ast::BinaryOp::kXor:
            // '^' only works on integer types, and in any such context, all other
            // integer operators also work.
            return true;
        case ast::BinaryOp::kShiftLeft:
        case ast::BinaryOp::kShiftRight:
            return IsSuitableForShift(lhs_type, rhs_type);
        default:
            return false;
    }
}

bool CanReplaceShiftLeftShiftRightWith(const sem::Type* lhs_type,
                                       const sem::Type* rhs_type,
                                       ast::BinaryOp new_operator) {
    switch (new_operator) {
        case ast::BinaryOp::kShiftLeft:
        case ast::BinaryOp::kShiftRight:
            // These operators are type-compatible.
            return true;
        case ast::BinaryOp::kAdd:
        case ast::BinaryOp::kSubtract:
        case ast::BinaryOp::kMultiply:
        case ast::BinaryOp::kDivide:
        case ast::BinaryOp::kModulo:
        case ast::BinaryOp::kAnd:
        case ast::BinaryOp::kOr:
        case ast::BinaryOp::kXor:
            // Shift operators allow mixing of signed and unsigned arguments, but in
            // the case where the arguments are homogeneous, they are type-compatible
            // with other numeric operators.
            return lhs_type == rhs_type;
        default:
            return false;
    }
}

bool CanReplaceEqualNotEqualWith(const sem::Type* lhs_type, ast::BinaryOp new_operator) {
    switch (new_operator) {
        case ast::BinaryOp::kEqual:
        case ast::BinaryOp::kNotEqual:
            // These operators are type-compatible.
            return true;
        case ast::BinaryOp::kLessThan:
        case ast::BinaryOp::kLessThanEqual:
        case ast::BinaryOp::kGreaterThan:
        case ast::BinaryOp::kGreaterThanEqual:
            // An equality comparison between numeric types can be changed to an
            // ordered comparison.
            return lhs_type->is_numeric_scalar_or_vector();
        case ast::BinaryOp::kLogicalAnd:
        case ast::BinaryOp::kLogicalOr:
            // An equality comparison between boolean scalars can be turned into a
            // logical operation.
            return lhs_type->Is<sem::Bool>();
        case ast::BinaryOp::kAnd:
        case ast::BinaryOp::kOr:
            // An equality comparison between boolean scalars or vectors can be turned
            // into a component-wise non-short-circuit logical operation.
            return lhs_type->is_bool_scalar_or_vector();
        default:
            return false;
    }
}

bool CanReplaceLessThanLessThanEqualGreaterThanGreaterThanEqualWith(ast::BinaryOp new_operator) {
    switch (new_operator) {
        case ast::BinaryOp::kEqual:
        case ast::BinaryOp::kNotEqual:
        case ast::BinaryOp::kLessThan:
        case ast::BinaryOp::kLessThanEqual:
        case ast::BinaryOp::kGreaterThan:
        case ast::BinaryOp::kGreaterThanEqual:
            // Ordered comparison operators can be interchanged, and equality
            // operators can be used in their place.
            return true;
        default:
            return false;
    }
}
}  // namespace

MutationChangeBinaryOperator::MutationChangeBinaryOperator(
    protobufs::MutationChangeBinaryOperator message)
    : message_(std::move(message)) {}

MutationChangeBinaryOperator::MutationChangeBinaryOperator(uint32_t binary_expr_id,
                                                           ast::BinaryOp new_operator) {
    message_.set_binary_expr_id(binary_expr_id);
    message_.set_new_operator(static_cast<uint32_t>(new_operator));
}

bool MutationChangeBinaryOperator::CanReplaceBinaryOperator(
    const Program& program,
    const ast::BinaryExpression& binary_expr,
    ast::BinaryOp new_operator) {
    if (new_operator == binary_expr.op) {
        // An operator should not be replaced with itself, as this would be a no-op.
        return false;
    }

    // Get the types of the operators.
    const auto* lhs_type = program.Sem().Get(binary_expr.lhs)->Type();
    const auto* rhs_type = program.Sem().Get(binary_expr.rhs)->Type();

    // If these are reference types, unwrap them to get the pointee type.
    const sem::Type* lhs_basic_type =
        lhs_type->Is<sem::Reference>() ? lhs_type->As<sem::Reference>()->StoreType() : lhs_type;
    const sem::Type* rhs_basic_type =
        rhs_type->Is<sem::Reference>() ? rhs_type->As<sem::Reference>()->StoreType() : rhs_type;

    switch (binary_expr.op) {
        case ast::BinaryOp::kAdd:
        case ast::BinaryOp::kSubtract:
            return CanReplaceAddSubtractWith(lhs_basic_type, rhs_basic_type, new_operator);
        case ast::BinaryOp::kMultiply:
            return CanReplaceMultiplyWith(lhs_basic_type, rhs_basic_type, new_operator);
        case ast::BinaryOp::kDivide:
        case ast::BinaryOp::kModulo:
            return CanReplaceDivideOrModuloWith(lhs_basic_type, rhs_basic_type, new_operator);
        case ast::BinaryOp::kAnd:
        case ast::BinaryOp::kOr:
            return CanReplaceAndOrWith(lhs_basic_type, rhs_basic_type, new_operator);
        case ast::BinaryOp::kXor:
            return CanReplaceXorWith(lhs_basic_type, rhs_basic_type, new_operator);
        case ast::BinaryOp::kShiftLeft:
        case ast::BinaryOp::kShiftRight:
            return CanReplaceShiftLeftShiftRightWith(lhs_basic_type, rhs_basic_type, new_operator);
        case ast::BinaryOp::kLogicalAnd:
        case ast::BinaryOp::kLogicalOr:
            return CanReplaceLogicalAndLogicalOrWith(new_operator);
        case ast::BinaryOp::kEqual:
        case ast::BinaryOp::kNotEqual:
            return CanReplaceEqualNotEqualWith(lhs_basic_type, new_operator);
        case ast::BinaryOp::kLessThan:
        case ast::BinaryOp::kLessThanEqual:
        case ast::BinaryOp::kGreaterThan:
        case ast::BinaryOp::kGreaterThanEqual:
        case ast::BinaryOp::kNone:
            return CanReplaceLessThanLessThanEqualGreaterThanGreaterThanEqualWith(new_operator);
            assert(false && "Unreachable");
            return false;
    }
}

bool MutationChangeBinaryOperator::IsApplicable(const Program& program,
                                                const NodeIdMap& node_id_map) const {
    const auto* binary_expr_node =
        As<ast::BinaryExpression>(node_id_map.GetNode(message_.binary_expr_id()));
    if (binary_expr_node == nullptr) {
        // Either the id does not exist, or does not correspond to a binary
        // expression.
        return false;
    }
    // Check whether the replacement is acceptable.
    const auto new_operator = static_cast<ast::BinaryOp>(message_.new_operator());
    return CanReplaceBinaryOperator(program, *binary_expr_node, new_operator);
}

void MutationChangeBinaryOperator::Apply(const NodeIdMap& node_id_map,
                                         CloneContext* clone_context,
                                         NodeIdMap* new_node_id_map) const {
    // Get the node whose operator is to be replaced.
    const auto* binary_expr_node =
        As<ast::BinaryExpression>(node_id_map.GetNode(message_.binary_expr_id()));

    // Clone the binary expression, with the appropriate new operator.
    const ast::BinaryExpression* cloned_replacement;
    switch (static_cast<ast::BinaryOp>(message_.new_operator())) {
        case ast::BinaryOp::kAnd:
            cloned_replacement =
                clone_context->dst->And(clone_context->Clone(binary_expr_node->lhs),
                                        clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kOr:
            cloned_replacement =
                clone_context->dst->Or(clone_context->Clone(binary_expr_node->lhs),
                                       clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kXor:
            cloned_replacement =
                clone_context->dst->Xor(clone_context->Clone(binary_expr_node->lhs),
                                        clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kLogicalAnd:
            cloned_replacement =
                clone_context->dst->LogicalAnd(clone_context->Clone(binary_expr_node->lhs),
                                               clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kLogicalOr:
            cloned_replacement =
                clone_context->dst->LogicalOr(clone_context->Clone(binary_expr_node->lhs),
                                              clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kEqual:
            cloned_replacement =
                clone_context->dst->Equal(clone_context->Clone(binary_expr_node->lhs),
                                          clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kNotEqual:
            cloned_replacement =
                clone_context->dst->NotEqual(clone_context->Clone(binary_expr_node->lhs),
                                             clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kLessThan:
            cloned_replacement =
                clone_context->dst->LessThan(clone_context->Clone(binary_expr_node->lhs),
                                             clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kGreaterThan:
            cloned_replacement =
                clone_context->dst->GreaterThan(clone_context->Clone(binary_expr_node->lhs),
                                                clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kLessThanEqual:
            cloned_replacement =
                clone_context->dst->LessThanEqual(clone_context->Clone(binary_expr_node->lhs),
                                                  clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kGreaterThanEqual:
            cloned_replacement =
                clone_context->dst->GreaterThanEqual(clone_context->Clone(binary_expr_node->lhs),
                                                     clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kShiftLeft:
            cloned_replacement =
                clone_context->dst->Shl(clone_context->Clone(binary_expr_node->lhs),
                                        clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kShiftRight:
            cloned_replacement =
                clone_context->dst->Shr(clone_context->Clone(binary_expr_node->lhs),
                                        clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kAdd:
            cloned_replacement =
                clone_context->dst->Add(clone_context->Clone(binary_expr_node->lhs),
                                        clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kSubtract:
            cloned_replacement =
                clone_context->dst->Sub(clone_context->Clone(binary_expr_node->lhs),
                                        clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kMultiply:
            cloned_replacement =
                clone_context->dst->Mul(clone_context->Clone(binary_expr_node->lhs),
                                        clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kDivide:
            cloned_replacement =
                clone_context->dst->Div(clone_context->Clone(binary_expr_node->lhs),
                                        clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kModulo:
            cloned_replacement =
                clone_context->dst->Mod(clone_context->Clone(binary_expr_node->lhs),
                                        clone_context->Clone(binary_expr_node->rhs));
            break;
        case ast::BinaryOp::kNone:
            cloned_replacement = nullptr;
            assert(false && "Unreachable");
    }
    // Set things up so that the original binary expression will be replaced with
    // its clone, and update the id mapping.
    clone_context->Replace(binary_expr_node, cloned_replacement);
    new_node_id_map->Add(cloned_replacement, message_.binary_expr_id());
}

protobufs::Mutation MutationChangeBinaryOperator::ToMessage() const {
    protobufs::Mutation mutation;
    *mutation.mutable_change_binary_operator() = message_;
    return mutation;
}

}  // namespace tint::fuzzers::ast_fuzzer