// Copyright (c) 2017 The Khronos Group Inc.

// Copyright (c) 2017 Valve Corporation

// Copyright (c) 2017 LunarG Inc.

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

#include "block_merge_util.h"

namespace spvtools {

namespace opt {

namespace blockmergeutil {

namespace {

// Returns true if |block| contains a merge instruction.

bool IsHeader(BasicBlock* block) { return block->GetMergeInst() != nullptr; }

// Returns true if |id| contains a merge instruction.

bool IsHeader(IRContext* context, uint32_t id) {

  return IsHeader(

      context->get_instr_block(context->get_def_use_mgr()->GetDef(id)));

}

// Returns true if |id| is the merge target of a merge instruction.

bool IsMerge(IRContext* context, uint32_t id) {

  return !context->get_def_use_mgr()->WhileEachUse(

      id, [](Instruction* user, uint32_t index) {

        spv::Op op = user->opcode();

        if ((op == spv::Op::OpLoopMerge || op == spv::Op::OpSelectionMerge) &&

            index == 0u) {

          return false;

        }

        return true;

      });

}

// Returns true if |block| is the merge target of a merge instruction.

bool IsMerge(IRContext* context, BasicBlock* block) {

  return IsMerge(context, block->id());

}

// Returns true if |id| is the continue target of a merge instruction.

bool IsContinue(IRContext* context, uint32_t id) {

  return !context->get_def_use_mgr()->WhileEachUse(

      id, [](Instruction* user, uint32_t index) {

        spv::Op op = user->opcode();

        if (op == spv::Op::OpLoopMerge && index == 1u) {

          return false;

        }

        return true;

      });

}

// Removes any OpPhi instructions in |block|, which should have exactly one

// predecessor, replacing uses of OpPhi ids with the ids associated with the

// predecessor.

void EliminateOpPhiInstructions(IRContext* context, BasicBlock* block) {

  block->ForEachPhiInst([context](Instruction* phi) {

    assert(2 == phi->NumInOperands() &&

           "A block can only have one predecessor for block merging to make "

           "sense.");

    context->ReplaceAllUsesWith(phi->result_id(),

                                phi->GetSingleWordInOperand(0));

    context->KillInst(phi);

  });

}

}  // Anonymous namespace

bool CanMergeWithSuccessor(IRContext* context, BasicBlock* block) {

  // Find block with single successor which has no other predecessors.

  auto ii = block->end();

  --ii;

  Instruction* br = &*ii;

  if (br->opcode() != spv::Op::OpBranch) {

    return false;

  }

  const uint32_t lab_id = br->GetSingleWordInOperand(0);

  if (context->cfg()->preds(lab_id).size() != 1) {

    return false;

  }

  bool pred_is_merge = IsMerge(context, block);

  bool succ_is_merge = IsMerge(context, lab_id);

  if (pred_is_merge && succ_is_merge) {

    // Cannot merge two merges together.

    return false;

  }

  // Note: This means that the instructions in a break block will execute as if

  // they were still diverged according to the loop iteration. This restricts

  // potential transformations an implementation may perform on the IR to match

  // shader author expectations. Similarly, instructions in the loop construct

  // cannot be moved into the continue construct unless it can be proven that

  // invocations are always converged.

  if (succ_is_merge && context->get_feature_mgr()->HasExtension(

                           kSPV_KHR_maximal_reconvergence)) {

    return false;

  }

  if (pred_is_merge && IsContinue(context, lab_id)) {

    // Cannot merge a continue target with a merge block.

    return false;

  }

  Instruction* merge_inst = block->GetMergeInst();

  const bool pred_is_header = IsHeader(block);

  if (pred_is_header && lab_id != merge_inst->GetSingleWordInOperand(0u)) {

    bool succ_is_header = IsHeader(context, lab_id);

    if (pred_is_header && succ_is_header) {

      // Cannot merge two headers together when the successor is not the merge

      // block of the predecessor.

      return false;

    }

    // If this is a header block and the successor is not its merge, we must

    // be careful about which blocks we are willing to merge together.

    // OpLoopMerge must be followed by a conditional or unconditional branch.

    // The merge must be a loop merge because a selection merge cannot be

    // followed by an unconditional branch.

    BasicBlock* succ_block = context->get_instr_block(lab_id);

    spv::Op succ_term_op = succ_block->terminator()->opcode();

    assert(merge_inst->opcode() == spv::Op::OpLoopMerge);

    if (succ_term_op != spv::Op::OpBranch &&

        succ_term_op != spv::Op::OpBranchConditional) {

      return false;

    }

  }

  if (succ_is_merge || IsContinue(context, lab_id)) {

    auto* struct_cfg = context->GetStructuredCFGAnalysis();

    auto switch_block_id = struct_cfg->ContainingSwitch(block->id());

    if (switch_block_id) {

      auto switch_merge_id = struct_cfg->SwitchMergeBlock(switch_block_id);

      const auto* switch_inst =

          &*block->GetParent()->FindBlock(switch_block_id)->tail();

      for (uint32_t i = 1; i < switch_inst->NumInOperands(); i += 2) {

        auto target_id = switch_inst->GetSingleWordInOperand(i);

        if (target_id == block->id() && target_id != switch_merge_id) {

          // Case constructs must be structurally dominated by the OpSwitch.

          // Since the successor is the merge/continue for another construct,

          // merging the blocks would break that requirement.

          return false;

        }

      }

    }

  }

  return true;

}

void MergeWithSuccessor(IRContext* context, Function* func,

                        Function::iterator bi) {

  assert(CanMergeWithSuccessor(context, &*bi) &&

         "Precondition failure for MergeWithSuccessor: it must be legal to "

         "merge the block and its successor.");

  auto ii = bi->end();

  --ii;

  Instruction* br = &*ii;

  const uint32_t lab_id = br->GetSingleWordInOperand(0);

  Instruction* merge_inst = bi->GetMergeInst();

  bool pred_is_header = IsHeader(&*bi);

  // Merge blocks.

  context->KillInst(br);

  auto sbi = bi;

  for (; sbi != func->end(); ++sbi)

    if (sbi->id() == lab_id) break;

  // If bi is sbi's only predecessor, it dominates sbi and thus

  // sbi must follow bi in func's ordering.

  assert(sbi != func->end());

  if (sbi->tail()->opcode() == spv::Op::OpSwitch &&

      sbi->MergeBlockIdIfAny() != 0) {

    context->InvalidateAnalyses(IRContext::Analysis::kAnalysisStructuredCFG);

  }

  // Update the inst-to-block mapping for the instructions in sbi.

  for (auto& inst : *sbi) {

    context->set_instr_block(&inst, &*bi);

  }

  EliminateOpPhiInstructions(context, &*sbi);

  // Now actually move the instructions.

  bi->AddInstructions(&*sbi);

  if (merge_inst) {

    if (pred_is_header && lab_id == merge_inst->GetSingleWordInOperand(0u)) {

      // Merging the header and merge blocks, so remove the structured control

      // flow declaration.

      context->KillInst(merge_inst);

    } else {

      // Move OpLine/OpNoLine information to merge_inst. This solves

      // the validation error that OpLine is placed between OpLoopMerge

      // and OpBranchConditional.

      auto terminator = bi->terminator();

      auto& vec = terminator->dbg_line_insts();

      if (vec.size() > 0) {

        merge_inst->ClearDbgLineInsts();

        auto& new_vec = merge_inst->dbg_line_insts();

        new_vec.insert(new_vec.end(), vec.begin(), vec.end());

        terminator->ClearDbgLineInsts();

        for (auto& l_inst : new_vec)

          context->get_def_use_mgr()->AnalyzeInstDefUse(&l_inst);

      }

      // Clear debug scope of terminator to avoid DebugScope

      // emitted between terminator and merge.

      terminator->SetDebugScope(DebugScope(kNoDebugScope, kNoInlinedAt));

      // Move the merge instruction to just before the terminator.

      merge_inst->InsertBefore(terminator);

    }

  }

  context->ReplaceAllUsesWith(lab_id, bi->id());

  context->KillInst(sbi->GetLabelInst());

  (void)sbi.Erase();

}

}  // namespace blockmergeutil

}  // namespace opt

}  // namespace spvtools