//===--- AllocBoxToStack.cpp - Promote alloc_box to alloc_stack -----------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2015 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See http://swift.org/LICENSE.txt for license information
// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "allocbox-to-stack"
#include "swift/SILPasses/Passes.h"
#include "swift/SIL/SILBuilder.h"
#include "swift/SIL/Dominance.h"
#include "swift/SILPasses/Utils/Local.h"
#include "swift/SILPasses/Transforms.h"
#include "swift/SILAnalysis/DominanceAnalysis.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Support/Debug.h"
using namespace swift;

STATISTIC(NumStackPromoted, "Number of alloc_box's promoted to the stack");
STATISTIC(NumEntryCycleBB, "Number of alloc_box's promotion disrupted due to "
          "entry cycle");

//===----------------------------------------------------------------------===//
//                           alloc_box Promotion
//===----------------------------------------------------------------------===//

/// \brief Checks if the single entry single exit region starting at \p Entry
/// and ending at \p Exit has cycles. This function returns True if there is
/// a path between \p  Entry to itself without going through \p Exit.
static bool regionHasCycles(SILBasicBlock *Entry, SILBasicBlock *Exit) {
  llvm::SmallVector<SILBasicBlock *, 16> Worklist;
  llvm::SmallPtrSet<SILBasicBlock *, 16> VisitedBBSet;

  // Start the scan at the entry block.
  Worklist.push_back(Entry);

  // For each BB in the worklist...
  while (!Worklist.empty()) {
    auto *BB = Worklist.pop_back_val();

    // Add the successors of the BB to the worklist...
    for (auto &Succ : BB->getSuccs()) {
      auto *BB = Succ.getBB();

      // If one of those successors is the entry block, we have a cycle.
      if (BB == Entry) {
        ++NumEntryCycleBB;
        return true;
      }

      // If the successor is the exit block, skip it.
      if (BB == Exit)
        continue;

      // Otherwise, if we have not visited the BB yet, add it to the worklist
      // for processing.
      if (VisitedBBSet.insert(BB))
        Worklist.push_back(BB);
    }
  }

  // The entry block is not reachable from itself, we are safe.
  return false;
}

/// getLastRelease - Determine if there is a single ReleaseInst or
/// DeallocBoxInst that post-dominates all of the uses of the specified
/// AllocBox.  If so, return it.  If not, return null.
static SILInstruction *getLastRelease(AllocBoxInst *ABI,
                                      SmallVectorImpl<SILInstruction*> &Users,
                                    SmallVectorImpl<SILInstruction*> &Releases,
                                      PostDominanceInfo *PDI) {
  // If there are no releases, then the box is leaked.  Don't transform it. This
  // can only happen in hand-written SIL code, not compiler generated code.
  if (Releases.empty())
    return nullptr;


  // If there is a single release, it must be the last release.  The calling
  // conventions used in SIL (at least in the case where the ABI doesn't escape)
  // are such that the value is live until it is explicitly released: there are
  // no calls in this case that pass ownership and release for us.
  if (Releases.size() == 1)
    return Releases.back();

  // If there are multiple releases of the value, we only support the case where
  // there is a single ultimate release that post-dominates all of the rest of
  // them.

  // Determine the most-post-dominating release by doing a linear scan over all
  // of the releases to find one that post-dominates them all.  Because we don't
  // want to do multiple scans of the block (which could be large) just keep
  // track of whether there are multiple releases in the ultimate block we find.
  bool MultipleReleasesInBlock = false;
  SILInstruction *LastRelease = Releases[0];

  for (unsigned i = 1, e = Releases.size(); i != e; ++i) {
    SILInstruction *RI = Releases[i];

    // If this release is in the same block as our candidate, keep track of the
    // multiple release nature of that block, but don't try to determine an
    // ordering between them yet.
    if (RI->getParent() == LastRelease->getParent()) {
      MultipleReleasesInBlock = true;
      continue;
    }

    if (PDI->properlyDominates(RI->getParent(), LastRelease->getParent())) {
      // RI post-dom's LastRelease, so it is our new LastRelease.
      LastRelease = RI;
      MultipleReleasesInBlock = false;
    }
  }

  // Okay, we found the most-post-dominating release. We are currently checking
  // for domination/post-domination but not control equivalence, a necessary
  // condition for a single entry/single-exit region. This check ensures
  // that we do not run into the issue by making sure that the allocbox is not
  // reachable from itself.
  if (regionHasCycles(ABI->getParent(), LastRelease->getParent()))
    return nullptr;

  // Make sure that the most-post-dominating release we have found postdom all
  // of our uses. If the release does not postdom a use, it would be unsafe to
  // perform the use.
  for (auto *User : Users) {
    if (User->getParent() == LastRelease->getParent())
      continue;

    if (!PDI->properlyDominates(LastRelease->getParent(), User->getParent()))
      return nullptr;
  }

  // Okay, the LastRelease block postdoms all users.  If there are multiple
  // releases in the block, make sure we're looking at the last one.
  if (MultipleReleasesInBlock) {
    for (auto MBBI = --LastRelease->getParent()->end(); ; --MBBI) {
      auto *RI = dyn_cast<StrongReleaseInst>(MBBI);
      if (RI == nullptr ||
          RI->getOperand() != SILValue(ABI, 0)) {
        assert(MBBI != LastRelease->getParent()->begin() &&
               "Didn't find any release in this block?");
        continue;
      }
      LastRelease = RI;
      break;
    }
  }

  return LastRelease;
}

/// \brief Returns True if the operand or one of its users is captured.
static bool useCaptured(Operand *UI) {
  auto *User = UI->getUser();

  // These instructions do not cause the box's address to escape.
  if (isa<CopyAddrInst>(User) ||
      isa<LoadInst>(User) ||
      isa<ProtocolMethodInst>(User) ||
      (isa<StoreInst>(User) && UI->getOperandNumber() == 1) ||
      (isa<AssignInst>(User) && UI->getOperandNumber() == 1)) {
    return false;
  }

  if ((isa<AddressToPointerInst>(User) || isa<PointerToAddressInst>(User)) &&
      User->hasOneUse()) {
    return useCaptured(*User->use_begin());
  }

  return true;
}

/// checkAllocBoxUses - Recursively scan all of the uses of the
/// specified alloc_box, returning true if a use escapes, or false
/// otherwise.
static bool checkAllocBoxUses(AllocBoxInst *ABI, SILValue V,
                              SmallVectorImpl<SILInstruction*> &Users) {
  for (auto UI : V.getUses()) {
    auto *User = UI->getUser();

    // These instructions do not cause the box's address to escape.
    if (!useCaptured(UI)) {
      Users.push_back(User);
      continue;
    }

    // These instructions only cause the alloc_box to escape if they are used in
    // a way that escapes.  Recursively check that the uses of the instruction
    // don't escape and collect all of the uses of the value.
    if (isa<StructElementAddrInst>(User) || isa<TupleElementAddrInst>(User) ||
        isa<ProjectExistentialInst>(User) || isa<OpenExistentialInst>(User) ||
        isa<MarkUninitializedInst>(User)) {
      Users.push_back(User);
      if (checkAllocBoxUses(ABI, User, Users))
        return true;
      continue;
    }

    // apply and partial_apply instructions do not capture the pointer when
    // it is passed through inout arguments or for indirect returns.
    if (auto apply = dyn_cast<ApplyInst>(User)) {
      if (apply->getSubstCalleeType()
            ->getInterfaceParameters()[UI->getOperandNumber()-1].isIndirect())
        continue;
    }
    if (auto partialApply = dyn_cast<PartialApplyInst>(User)) {
      if (partialApply->getSubstCalleeType()
            ->getInterfaceParameters()[UI->getOperandNumber()-1].isIndirect())
        continue;

    }

    // Otherwise, this looks like it escapes.
    DEBUG(llvm::dbgs() << "*** Failed to promote alloc_box in @"
          << ABI->getFunction()->getName() << ": " << *ABI
          << "    Due to user: " << *User << "\n");

    return true;
  }

  return false;
}


/// optimizeAllocBox - Try to promote an alloc_box instruction to an
/// alloc_stack.  On success, this updates the IR and returns true, but does not
/// remove the alloc_box itself.
static bool optimizeAllocBox(AllocBoxInst *ABI, PostDominanceInfo *PDI) {

  // Scan all of the uses of the alloc_box to see if any of them cause the
  // allocated memory to escape.  If so, we can't promote it to the stack.  If
  // not, we can turn it into an alloc_stack.
  SmallVector<SILInstruction*, 32> Users;
  if (checkAllocBoxUses(ABI, SILValue(ABI, 1), Users))
    return false;

  // Scan all of the uses of the retain count value, collecting all the releases
  // and validating that we don't have an unexpected user.
  SmallVector<SILInstruction*, 4> Releases;
  for (auto UI : SILValue(ABI, 0).getUses()) {
    auto *User = UI->getUser();

    if (isa<StrongRetainInst>(User))
      continue;

    // Release doesn't either, but we want to keep track of where this value
    // gets released.
    if (isa<StrongReleaseInst>(User) || isa<DeallocBoxInst>(User)) {
      Releases.push_back(User);
      Users.push_back(User);
      continue;
    }

    // Otherwise, this looks like it escapes.
    DEBUG(llvm::dbgs() << "*** Failed to promote alloc_box in @"
          << ABI->getFunction()->getName() << ": " << *ABI
          << "    Due to user: " << *User << "\n");

    return false;
  }


  // Okay, the value doesn't escape.  Determine where the last release is.  This
  // code only handles the case where there is a single "last release".  This
  // should work for us, because we don't expect code duplication that can
  // introduce different releases for different codepaths.  If this ends up
  // mattering in the future, this can be generalized.
  SILInstruction *LastRelease = getLastRelease(ABI, Users, Releases, PDI);

  // FIXME: If there's no last release, we can't balance the alloc_stack.
  if (!LastRelease)
    return false;

  auto &lowering =
    ABI->getModule().Types.getTypeLowering(ABI->getElementType());
  if (LastRelease == nullptr && !lowering.isTrivial()) {
    // If we can't tell where the last release is, we don't know where to insert
    // the destroy_addr for this box.
    DEBUG(llvm::dbgs() << "*** Failed to promote alloc_box: " << *ABI
          << "    Cannot determine location of the last release!\n"
          << *ABI->getParent()->getParent() << "\n");
    return false;
  }

  DEBUG(llvm::dbgs() << "*** Promoting alloc_box to stack: " << *ABI);

  // Okay, it looks like this value doesn't escape.  Promote it to an
  // alloc_stack.  Start by inserting the alloc stack after the alloc_box.
  SILBuilder B1(ABI->getParent(), ++SILBasicBlock::iterator(ABI));
  auto *ASI = B1.createAllocStack(ABI->getLoc(), ABI->getElementType());
  ASI->setDebugScope(ABI->getDebugScope());

  // Replace all uses of the pointer operand with the spiffy new AllocStack.
  SILValue(ABI, 1).replaceAllUsesWith(ASI->getAddressResult());

  // If we found a 'last release', insert a dealloc_stack instruction and a
  // destroy_addr if its type is non-trivial.
  if (LastRelease) {
    SILBuilder B2(LastRelease);

    if (!lowering.isTrivial() && !isa<DeallocBoxInst>(LastRelease))
      B2.emitDestroyAddr(CleanupLocation::getCleanupLocation(ABI->getLoc()),
                         ASI->getAddressResult());

    // Reset the insertion point in case the destroy address expanded to
    // multiple blocks.
    B2.setInsertionPoint(LastRelease);
    B2.createDeallocStack(LastRelease->getLoc(), ASI->getContainerResult());
  }

  // Remove any retain and release instructions.  Since all uses of result #1
  // are gone, this only walks through uses of result #0 (the retain count
  // pointer).
  while (!ABI->use_empty()) {
    auto *User = (*ABI->use_begin())->getUser();
    assert(isa<StrongReleaseInst>(User) || isa<StrongRetainInst>(User) ||
           isa<DeallocBoxInst>(User));

    User->eraseFromParent();
  }

  return true;
}

class AllocBoxToStack : public SILFunctionTransform {
  /// The entry point to the transformation.
  void run() {
    DominanceAnalysis* DA = getAnalysis<DominanceAnalysis>();
    PostDominanceInfo *PDI = DA->getPostDomInfo(getFunction());
    bool Changed = false;

    for (auto &BB : *getFunction()) {
      auto I = BB.begin(), E = BB.end();
      while (I != E) {
        if (auto *ABI = dyn_cast<AllocBoxInst>(I))
          if (optimizeAllocBox(ABI, PDI)) {
            ++NumStackPromoted;
            // Carefully move iterator to avoid invalidation problems.
            ++I;
            ABI->eraseFromParent();
            Changed = true;
            continue;
          }

        ++I;
      }
    }
    if (Changed)
      invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
  }

  StringRef getName() override { return "AllocBox-To-Stack Optimization"; }
};

SILTransform *swift::createAllocBoxToStack() {
  return new AllocBoxToStack();
}