//===- PackMetadataMarkerInserter.cpp - Add markers for metadata de/allocs ===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// IRGen initially assembles metadata and wtable packs on the stack.  It then
// heapifies those packs.
//
// As an optimization, that heapification can be avoided some of the time:
// whenever the instructions which may result in such allocations are identified
// in functions which do not have instructions that obstruct StackNesting.
//
// Finds all instructions on behalf of which IRGen may emit on-stack pack
// metadata, looking for instructions which obstruct that optimization
//
// No code motion may occur after this pass: alloc_pack_metadata must directly
// precede the instruction on behalf of which metadata will actually be emitted
// (e.g. apply).
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "pack-metadata-dealloc-inserter"

#include "swift/AST/TypeWalker.h"
#include "swift/IRGen/IRGenSILPasses.h"
#include "swift/SIL/ApplySite.h"
#include "swift/SIL/Dominance.h"
#include "swift/SIL/SILBasicBlock.h"
#include "swift/SIL/SILBuilder.h"
#include "swift/SIL/SILInstruction.h"
#include "swift/SIL/SILNode.h"
#include "swift/SIL/StackList.h"
#include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
#include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
#include "swift/SILOptimizer/PassManager/Transforms.h"
#include "swift/SILOptimizer/Utils/CFGOptUtils.h"
#include "swift/SILOptimizer/Utils/StackNesting.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/ErrorHandling.h"

using namespace swift;

namespace {

/// Finds and inserts marker instructions indicating where on-stack pack
/// metadata and wtables should be cleaned up.
///
/// Inserts alloc_pack_metadata before such instructions, and inserts
/// dealloc_pack_metadata on the dominance frontier.  The client is responsible
/// for ensuring that critical edges have been split before inserting markers.
class Inserter {
  SILFunction *function;
  StackList<SILInstruction *> instructionsToCleanup;

public:
  Inserter(SILFunction *function)
      : function(function), instructionsToCleanup(function) {}

  enum class FindResult {
    None,
    Some,
    Unhandleable,
  };

  Inserter::FindResult find();

  void insert(DominanceInfo *tree, DeadEndBlocks *deBlocks);

private:
  FindResult shouldInsertMarkersForInstruction(SILInstruction *inst);
  void insertMarkers(SILInstruction *instruction, DominanceInfo *tree,
                     DeadEndBlocks *deBlocks);
};

Inserter::FindResult
Inserter::shouldInsertMarkersForInstruction(SILInstruction *inst) {
  switch (inst->getKind()) {
  case SILInstructionKind::BuiltinInst: {
    auto *bi = cast<BuiltinInst>(inst);
    // Functions with async lets may use the stack, but that is not encoded in
    // SIL.  As a result, stack nesting is unable to properly nest async lets
    // with on-stack pack metadata.  rdar://109850951
    if (bi->getBuiltinKind() ==
            BuiltinValueKind::StartAsyncLetWithLocalBuffer ||
        bi->getBuiltinKind() == BuiltinValueKind::StartAsyncLet)
      return Inserter::FindResult::Unhandleable;
    LLVM_FALLTHROUGH;
  }
  default:
    return inst->mayRequirePackMetadata(*inst->getFunction())
               ? FindResult::Some
               : FindResult::None;
  }
}

void Inserter::insertMarkers(SILInstruction *instruction, DominanceInfo *tree,
                             DeadEndBlocks *deBlocks) {
  SILBuilderWithScope builder(instruction);
  auto *apmi = builder.createAllocPackMetadata(CleanupLocation(
      RegularLocation::getAutoGeneratedLocation(instruction->getLoc())));
  SmallVector<SILBasicBlock *, 4> boundary;
  auto *block = apmi->getParent();
  computeDominatedBoundaryBlocks(block, tree, boundary);
  for (auto *block : boundary) {
    if (deBlocks->isDeadEnd(block))
      continue;
    SILBuilderWithScope builder(&block->back());
    builder.createDeallocPackMetadata(
        CleanupLocation(RegularLocation::getAutoGeneratedLocation()), apmi);
  }
}

/// Find instructions that may entail materializing a metadata pack and those
/// that will prevent stack allocation of metadata.
///
/// Found instructions are recorded in instructionsToCleanup.
Inserter::FindResult Inserter::find() {
  auto collected = false;
  for (auto &block : *function) {
    for (auto &instruction : block) {
      auto result = shouldInsertMarkersForInstruction(&instruction);
      switch (result) {
      case FindResult::None:
        continue;
      case FindResult::Some:
        instructionsToCleanup.push_back(&instruction);
        collected = true;
        continue;
      case FindResult::Unhandleable:
        return FindResult::Unhandleable;
      }
      llvm_unreachable("covered switch");
    }
  }
  return collected ? FindResult::Some : FindResult::None;
}

/// Insert the marker instructions around the instructions previously found.
void Inserter::insert(DominanceInfo *tree, DeadEndBlocks *deBlocks) {
  for (auto *instruction : instructionsToCleanup) {
    insertMarkers(instruction, tree, deBlocks);
  }
}

class PackMetadataMarkerInserter : public SILFunctionTransform {
  void run() override {
    auto *function = getFunction();

    // If on-stack metadata is disabled, this pass does nothing.
    auto options = function->getModule().getOptions();
    if (!options.EnablePackMetadataStackPromotion)
      return;

    Inserter inserter(function);
    auto found = inserter.find();
    switch (found) {
    case Inserter::FindResult::None:
      return;
    case Inserter::FindResult::Unhandleable:
      LLVM_DEBUG(llvm::dbgs() << "Found obstructive instructions in "
                              << function->getName() << ".\n");
      /// The function contains some unhandleable instruction.  Record that on
      /// the function.
      function->setUseStackForPackMetadata(DoNotUseStackForPackMetadata);
      return;
    case Inserter::FindResult::Some:
      LLVM_DEBUG(llvm::dbgs() << "Found potential pack metadata emitters in "
                              << function->getName() << ".\n");
      break;
    }

    auto *dominance = getAnalysis<DominanceAnalysis>();
    auto *tree = dominance->get(function);
    auto split = splitAllCriticalEdges(*function, /*domInfo=*/tree,
                                       /*loopInfo=*/nullptr);
    auto *deadEnds = getAnalysis<DeadEndBlocksAnalysis>();
    if (split) {
      deadEnds->invalidateFunction(function);
    }
    auto *deBlocks = deadEnds->get(function);
    inserter.insert(tree, deBlocks);
    auto changes = StackNesting::fixNesting(function);
    invalidateAnalysis(
        (split || changes == StackNesting::Changes::CFG)
            ? SILAnalysis::InvalidationKind::BranchesAndInstructions
            : SILAnalysis::InvalidationKind::Instructions);
  }
};

} // end anonymous namespace

SILTransform *irgen::createPackMetadataMarkerInserter() {
  return new PackMetadataMarkerInserter();
}