// Copyright (c) 2022 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 SOURCE_OPT_INTERFACE_VAR_SROA_H_
#define SOURCE_OPT_INTERFACE_VAR_SROA_H_

#include <optional>
#include <unordered_set>

#include "source/opt/pass.h"

namespace spvtools {
namespace opt {

// See optimizer.hpp for documentation.
//
// Note that the current implementation of this pass covers only store, load,
// access chain instructions for the interface variables. Supporting other types
// of instructions is a future work.
class InterfaceVariableScalarReplacement : public Pass {
 public:
  InterfaceVariableScalarReplacement() {}

  const char* name() const override {
    return "interface-variable-scalar-replacement";
  }
  Status Process() override;

  IRContext::Analysis GetPreservedAnalyses() override {
    return IRContext::kAnalysisDecorations | IRContext::kAnalysisDefUse |
           IRContext::kAnalysisConstants | IRContext::kAnalysisTypes;
  }

 private:
  // A struct containing components of a composite variable. If the composite
  // consists of multiple or recursive components, |component_variable| is
  // nullptr and |nested_composite_components| keeps the components. If it has a
  // single component, |nested_composite_components| is empty and
  // |component_variable| is the component. Note that each element of
  // |nested_composite_components| has the NestedCompositeComponents struct as
  // its type that can recursively keep the components.
  struct NestedCompositeComponents {
    NestedCompositeComponents() : component_variable(nullptr) {}

    bool HasMultipleComponents() const {
      return !nested_composite_components.empty();
    }

    const std::vector<NestedCompositeComponents>& GetComponents() const {
      return nested_composite_components;
    }

    void AddComponent(const NestedCompositeComponents& component) {
      nested_composite_components.push_back(component);
    }

    Instruction* GetComponentVariable() const { return component_variable; }

    void SetSingleComponentVariable(Instruction* var) {
      component_variable = var;
    }

   private:
    std::vector<NestedCompositeComponents> nested_composite_components;
    Instruction* component_variable;
  };

  // Collects all interface variables used by the |entry_point|.
  std::vector<Instruction*> CollectInterfaceVariables(Instruction& entry_point);

  // Returns whether |var| has the extra arrayness for the entry point
  // |entry_point| or not.
  bool HasExtraArrayness(Instruction& entry_point, Instruction* var);

  // Finds a Location BuiltIn decoration of |var| and returns it via
  // |location|. Returns true whether the location exists or not.
  bool GetVariableLocation(Instruction* var, uint32_t* location);

  // Finds a Component BuiltIn decoration of |var| and returns it via
  // |component|. Returns true whether the component exists or not.
  bool GetVariableComponent(Instruction* var, uint32_t* component);

  // Returns the type of |var| as an instruction.
  Instruction* GetTypeOfVariable(Instruction* var);

  // Replaces an interface variable |interface_var| whose type is
  // |interface_var_type| with scalars and returns whether it succeeds or not.
  // |location| is the value of Location Decoration for |interface_var|.
  // |component| is the value of Component Decoration for |interface_var|.
  // If |extra_array_length| is 0, it means |interface_var| has a Patch
  // decoration. Otherwise, |extra_array_length| denotes the length of the extra
  // array of |interface_var|.
  Status ReplaceInterfaceVariableWithScalars(Instruction* interface_var,
                                             Instruction* interface_var_type,
                                             uint32_t location,
                                             uint32_t component,
                                             uint32_t extra_array_length);

  // Creates scalar variables with the storage classe |storage_class| to replace
  // an interface variable whose type is |interface_var_type|. If
  // |extra_array_length| is not zero, adds the extra arrayness to the created
  // scalar variables.
  std::optional<NestedCompositeComponents>
  CreateScalarInterfaceVarsForReplacement(Instruction* interface_var_type,
                                          spv::StorageClass storage_class,
                                          uint32_t extra_array_length);

  // Creates scalar variables with the storage classe |storage_class| to replace
  // the interface variable whose type is OpTypeArray |interface_var_type| with.
  // If |extra_array_length| is not zero, adds the extra arrayness to all the
  // scalar variables.
  std::optional<NestedCompositeComponents> CreateScalarInterfaceVarsForArray(
      Instruction* interface_var_type, spv::StorageClass storage_class,
      uint32_t extra_array_length);

  // Creates scalar variables with the storage classe |storage_class| to replace
  // the interface variable whose type is OpTypeMatrix |interface_var_type|
  // with. If |extra_array_length| is not zero, adds the extra arrayness to all
  // the scalar variables.
  std::optional<NestedCompositeComponents> CreateScalarInterfaceVarsForMatrix(
      Instruction* interface_var_type, spv::StorageClass storage_class,
      uint32_t extra_array_length);

  // Recursively adds Location and Component decorations to variables in
  // |vars| with |location| and |component|. Increases |location| by one after
  // it actually adds Location and Component decorations for a variable.
  void AddLocationAndComponentDecorations(const NestedCompositeComponents& vars,
                                          uint32_t* location,
                                          uint32_t component);

  // Replaces the interface variable |interface_var| with
  // |scalar_interface_vars| and returns whether it succeeds or not.
  // |extra_arrayness| is the extra arrayness of the interface variable.
  // |scalar_interface_vars| contains the nested variables to replace the
  // interface variable with.
  Status ReplaceInterfaceVarWith(
      Instruction* interface_var, uint32_t extra_arrayness,
      const NestedCompositeComponents& scalar_interface_vars);

  // Replaces |interface_var| in the operands of instructions
  // |interface_var_users| with |scalar_interface_vars|. This is a recursive
  // method and |interface_var_component_indices| is used to specify which
  // recursive component of |interface_var| is replaced. Returns composite
  // construct instructions to be replaced with load instructions of
  // |interface_var_users| via |loads_to_composites|. Returns composite
  // construct instructions to be replaced with load instructions of access
  // chain instructions in |interface_var_users| via
  // |loads_for_access_chain_to_composites|.
  Status ReplaceComponentsOfInterfaceVarWith(
      Instruction* interface_var,
      const std::vector<Instruction*>& interface_var_users,
      const NestedCompositeComponents& scalar_interface_vars,
      std::vector<uint32_t>& interface_var_component_indices,
      const uint32_t* extra_array_index,
      std::unordered_map<Instruction*, Instruction*>* loads_to_composites,
      std::unordered_map<Instruction*, Instruction*>*
          loads_for_access_chain_to_composites);

  // Replaces |interface_var| in the operands of instructions
  // |interface_var_users| with |components| that is a vector of components for
  // the interface variable |interface_var|. This is a recursive method and
  // |interface_var_component_indices| is used to specify which recursive
  // component of |interface_var| is replaced. Returns composite construct
  // instructions to be replaced with load instructions of |interface_var_users|
  // via |loads_to_composites|. Returns composite construct instructions to be
  // replaced with load instructions of access chain instructions in
  // |interface_var_users| via |loads_for_access_chain_to_composites|.
  Status ReplaceMultipleComponentsOfInterfaceVarWith(
      Instruction* interface_var,
      const std::vector<Instruction*>& interface_var_users,
      const std::vector<NestedCompositeComponents>& components,
      std::vector<uint32_t>& interface_var_component_indices,
      const uint32_t* extra_array_index,
      std::unordered_map<Instruction*, Instruction*>* loads_to_composites,
      std::unordered_map<Instruction*, Instruction*>*
          loads_for_access_chain_to_composites);

  // Replaces a component of |interface_var| that is used as an operand of
  // instruction |interface_var_user| with |scalar_var|.
  // |interface_var_component_indices| is a vector of recursive indices for
  // which recursive component of |interface_var| is replaced. If
  // |interface_var_user| is a load, returns the component value via
  // |loads_to_component_values|. If |interface_var_user| is an access chain,
  // returns the component value for loads of |interface_var_user| via
  // |loads_for_access_chain_to_component_values|.
  Status ReplaceComponentOfInterfaceVarWith(
      Instruction* interface_var, Instruction* interface_var_user,
      Instruction* scalar_var,
      const std::vector<uint32_t>& interface_var_component_indices,
      const uint32_t* extra_array_index,
      std::unordered_map<Instruction*, Instruction*>* loads_to_component_values,
      std::unordered_map<Instruction*, Instruction*>*
          loads_for_access_chain_to_component_values);

  // Creates instructions to load |scalar_var| and inserts them before
  // |insert_before|. If |extra_array_index| is not null, they load
  // |extra_array_index| th component of |scalar_var| instead of |scalar_var|
  // itself.
  Instruction* LoadScalarVar(Instruction* scalar_var,
                             const uint32_t* extra_array_index,
                             Instruction* insert_before);

  // Creates instructions to load an access chain to |var| and inserts them
  // before |insert_before|. |Indexes| will be Indexes operand of the access
  // chain.
  Instruction* LoadAccessChainToVar(Instruction* var,
                                    const std::vector<uint32_t>& indexes,
                                    Instruction* insert_before);

  // Creates instructions to store a component of an aggregate whose id is
  // |value_id| to an access chain to |scalar_var| and inserts the created
  // instructions before |insert_before|. To get the component, recursively
  // traverses the aggregate with |component_indices| as indexes.
  // Numbers in |access_chain_indices| are the Indexes operand of the access
  // chain to |scalar_var|
  void StoreComponentOfValueToAccessChainToScalarVar(
      uint32_t value_id, const std::vector<uint32_t>& component_indices,
      Instruction* scalar_var,
      const std::vector<uint32_t>& access_chain_indices,
      Instruction* insert_before);

  // Creates instructions to store a component of an aggregate whose id is
  // |value_id| to |scalar_var| and inserts the created instructions before
  // |insert_before|. To get the component, recursively traverses the aggregate
  // using |extra_array_index| and |component_indices| as indexes.
  void StoreComponentOfValueToScalarVar(
      uint32_t value_id, const std::vector<uint32_t>& component_indices,
      Instruction* scalar_var, const uint32_t* extra_array_index,
      Instruction* insert_before);

  // Creates instructions to store a component of an aggregate whose id is
  // |value_id| to |ptr| and inserts the created instructions before
  // |insert_before|. To get the component, recursively traverses the aggregate
  // using |extra_array_index| and |component_indices| as indexes.
  // |component_type_id| is the id of the type instruction of the component.
  void StoreComponentOfValueTo(uint32_t component_type_id, uint32_t value_id,
                               const std::vector<uint32_t>& component_indices,
                               Instruction* ptr,
                               const uint32_t* extra_array_index,
                               Instruction* insert_before);

  // Creates new OpCompositeExtract with |type_id| for Result Type,
  // |composite_id| for Composite operand, and |indexes| for Indexes operands.
  // If |extra_first_index| is not nullptr, uses it as the first Indexes
  // operand.
  Instruction* CreateCompositeExtract(uint32_t type_id, uint32_t composite_id,
                                      const std::vector<uint32_t>& indexes,
                                      const uint32_t* extra_first_index);

  // Creates a new OpLoad whose Result Type is |type_id| and Pointer operand is
  // |ptr|. Inserts the new instruction before |insert_before|.
  Instruction* CreateLoad(uint32_t type_id, Instruction* ptr,
                          Instruction* insert_before);

  // Clones an annotation instruction |annotation_inst| and sets the target
  // operand of the new annotation instruction as |var_id|.
  void CloneAnnotationForVariable(Instruction* annotation_inst,
                                  uint32_t var_id);

  // Replaces the interface variable |interface_var| in the operands of the
  // entry point |entry_point| with |scalar_var_id|. If it cannot find
  // |interface_var| from the operands of the entry point |entry_point|, adds
  // |scalar_var_id| as an operand of the entry point |entry_point|.
  bool ReplaceInterfaceVarInEntryPoint(Instruction* interface_var,
                                       Instruction* entry_point,
                                       uint32_t scalar_var_id);

  // Creates an access chain instruction whose Base operand is |var| and Indexes
  // operand is |index|. |component_type_id| is the id of the type instruction
  // that is the type of component. Inserts the new access chain before
  // |insert_before|.
  Instruction* CreateAccessChainWithIndex(uint32_t component_type_id,
                                          Instruction* var, uint32_t index,
                                          Instruction* insert_before);

  // Returns the pointee type of the type of variable |var|.
  uint32_t GetPointeeTypeIdOfVar(Instruction* var);

  // Replaces the access chain |access_chain| and its users with a new access
  // chain that points |scalar_var| as the Base operand having
  // |interface_var_component_indices| as Indexes operands and users of the new
  // access chain. When some of the users are load instructions, returns the
  // original load instruction to the new instruction that loads a component of
  // the original load value via |loads_to_component_values|.
  void ReplaceAccessChainWith(
      Instruction* access_chain,
      const std::vector<uint32_t>& interface_var_component_indices,
      Instruction* scalar_var,
      std::unordered_map<Instruction*, Instruction*>*
          loads_to_component_values);

  // Assuming that |access_chain| is an access chain instruction whose Base
  // operand is |base_access_chain|, replaces the operands of |access_chain|
  // with operands of |base_access_chain| and Indexes operands of
  // |access_chain|.
  void UseBaseAccessChainForAccessChain(Instruction* access_chain,
                                        Instruction* base_access_chain);

  // Creates composite construct instructions for load instructions that are the
  // keys of |loads_to_component_values| if no such composite construct
  // instructions exist. Adds a component of the composite as an operand of the
  // created composite construct instruction. Each value of
  // |loads_to_component_values| is the component. Returns the created composite
  // construct instructions using |loads_to_composites|. |depth_to_component| is
  // the number of recursive access steps to get the component from the
  // composite.
  void AddComponentsToCompositesForLoads(
      const std::unordered_map<Instruction*, Instruction*>&
          loads_to_component_values,
      std::unordered_map<Instruction*, Instruction*>* loads_to_composites,
      uint32_t depth_to_component);

  // Creates a composite construct instruction for a component of the value of
  // instruction |load| in |depth_to_component| th recursive depth and inserts
  // it after |load|.
  Instruction* CreateCompositeConstructForComponentOfLoad(
      Instruction* load, uint32_t depth_to_component);

  // Creates a new access chain instruction that points to variable |var| whose
  // type is the instruction with |var_type_id| and inserts it before
  // |insert_before|. The new access chain will have |index_ids| for Indexes
  // operands. Returns the type id of the component that is pointed by the new
  // access chain via |component_type_id|.
  Instruction* CreateAccessChainToVar(uint32_t var_type_id, Instruction* var,
                                      const std::vector<uint32_t>& index_ids,
                                      Instruction* insert_before,
                                      uint32_t* component_type_id);

  // Returns the result id of OpTypeArray instrunction whose Element Type
  // operand is |elem_type_id| and Length operand is |array_length|.
  uint32_t GetArrayType(uint32_t elem_type_id, uint32_t array_length);

  // Returns the result id of OpTypePointer instrunction whose Type
  // operand is |type_id| and Storage Class operand is |storage_class|.
  uint32_t GetPointerType(uint32_t type_id, spv::StorageClass storage_class);

  // Kills an instrunction |inst| and its users.
  void KillInstructionAndUsers(Instruction* inst);

  // Kills a vector of instrunctions |insts| and their users.
  void KillInstructionsAndUsers(const std::vector<Instruction*>& insts);

  // Kills all OpDecorate instructions for Location and Component of the
  // variable whose id is |var_id|.
  void KillLocationAndComponentDecorations(uint32_t var_id);

  // If |var| has the extra arrayness for an entry point, reports an error and
  // returns true. Otherwise, returns false.
  bool ReportErrorIfHasExtraArraynessForOtherEntry(Instruction* var);

  // If |var| does not have the extra arrayness for an entry point, reports an
  // error and returns true. Otherwise, returns false.
  bool ReportErrorIfHasNoExtraArraynessForOtherEntry(Instruction* var);

  // If |interface_var| has the extra arrayness for an entry point but it does
  // not have one for another entry point, reports an error and returns false.
  // Otherwise, returns true. |has_extra_arrayness| denotes whether it has an
  // extra arrayness for an entry point or not.
  bool CheckExtraArraynessConflictBetweenEntries(Instruction* interface_var,
                                                 bool has_extra_arrayness);

  // Conducts the scalar replacement for the interface variables used by the
  // |entry_point|.
  Pass::Status ReplaceInterfaceVarsWithScalars(Instruction& entry_point);

  // A set of interface variable ids that were already removed from operands of
  // the entry point.
  std::unordered_set<uint32_t>
      interface_vars_removed_from_entry_point_operands_;

  // A mapping from ids of new composite construct instructions that load
  // instructions are replaced with to the recursive depth of the component of
  // load that the new component construct instruction is used for.
  std::unordered_map<uint32_t, uint32_t> composite_ids_to_component_depths;

  // A set of interface variables with the extra arrayness for any of the entry
  // points.
  std::unordered_set<Instruction*> vars_with_extra_arrayness;

  // A set of interface variables without the extra arrayness for any of the
  // entry points.
  std::unordered_set<Instruction*> vars_without_extra_arrayness;

  // Returns the next available id, or 0 if the id overflows.
  uint32_t TakeNextId() { return context()->TakeNextId(); }
};

}  // namespace opt
}  // namespace spvtools

#endif  // SOURCE_OPT_INTERFACE_VAR_SROA_H_