//===--- OSSALifetimeCompletion.cpp ---------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
///
/// OSSA lifetime completion adds lifetime ending instructions to make
/// linear lifetimes complete.
///
/// Interior liveness handles the following cases naturally:
///
/// When completing the lifetime of the initial value, %v1, transitively
/// include all uses of dominated reborrows as, such as %phi1 in this example:
///
///     %v1 = ...
///     cond_br bb1, bb2
///   bb1:
///     %b1 = begin_borrow %v1
///     br bb3(%b1)
///   bb2:
///     %b2 = begin_borrow %v1
///     br bb3(%b2)
///   bb3(%phi1):
///     %u1 = %phi1
///     end_borrow %phi1
///     %k1 = destroy_value %v1 // must be below end_borrow %phi1
///
/// When completing the lifetime for a phi (%phi2) transitively include all
/// uses of inner adjacent reborrows, such as %phi1 in this example:
///
///   bb1:
///     %v1 = ...
///     %b1 = begin_borrow %v1
///     br bb3(%b1, %v1)
///   bb2:
///     %v2 = ...
///     %b2 = begin_borrow %v2
///     br bb3(%b2, %v2)
///   bb3(%phi1, %phi2):
///     %u1 = %phi1
///     end_borrow %phi1
///     %k1 = destroy_value %phi1
///
//===----------------------------------------------------------------------===//

#include "swift/SIL/OSSALifetimeCompletion.h"
#include "swift/SIL/SILBuilder.h"
#include "swift/SIL/SILInstruction.h"
#include "swift/SIL/Test.h"

using namespace swift;

static SILInstruction *endOSSALifetime(SILValue value, SILBuilder &builder) {
  auto loc =
    RegularLocation::getAutoGeneratedLocation(builder.getInsertionPointLoc());
  if (value->getOwnershipKind() == OwnershipKind::Owned) {
    return builder.createDestroyValue(loc, value);
  }
  return builder.createEndBorrow(loc, value);
}

static bool endLifetimeAtBoundary(SILValue value,
                                  const SSAPrunedLiveness &liveness) {
  PrunedLivenessBoundary boundary;
  liveness.computeBoundary(boundary);

  bool changed = false;
  for (SILInstruction *lastUser : boundary.lastUsers) {
    if (liveness.isInterestingUser(lastUser)
        != PrunedLiveness::LifetimeEndingUse) {
      changed = true;
      SILBuilderWithScope::insertAfter(lastUser, [value](SILBuilder &builder) {
        endOSSALifetime(value, builder);
      });
    }
  }
  for (SILBasicBlock *edge : boundary.boundaryEdges) {
    changed = true;
    SILBuilderWithScope builder(edge->begin());
    endOSSALifetime(value, builder);
  }
  for (SILNode *deadDef : boundary.deadDefs) {
    SILInstruction *next = nullptr;
    if (auto *deadInst = dyn_cast<SILInstruction>(deadDef)) {
      next = deadInst->getNextInstruction();
    } else {
      next = cast<ValueBase>(deadDef)->getNextInstruction();
    }
    changed = true;
    SILBuilderWithScope builder(next);
    endOSSALifetime(value, builder);
  }
  return changed;
}

static bool endLifetimeAtUnreachableBlocks(SILValue value,
                                           const SSAPrunedLiveness &liveness) {
  PrunedLivenessBoundary boundary;
  liveness.computeBoundary(boundary);

  BasicBlockWorklist deadEndBlocks(value->getFunction());
  for (SILInstruction *lastUser : boundary.lastUsers) {
    if (liveness.isInterestingUser(lastUser)
        != PrunedLiveness::LifetimeEndingUse) {
      deadEndBlocks.push(lastUser->getParent());
    }
  }
  for (SILBasicBlock *edge : boundary.boundaryEdges) {
    deadEndBlocks.push(edge);
  }
  for (SILNode *deadDef : boundary.deadDefs) {
    deadEndBlocks.push(deadDef->getParentBlock());
  }
  // Forward CFG walk from the non-lifetime-ending boundary to the unreachable
  // instructions.
  bool changed = false;
  while (auto *block = deadEndBlocks.pop()) {
    if (block->succ_empty()) {
      // This assert will fail unless there are already lifetime-ending
      // instruction on all paths to normal function exits.
      auto *unreachable = cast<UnreachableInst>(block->getTerminator());
      SILBuilderWithScope builder(unreachable);
      endOSSALifetime(value, builder);
      changed = true;
    }
    for (auto *successor : block->getSuccessorBlocks()) {
      deadEndBlocks.pushIfNotVisited(successor);
    }
  }
  return changed;
}

/// End the lifetime of \p value at unreachable instructions.
///
/// Returns true if any new instructions were created to complete the lifetime.
///
/// This is only meant to cleanup lifetimes that lead to dead-end blocks. After
/// recursively completing all nested scopes, it then simply ends the lifetime
/// at the Unreachable instruction.
bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(
    SILValue value, bool forceBoundaryCompletion) {
  // Called for inner borrows, inner adjacent reborrows, inner reborrows, and
  // scoped addresses.
  auto handleInnerScope = [this](SILValue innerBorrowedValue) {
    completeOSSALifetime(innerBorrowedValue);
  };
  InteriorLiveness liveness(value);
  liveness.compute(domInfo, handleInnerScope);

  bool changed = false;
  if (value->isLexical() && !forceBoundaryCompletion) {
    changed |= endLifetimeAtUnreachableBlocks(value, liveness.getLiveness());
  } else {
    changed |= endLifetimeAtBoundary(value, liveness.getLiveness());
  }
  // TODO: Rebuild outer adjacent phis on demand (SILGen does not currently
  // produce guaranteed phis). See FindEnclosingDefs &
  // findSuccessorDefsFromPredDefs. If no enclosing phi is found, we can create
  // it here and use updateSSA to recursively populate phis.
  assert(liveness.getUnenclosedPhis().empty());
  return changed;
}

namespace swift::test {
// Arguments:
// - SILValue: value
// Dumps:
// - function
static FunctionTest OSSALifetimeCompletionTest(
    "ossa-lifetime-completion",
    [](auto &function, auto &arguments, auto &test) {
      SILValue value = arguments.takeValue();
      llvm::dbgs() << "OSSA lifetime completion: " << value;
      OSSALifetimeCompletion completion(&function, /*domInfo*/ nullptr);
      completion.completeOSSALifetime(value);
      function.dump();
    });
} // end namespace swift::test

// TODO: create a fast check for 'mayEndLifetime(SILInstruction *)'. Verify that
// it returns true for every instruction that has a lifetime-ending operand.
void UnreachableLifetimeCompletion::visitUnreachableInst(
    SILInstruction *instruction) {
  auto *block = instruction->getParent();
  bool inReachableBlock = !unreachableBlocks.contains(block);
  // If this instruction's block is already marked unreachable, and
  // updatingLifetimes is not yet set, then this instruction will be visited
  // again later when propagating unreachable blocks.
  if (!inReachableBlock && !updatingLifetimes)
    return;

  for (Operand &operand : instruction->getAllOperands()) {
    if (!operand.isLifetimeEnding())
      continue;

    SILValue value = operand.get();
    SILBasicBlock *defBlock = value->getParentBlock();
    if (unreachableBlocks.contains(defBlock))
      continue;

    auto *def = value->getDefiningInstruction();
    if (def && unreachableInsts.contains(def))
      continue;

    // The operand's definition is still reachable and its lifetime ends on a
    // newly unreachable path.
    //
    // Note: The arguments of a no-return try_apply may still appear reachable
    // here because the try_apply itself is never visited as unreachable, hence
    // its successor blocks are not marked . But it
    // seems harmless to recompute their lifetimes.

    // Insert this unreachable instruction in unreachableInsts if its parent
    // block is not already marked unreachable.
    if (inReachableBlock) {
      unreachableInsts.insert(instruction);
    }
    incompleteValues.insert(value);

    // Add unreachable successors to the forward traversal worklist.
    if (auto *term = dyn_cast<TermInst>(instruction)) {
      for (auto *succBlock : term->getSuccessorBlocks()) {
        if (llvm::all_of(succBlock->getPredecessorBlocks(),
                         [&](SILBasicBlock *predBlock) {
                           if (predBlock == block)
                             return true;

                           return unreachableBlocks.contains(predBlock);
                         })) {
          unreachableBlocks.insert(succBlock);
        }
      }
    }
  }
}

bool UnreachableLifetimeCompletion::completeLifetimes() {
  assert(!updatingLifetimes && "don't call this more than once");
  updatingLifetimes = true;

  // Now that all unreachable terminator instructions have been visited,
  // propagate unreachable blocks.
  for (auto blockIt = unreachableBlocks.begin();
       blockIt != unreachableBlocks.end(); ++blockIt) {
    auto *block = *blockIt;
    for (auto &instruction : *block) {
      visitUnreachableInst(&instruction);
    }
  }

  OSSALifetimeCompletion completion(function, domInfo);

  bool changed = false;
  for (auto value : incompleteValues) {
    if (completion.completeOSSALifetime(value)
        == LifetimeCompletion::WasCompleted) {
      changed = true;
    }
  }
  return changed;
}