// Copyright (c) 2018 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.

#ifndef LIBSPIRV_OPT_UPGRADE_MEMORY_MODEL_H_
#define LIBSPIRV_OPT_UPGRADE_MEMORY_MODEL_H_

#include <functional>
#include <tuple>

#include "pass.h"

namespace spvtools {
namespace opt {

// Hashing functor for the memoized result store.
struct CacheHash {
  size_t operator()(
      const std::pair<uint32_t, std::vector<uint32_t>>& item) const {
    std::u32string to_hash;
    to_hash.push_back(item.first);
    for (auto i : item.second) to_hash.push_back(i);
    return std::hash<std::u32string>()(to_hash);
  }
};

// Upgrades the memory model from Logical GLSL450 to Logical VulkanKHR.
//
// This pass remove deprecated decorations (Volatile and Coherent) and replaces
// them with new flags on individual instructions. It adds the Output storage
// class semantic to control barriers in tessellation control shaders that have
// an access to Output memory.
class UpgradeMemoryModel : public Pass {
 public:
  const char* name() const override { return "upgrade-memory-model"; }
  Status Process() override;

 private:
  // Used to indicate whether the operation performs an availability or
  // visibility operation.
  enum OperationType { kVisibility, kAvailability };

  // Used to indicate whether the instruction is a memory or image instruction.
  enum InstructionType { kMemory, kImage };

  // Modifies the OpMemoryModel to use VulkanKHR. Adds the Vulkan memory model
  // capability and extension.
  void UpgradeMemoryModelInstruction();

  // Upgrades memory, image and atomic instructions.
  // Memory and image instructions convert coherent and volatile decorations
  // into flags on the instruction.
  // Atomic memory semantics convert volatile decoration into flags on the
  // instruction.
  void UpgradeInstructions();

  // Upgrades memory and image operands for instructions that have them.
  void UpgradeMemoryAndImages();

  // Adds the volatile memory semantic if necessary.
  void UpgradeAtomics();

  // Returns whether |id| is coherent and/or volatile.
  std::tuple<bool, bool, spv::Scope> GetInstructionAttributes(uint32_t id);

  // Traces |inst| to determine if it is coherent and/or volatile.
  // |indices| tracks the access chain indices seen so far.
  std::pair<bool, bool> TraceInstruction(Instruction* inst,
                                         std::vector<uint32_t> indices,
                                         std::unordered_set<uint32_t>* visited);

  // Return true if |inst| is decorated with |decoration|.
  // If |inst| is decorated by member decorations then either |value| must
  // match the index or |value| must be a maximum allowable value. The max
  // value allows any element to match.
  bool HasDecoration(const Instruction* inst, uint32_t value,
                     spv::Decoration decoration);

  // Returns whether |type_id| indexed via |indices| is coherent and/or
  // volatile.
  std::pair<bool, bool> CheckType(uint32_t type_id,
                                  const std::vector<uint32_t>& indices);

  // Returns whether any type/element under |inst| is coherent and/or volatile.
  std::pair<bool, bool> CheckAllTypes(const Instruction* inst);

  // Modifies the flags of |inst| to include the new flags for the Vulkan
  // memory model. |operation_type| indicates whether flags should use
  // MakeVisible or MakeAvailable variants. |inst_type| indicates whether the
  // Pointer or Texel variants of flags should be used.
  void UpgradeFlags(Instruction* inst, uint32_t in_operand, bool is_coherent,
                    bool is_volatile, OperationType operation_type,
                    InstructionType inst_type);

  // Modifies the semantics at |in_operand| of |inst| to include the volatile
  // bit if |is_volatile| is true.
  void UpgradeSemantics(Instruction* inst, uint32_t in_operand,
                        bool is_volatile);

  // Returns the result id for a constant for |scope|.
  uint32_t GetScopeConstant(spv::Scope scope);

  // Returns the value of |index_inst|. |index_inst| must be an OpConstant of
  // integer type.g
  uint64_t GetIndexValue(Instruction* index_inst);

  // Removes coherent and volatile decorations.
  void CleanupDecorations();

  // For all tessellation control entry points, if there is an operation on
  // Output storage class, then all barriers are modified to include the
  // OutputMemoryKHR semantic.
  void UpgradeBarriers();

  // If the Vulkan memory model is specified, device scope actually means
  // device scope. The memory scope must be modified to be QueueFamilyKHR
  // scope.
  void UpgradeMemoryScope();

  // Returns true if |scope_id| is spv::Scope::Device.
  bool IsDeviceScope(uint32_t scope_id);

  // Upgrades GLSL.std.450 modf and frexp. Both instructions are replaced with
  // their struct versions. New extracts and a store are added in order to
  // facilitate adding memory model flags.
  void UpgradeExtInst(Instruction* modf);

  // Returns the number of words taken up by a memory access argument and its
  // implied operands.
  uint32_t MemoryAccessNumWords(uint32_t mask);

  // Caches the result of TraceInstruction. For a given result id and set of
  // indices, stores whether that combination is coherent and/or volatile.
  std::unordered_map<std::pair<uint32_t, std::vector<uint32_t>>,
                     std::pair<bool, bool>, CacheHash>
      cache_;
};
}  // namespace opt
}  // namespace spvtools
#endif  // LIBSPIRV_OPT_UPGRADE_MEMORY_MODEL_H_