// 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 <string>
#include <unordered_set>
#include <vector>

#include "gtest/gtest.h"

#include "src/tint/fuzzers/tint_ast_fuzzer/mutator.h"
#include "src/tint/fuzzers/tint_ast_fuzzer/node_id_map.h"
#include "src/tint/program_builder.h"
#include "src/tint/reader/wgsl/parser.h"
#include "src/tint/writer/wgsl/generator.h"

namespace tint::fuzzers::ast_fuzzer {
namespace {

std::string OpToString(ast::BinaryOp op) {
  switch (op) {
    case ast::BinaryOp::kNone:
      assert(false && "Unreachable");
      return "";
    case ast::BinaryOp::kAnd:
      return "&";
    case ast::BinaryOp::kOr:
      return "|";
    case ast::BinaryOp::kXor:
      return "^";
    case ast::BinaryOp::kLogicalAnd:
      return "&&";
    case ast::BinaryOp::kLogicalOr:
      return "||";
    case ast::BinaryOp::kEqual:
      return "==";
    case ast::BinaryOp::kNotEqual:
      return "!=";
    case ast::BinaryOp::kLessThan:
      return "<";
    case ast::BinaryOp::kGreaterThan:
      return ">";
    case ast::BinaryOp::kLessThanEqual:
      return "<=";
    case ast::BinaryOp::kGreaterThanEqual:
      return ">=";
    case ast::BinaryOp::kShiftLeft:
      return "<<";
    case ast::BinaryOp::kShiftRight:
      return ">>";
    case ast::BinaryOp::kAdd:
      return "+";
    case ast::BinaryOp::kSubtract:
      return "-";
    case ast::BinaryOp::kMultiply:
      return "*";
    case ast::BinaryOp::kDivide:
      return "/";
    case ast::BinaryOp::kModulo:
      return "%";
  }
}

TEST(ChangeBinaryOperatorTest, NotApplicable_Simple) {
  std::string content = R"(
    fn main() {
      let a : i32 = 1 + 2;
    }
  )";
  Source::File file("test.wgsl", content);
  auto program = reader::wgsl::Parse(&file);
  ASSERT_TRUE(program.IsValid()) << program.Diagnostics().str();

  NodeIdMap node_id_map(program);

  const auto& main_fn_stmts = program.AST().Functions()[0]->body->statements;

  const auto* a_var =
      main_fn_stmts[0]->As<ast::VariableDeclStatement>()->variable;
  ASSERT_NE(a_var, nullptr);

  auto a_var_id = node_id_map.GetId(a_var);

  const auto* sum_expr = a_var->constructor->As<ast::BinaryExpression>();
  ASSERT_NE(sum_expr, nullptr);

  auto sum_expr_id = node_id_map.GetId(sum_expr);
  ASSERT_NE(sum_expr_id, 0);

  // binary_expr_id is invalid.
  EXPECT_FALSE(MutationChangeBinaryOperator(0, ast::BinaryOp::kSubtract)
                   .IsApplicable(program, node_id_map));

  // binary_expr_id is not a binary expression.
  EXPECT_FALSE(MutationChangeBinaryOperator(a_var_id, ast::BinaryOp::kSubtract)
                   .IsApplicable(program, node_id_map));

  // new_operator is applicable to the argument types.
  EXPECT_FALSE(MutationChangeBinaryOperator(0, ast::BinaryOp::kLogicalAnd)
                   .IsApplicable(program, node_id_map));

  // new_operator does not have the right result type.
  EXPECT_FALSE(MutationChangeBinaryOperator(0, ast::BinaryOp::kLessThan)
                   .IsApplicable(program, node_id_map));
}

TEST(ChangeBinaryOperatorTest, Applicable_Simple) {
  std::string shader = R"(fn main() {
  let a : i32 = (1 + 2);
}
)";
  Source::File file("test.wgsl", shader);
  auto program = reader::wgsl::Parse(&file);
  ASSERT_TRUE(program.IsValid()) << program.Diagnostics().str();

  NodeIdMap node_id_map(program);

  const auto& main_fn_stmts = program.AST().Functions()[0]->body->statements;

  const auto* a_var =
      main_fn_stmts[0]->As<ast::VariableDeclStatement>()->variable;
  ASSERT_NE(a_var, nullptr);

  const auto* sum_expr = a_var->constructor->As<ast::BinaryExpression>();
  ASSERT_NE(sum_expr, nullptr);

  auto sum_expr_id = node_id_map.GetId(sum_expr);
  ASSERT_NE(sum_expr_id, 0);

  ASSERT_TRUE(MaybeApplyMutation(
      program,
      MutationChangeBinaryOperator(sum_expr_id, ast::BinaryOp::kSubtract),
      node_id_map, &program, &node_id_map, nullptr));
  ASSERT_TRUE(program.IsValid()) << program.Diagnostics().str();

  writer::wgsl::Options options;
  auto result = writer::wgsl::Generate(&program, options);
  ASSERT_TRUE(result.success) << result.error;

  std::string expected_shader = R"(fn main() {
  let a : i32 = (1 - 2);
}
)";
  ASSERT_EQ(expected_shader, result.wgsl);
}

void CheckMutations(
    const std::string& lhs_type,
    const std::string& rhs_type,
    const std::string& result_type,
    ast::BinaryOp original_operator,
    const std::unordered_set<ast::BinaryOp>& allowed_replacement_operators) {
  std::stringstream shader;
  shader << "fn foo(a : " << lhs_type << ", b : " << rhs_type + ") {\n"
         << "  let r : " << result_type
         << " = (a " + OpToString(original_operator) << " b);\n}\n";

  const std::vector<ast::BinaryOp> all_operators = {
      ast::BinaryOp::kAnd,
      ast::BinaryOp::kOr,
      ast::BinaryOp::kXor,
      ast::BinaryOp::kLogicalAnd,
      ast::BinaryOp::kLogicalOr,
      ast::BinaryOp::kEqual,
      ast::BinaryOp::kNotEqual,
      ast::BinaryOp::kLessThan,
      ast::BinaryOp::kGreaterThan,
      ast::BinaryOp::kLessThanEqual,
      ast::BinaryOp::kGreaterThanEqual,
      ast::BinaryOp::kShiftLeft,
      ast::BinaryOp::kShiftRight,
      ast::BinaryOp::kAdd,
      ast::BinaryOp::kSubtract,
      ast::BinaryOp::kMultiply,
      ast::BinaryOp::kDivide,
      ast::BinaryOp::kModulo};

  for (auto new_operator : all_operators) {
    Source::File file("test.wgsl", shader.str());
    auto program = reader::wgsl::Parse(&file);
    ASSERT_TRUE(program.IsValid()) << program.Diagnostics().str();

    NodeIdMap node_id_map(program);

    const auto& stmts = program.AST().Functions()[0]->body->statements;

    const auto* r_var = stmts[0]->As<ast::VariableDeclStatement>()->variable;
    ASSERT_NE(r_var, nullptr);

    const auto* binary_expr = r_var->constructor->As<ast::BinaryExpression>();
    ASSERT_NE(binary_expr, nullptr);

    auto binary_expr_id = node_id_map.GetId(binary_expr);
    ASSERT_NE(binary_expr_id, 0);

    MutationChangeBinaryOperator mutation(binary_expr_id, new_operator);

    std::stringstream expected_shader;
    expected_shader << "fn foo(a : " << lhs_type << ", b : " << rhs_type
                    << ") {\n"
                    << "  let r : " << result_type << " = (a "
                    << OpToString(new_operator) << " b);\n}\n";

    if (allowed_replacement_operators.count(new_operator) == 0) {
      ASSERT_FALSE(mutation.IsApplicable(program, node_id_map));
      if (new_operator != binary_expr->op) {
        Source::File invalid_file("test.wgsl", expected_shader.str());
        auto invalid_program = reader::wgsl::Parse(&invalid_file);
        ASSERT_FALSE(invalid_program.IsValid()) << program.Diagnostics().str();
      }
    } else {
      ASSERT_TRUE(MaybeApplyMutation(program, mutation, node_id_map, &program,
                                     &node_id_map, nullptr));
      ASSERT_TRUE(program.IsValid()) << program.Diagnostics().str();

      writer::wgsl::Options options;
      auto result = writer::wgsl::Generate(&program, options);
      ASSERT_TRUE(result.success) << result.error;

      ASSERT_EQ(expected_shader.str(), result.wgsl);
    }
  }
}

TEST(ChangeBinaryOperatorTest, AddSubtract) {
  for (auto op : {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract}) {
    const ast::BinaryOp other_op = op == ast::BinaryOp::kAdd
                                       ? ast::BinaryOp::kSubtract
                                       : ast::BinaryOp::kAdd;
    for (std::string type : {"i32", "vec2<i32>", "vec3<i32>", "vec4<i32>"}) {
      CheckMutations(
          type, type, type, op,
          {other_op, ast::BinaryOp::kMultiply, ast::BinaryOp::kDivide,
           ast::BinaryOp::kModulo, ast::BinaryOp::kAnd, ast::BinaryOp::kOr,
           ast::BinaryOp::kXor});
    }
    for (std::string type : {"u32", "vec2<u32>", "vec3<u32>", "vec4<u32>"}) {
      CheckMutations(
          type, type, type, op,
          {other_op, ast::BinaryOp::kMultiply, ast::BinaryOp::kDivide,
           ast::BinaryOp::kModulo, ast::BinaryOp::kAnd, ast::BinaryOp::kOr,
           ast::BinaryOp::kXor, ast::BinaryOp::kShiftLeft,
           ast::BinaryOp::kShiftRight});
    }
    for (std::string type : {"f32", "vec2<f32>", "vec3<f32>", "vec4<f32>"}) {
      CheckMutations(type, type, type, op,
                     {other_op, ast::BinaryOp::kMultiply,
                      ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
    }
    for (std::string vector_type : {"vec2<i32>", "vec3<i32>", "vec4<i32>"}) {
      std::string scalar_type = "i32";
      CheckMutations(vector_type, scalar_type, vector_type, op,
                     {other_op, ast::BinaryOp::kMultiply,
                      ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
      CheckMutations(scalar_type, vector_type, vector_type, op,
                     {other_op, ast::BinaryOp::kMultiply,
                      ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
    }
    for (std::string vector_type : {"vec2<u32>", "vec3<u32>", "vec4<u32>"}) {
      std::string scalar_type = "u32";
      CheckMutations(vector_type, scalar_type, vector_type, op,
                     {other_op, ast::BinaryOp::kMultiply,
                      ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
      CheckMutations(scalar_type, vector_type, vector_type, op,
                     {other_op, ast::BinaryOp::kMultiply,
                      ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
    }
    for (std::string vector_type : {"vec2<f32>", "vec3<f32>", "vec4<f32>"}) {
      std::string scalar_type = "f32";
      CheckMutations(vector_type, scalar_type, vector_type, op,
                     {other_op, ast::BinaryOp::kMultiply,
                      ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
      CheckMutations(scalar_type, vector_type, vector_type, op,
                     {other_op, ast::BinaryOp::kMultiply,
                      ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
    }
    for (std::string square_matrix_type :
         {"mat2x2<f32>", "mat3x3<f32>", "mat4x4<f32>"}) {
      CheckMutations(square_matrix_type, square_matrix_type, square_matrix_type,
                     op, {other_op, ast::BinaryOp::kMultiply});
    }
    for (std::string non_square_matrix_type :
         {"mat2x3<f32>", "mat2x4<f32>", "mat3x2<f32>", "mat3x4<f32>",
          "mat4x2<f32>", "mat4x3<f32>"}) {
      CheckMutations(non_square_matrix_type, non_square_matrix_type,
                     non_square_matrix_type, op, {other_op});
    }
  }
}

TEST(ChangeBinaryOperatorTest, Mul) {
  for (std::string type : {"i32", "vec2<i32>", "vec3<i32>", "vec4<i32>"}) {
    CheckMutations(
        type, type, type, ast::BinaryOp::kMultiply,
        {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract, ast::BinaryOp::kDivide,
         ast::BinaryOp::kModulo, ast::BinaryOp::kAnd, ast::BinaryOp::kOr,
         ast::BinaryOp::kXor});
  }
  for (std::string type : {"u32", "vec2<u32>", "vec3<u32>", "vec4<u32>"}) {
    CheckMutations(
        type, type, type, ast::BinaryOp::kMultiply,
        {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract, ast::BinaryOp::kDivide,
         ast::BinaryOp::kModulo, ast::BinaryOp::kAnd, ast::BinaryOp::kOr,
         ast::BinaryOp::kXor, ast::BinaryOp::kShiftLeft,
         ast::BinaryOp::kShiftRight});
  }
  for (std::string type : {"f32", "vec2<f32>", "vec3<f32>", "vec4<f32>"}) {
    CheckMutations(type, type, type, ast::BinaryOp::kMultiply,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
  }
  for (std::string vector_type : {"vec2<i32>", "vec3<i32>", "vec4<i32>"}) {
    std::string scalar_type = "i32";
    CheckMutations(vector_type, scalar_type, vector_type,
                   ast::BinaryOp::kMultiply,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
    CheckMutations(scalar_type, vector_type, vector_type,
                   ast::BinaryOp::kMultiply,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
  }
  for (std::string vector_type : {"vec2<u32>", "vec3<u32>", "vec4<u32>"}) {
    std::string scalar_type = "u32";
    CheckMutations(vector_type, scalar_type, vector_type,
                   ast::BinaryOp::kMultiply,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
    CheckMutations(scalar_type, vector_type, vector_type,
                   ast::BinaryOp::kMultiply,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
  }
  for (std::string vector_type : {"vec2<f32>", "vec3<f32>", "vec4<f32>"}) {
    std::string scalar_type = "f32";
    CheckMutations(vector_type, scalar_type, vector_type,
                   ast::BinaryOp::kMultiply,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
    CheckMutations(scalar_type, vector_type, vector_type,
                   ast::BinaryOp::kMultiply,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kDivide, ast::BinaryOp::kModulo});
  }
  for (std::string square_matrix_type :
       {"mat2x2<f32>", "mat3x3<f32>", "mat4x4<f32>"}) {
    CheckMutations(square_matrix_type, square_matrix_type, square_matrix_type,
                   ast::BinaryOp::kMultiply,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract});
  }

  CheckMutations("vec2<f32>", "mat2x2<f32>", "vec2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("vec2<f32>", "mat3x2<f32>", "vec3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("vec2<f32>", "mat4x2<f32>", "vec4<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("mat2x2<f32>", "vec2<f32>", "vec2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat2x2<f32>", "mat3x2<f32>", "mat3x2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat2x2<f32>", "mat4x2<f32>", "mat4x2<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("mat2x3<f32>", "vec2<f32>", "vec3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat2x3<f32>", "mat2x2<f32>", "mat2x3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat2x3<f32>", "mat3x2<f32>", "mat3x3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat2x3<f32>", "mat4x2<f32>", "mat4x3<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("mat2x4<f32>", "vec2<f32>", "vec4<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat2x4<f32>", "mat2x2<f32>", "mat2x4<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat2x4<f32>", "mat3x2<f32>", "mat3x4<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat2x4<f32>", "mat4x2<f32>", "mat4x4<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("vec3<f32>", "mat2x3<f32>", "vec2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("vec3<f32>", "mat3x3<f32>", "vec3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("vec3<f32>", "mat4x3<f32>", "vec4<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("mat3x2<f32>", "vec3<f32>", "vec2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat3x2<f32>", "mat2x3<f32>", "mat2x2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat3x2<f32>", "mat3x3<f32>", "mat3x2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat3x2<f32>", "mat4x3<f32>", "mat4x2<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("mat3x3<f32>", "vec3<f32>", "vec3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat3x3<f32>", "mat2x3<f32>", "mat2x3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat3x3<f32>", "mat4x3<f32>", "mat4x3<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("mat3x4<f32>", "vec3<f32>", "vec4<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat3x4<f32>", "mat2x3<f32>", "mat2x4<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat3x4<f32>", "mat3x3<f32>", "mat3x4<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat3x4<f32>", "mat4x3<f32>", "mat4x4<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("vec4<f32>", "mat2x4<f32>", "vec2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("vec4<f32>", "mat3x4<f32>", "vec3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("vec4<f32>", "mat4x4<f32>", "vec4<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("mat4x2<f32>", "vec4<f32>", "vec2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat4x2<f32>", "mat2x4<f32>", "mat2x2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat4x2<f32>", "mat3x4<f32>", "mat3x2<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat4x2<f32>", "mat4x4<f32>", "mat4x2<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("mat4x3<f32>", "vec4<f32>", "vec3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat4x3<f32>", "mat2x4<f32>", "mat2x3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat4x3<f32>", "mat3x4<f32>", "mat3x3<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat4x3<f32>", "mat4x4<f32>", "mat4x3<f32>",
                 ast::BinaryOp::kMultiply, {});

  CheckMutations("mat4x4<f32>", "vec4<f32>", "vec4<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat4x4<f32>", "mat2x4<f32>", "mat2x4<f32>",
                 ast::BinaryOp::kMultiply, {});
  CheckMutations("mat4x4<f32>", "mat3x4<f32>", "mat3x4<f32>",
                 ast::BinaryOp::kMultiply, {});
}

TEST(ChangeBinaryOperatorTest, DivideAndModulo) {
  for (std::string type : {"i32", "vec2<i32>", "vec3<i32>", "vec4<i32>"}) {
    CheckMutations(
        type, type, type, ast::BinaryOp::kDivide,
        {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
         ast::BinaryOp::kMultiply, ast::BinaryOp::kModulo, ast::BinaryOp::kAnd,
         ast::BinaryOp::kOr, ast::BinaryOp::kXor});
  }
  for (std::string type : {"u32", "vec2<u32>", "vec3<u32>", "vec4<u32>"}) {
    CheckMutations(
        type, type, type, ast::BinaryOp::kDivide,
        {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
         ast::BinaryOp::kMultiply, ast::BinaryOp::kModulo, ast::BinaryOp::kAnd,
         ast::BinaryOp::kOr, ast::BinaryOp::kXor, ast::BinaryOp::kShiftLeft,
         ast::BinaryOp::kShiftRight});
  }
  for (std::string type : {"f32", "vec2<f32>", "vec3<f32>", "vec4<f32>"}) {
    CheckMutations(type, type, type, ast::BinaryOp::kDivide,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kModulo});
  }
  for (std::string vector_type : {"vec2<i32>", "vec3<i32>", "vec4<i32>"}) {
    std::string scalar_type = "i32";
    CheckMutations(vector_type, scalar_type, vector_type,
                   ast::BinaryOp::kDivide,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kModulo});
    CheckMutations(scalar_type, vector_type, vector_type,
                   ast::BinaryOp::kDivide,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kModulo});
  }
  for (std::string vector_type : {"vec2<u32>", "vec3<u32>", "vec4<u32>"}) {
    std::string scalar_type = "u32";
    CheckMutations(vector_type, scalar_type, vector_type,
                   ast::BinaryOp::kDivide,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kModulo});
    CheckMutations(scalar_type, vector_type, vector_type,
                   ast::BinaryOp::kDivide,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kModulo});
  }
  for (std::string vector_type : {"vec2<f32>", "vec3<f32>", "vec4<f32>"}) {
    std::string scalar_type = "f32";
    CheckMutations(vector_type, scalar_type, vector_type,
                   ast::BinaryOp::kDivide,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kModulo});
    CheckMutations(scalar_type, vector_type, vector_type,
                   ast::BinaryOp::kDivide,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kModulo});
  }
  for (std::string type : {"i32", "vec2<i32>", "vec3<i32>", "vec4<i32>"}) {
    CheckMutations(
        type, type, type, ast::BinaryOp::kModulo,
        {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
         ast::BinaryOp::kMultiply, ast::BinaryOp::kDivide, ast::BinaryOp::kAnd,
         ast::BinaryOp::kOr, ast::BinaryOp::kXor});
  }
  for (std::string type : {"u32", "vec2<u32>", "vec3<u32>", "vec4<u32>"}) {
    CheckMutations(
        type, type, type, ast::BinaryOp::kModulo,
        {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
         ast::BinaryOp::kMultiply, ast::BinaryOp::kDivide, ast::BinaryOp::kAnd,
         ast::BinaryOp::kOr, ast::BinaryOp::kXor, ast::BinaryOp::kShiftLeft,
         ast::BinaryOp::kShiftRight});
  }
  for (std::string type : {"f32", "vec2<f32>", "vec3<f32>", "vec4<f32>"}) {
    CheckMutations(type, type, type, ast::BinaryOp::kModulo,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kDivide});
  }
  for (std::string vector_type : {"vec2<i32>", "vec3<i32>", "vec4<i32>"}) {
    std::string scalar_type = "i32";
    CheckMutations(vector_type, scalar_type, vector_type,
                   ast::BinaryOp::kModulo,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kDivide});
    CheckMutations(scalar_type, vector_type, vector_type,
                   ast::BinaryOp::kModulo,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kDivide});
  }
  for (std::string vector_type : {"vec2<u32>", "vec3<u32>", "vec4<u32>"}) {
    std::string scalar_type = "u32";
    CheckMutations(vector_type, scalar_type, vector_type,
                   ast::BinaryOp::kModulo,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kDivide});
    CheckMutations(scalar_type, vector_type, vector_type,
                   ast::BinaryOp::kModulo,
                   {ast::BinaryOp::kAdd, ast::BinaryOp::kSubtract,
                    ast::BinaryOp::kMultiply, ast::BinaryOp::kDivide});
  }
}

TEST(ChangeBinaryOperatorTest, AndOrXor) {
  for (auto op :
       {ast::BinaryOp::kAnd, ast::BinaryOp::kOr, ast::BinaryOp::kXor}) {
    std::unordered_set<ast::BinaryOp> allowed_replacement_operators_signed{
        ast::BinaryOp::kAdd,      ast::BinaryOp::kSubtract,
        ast::BinaryOp::kMultiply, ast::BinaryOp::kDivide,
        ast::BinaryOp::kModulo,   ast::BinaryOp::kAnd,
        ast::BinaryOp::kOr,       ast::BinaryOp::kXor};
    allowed_replacement_operators_signed.erase(op);
    for (std::string type : {"i32", "vec2<i32>", "vec3<i32>", "vec4<i32>"}) {
      CheckMutations(type, type, type, op,
                     allowed_replacement_operators_signed);
    }
    std::unordered_set<ast::BinaryOp> allowed_replacement_operators_unsigned{
        ast::BinaryOp::kAdd,        ast::BinaryOp::kSubtract,
        ast::BinaryOp::kMultiply,   ast::BinaryOp::kDivide,
        ast::BinaryOp::kModulo,     ast::BinaryOp::kShiftLeft,
        ast::BinaryOp::kShiftRight, ast::BinaryOp::kAnd,
        ast::BinaryOp::kOr,         ast::BinaryOp::kXor};
    allowed_replacement_operators_unsigned.erase(op);
    for (std::string type : {"u32", "vec2<u32>", "vec3<u32>", "vec4<u32>"}) {
      CheckMutations(type, type, type, op,
                     allowed_replacement_operators_unsigned);
    }
    if (op != ast::BinaryOp::kXor) {
      for (std::string type :
           {"bool", "vec2<bool>", "vec3<bool>", "vec4<bool>"}) {
        std::unordered_set<ast::BinaryOp> allowed_replacement_operators_bool{
            ast::BinaryOp::kAnd, ast::BinaryOp::kOr, ast::BinaryOp::kEqual,
            ast::BinaryOp::kNotEqual};
        allowed_replacement_operators_bool.erase(op);
        if (type == "bool") {
          allowed_replacement_operators_bool.insert(ast::BinaryOp::kLogicalAnd);
          allowed_replacement_operators_bool.insert(ast::BinaryOp::kLogicalOr);
        }
        CheckMutations(type, type, type, op,
                       allowed_replacement_operators_bool);
      }
    }
  }
}

TEST(ChangeBinaryOperatorTest, EqualNotEqual) {
  for (auto op : {ast::BinaryOp::kEqual, ast::BinaryOp::kNotEqual}) {
    for (std::string element_type : {"i32", "u32", "f32"}) {
      for (size_t element_count = 1; element_count <= 4; element_count++) {
        std::stringstream argument_type;
        std::stringstream result_type;
        if (element_count == 1) {
          argument_type << element_type;
          result_type << "bool";
        } else {
          argument_type << "vec" << element_count << "<" << element_type << ">";
          result_type << "vec" << element_count << "<bool>";
        }
        std::unordered_set<ast::BinaryOp> allowed_replacement_operators{
            ast::BinaryOp::kLessThan,    ast::BinaryOp::kLessThanEqual,
            ast::BinaryOp::kGreaterThan, ast::BinaryOp::kGreaterThanEqual,
            ast::BinaryOp::kEqual,       ast::BinaryOp::kNotEqual};
        allowed_replacement_operators.erase(op);
        CheckMutations(argument_type.str(), argument_type.str(),
                       result_type.str(), op, allowed_replacement_operators);
      }
    }
    {
      std::unordered_set<ast::BinaryOp> allowed_replacement_operators{
          ast::BinaryOp::kLogicalAnd, ast::BinaryOp::kLogicalOr,
          ast::BinaryOp::kAnd,        ast::BinaryOp::kOr,
          ast::BinaryOp::kEqual,      ast::BinaryOp::kNotEqual};
      allowed_replacement_operators.erase(op);
      CheckMutations("bool", "bool", "bool", op, allowed_replacement_operators);
    }
    for (size_t element_count = 2; element_count <= 4; element_count++) {
      std::stringstream argument_and_result_type;
      argument_and_result_type << "vec" << element_count << "<bool>";
      std::unordered_set<ast::BinaryOp> allowed_replacement_operators{
          ast::BinaryOp::kAnd, ast::BinaryOp::kOr, ast::BinaryOp::kEqual,
          ast::BinaryOp::kNotEqual};
      allowed_replacement_operators.erase(op);
      CheckMutations(
          argument_and_result_type.str(), argument_and_result_type.str(),
          argument_and_result_type.str(), op, allowed_replacement_operators);
    }
  }
}

TEST(ChangeBinaryOperatorTest,
     LessThanLessThanEqualGreaterThanGreaterThanEqual) {
  for (auto op :
       {ast::BinaryOp::kLessThan, ast::BinaryOp::kLessThanEqual,
        ast::BinaryOp::kGreaterThan, ast::BinaryOp::kGreaterThanEqual}) {
    for (std::string element_type : {"i32", "u32", "f32"}) {
      for (size_t element_count = 1; element_count <= 4; element_count++) {
        std::stringstream argument_type;
        std::stringstream result_type;
        if (element_count == 1) {
          argument_type << element_type;
          result_type << "bool";
        } else {
          argument_type << "vec" << element_count << "<" << element_type << ">";
          result_type << "vec" << element_count << "<bool>";
        }
        std::unordered_set<ast::BinaryOp> allowed_replacement_operators{
            ast::BinaryOp::kLessThan,    ast::BinaryOp::kLessThanEqual,
            ast::BinaryOp::kGreaterThan, ast::BinaryOp::kGreaterThanEqual,
            ast::BinaryOp::kEqual,       ast::BinaryOp::kNotEqual};
        allowed_replacement_operators.erase(op);
        CheckMutations(argument_type.str(), argument_type.str(),
                       result_type.str(), op, allowed_replacement_operators);
      }
    }
  }
}

TEST(ChangeBinaryOperatorTest, LogicalAndLogicalOr) {
  for (auto op : {ast::BinaryOp::kLogicalAnd, ast::BinaryOp::kLogicalOr}) {
    std::unordered_set<ast::BinaryOp> allowed_replacement_operators{
        ast::BinaryOp::kLogicalAnd, ast::BinaryOp::kLogicalOr,
        ast::BinaryOp::kAnd,        ast::BinaryOp::kOr,
        ast::BinaryOp::kEqual,      ast::BinaryOp::kNotEqual};
    allowed_replacement_operators.erase(op);
    CheckMutations("bool", "bool", "bool", op, allowed_replacement_operators);
  }
}

TEST(ChangeBinaryOperatorTest, ShiftLeftShiftRight) {
  for (auto op : {ast::BinaryOp::kShiftLeft, ast::BinaryOp::kShiftRight}) {
    for (std::string lhs_element_type : {"i32", "u32"}) {
      for (size_t element_count = 1; element_count <= 4; element_count++) {
        std::stringstream lhs_and_result_type;
        std::stringstream rhs_type;
        if (element_count == 1) {
          lhs_and_result_type << lhs_element_type;
          rhs_type << "u32";
        } else {
          lhs_and_result_type << "vec" << element_count << "<"
                              << lhs_element_type << ">";
          rhs_type << "vec" << element_count << "<u32>";
        }
        std::unordered_set<ast::BinaryOp> allowed_replacement_operators{
            ast::BinaryOp::kShiftLeft, ast::BinaryOp::kShiftRight};
        allowed_replacement_operators.erase(op);
        if (lhs_element_type == "u32") {
          allowed_replacement_operators.insert(ast::BinaryOp::kAdd);
          allowed_replacement_operators.insert(ast::BinaryOp::kSubtract);
          allowed_replacement_operators.insert(ast::BinaryOp::kMultiply);
          allowed_replacement_operators.insert(ast::BinaryOp::kDivide);
          allowed_replacement_operators.insert(ast::BinaryOp::kModulo);
          allowed_replacement_operators.insert(ast::BinaryOp::kAnd);
          allowed_replacement_operators.insert(ast::BinaryOp::kOr);
          allowed_replacement_operators.insert(ast::BinaryOp::kXor);
        }
        CheckMutations(lhs_and_result_type.str(), rhs_type.str(),
                       lhs_and_result_type.str(), op,
                       allowed_replacement_operators);
      }
    }
  }
}

}  // namespace
}  // namespace tint::fuzzers::ast_fuzzer