// 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 SOURCE_OPT_SSA_REWRITE_PASS_H_
#define SOURCE_OPT_SSA_REWRITE_PASS_H_

#include <queue>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include "source/opt/basic_block.h"
#include "source/opt/ir_context.h"
#include "source/opt/mem_pass.h"

namespace spvtools {
namespace opt {

// Utility class for passes that need to rewrite a function into SSA.  This
// converts load/store operations on function-local variables into SSA IDs,
// which allows them to be the target of optimizing transformations.
//
// Store and load operations to these variables are converted into
// operations on SSA IDs.  Phi instructions are added when needed.  See the
// SSA construction paper for algorithmic details
// (https://link.springer.com/chapter/10.1007/978-3-642-37051-9_6)
class SSARewriter {
 public:
  SSARewriter(MemPass* pass)
      : pass_(pass), first_phi_id_(pass_->get_module()->IdBound()) {}

  // Rewrites SSA-target variables in function |fp| into SSA.  This is the
  // entry point for the SSA rewrite algorithm.  SSA-target variables are
  // locally defined variables that meet the criteria set by IsSSATargetVar.
  //
  // It returns true if function |fp| was modified.  Otherwise, it returns
  // false.
  bool RewriteFunctionIntoSSA(Function* fp);

 private:
  class PhiCandidate {
   public:
    explicit PhiCandidate(uint32_t var, uint32_t result, BasicBlock* block)
        : var_id_(var),
          result_id_(result),
          bb_(block),
          phi_args_(),
          copy_of_(0),
          is_complete_(false),
          users_() {}

    uint32_t var_id() const { return var_id_; }
    uint32_t result_id() const { return result_id_; }
    BasicBlock* bb() const { return bb_; }
    std::vector<uint32_t>& phi_args() { return phi_args_; }
    const std::vector<uint32_t>& phi_args() const { return phi_args_; }
    uint32_t copy_of() const { return copy_of_; }
    bool is_complete() const { return is_complete_; }
    std::vector<uint32_t>& users() { return users_; }
    const std::vector<uint32_t>& users() const { return users_; }

    // Marks this phi candidate as a trivial copy of |orig_id|.
    void MarkCopyOf(uint32_t orig_id) { copy_of_ = orig_id; }

    // Marks this phi candidate as incomplete.
    void MarkIncomplete() { is_complete_ = false; }

    // Marks this phi candidate as complete.
    void MarkComplete() { is_complete_ = true; }

    // Returns true if this Phi candidate is ready to be emitted.
    bool IsReady() const { return is_complete() && copy_of() == 0; }

    // Pretty prints this Phi candidate into a string and returns it. |cfg| is
    // needed to lookup basic block predecessors.
    std::string PrettyPrint(const CFG* cfg) const;

    // Registers |operand_id| as a user of this Phi candidate.
    void AddUser(uint32_t operand_id) { users_.push_back(operand_id); }

   private:
    // Variable ID that this Phi is merging.
    uint32_t var_id_;

    // SSA ID generated by this Phi (i.e., this is the result ID of the eventual
    // Phi instruction).
    uint32_t result_id_;

    // Basic block to hold this Phi.
    BasicBlock* bb_;

    // Vector of operands for every predecessor block of |bb|.  This vector is
    // organized so that the Ith slot contains the argument coming from the Ith
    // predecessor of |bb|.
    std::vector<uint32_t> phi_args_;

    // If this Phi is a trivial copy of another Phi, this is the ID of the
    // original. If this is 0, it means that this is not a trivial Phi.
    uint32_t copy_of_;

    // False, if this Phi candidate has no arguments or at least one argument is
    // %0.
    bool is_complete_;

    // List of all users for this Phi instruction. Each element is the result ID
    // of the load instruction replaced by this Phi, or the result ID of a Phi
    // candidate that has this Phi in its list of operands.
    std::vector<uint32_t> users_;
  };

  // Type used to keep track of store operations in each basic block.
  typedef std::unordered_map<BasicBlock*,
                             std::unordered_map<uint32_t, uint32_t>>
      BlockDefsMap;

  // Generates all the SSA rewriting decisions for basic block |bb|.  This
  // populates the Phi candidate table (|phi_candidate_|) and the load
  // replacement table (|load_replacement_).
  void GenerateSSAReplacements(BasicBlock* bb);

  // Seals block |bb|.  Sealing a basic block means |bb| and all its
  // predecessors of |bb| have been scanned for loads/stores.
  void SealBlock(BasicBlock* bb);

  // Returns true if |bb| has been sealed.
  bool IsBlockSealed(BasicBlock* bb) { return sealed_blocks_.count(bb) != 0; }

  // Returns the Phi candidate with result ID |id| if it exists in the table
  // |phi_candidates_|. If no such Phi candidate exists, it returns nullptr.
  PhiCandidate* GetPhiCandidate(uint32_t id) {
    auto it = phi_candidates_.find(id);
    return (it != phi_candidates_.end()) ? &it->second : nullptr;
  }

  // Replaces all the users of Phi candidate |phi_cand| to be users of
  // |repl_id|.
  void ReplacePhiUsersWith(const PhiCandidate& phi_cand, uint32_t repl_id);

  // Returns the value ID that should replace the load ID in the given
  // replacement pair |repl|.  The replacement is a pair (|load_id|, |val_id|).
  // If |val_id| is itself replaced by another value in the table, this function
  // will look the replacement for |val_id| until it finds one that is not
  // itself replaced.  For instance, given:
  //
  //            %34 = OpLoad %float %f1
  //            OpStore %t %34
  //            %36 = OpLoad %float %t
  //
  // Assume that %f1 is reached by a Phi candidate %42, the load
  // replacement table will have the following entries:
  //
  //            %34 -> %42
  //            %36 -> %34
  //
  // So, when looking for the replacement for %36, we should not use
  // %34. Rather, we should use %42.  To do this, the chain of
  // replacements must be followed until we reach an element that has
  // no replacement.
  uint32_t GetReplacement(std::pair<uint32_t, uint32_t> repl);

  // Returns the argument at index |ix| from |phi_candidate|. If argument |ix|
  // comes from a trivial Phi, it follows the copy-of chain from that trivial
  // Phi until it finds the original Phi candidate.
  //
  // This is only valid after all Phi candidates have been completed. It can
  // only be called when generating the IR for these Phis.
  uint32_t GetPhiArgument(const PhiCandidate* phi_candidate, uint32_t ix);

  // Applies all the SSA replacement decisions.  This replaces loads/stores to
  // SSA target variables with their corresponding SSA IDs, and inserts Phi
  // instructions for them.
  bool ApplyReplacements();

  // Registers a definition for variable |var_id| in basic block |bb| with
  // value |val_id|.
  void WriteVariable(uint32_t var_id, BasicBlock* bb, uint32_t val_id) {
    defs_at_block_[bb][var_id] = val_id;
    if (auto* pc = GetPhiCandidate(val_id)) {
      pc->AddUser(bb->id());
    }
  }

  // Processes the store operation |inst| in basic block |bb|. This extracts
  // the variable ID being stored into, determines whether the variable is an
  // SSA-target variable, and, if it is, it stores its value in the
  // |defs_at_block_| map.
  void ProcessStore(Instruction* inst, BasicBlock* bb);

  // Processes the load operation |inst| in basic block |bb|. This extracts
  // the variable ID being stored into, determines whether the variable is an
  // SSA-target variable, and, if it is, it reads its reaching definition by
  // calling |GetReachingDef|.
  void ProcessLoad(Instruction* inst, BasicBlock* bb);

  // Reads the current definition for variable |var_id| in basic block |bb|.
  // If |var_id| is not defined in block |bb| it walks up the predecessors of
  // |bb|, creating new Phi candidates along the way, if needed.
  //
  // It returns the value for |var_id| from the RHS of the current reaching
  // definition for |var_id|.
  uint32_t GetReachingDef(uint32_t var_id, BasicBlock* bb);

  // Adds arguments to |phi_candidate| by getting the reaching definition of
  // |phi_candidate|'s variable on each of the predecessors of its basic
  // block. After populating the argument list, it determines whether all its
  // arguments are the same.  If so, it returns the ID of the argument that
  // this Phi copies.
  uint32_t AddPhiOperands(PhiCandidate* phi_candidate);

  // Creates a Phi candidate instruction for variable |var_id| in basic block
  // |bb|.
  //
  // Since the rewriting algorithm may remove Phi candidates when it finds
  // them to be trivial, we avoid the expense of creating actual Phi
  // instructions by keeping a pool of Phi candidates (|phi_candidates_|)
  // during rewriting.
  //
  // Once the candidate Phi is created, it returns its ID.
  PhiCandidate& CreatePhiCandidate(uint32_t var_id, BasicBlock* bb);

  // Attempts to remove a trivial Phi candidate |phi_cand|. Trivial Phis are
  // those that only reference themselves and one other value |val| any number
  // of times. This will try to remove any other Phis that become trivial
  // after |phi_cand| is removed.
  //
  // If |phi_cand| is trivial, it returns the SSA ID for the value that should
  // replace it.  Otherwise, it returns the SSA ID for |phi_cand|.
  uint32_t TryRemoveTrivialPhi(PhiCandidate* phi_cand);

  // Finalizes |phi_candidate| by replacing every argument that is still %0
  // with its reaching definition.
  void FinalizePhiCandidate(PhiCandidate* phi_candidate);

  // Finalizes processing of Phi candidates.  Once the whole function has been
  // scanned for loads and stores, the CFG will still have some incomplete and
  // trivial Phis.  This will add missing arguments and remove trivial Phi
  // candidates.
  void FinalizePhiCandidates();

  // Prints the table of Phi candidates to std::cerr.
  void PrintPhiCandidates() const;

  // Prints the load replacement table to std::cerr.
  void PrintReplacementTable() const;

  // Map holding the value of every SSA-target variable at every basic block
  // where the variable is stored. defs_at_block_[block][var_id] = val_id
  // means that there is a store or Phi instruction for variable |var_id| at
  // basic block |block| with value |val_id|.
  BlockDefsMap defs_at_block_;

  // Map, indexed by Phi ID, holding all the Phi candidates created during SSA
  // rewriting.  |phi_candidates_[id]| returns the Phi candidate whose result
  // is |id|.
  std::unordered_map<uint32_t, PhiCandidate> phi_candidates_;

  // Queue of incomplete Phi candidates. These are Phi candidates created at
  // unsealed blocks. They need to be completed before they are instantiated
  // in ApplyReplacements.
  std::queue<PhiCandidate*> incomplete_phis_;

  // List of completed Phi candidates.  These are the only candidates that
  // will become real Phi instructions.
  std::vector<PhiCandidate*> phis_to_generate_;

  // SSA replacement table.  This maps variable IDs, resulting from a load
  // operation, to the value IDs that will replace them after SSA rewriting.
  // After all the rewriting decisions are made, a final scan through the IR
  // is done to replace all uses of the original load ID with the value ID.
  std::unordered_map<uint32_t, uint32_t> load_replacement_;

  // Set of blocks that have been sealed already.
  std::unordered_set<BasicBlock*> sealed_blocks_;

  // Memory pass requesting the SSA rewriter.
  MemPass* pass_;

  // ID of the first Phi created by the SSA rewriter.  During rewriting, any
  // ID bigger than this corresponds to a Phi candidate.
  uint32_t first_phi_id_;
};

class SSARewritePass : public MemPass {
 public:
  SSARewritePass() = default;

  const char* name() const override { return "ssa-rewrite"; }
  Status Process() override;
};

}  // namespace opt
}  // namespace spvtools

#endif  // SOURCE_OPT_SSA_REWRITE_PASS_H_