// Copyright 2021 The Dawn & Tint Authors
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
//    contributors may be used to endorse or promote products derived from
//    this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "src/tint/lang/wgsl/ast/transform/binding_remapper.h"

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

#include "src/tint/lang/core/fluent_types.h"
#include "src/tint/lang/wgsl/ast/disable_validation_attribute.h"
#include "src/tint/lang/wgsl/program/clone_context.h"
#include "src/tint/lang/wgsl/program/program_builder.h"
#include "src/tint/lang/wgsl/resolver/resolve.h"
#include "src/tint/lang/wgsl/sem/function.h"
#include "src/tint/lang/wgsl/sem/variable.h"
#include "src/tint/utils/text/string.h"

TINT_INSTANTIATE_TYPEINFO(tint::ast::transform::BindingRemapper);
TINT_INSTANTIATE_TYPEINFO(tint::ast::transform::BindingRemapper::Remappings);

using namespace tint::core::fluent_types;  // NOLINT

namespace tint::ast::transform {

BindingRemapper::Remappings::Remappings(BindingPoints bp, AccessControls ac, bool may_collide)
    : binding_points(std::move(bp)),
      access_controls(std::move(ac)),
      allow_collisions(may_collide) {}

BindingRemapper::Remappings::Remappings(const Remappings&) = default;
BindingRemapper::Remappings::~Remappings() = default;

BindingRemapper::BindingRemapper() = default;
BindingRemapper::~BindingRemapper() = default;

Transform::ApplyResult BindingRemapper::Apply(const Program& src,
                                              const DataMap& inputs,
                                              DataMap&) const {
    ProgramBuilder b;
    program::CloneContext ctx{&b, &src, /* auto_clone_symbols */ true};

    auto* remappings = inputs.Get<Remappings>();
    if (!remappings) {
        b.Diagnostics().add_error(diag::System::Transform,
                                  "missing transform data for " + std::string(TypeInfo().name));
        return resolver::Resolve(b);
    }

    if (remappings->binding_points.empty() && remappings->access_controls.empty()) {
        return SkipTransform;
    }

    // A set of post-remapped binding points that need to be decorated with a
    // DisableValidationAttribute to disable binding-point-collision validation
    std::unordered_set<BindingPoint> add_collision_attr;

    if (remappings->allow_collisions) {
        // Scan for binding point collisions generated by this transform.
        // Populate all collisions in the `add_collision_attr` set.
        for (auto* func_ast : src.AST().Functions()) {
            if (!func_ast->IsEntryPoint()) {
                continue;
            }
            auto* func = src.Sem().Get(func_ast);
            std::unordered_map<BindingPoint, int> binding_point_counts;
            for (auto* global : func->TransitivelyReferencedGlobals()) {
                if (auto from = global->BindingPoint()) {
                    auto bp_it = remappings->binding_points.find(*from);
                    if (bp_it != remappings->binding_points.end()) {
                        // Remapped
                        BindingPoint to = bp_it->second;
                        if (binding_point_counts[to]++) {
                            add_collision_attr.emplace(to);
                        }
                    } else {
                        // No remapping
                        if (binding_point_counts[*from]++) {
                            add_collision_attr.emplace(*from);
                        }
                    }
                }
            }
        }
    }

    for (auto* var : src.AST().Globals<Var>()) {
        if (var->HasBindingPoint()) {
            auto* global_sem = src.Sem().Get<sem::GlobalVariable>(var);

            // The original binding point
            BindingPoint from = *global_sem->BindingPoint();

            // The binding point after remapping
            BindingPoint bp = from;

            // Replace any group or binding attributes.
            // Note: This has to be performed *before* remapping access controls, as
            // `ctx.Clone(var->attributes)` depend on these replacements.
            auto bp_it = remappings->binding_points.find(from);
            if (bp_it != remappings->binding_points.end()) {
                BindingPoint to = bp_it->second;
                auto* new_group = b.Group(AInt(to.group));
                auto* new_binding = b.Binding(AInt(to.binding));

                auto* old_group = GetAttribute<GroupAttribute>(var->attributes);
                auto* old_binding = GetAttribute<BindingAttribute>(var->attributes);

                ctx.Replace(old_group, new_group);
                ctx.Replace(old_binding, new_binding);
                bp = to;
            }

            // Replace any access controls.
            auto ac_it = remappings->access_controls.find(from);
            if (ac_it != remappings->access_controls.end()) {
                core::Access access = ac_it->second;
                if (access == core::Access::kUndefined) {
                    b.Diagnostics().add_error(diag::System::Transform,
                                              "invalid access mode (" +
                                                  std::to_string(static_cast<uint32_t>(access)) +
                                                  ")");
                    return resolver::Resolve(b);
                }
                auto* sem = src.Sem().Get(var);
                if (sem->AddressSpace() != core::AddressSpace::kStorage) {
                    b.Diagnostics().add_error(
                        diag::System::Transform,
                        "cannot apply access control to variable with address space " +
                            std::string(tint::ToString(sem->AddressSpace())));
                    return resolver::Resolve(b);
                }
                auto* ty = sem->Type()->UnwrapRef();
                auto inner_ty = CreateASTTypeFor(ctx, ty);
                auto* new_var =
                    b.create<Var>(ctx.Clone(var->source),                  // source
                                  b.Ident(ctx.Clone(var->name->symbol)),   // name
                                  inner_ty,                                // type
                                  ctx.Clone(var->declared_address_space),  // address space
                                  b.Expr(access),                          // access
                                  ctx.Clone(var->initializer),             // initializer
                                  ctx.Clone(var->attributes));             // attributes
                ctx.Replace(var, new_var);
            }

            // Add `DisableValidationAttribute`s if required
            if (add_collision_attr.count(bp)) {
                auto* attribute = b.Disable(DisabledValidation::kBindingPointCollision);
                ctx.InsertBefore(var->attributes, *var->attributes.begin(), attribute);
            }
        }
    }

    ctx.Clone();
    return resolver::Resolve(b);
}

}  // namespace tint::ast::transform