//===--- ArrayElementValuePropagation.cpp - Propagate values of arrays ----===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2016 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
//
//===----------------------------------------------------------------------===//
#define DEBUG_TYPE "array-element-propagation"

#include "llvm/ADT/SetVector.h"
#include "swift/AST/NameLookup.h"
#include "swift/AST/ParameterList.h"
#include "swift/AST/GenericEnvironment.h"
#include "swift/SIL/SILBasicBlock.h"
#include "swift/SIL/SILInstruction.h"
#include "swift/SIL/SILFunction.h"
#include "swift/SIL/DebugUtils.h"
#include "swift/SILOptimizer/Analysis/ArraySemantic.h"
#include "swift/SILOptimizer/PassManager/Passes.h"
#include "swift/SILOptimizer/PassManager/Transforms.h"
#include "llvm/ADT/SmallVector.h"

using namespace swift;

// After approx. this many elements, it's faster to use append(contentsOf:)
constexpr unsigned APPEND_CONTENTSOF_REPLACEMENT_VALUES_MAX = 6;

typedef std::pair<ApplyInst *, SILValue> ValueReplacementPair;
typedef std::tuple<ApplyInst *, llvm::SmallVector<SILValue, 4>,
                   ArrayRef<Substitution>>
    ValueReplacementsPair;

/// Propagate the elements of array values to calls of the array's get_element
/// method, and replace calls of append(contentsOf:) with append(element:).
///
/// Array literal construction and array initialization of array values
/// associates element values with the array value. These values can be
/// propagated to the get_element method if we can prove that the array value
/// has not changed until reading the array value's element. These values can
/// also be used to replace append(contentsOf:) with multiple append(element:)
/// calls.
///
/// Propagation of the elements of one array allocation.
///
/// We propagate the elements associated with calls of
///
/// * Array.init(count:repeatedValue:)
///   The 'repeatedValue'.
///   TODO: this is not yet implemented.
///
/// * Array._adoptStorage(storage:count:)
///   The stores on the returned array element buffer pointer.
///
namespace {
class ArrayAllocation {
  /// The array allocation call.
  ApplyInst *Alloc;
  /// The array value returned by the allocation call.
  SILValue ArrayValue;

  /// The pointer to the returned array element buffer pointer.
  SILValue ElementBuffer;

  /// The calls to Array get_element that use this array allocation.
  llvm::SmallSetVector<ApplyInst *, 16> GetElementCalls;

  /// The calls to Array append_contentsOf that use this array allocation.
  llvm::SmallVector<ApplyInst *, 4> AppendContentsOfCalls;

  /// A map of Array indices to element values
  llvm::DenseMap<uint64_t, SILValue> ElementValueMap;

  /// Array get_element calls and their matching array element value for later
  /// replacement.
  llvm::SmallVectorImpl<ValueReplacementPair> &GetElementReplacementMap;

  /// Array append_contentsOf calls and their matching array values for later
  /// replacement.
  llvm::SmallVectorImpl<ValueReplacementsPair> &AppendContentsOfReplacementMap;

  /// Substitutions to be used when replacing append_contentsOf calls.
  ArrayRef<Substitution> Substitutions;

  ArrayAllocation(
      ApplyInst *AI,
      llvm::SmallVectorImpl<ValueReplacementPair> &GetElementReplacements,
      llvm::SmallVectorImpl<ValueReplacementsPair> &AppendContentsOfReplacements)
      : Alloc(AI), GetElementReplacementMap(GetElementReplacements),
        AppendContentsOfReplacementMap(AppendContentsOfReplacements) {}

  bool findReplacements();
  bool isInitializationWithKnownElements();
  bool mapInitializationStores();
  bool analyzeArrayValueUses();
  bool recursivelyCollectUses(ValueBase *Def);
  bool collectForwardableValues();
  bool replacementsAreValid();

public:

  /// Find a set of get_element calls that can be replace by the initialization
  /// value of the array allocation call.
  ///
  /// Returns true if an access can be replaced. The replacements are stored in
  /// the \p ReplacementMap.
  static bool findReplacements(
      ApplyInst *Inst,
      llvm::SmallVectorImpl<ValueReplacementPair> &GetElementReplacements,
      llvm::SmallVectorImpl<ValueReplacementsPair> &AppendContentsOfReplacements) {
    return ArrayAllocation(Inst, GetElementReplacements,
                           AppendContentsOfReplacements)
        .findReplacements();
  }
};
}


/// Map the indices of array element initialization stores to their values.
bool ArrayAllocation::mapInitializationStores() {
  assert(ElementBuffer &&
         "Must have identified an array element storage pointer");

  // Match initialization stores.
  // %83 = struct_extract %element_buffer : $UnsafeMutablePointer<Int>
  // %84 = pointer_to_address %83 : $Builtin.RawPointer to strict $*Int
  // store %85 to %84 : $*Int
  // %87 = integer_literal $Builtin.Word, 1
  // %88 = index_addr %84 : $*Int, %87 : $Builtin.Word
  // store %some_value to %88 : $*Int

  auto *UnsafeMutablePointerExtract =
      dyn_cast_or_null<StructExtractInst>(getSingleNonDebugUser(ElementBuffer));
  if (!UnsafeMutablePointerExtract)
    return false;
  auto *PointerToAddress = dyn_cast_or_null<PointerToAddressInst>(
      getSingleNonDebugUser(UnsafeMutablePointerExtract));
  if (!PointerToAddress)
    return false;

  // Match the stores. We can have either a store directly to the address or
  // to an index_addr projection.
  for (auto *Op : PointerToAddress->getUses()) {
    auto *Inst = Op->getUser();

    // Store to the base.
    auto *SI = dyn_cast<StoreInst>(Inst);
    if (SI && SI->getDest() == PointerToAddress) {
      // We have already seen an entry for this index bail.
      if (ElementValueMap.count(0))
        return false;
      ElementValueMap[0] = SI->getSrc();
      continue;
    } else if (SI)
      return false;

    // Store an index_addr projection.
    auto *IndexAddr = dyn_cast<IndexAddrInst>(Inst);
    if (!IndexAddr)
      return false;
    SI = dyn_cast_or_null<StoreInst>(getSingleNonDebugUser(IndexAddr));
    if (!SI || SI->getDest() != IndexAddr)
      return false;
    auto *Index = dyn_cast<IntegerLiteralInst>(IndexAddr->getIndex());
    if (!Index)
      return false;
    auto IndexVal = Index->getValue();
    // Let's not blow up our map.
    if (IndexVal.getActiveBits() > 16)
      return false;
    // Already saw an entry.
    if (ElementValueMap.count(IndexVal.getZExtValue()))
      return false;

    ElementValueMap[IndexVal.getZExtValue()] = SI->getSrc();
  }

  return !ElementValueMap.empty();
}

/// Check that we have an array initialization call with known elements.
///
/// The returned array value is known not to be aliased since it was just
/// allocated.
bool ArrayAllocation::isInitializationWithKnownElements() {
  ArraySemanticsCall Uninitialized(Alloc, "array.uninitialized");

  if (!Uninitialized)
    return false;

  ArrayValue = Uninitialized.getArrayValue();
  if (!ArrayValue)
    return false;

  ElementBuffer = Uninitialized.getArrayElementStoragePointer();
  if (!ElementBuffer)
    return false;

  Substitutions =
      ArrayValue->getType().gatherAllSubstitutions(Alloc->getModule());

  return mapInitializationStores();
}

bool ArrayAllocation::replacementsAreValid() {
  unsigned ElementCount = ElementValueMap.size();

  if (ElementCount > APPEND_CONTENTSOF_REPLACEMENT_VALUES_MAX)
    return false;

  // Bail if elements aren't contiguous
  for (unsigned i = 0; i < ElementCount; ++i)
    if (!ElementValueMap.count(i))
      return false;

  return true;
}

/// Propagate the elements of an array literal to get_element method calls on
/// the same array.
///
/// We have to prove that the array value is not changed in between the
/// creation and the method call to get_element.
bool ArrayAllocation::findReplacements() {
  if (!isInitializationWithKnownElements())
    return false;

  // The array value was stored or has escaped.
  if (!analyzeArrayValueUses())
    return false;

  if (GetElementCalls.empty() && AppendContentsOfCalls.empty())
    return false;

  if (AppendContentsOfCalls.empty())
    return collectForwardableValues();

  // Find the substitutions
  if (!replacementsAreValid())
    return collectForwardableValues();

  llvm::SmallVector<SILValue, 4> ElementValueVector;
  for (unsigned i = 0; i < ElementValueMap.size(); ++i)
    ElementValueVector.push_back(ElementValueMap[i]);

  for (auto *Call : AppendContentsOfCalls)
    AppendContentsOfReplacementMap.push_back(
        std::make_tuple(Call, ElementValueVector, Substitutions));

  return collectForwardableValues();
}

/// Collect all get_element users and check that there are no escapes or uses
/// that could change the array value.
bool ArrayAllocation::analyzeArrayValueUses() {
  return recursivelyCollectUses(ArrayValue);
}

/// Recursively look at all uses of this definition. Abort if the array value
/// could escape or be changed. Collect all uses that are calls to array.count.
bool ArrayAllocation::recursivelyCollectUses(ValueBase *Def) {
  for (auto *Opd : Def->getUses()) {
    auto *User = Opd->getUser();
    // Ignore reference counting and debug instructions.
    if (isa<RefCountingInst>(User) || isa<DebugValueInst>(User))
      continue;

    // Array value projection.
    if (auto *SEI = dyn_cast<StructExtractInst>(User)) {
      if (!recursivelyCollectUses(SEI))
        return false;
      continue;
    }

    // Check array semantic calls.
    ArraySemanticsCall ArrayOp(User);
    if (ArrayOp) {
      if (ArrayOp.getKind() == ArrayCallKind::kAppendContentsOf) {
        AppendContentsOfCalls.push_back(ArrayOp);
        continue;
      } else if (ArrayOp.getKind() == ArrayCallKind::kGetElement) {
        GetElementCalls.insert(ArrayOp);
        continue;
      } else if (ArrayOp.doesNotChangeArray()) {
        continue;
      }
    }

    // An operation that escapes or modifies the array value.
    return false;
  }
  return true;
}

/// Look at the get_element calls and match them to values by index.
bool ArrayAllocation::collectForwardableValues() {
  bool FoundForwardableValue = false;
  for (auto *GetElementCall : GetElementCalls) {
    ArraySemanticsCall GetElement(GetElementCall);
    assert(GetElement.getKind() == ArrayCallKind::kGetElement);

    auto ConstantIndex = GetElement.getConstantIndex();
    if (ConstantIndex == None)
      continue;

    assert(*ConstantIndex >= 0 && "Must have a positive index");

    auto EltValueIt = ElementValueMap.find(*ConstantIndex);
    if (EltValueIt == ElementValueMap.end())
      continue;

    GetElementReplacementMap.push_back(
        std::make_pair(GetElementCall, EltValueIt->second));
    FoundForwardableValue = true;
  }
  return FoundForwardableValue;
}

// =============================================================================
//                                 Driver
// =============================================================================

namespace {

class ArrayElementPropagation : public SILFunctionTransform {

public:
  ArrayElementPropagation() {}

  StringRef getName() override {
    return "Array Element Propagation";
  }

  void replaceAppendCalls(llvm::SmallVector<ValueReplacementsPair, 4> Repls) {
    auto &Fn = *getFunction();
    auto &M = Fn.getModule();
    auto &Ctx = M.getASTContext();

    if (Repls.empty())
      return;

    DEBUG(llvm::dbgs() << "Array append contentsOf calls replaced in "
                       << Fn.getName() << " (" << Repls.size() << ")\n");

    auto *AppendFnDecl = Ctx.getArrayAppendElementDecl();
    if (!AppendFnDecl)
      return;

    auto Mangled = SILDeclRef(AppendFnDecl, SILDeclRef::Kind::Func).mangle();
    auto *AppendFn = M.hasFunction(Mangled, SILLinkage::PublicExternal);
    if (!AppendFn)
      return;

    for (auto &Repl : Repls) {
      ArraySemanticsCall AppendContentsOf(std::get<0>(Repl));
      assert(AppendContentsOf && "Must be AppendContentsOf call");
      AppendContentsOf.replaceByAppendingValues(M, AppendFn, std::get<1>(Repl),
                                                std::get<2>(Repl));
    }
  }

  void run() override {
    auto &Fn = *getFunction();

    bool Changed = false;

    // Propagate the elements an of array value to its users.
    llvm::SmallVector<ValueReplacementPair, 16> GetElementReplacements;
    llvm::SmallVector<ValueReplacementsPair, 4> AppendContentsOfReplacements;
    for (auto &BB :Fn) {
      for (auto &Inst : BB) {
        if (auto *Apply = dyn_cast<ApplyInst>(&Inst))
          Changed |= ArrayAllocation::findReplacements(
              Apply, GetElementReplacements, AppendContentsOfReplacements);
      }
    }

    DEBUG(if (!GetElementReplacements.empty()) {
      llvm::dbgs() << "Array elements replaced in " << Fn.getName() << " ("
                   << GetElementReplacements.size() << ")\n";
    });
    // Perform the actual replacement of the get_element call by its value.
    for (auto &Repl : GetElementReplacements) {
      ArraySemanticsCall GetElement(Repl.first);
      GetElement.replaceByValue(Repl.second);
    }

    replaceAppendCalls(AppendContentsOfReplacements);

    if (Changed) {
      PM->invalidateAnalysis(
          &Fn, SILAnalysis::InvalidationKind::CallsAndInstructions);
    }
  }
};
} // End anonymous namespace.

SILTransform *swift::createArrayElementPropagation() {
  return new ArrayElementPropagation();
}