// Copyright 2021 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 "fuzzers/tint_spirv_tools_fuzzer/spirv_opt_mutator.h"

#include <fstream>
#include <iostream>
#include <unordered_set>
#include <utility>

#include "fuzzers/tint_spirv_tools_fuzzer/util.h"
#include "spirv-tools/optimizer.hpp"

namespace tint {
namespace fuzzers {
namespace spvtools_fuzzer {

SpirvOptMutator::SpirvOptMutator(spv_target_env target_env,
                                 uint32_t seed,
                                 std::vector<uint32_t> binary,
                                 bool validate_after_each_opt,
                                 uint32_t opt_batch_size)
    : num_executions_(0),
      is_valid_(true),
      target_env_(target_env),
      original_binary_(std::move(binary)),
      seed_(seed),
      opt_passes_({"--combine-access-chains",
                   "--loop-unroll",
                   "--merge-blocks",
                   "--cfg-cleanup",
                   "--eliminate-dead-functions",
                   "--merge-return",
                   "--wrap-opkill",
                   "--eliminate-dead-code-aggressive",
                   "--if-conversion",
                   "--eliminate-local-single-store",
                   "--eliminate-local-single-block",
                   "--eliminate-dead-branches",
                   "--scalar-replacement=0",
                   "--eliminate-dead-inserts",
                   "--eliminate-dead-members",
                   "--simplify-instructions",
                   "--private-to-local",
                   "--ssa-rewrite",
                   "--ccp",
                   "--reduce-load-size",
                   "--vector-dce",
                   "--scalar-replacement=100",
                   "--inline-entry-points-exhaustive",
                   "--redundancy-elimination",
                   "--convert-local-access-chains",
                   "--copy-propagate-arrays",
                   "--fix-storage-class"}),
      optimized_binary_(),
      validate_after_each_opt_(validate_after_each_opt),
      opt_batch_size_(opt_batch_size),
      rng_(seed) {
  assert(spvtools::SpirvTools(target_env).Validate(original_binary_) &&
         "Initial binary is invalid");
  assert(!opt_passes_.empty() && "Must be at least one pass");
}

SpirvOptMutator::Result SpirvOptMutator::Mutate() {
  assert(is_valid_ && "The optimizer is not longer valid");

  const uint32_t kMaxNumExecutions = 100;
  const uint32_t kMaxNumStuck = 10;

  if (num_executions_ == kMaxNumExecutions) {
    // We've applied this mutator many times already. Indicate to the user that
    // it might be better to try a different mutator.
    return {Status::kLimitReached, false};
  }

  num_executions_++;

  // Get the input binary. If this is the first time we run this mutator, use
  // the `original_binary_`. Otherwise, one of the following will be true:
  // - the `optimized_binary_` is not empty.
  // - the previous call to the `Mutate` method returned `kStuck`.
  auto binary = num_executions_ == 1 ? original_binary_ : optimized_binary_;
  optimized_binary_.clear();

  assert(!binary.empty() && "Can't run the optimizer on an empty binary");

  // Number of times spirv-opt wasn't able to produce any new result.
  uint32_t num_stuck = 0;
  do {
    // Randomly select `opt_batch_size` optimization passes. If `opt_batch_size`
    // is equal to 0, we will use the number of passes equal to the number of
    // all available passes.
    auto num_of_passes = opt_batch_size_ ? opt_batch_size_ : opt_passes_.size();
    std::vector<std::string> passes;

    while (passes.size() < num_of_passes) {
      auto idx = std::uniform_int_distribution<size_t>(
          0, opt_passes_.size() - 1)(rng_);
      passes.push_back(opt_passes_[idx]);
    }

    // Run the `binary` into the `optimized_binary_`.
    spvtools::Optimizer optimizer(target_env_);
    optimizer.SetMessageConsumer(util::GetBufferMessageConsumer(&errors_));
    optimizer.SetValidateAfterAll(validate_after_each_opt_);
    optimizer.RegisterPassesFromFlags(passes);
    if (!optimizer.Run(binary.data(), binary.size(), &optimized_binary_)) {
      is_valid_ = false;
      return {Status::kInvalid, true};
    }
  } while (optimized_binary_.empty() && ++num_stuck < kMaxNumStuck);

  return {optimized_binary_.empty() ? Status::kStuck : Status::kComplete,
          !optimized_binary_.empty()};
}

std::vector<uint32_t> SpirvOptMutator::GetBinary() const {
  return optimized_binary_;
}

std::string SpirvOptMutator::GetErrors() const {
  return errors_.str();
}

void SpirvOptMutator::LogErrors(const std::string* path, uint32_t count) const {
  auto message = GetErrors();
  std::cout << count << " | SpirvOptMutator (seed: " << seed_ << ")"
            << std::endl;
  std::cout << message << std::endl;

  if (path) {
    auto prefix = *path + std::to_string(count);

    // Write errors to file.
    std::ofstream(prefix + ".opt.log") << "seed: " << seed_ << std::endl
                                       << message << std::endl;

    // Write the invalid SPIR-V binary.
    util::WriteBinary(prefix + ".opt.invalid.spv", optimized_binary_);

    // Write the original SPIR-V binary.
    util::WriteBinary(prefix + ".opt.original.spv", original_binary_);
  }
}

}  // namespace spvtools_fuzzer
}  // namespace fuzzers
}  // namespace tint