// 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 "src/tint/transform/num_workgroups_from_uniform.h"

#include <memory>
#include <string>
#include <unordered_set>
#include <utility>

#include "src/tint/program_builder.h"
#include "src/tint/sem/function.h"
#include "src/tint/transform/canonicalize_entry_point_io.h"
#include "src/tint/utils/hash.h"

TINT_INSTANTIATE_TYPEINFO(tint::transform::NumWorkgroupsFromUniform);
TINT_INSTANTIATE_TYPEINFO(tint::transform::NumWorkgroupsFromUniform::Config);

namespace tint::transform {
namespace {
/// Accessor describes the identifiers used in a member accessor that is being
/// used to retrieve the num_workgroups builtin from a parameter.
struct Accessor {
  Symbol param;
  Symbol member;

  /// Equality operator
  bool operator==(const Accessor& other) const {
    return param == other.param && member == other.member;
  }
  /// Hash function
  struct Hasher {
    size_t operator()(const Accessor& a) const {
      return utils::Hash(a.param, a.member);
    }
  };
};
}  // namespace

NumWorkgroupsFromUniform::NumWorkgroupsFromUniform() = default;
NumWorkgroupsFromUniform::~NumWorkgroupsFromUniform() = default;

bool NumWorkgroupsFromUniform::ShouldRun(const Program* program,
                                         const DataMap&) const {
  for (auto* node : program->ASTNodes().Objects()) {
    if (auto* attr = node->As<ast::BuiltinAttribute>()) {
      if (attr->builtin == ast::Builtin::kNumWorkgroups) {
        return true;
      }
    }
  }
  return false;
}

void NumWorkgroupsFromUniform::Run(CloneContext& ctx,
                                   const DataMap& inputs,
                                   DataMap&) const {
  auto* cfg = inputs.Get<Config>();
  if (cfg == nullptr) {
    ctx.dst->Diagnostics().add_error(
        diag::System::Transform,
        "missing transform data for " + std::string(TypeInfo().name));
    return;
  }

  const char* kNumWorkgroupsMemberName = "num_workgroups";

  // Find all entry point parameters that declare the num_workgroups builtin.
  std::unordered_set<Accessor, Accessor::Hasher> to_replace;
  for (auto* func : ctx.src->AST().Functions()) {
    // num_workgroups is only valid for compute stages.
    if (func->PipelineStage() != ast::PipelineStage::kCompute) {
      continue;
    }

    for (auto* param : ctx.src->Sem().Get(func)->Parameters()) {
      // Because the CanonicalizeEntryPointIO transform has been run, builtins
      // will only appear as struct members.
      auto* str = param->Type()->As<sem::Struct>();
      if (!str) {
        continue;
      }

      for (auto* member : str->Members()) {
        auto* builtin = ast::GetAttribute<ast::BuiltinAttribute>(
            member->Declaration()->attributes);
        if (!builtin || builtin->builtin != ast::Builtin::kNumWorkgroups) {
          continue;
        }

        // Capture the symbols that would be used to access this member, which
        // we will replace later. We currently have no way to get from the
        // parameter directly to the member accessor expressions that use it.
        to_replace.insert(
            {param->Declaration()->symbol, member->Declaration()->symbol});

        // Remove the struct member.
        // The CanonicalizeEntryPointIO transform will have generated this
        // struct uniquely for this particular entry point, so we know that
        // there will be no other uses of this struct in the module and that we
        // can safely modify it here.
        ctx.Remove(str->Declaration()->members, member->Declaration());

        // If this is the only member, remove the struct and parameter too.
        if (str->Members().size() == 1) {
          ctx.Remove(func->params, param->Declaration());
          ctx.Remove(ctx.src->AST().GlobalDeclarations(), str->Declaration());
        }
      }
    }
  }

  // Get (or create, on first call) the uniform buffer that will receive the
  // number of workgroups.
  const ast::Variable* num_workgroups_ubo = nullptr;
  auto get_ubo = [&]() {
    if (!num_workgroups_ubo) {
      auto* num_workgroups_struct = ctx.dst->Structure(
          ctx.dst->Sym(),
          {ctx.dst->Member(kNumWorkgroupsMemberName,
                           ctx.dst->ty.vec3(ctx.dst->ty.u32()))});
      num_workgroups_ubo = ctx.dst->Global(
          ctx.dst->Sym(), ctx.dst->ty.Of(num_workgroups_struct),
          ast::StorageClass::kUniform,
          ast::AttributeList{ctx.dst->GroupAndBinding(
              cfg->ubo_binding.group, cfg->ubo_binding.binding)});
    }
    return num_workgroups_ubo;
  };

  // Now replace all the places where the builtins are accessed with the value
  // loaded from the uniform buffer.
  for (auto* node : ctx.src->ASTNodes().Objects()) {
    auto* accessor = node->As<ast::MemberAccessorExpression>();
    if (!accessor) {
      continue;
    }
    auto* ident = accessor->structure->As<ast::IdentifierExpression>();
    if (!ident) {
      continue;
    }

    if (to_replace.count({ident->symbol, accessor->member->symbol})) {
      ctx.Replace(accessor, ctx.dst->MemberAccessor(get_ubo()->symbol,
                                                    kNumWorkgroupsMemberName));
    }
  }

  ctx.Clone();
}

NumWorkgroupsFromUniform::Config::Config(sem::BindingPoint ubo_bp)
    : ubo_binding(ubo_bp) {}
NumWorkgroupsFromUniform::Config::Config(const Config&) = default;
NumWorkgroupsFromUniform::Config::~Config() = default;

}  // namespace tint::transform