// Copyright (c) 2025 LunarG Inc.
//
// 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 <algorithm>
#include <map>
#include <set>
#include <vector>

#include "source/opt/pass.h"

namespace spvtools {
namespace opt {

// The canonicalize IDs pass is an optimization to improve compression of SPIR-V
// binary files via entropy reduction. It transforms SPIR-V to SPIR-V, remapping
// IDs. The resulting modules have an increased ID range (IDs are not as tightly
// packed around zero), but will compress better when multiple modules are
// compressed together, since the compressor's dictionary can find better cross
// module commonality. Remapping is accomplished via canonicalization. Thus,
// modules can be compressed one at a time with no loss of quality relative to
// operating on many modules at once.

// This pass should be run after most optimization passes except for
// --strip-debug because this pass will use OpName to canonicalize IDs. i.e. Run
// --strip-debug after this pass.

// This is a port of remap utility in glslang. There are great deal of magic
// numbers that are present throughout this code. The general goal is to replace
// the IDs with a hash value such that the distribution of IDs is deterministic
// and minimizes collisions. The magic numbers in the glslang version were
// chosen semi-arbitrarily and have been preserved in this port in order to
// maintain backward compatibility.

class CanonicalizeIdsPass : public Pass {
 public:
  CanonicalizeIdsPass() = default;
  virtual ~CanonicalizeIdsPass() = default;

  Pass::Status Process() override;

  const char* name() const override { return "canonicalize-ids"; }

 private:
  // Special values for IDs.
  static constexpr spv::Id unmapped_{spv::Id(-10000)};
  static constexpr spv::Id unused_{spv::Id(-10001)};

  // Scans the module for IDs and sets them to unmapped_.
  void ScanIds();

  // Functions to compute new IDs.
  void CanonicalizeTypeAndConst();
  spv::Id HashTypeAndConst(
      spv::Id const id) const;  // Helper for CanonicalizeTypeAndConst.
  void CanonicalizeNames();
  void CanonicalizeFunctions();
  spv::Id HashOpCode(Instruction const* const inst)
      const;  // Helper for CanonicalizeFunctions.
  void CanonicalizeRemainders();

  // Applies the new IDs.
  bool ApplyMap();

  // Methods to manage the bound field in header.
  spv::Id GetBound() const;  // All IDs must satisfy 0 < ID < bound.
  void UpdateBound();

  // Methods to map from old IDs to new IDs.
  spv::Id GetNewId(spv::Id const old_id) const { return new_id_[old_id]; }
  spv::Id SetNewId(spv::Id const old_id, spv::Id new_id);

  // Methods to manage claimed IDs.
  spv::Id ClaimNewId(spv::Id new_id);
  bool IsNewIdClaimed(spv::Id const new_id) const {
    return claimed_new_ids_.find(new_id) != claimed_new_ids_.end();
  }

  // Queries for old IDs.
  bool IsOldIdUnmapped(spv::Id const old_id) const {
    return GetNewId(old_id) == unmapped_;
  }
  bool IsOldIdUnused(spv::Id const old_id) const {
    return GetNewId(old_id) == unused_;
  }

  // Container to map old IDs to new IDs. e.g. new_id_[old_id] = new_id
  std::vector<spv::Id> new_id_;

  // IDs from the new ID space that have been claimed (faster than searching
  // through new_id_).
  std::set<spv::Id> claimed_new_ids_;

  // Helper functions for printing IDs (useful for debugging).
  std::string IdAsString(spv::Id const id) const;
  void PrintNewIds() const;

  // Containers to track IDs we want to canonicalize.
  std::vector<spv::Id> type_and_const_ids_;
  std::map<std::string, spv::Id> name_ids_;
  std::vector<spv::Id> function_ids_;
};

}  // namespace opt
}  // namespace spvtools