//===--- 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/SILFunction.h"
#include "swift/SIL/SILInstruction.h"
#include "swift/SIL/Test.h"
#include "llvm/ADT/STLExtras.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;
}

namespace {
/// Implements OSSALifetimeCompletion::visitUnreachableLifetimeEnds.  Finds
/// positions as near as possible to unreachables at which `value`'s lifetime
/// is available.
///
/// Finding these positions is a three step process:
/// 1) computeRegion: Forward CFG walk from non-lifetime-ending boundary to find
///                   the dead-end region in which the value might be available.
/// 2) propagateAvailability: Forward iterative dataflow within the region to
///                           determine which blocks the value is available in.
/// 3) visitAvailabilityBoundary: Visits the final blocks in the region where
///                               the value is available--these are the blocks
///                               without successors or with at least one
///                               unavailable successor.
class VisitUnreachableLifetimeEnds {
  /// The value whose dead-end block lifetime ends are to be visited.
  SILValue value;

  /// The non-lifetime-ending boundary of `value`.
  BasicBlockSet starts;
  /// The region between (inclusive) the `starts` and the unreachable blocks.
  BasicBlockSetVector region;

public:
  VisitUnreachableLifetimeEnds(SILValue value)
      : value(value), starts(value->getFunction()),
        region(value->getFunction()) {}

  /// Region discovery.
  ///
  /// Forward CFG walk from non-lifetime-ending boundary to unreachable
  /// instructions.
  void computeRegion(const SSAPrunedLiveness &liveness);

  struct Result;

  /// Iterative dataflow to determine availability for each block in `region`.
  void propagateAvailablity(Result &result);

  /// Visit the terminators of blocks on the boundary of availability.
  void
  visitAvailabilityBoundary(Result const &result,
                            llvm::function_ref<void(SILInstruction *)> visit);

  struct State {
    enum class Value : uint8_t {
      Unavailable = 0,
      Available,
      Unknown,
    };
    Value value;

    State(Value value) : value(value){};
    operator Value() const { return value; }
    State meet(State const other) const {
      return *this < other ? *this : other;
    }

    static State Unavailable() { return {Value::Unavailable}; }
    static State Available() { return {Value::Available}; }
    static State Unknown() { return {Value::Unknown}; }
  };

  struct Result {
    BasicBlockBitfield states;

    Result(SILFunction *function) : states(function, 2) {}

    State getState(SILBasicBlock *block) const {
      return {(State::Value)states.get(block)};
    }

    void setState(SILBasicBlock *block, State newState) {
      states.set(block, (unsigned)newState.value);
    }

    /// Propagate predecessors' state into `block`.
    ///
    /// states[block] ∧= state[predecessor_1] ∧ ... ∧ state[predecessor_n]
    bool updateState(SILBasicBlock *block) {
      auto oldState = getState(block);
      auto state = oldState;
      for (auto *predecessor : block->getPredecessorBlocks()) {
        state = state.meet(getState(predecessor));
      }
      setState(block, state);
      return state != oldState;
    }
  };
};

void VisitUnreachableLifetimeEnds::computeRegion(
    const SSAPrunedLiveness &liveness) {
  // Find the non-lifetime-ending boundary of `value`.
  PrunedLivenessBoundary boundary;
  liveness.computeBoundary(boundary);

  for (SILInstruction *lastUser : boundary.lastUsers) {
    if (liveness.isInterestingUser(lastUser)
        != PrunedLiveness::LifetimeEndingUse) {
      region.insert(lastUser->getParent());
      starts.insert(lastUser->getParent());
    }
  }
  for (SILBasicBlock *edge : boundary.boundaryEdges) {
    region.insert(edge);
    starts.insert(edge);
  }
  for (SILNode *deadDef : boundary.deadDefs) {
    region.insert(deadDef->getParentBlock());
    starts.insert(deadDef->getParentBlock());
  }

  // Forward walk to find the region in which `value` might be available.
  BasicBlockWorklist regionWorklist(value->getFunction());
  // Start the forward walk from the non-lifetime-ending boundary.
  for (auto *start : region) {
    regionWorklist.push(start);
  }
  while (auto *block = regionWorklist.pop()) {
    if (block->succ_empty()) {
      // This assert will fail unless there are already lifetime-ending
      // instruction on all paths to normal function exits.
      assert(isa<UnreachableInst>(block->getTerminator()));
    }
    for (auto *successor : block->getSuccessorBlocks()) {
      regionWorklist.pushIfNotVisited(successor);
      region.insert(successor);
    }
  }
}

void VisitUnreachableLifetimeEnds::propagateAvailablity(Result &result) {
  // Initialize per-block state.
  // - all blocks outside of the region are ::Unavailable (automatically
  //   initialized)
  // - non-initial in-region blocks are Unknown
  // - start blocks are ::Available
  for (auto *block : region) {
    if (starts.contains(block))
      result.setState(block, State::Available());
    else
      result.setState(block, State::Unknown());
  }

  BasicBlockWorklist worklist(value->getFunction());

  // Initialize worklist with all participating blocks.
  //
  // Only perform dataflow in the non-initial region.  Every initial block is
  // by definition ::Available.
  for (auto *block : region) {
    if (starts.contains(block))
      continue;
    worklist.push(block);
  }

  // Iterate over blocks which are successors of blocks whose state changed.
  while (auto *block = worklist.popAndForget()) {
    // Only propagate availability in non-initial, in-region blocks.
    if (!region.contains(block) || starts.contains(block))
      continue;
    auto changed = result.updateState(block);
    if (!changed) {
      continue;
    }
    // The state has changed.  Propagate the new state into successors.
    for (auto *successor : block->getSuccessorBlocks()) {
      worklist.pushIfNotVisited(successor);
    }
  }
}

void VisitUnreachableLifetimeEnds::visitAvailabilityBoundary(
    Result const &result, llvm::function_ref<void(SILInstruction *)> visit) {
  for (auto *block : region) {
    auto available = result.getState(block) == State::Available();
    if (!available) {
      continue;
    }
    auto hasUnreachableSuccessor = [&]() {
      // Use a lambda to avoid checking if possible.
      return llvm::any_of(block->getSuccessorBlocks(), [&result](auto *block) {
        return result.getState(block) == State::Unavailable();
      });
    };
    if (!block->succ_empty() && !hasUnreachableSuccessor()) {
      continue;
    }
    assert(hasUnreachableSuccessor() ||
           isa<UnreachableInst>(block->getTerminator()));
    visit(block->getTerminator());
  }
}
} // end anonymous namespace

void OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
    SILValue value, const SSAPrunedLiveness &liveness,
    llvm::function_ref<void(SILInstruction *)> visit) {

  VisitUnreachableLifetimeEnds visitor(value);

  visitor.computeRegion(liveness);

  VisitUnreachableLifetimeEnds::Result result(value->getFunction());

  visitor.propagateAvailablity(result);

  visitor.visitAvailabilityBoundary(result, visit);
}

static bool endLifetimeAtUnreachableBlocks(SILValue value,
                                           const SSAPrunedLiveness &liveness) {
  bool changed = false;
  OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
      value, liveness, [&](auto *unreachable) {
        SILBuilderWithScope builder(unreachable);
        endOSSALifetime(value, builder);
        changed = true;
      });
  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::outs() << "OSSA lifetime completion: " << value;
      OSSALifetimeCompletion completion(&function, /*domInfo*/ nullptr);
      completion.completeOSSALifetime(value);
      function.print(llvm::outs());
    });
} // 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;
}