// Copyright (c) 2025 Google LLC

//

// 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 "source/opt/resolve_binding_conflicts_pass.h"

#include <algorithm>

#include <unordered_map>

#include <unordered_set>

#include <vector>

#include "source/opt/decoration_manager.h"

#include "source/opt/def_use_manager.h"

#include "source/opt/instruction.h"

#include "source/opt/ir_builder.h"

#include "source/opt/ir_context.h"

#include "spirv/unified1/spirv.h"

namespace spvtools {

namespace opt {

// A VarBindingInfo contains the binding information for a single resource

// variable.

//

// Exactly one such object is created per resource variable in the

// module. In particular, when a resource variable is statically used by

// more than one entry point, those entry points share the same VarBindingInfo

// object for that variable.

struct VarBindingInfo {

  const Instruction* const var;

  const uint32_t descriptor_set;

  Instruction* const binding_decoration;

  // Returns the binding number.

  uint32_t binding() const {

    return binding_decoration->GetSingleWordInOperand(2);

  }

  // Sets the binding number to 'b'.

  void updateBinding(uint32_t b) { binding_decoration->SetOperand(2, {b}); }

};

// The bindings in the same descriptor set that are used by an entry point.

using BindingList = std::vector<VarBindingInfo*>;

// A map from descriptor set number to the list of bindings in that descriptor

// set, as used by a particular entry point.

using DescriptorSets = std::unordered_map<uint32_t, BindingList>;

IRContext::Analysis ResolveBindingConflictsPass::GetPreservedAnalyses() {

  // All analyses are kept up to date.

  // At most this modifies the Binding numbers on variables.

  return IRContext::kAnalysisDefUse | IRContext::kAnalysisInstrToBlockMapping |

         IRContext::kAnalysisDecorations | IRContext::kAnalysisCombinators |

         IRContext::kAnalysisCFG | IRContext::kAnalysisDominatorAnalysis |

         IRContext::kAnalysisLoopAnalysis | IRContext::kAnalysisNameMap |

         IRContext::kAnalysisScalarEvolution |

         IRContext::kAnalysisRegisterPressure |

         IRContext::kAnalysisValueNumberTable |

         IRContext::kAnalysisStructuredCFG | IRContext::kAnalysisBuiltinVarId |

         IRContext::kAnalysisIdToFuncMapping | IRContext::kAnalysisConstants |

         IRContext::kAnalysisTypes | IRContext::kAnalysisDebugInfo |

         IRContext::kAnalysisLiveness;

}

// Orders variable binding info objects.

// * The binding number is most signficant;

// * Then a sampler-like object compares greater than non-sampler like object.

// * Otherwise compare based on variable ID.

// This provides a total order among bindings in a descriptor set for a valid

// Vulkan module.

bool Less(const VarBindingInfo* const lhs, const VarBindingInfo* const rhs) {

  if (lhs->binding() < rhs->binding()) return true;

  if (lhs->binding() > rhs->binding()) return false;

  // Examine types.

  // In valid Vulkan the only conflict can occur between

  // images and samplers.  We only care about a specific

  // comparison when one is a image-like thing and the other

  // is a sampler-like thing of the same shape.  So unwrap

  // types until we hit one of those two.

  auto* def_use_mgr = lhs->var->context()->get_def_use_mgr();

  // Returns the type found by iteratively following pointer pointee type,

  // or array element type.

  auto unwrap = [&def_use_mgr](Instruction* ty) {

    bool keep_going = true;

    do {

      switch (ty->opcode()) {

        case spv::Op::OpTypePointer:

          ty = def_use_mgr->GetDef(ty->GetSingleWordInOperand(1));

          break;

        case spv::Op::OpTypeArray:

        case spv::Op::OpTypeRuntimeArray:

          ty = def_use_mgr->GetDef(ty->GetSingleWordInOperand(0));

          break;

        default:

          keep_going = false;

          break;

      }

    } while (keep_going);

    return ty;

  };

  auto* lhs_ty = unwrap(def_use_mgr->GetDef(lhs->var->type_id()));

  auto* rhs_ty = unwrap(def_use_mgr->GetDef(rhs->var->type_id()));

  if (lhs_ty->opcode() == rhs_ty->opcode()) {

    // Pick based on variable ID.

    return lhs->var->result_id() < rhs->var->result_id();

  }

  // A sampler is always greater than an image.

  if (lhs_ty->opcode() == spv::Op::OpTypeSampler) {

    return false;

  }

  if (rhs_ty->opcode() == spv::Op::OpTypeSampler) {

    return true;

  }

  // Pick based on variable ID.

  return lhs->var->result_id() < rhs->var->result_id();

}

// Summarizes the caller-callee relationships between functions in a module.

class CallGraph {

 public:

  // Returns the list of all functions statically reachable from entry points,

  // where callees precede callers.

  const std::vector<uint32_t>& CalleesBeforeCallers() const {

    return visit_order_;

  }

  // Returns the list functions called from a given function.

  const std::unordered_set<uint32_t>& Callees(uint32_t caller) {

    return calls_[caller];

  }

  CallGraph(IRContext& context) {

    // Populate calls_.

    std::queue<uint32_t> callee_queue;

    for (const auto& fn : *context.module()) {

      auto& callees = calls_[fn.result_id()];

      context.AddCalls(&fn, &callee_queue);

      while (!callee_queue.empty()) {

        callees.insert(callee_queue.front());

        callee_queue.pop();

      }

    }

    // Perform depth-first search, starting from each entry point.

    // Populates visit_order_.

    for (const auto& ep : context.module()->entry_points()) {

      Visit(ep.GetSingleWordInOperand(1));

    }

  }

 private:

  // Visits a function, recursively visiting its callees. Adds this ID

  // to the visit_order after all callees have been visited.

  void Visit(uint32_t func_id) {

    if (visited_.count(func_id)) {

      return;

    }

    visited_.insert(func_id);

    for (auto callee_id : calls_[func_id]) {

      Visit(callee_id);

    }

    visit_order_.push_back(func_id);

  }

  // Maps the ID of a function to the IDs of functions it calls.

  std::unordered_map<uint32_t, std::unordered_set<uint32_t>> calls_;

  // IDs of visited functions;

  std::unordered_set<uint32_t> visited_;

  // IDs of functions, where callees precede callers.

  std::vector<uint32_t> visit_order_;

};

// Returns vector binding info for all resource variables in the module.

auto GetVarBindings(IRContext& context) {

  std::vector<VarBindingInfo> vars;

  auto* deco_mgr = context.get_decoration_mgr();

  for (auto& inst : context.module()->types_values()) {

    if (inst.opcode() == spv::Op::OpVariable) {

      Instruction* descriptor_set_deco = nullptr;

      Instruction* binding_deco = nullptr;

      for (auto* deco : deco_mgr->GetDecorationsFor(inst.result_id(), false)) {

        switch (static_cast<spv::Decoration>(deco->GetSingleWordInOperand(1))) {

          case spv::Decoration::DescriptorSet:

            assert(!descriptor_set_deco);

            descriptor_set_deco = deco;

            break;

          case spv::Decoration::Binding:

            assert(!binding_deco);

            binding_deco = deco;

            break;

          default:

            break;

        }

      }

      if (descriptor_set_deco && binding_deco) {

        vars.push_back({&inst, descriptor_set_deco->GetSingleWordInOperand(2),

                        binding_deco});

      }

    }

  }

  return vars;

}

// Merges the bindings from source into sink. Maintains order and uniqueness

// within a list of bindings.

void Merge(DescriptorSets& sink, const DescriptorSets& source) {

  for (auto index_and_bindings : source) {

    const uint32_t index = index_and_bindings.first;

    const BindingList& src1 = index_and_bindings.second;

    const BindingList& src2 = sink[index];

    BindingList merged;

    merged.resize(src1.size() + src2.size());

    auto merged_end = std::merge(src1.begin(), src1.end(), src2.begin(),

                                 src2.end(), merged.begin(), Less);

    auto unique_end = std::unique(merged.begin(), merged_end);

    merged.resize(unique_end - merged.begin());

    sink[index] = std::move(merged);

  }

}

// Resolves conflicts within this binding list, so the binding number on an

// item is at least one more than the binding number on the previous item.

// When this does not yet hold, increase the binding number on the second

// item in the pair. Returns true if any changes were applied.

bool ResolveConflicts(BindingList& bl) {

  bool changed = false;

  for (size_t i = 1; i < bl.size(); i++) {

    const auto prev_num = bl[i - 1]->binding();

    if (prev_num >= bl[i]->binding()) {

      bl[i]->updateBinding(prev_num + 1);

      changed = true;

    }

  }

  return changed;

}

Pass::Status ResolveBindingConflictsPass::Process() {

  // Assumes the descriptor set and binding decorations are not provided

  // via decoration groups.  Decoration groups were deprecated in SPIR-V 1.3

  // Revision 6.  I have not seen any compiler generate them. --dneto

  auto vars = GetVarBindings(*context());

  // Maps a function ID to the variables used directly or indirectly by the

  // function, organized into descriptor sets. Each descriptor set

  // consists of a BindingList of distinct variables.

  std::unordered_map<uint32_t, DescriptorSets> used_vars;

  // Determine variables directly used by functions.

  auto* def_use_mgr = context()->get_def_use_mgr();

  for (auto& var : vars) {

    std::unordered_set<uint32_t> visited_functions_for_var;

    def_use_mgr->ForEachUser(var.var, [&](Instruction* user) {

      if (auto* block = context()->get_instr_block(user)) {

        auto* fn = block->GetParent();

        assert(fn);

        const auto fn_id = fn->result_id();

        if (visited_functions_for_var.insert(fn_id).second) {

          used_vars[fn_id][var.descriptor_set].push_back(&var);

        }

      }

    });

  }

  // Sort within a descriptor set by binding number.

  for (auto& sets_for_fn : used_vars) {

    for (auto& ds : sets_for_fn.second) {

      BindingList& bindings = ds.second;

      std::stable_sort(bindings.begin(), bindings.end(), Less);

    }

  }

  // Propagate from callees to callers.

  CallGraph call_graph(*context());

  for (const uint32_t caller : call_graph.CalleesBeforeCallers()) {

    DescriptorSets& caller_ds = used_vars[caller];

    for (const uint32_t callee : call_graph.Callees(caller)) {

      Merge(caller_ds, used_vars[callee]);

    }

  }

  // At this point, the descriptor sets associated with each entry point

  // capture exactly the set of resource variables statically used

  // by the static call tree of that entry point.

  // Resolve conflicts.

  // VarBindingInfo objects may be shared between the bindings lists.

  // Updating a binding in one list can require updating another list later.

  // So repeat updates until settling.

  // The union of BindingLists across all entry points.

  std::vector<BindingList*> ep_bindings;

  for (auto& ep : context()->module()->entry_points()) {

    for (auto& ds : used_vars[ep.GetSingleWordInOperand(1)]) {

      BindingList& bindings = ds.second;

      ep_bindings.push_back(&bindings);

    }

  }

  bool modified = false;

  bool found_conflict;

  do {

    found_conflict = false;

    for (BindingList* bl : ep_bindings) {

      found_conflict |= ResolveConflicts(*bl);

    }

    modified |= found_conflict;

  } while (found_conflict);

  return modified ? Pass::Status::SuccessWithChange

                  : Pass::Status::SuccessWithoutChange;

}

}  // namespace opt

}  // namespace spvtools