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

#ifndef LIBSPIRV_OPT_MERGE_RETURN_PASS_H_
#define LIBSPIRV_OPT_MERGE_RETURN_PASS_H_

#include "basic_block.h"
#include "function.h"
#include "mem_pass.h"

#include <unordered_set>
#include <vector>

namespace spvtools {
namespace opt {

/*******************************************************************************
 *
 * Handling Structured Control Flow:
 *
 * Structured control flow guarantees that the CFG will reconverge at a given
 * point (the merge block). Within structured control flow, all blocks must be
 * post-dominated by the merge block, except return blocks and break blocks.
 * A break block is a block that branches to the innermost loop's merge block.
 *
 * Beyond this, we further assume that all unreachable blocks have been
 * cleanedup.  This means that the only unreachable blocks are those necessary
 * for valid structured control flow.
 *
 * Algorithm:
 *
 * If a return is encountered, it should record that: i) the function has
 * "returned" and ii) the value of the return. The return should be replaced
 * with a branch. If current block is not within structured control flow, this
 * is the final return. This block should branch to the new return block (its
 * direct successor). If the current block is within structured control flow,
 * the branch destination should be the innermost loop's merge (if it exists)
 * or the merge block of the immediate structured control flow. If the merge
 * block produces any live values it will need to be predicated. While the merge
 * is nested in structured control flow, the predication path should branch to
 * the next best merge block available. Once structured control flow has been
 * exited, remaining blocks must be predicated with new structured control flow
 * (OpSelectionMerge). These should be nested correctly in case of straight line
 * branching to reach the final return block.
 *
 * In the final return block, the return value should be loaded and returned.
 * Memory promotion passes should be able to promote the newly introduced
 * variables ("has returned" and "return value").
 *
 * Predicating the Final Merge:
 *
 * At each merge block predication needs to be introduced (optimization: only if
 * that block produces value live beyond it). This needs to be done carefully.
 * The merge block should be split into multiple blocks.
 *
 *          1 (header)
 *        /   \
 * (ret) 2     3 (merge)
 *
 *         ||
 *         \/
 *
 *          1 (header)
 *        /   \
 *       2     |
 *        \   /
 *          3 (merge for 1, new header)
 *        /   \
 *       |     3 (old body)
 *        \   /
 *    (ret) 4 (new merge)
 *
 * In the above (simple) example, the return originally in |2| is passed through
 * the merge. That merge is predicated such that the old body of the block is
 * the else branch. The branch condition is based on the value of the "has
 * returned" variable. In more complicated examples (blocks between |1| and
 * |3|), the SSA would need to fixed up due the newly reconvergent path at the
 * merge for |1|. Assuming |3| originally was also a return block, the old body
 * of |3| should also store the return value for that case. The return value in
 * |4| just requires loading the return value variable.
 *
 ******************************************************************************/

// Documented in optimizer.hpp
class MergeReturnPass : public MemPass {
 public:
  MergeReturnPass()
      : function_(nullptr),
        return_flag_(nullptr),
        return_value_(nullptr),
        constant_true_(nullptr),
        final_return_block_(nullptr) {}

  const char* name() const override { return "merge-return"; }
  Status Process(ir::IRContext*) override;

  ir::IRContext::Analysis GetPreservedAnalyses() override {
    // return ir::IRContext::kAnalysisDefUse;
    return ir::IRContext::kAnalysisNone;
  }

 private:
  // This class is used to store the a loop merge instruction and a selection
  // merge instruction.  The intended use is that is represent the inner most
  // contain selection construct and the inner most loop construct.
  class StructuredControlState {
   public:
    StructuredControlState(ir::Instruction* loop, ir::Instruction* merge)
        : loop_merge_(loop), current_merge_(merge) {}

    StructuredControlState(const StructuredControlState&) = default;

    bool InLoop() const { return loop_merge_; }
    bool InStructuredFlow() const { return CurrentMergeId() != 0; }

    uint32_t CurrentMergeId() const {
      return current_merge_ ? current_merge_->GetSingleWordInOperand(0u) : 0u;
    }

    uint32_t CurrentMergeHeader() const {
      return current_merge_ ? current_merge_->context()
                                  ->get_instr_block(current_merge_)
                                  ->id()
                            : 0;
    }

    uint32_t LoopMergeId() const {
      return loop_merge_ ? loop_merge_->GetSingleWordInOperand(0u) : 0u;
    }

    uint32_t CurrentLoopHeader() const {
      return loop_merge_
                 ? loop_merge_->context()->get_instr_block(loop_merge_)->id()
                 : 0;
    }

    ir::Instruction* LoopMergeInst() const { return loop_merge_; }

   private:
    ir::Instruction* loop_merge_;
    ir::Instruction* current_merge_;
  };

  // Returns all BasicBlocks terminated by OpReturn or OpReturnValue in
  // |function|.
  std::vector<ir::BasicBlock*> CollectReturnBlocks(ir::Function* function);

  // Creates a new basic block with a single return. If |function| returns a
  // value, a phi node is created to select the correct value to return.
  // Replaces old returns with an unconditional branch to the new block.
  void MergeReturnBlocks(ir::Function* function,
                         const std::vector<ir::BasicBlock*>& returnBlocks);

  // Merges the return instruction in |function| so that it has a single return
  // statement.  It is assumed that |function| has structured control flow, and
  // that |return_blocks| is a list of all of the basic blocks in |function|
  // that have a return.
  void ProcessStructured(ir::Function* function,
                         const std::vector<ir::BasicBlock*>& return_blocks);

  // Changes an OpReturn* or OpUnreachable instruction at the end of |block|
  // into a store to |return_flag_|, a store to |return_value_| (if necessary),
  // and a branch to the appropriate merge block.
  //
  // Is is assumed that |AddReturnValue| have already been called to created the
  // variable to store a return value if there is one.
  //
  // Note this will break the semantics.  To fix this, PredicateBlock will have
  // to be called on the merge block the branch targets.
  void ProcessStructuredBlock(ir::BasicBlock* block);

  // Creates a variable used to store whether or not the control flow has
  // traversed a block that used to have a return.  A pointer to the instruction
  // declaring the variable is stored in |return_flag_|.
  void AddReturnFlag();

  // Creates the variable used to store the return value when passing through
  // a block that use to contain an OpReturnValue.
  void AddReturnValue();

  // Adds a store that stores true to |return_flag_| immediately before the
  // terminator of |block|. It is assumed that |AddReturnFlag| has already been
  // called.
  void RecordReturned(ir::BasicBlock* block);

  // Adds an instruction that stores the value being returned in the
  // OpReturnValue in |block|.  The value is stored to |return_value_|, and the
  // store is placed before the OpReturnValue.
  //
  // If |block| does not contain an OpReturnValue, then this function has no
  // effect. If |block| contains an OpReturnValue, then |AddReturnValue| must
  // have already been called to create the variable to store to.
  void RecordReturnValue(ir::BasicBlock* block);

  // Adds an unconditional branch in |block| that branches to |target|.  It also
  // adds stores to |return_flag_| and |return_value_| as needed.
  // |AddReturnFlag| and |AddReturnValue| must have already been called.
  void BranchToBlock(ir::BasicBlock* block, uint32_t target);

  // Returns true if we need to pridicate |block| where |tail_block| is the
  // merge point.  (See |PredicateBlocks|).  There is no need to predicate if
  // there is no code that could be executed.
  bool RequiresPredication(const ir::BasicBlock* block,
                           const ir::BasicBlock* tail_block) const;

  // For every basic block that is reachable from a basic block in
  // |return_blocks|, extra code is added to jump around any code that should
  // not be executed because the original code would have already returned. This
  // involves adding new selections constructs to jump around these
  // instructions.
  void PredicateBlocks(const std::vector<ir::BasicBlock*>& return_blocks);

  // Add the predication code (see |PredicateBlocks|) to |tail_block| if it
  // requires predication.  |tail_block| and any new blocks that are known to
  // not require predication will be added to |predicated|.
  void PredicateBlock(ir::BasicBlock* block, ir::BasicBlock* tail_block,
                      std::unordered_set<ir::BasicBlock*>* predicated);

  // Add an |OpReturn| or |OpReturnValue| to the end of |block|.  If an
  // |OpReturnValue| is needed, the return value is loaded from |return_value_|.
  void CreateReturn(ir::BasicBlock* block);

  // Creates a block at the end of the function that will become the single
  // return block at the end of the pass.
  void CreateReturnBlock();

  // Creates a Phi node in |merge_block| for the result of |inst| coming from
  // |predecessor|.  Any uses of the result of |inst| that are no longer
  // dominated by |inst|, are replaced with the result of the new |OpPhi|
  // instruction.
  void CreatePhiNodesForInst(ir::BasicBlock* merge_block, uint32_t predecessor,
                             ir::Instruction& inst);

  // Traverse the nodes in |new_merge_nodes_|, and adds the OpPhi instructions
  // that are needed to make the code correct.  It is assumed that at this point
  // there are no unreachable blocks in the control flow graph.
  void AddNewPhiNodes();

  // Creates any new phi nodes that are needed in |bb| now that |pred| is no
  // longer the only block that preceedes |bb|.  |header_id| is the id of the
  // basic block for the loop or selection construct that merges at |bb|.
  void AddNewPhiNodes(ir::BasicBlock* bb, ir::BasicBlock* pred,
                      uint32_t header_id);

  // Saves |block| to a list of basic block that will require OpPhi nodes to be
  // added by calling |AddNewPhiNodes|.  It is assumed that |block| used to have
  // a single predecessor, |single_original_pred|, but now has more.
  void MarkForNewPhiNodes(ir::BasicBlock* block,
                          ir::BasicBlock* single_original_pred);

  // Return the original single predcessor of |block| if it was flagged as
  // having a single predecessor.  |nullptr| is returned otherwise.
  ir::BasicBlock* MarkedSinglePred(ir::BasicBlock* block) {
    auto it = new_merge_nodes_.find(block);
    if (it != new_merge_nodes_.end()) {
      return it->second;
    } else {
      return nullptr;
    }
  }

  StructuredControlState& CurrentState() { return state_.back(); }

  // A stack used to keep track of the innermost contain loop and selection
  // constructs.
  std::vector<StructuredControlState> state_;

  // The current function being transformed.
  ir::Function* function_;

  // The |OpVariable| instruction defining a boolean variable used to keep track
  // of whether or not the function is trying to return.
  ir::Instruction* return_flag_;

  // The |OpVariable| instruction defining a variabled to used to keep track of
  // the value that was returned when passing through a block that use to
  // contain an |OpReturnValue|.
  ir::Instruction* return_value_;

  // The instruction defining the boolean constant true.
  ir::Instruction* constant_true_;

  // The basic block that is suppose to become the contain the only return value
  // after processing the current function.
  ir::BasicBlock* final_return_block_;
  // This map contains the set of nodes that use to have a single predcessor,
  // but now have more.  They will need new OpPhi nodes.  For each of the nodes,
  // it is mapped to it original single predcessor.  It is assumed there are no
  // values that will need a phi on the new edges.
  std::unordered_map<ir::BasicBlock*, ir::BasicBlock*> new_merge_nodes_;
};

}  // namespace opt
}  // namespace spvtools

#endif  // LIBSPIRV_OPT_MERGE_RETURN_PASS_H_