//===--- ObjectOutliner.swift ----------------------------------------------==//
//
// 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
//
//===----------------------------------------------------------------------===//

import SIL

/// Outlines COW objects from functions into statically initialized global variables.
/// This is currently only done for Arrays.
/// If a function constructs an Array literal with constant elements (done by storing
/// the element values into the array buffer), a new global variable is created which
/// contains the constant elements in its static initializer.
/// For example:
/// ```
///   public func arrayLookup(_ i: Int) -> Int {
///     let lookupTable = [10, 11, 12]
///     return lookupTable[i]
///   }
/// ```
/// is turned into
/// ```
///   private let outlinedVariable_from_arrayLookup = [10, 11, 12]  // statically initialized
///
///   public func arrayLookup(_ i: Int) -> Int {
///     return outlinedVariable_from_arrayLookup[i]
///   }
/// ```
///
/// As a second optimization, if the array is a string literal which is a parameter to the
/// `_findStringSwitchCase` library function and the array has many elements (> 16), the
/// call is redirected to `_findStringSwitchCaseWithCache`. This function builds a cache
/// (e.g. a Dictionary) and stores it into a global variable.
/// Then subsequent calls to this function can do a fast lookup using the cache.
///
let objectOutliner = FunctionPass(name: "object-outliner") {
  (function: Function, context: FunctionPassContext) in
  for inst in function.instructions {
    if let ari = inst as? AllocRefInstBase {
      if let globalValue = optimizeObjectAllocation(allocRef: ari, context) {
        optimizeFindStringCall(stringArray: globalValue, context)
      }
    }
  }
}

private func optimizeObjectAllocation(allocRef: AllocRefInstBase, _ context: FunctionPassContext) -> GlobalValueInst? {
  if !allocRef.fieldsKnownStatically {
    return nil
  }

  // The presence of an end_cow_mutation guarantees that the originally initialized
  // object is not mutated (because it must be copied before mutation).
  guard let endCOW = findEndCOWMutation(of: allocRef),
        !endCOW.doKeepUnique else {
    return nil
  }

  guard let (storesToClassFields, storesToTailElements) = getInitialization(of: allocRef) else {
    return nil
  }

  let outlinedGlobal = context.createGlobalVariable(
        name: context.mangleOutlinedVariable(from: allocRef.parentFunction),
        type: allocRef.type, isPrivate: true)

  constructObject(of: allocRef, inInitializerOf: outlinedGlobal, storesToClassFields, storesToTailElements, context)
  context.erase(instructions: storesToClassFields)
  context.erase(instructions: storesToTailElements)

  return replace(object: allocRef, with: outlinedGlobal, context)
}

private func findEndCOWMutation(of object: Value) -> EndCOWMutationInst? {
  for use in object.uses {
    switch use.instruction {
    case let uci as UpcastInst:
      if let ecm = findEndCOWMutation(of: uci) {
        return ecm
      }
    case let urci as UncheckedRefCastInst:
      if let ecm = findEndCOWMutation(of: urci) {
        return ecm
      }
    case let mv as MoveValueInst:
      if let ecm = findEndCOWMutation(of: mv) {
        return ecm
      }
    case let ecm as EndCOWMutationInst:
      return ecm
    default:
      break
    }
  }
  return nil
}

private func getInitialization(of allocRef: AllocRefInstBase) -> (storesToClassFields: [StoreInst],
                                                                  storesToTailElements: [StoreInst])? {
  guard let numTailElements = allocRef.numTailElements else {
    return nil
  }
  var fieldStores = Array<StoreInst?>(repeating: nil, count: allocRef.numClassFields)

  // If the tail element is a tuple, then its tuple elements are initialized with separate stores.
  // E.g:
  //   %2 = ref_tail_addr
  //   %3 = tuple_element_addr %2, 0
  //   store %0 to %3
  //   %4 = tuple_element_addr %2, 1
  //   store %1 to %4
  var tailStores = Array<StoreInst?>(repeating: nil, count: numTailElements * allocRef.numStoresPerTailElement)

  if !findInitStores(of: allocRef, &fieldStores, &tailStores) {
    return nil
  }

  // Check that all fields and tail elements are initialized.
  if fieldStores.contains(nil) || tailStores.contains(nil) {
    return nil
  }
  return (fieldStores.map { $0! }, tailStores.map { $0! })
}

private func findInitStores(of object: Value,
                            _ fieldStores: inout [StoreInst?],
                            _ tailStores: inout [StoreInst?]) -> Bool {
  for use in object.uses {
    switch use.instruction {
    case let uci as UpcastInst:
      if !findInitStores(of: uci, &fieldStores, &tailStores) {
        return false
      }
    case let urci as UncheckedRefCastInst:
      if !findInitStores(of: urci, &fieldStores, &tailStores) {
        return false
      }
    case let mvi as MoveValueInst:
      if !findInitStores(of: mvi, &fieldStores, &tailStores) {
        return false
      }
    case let rea as RefElementAddrInst:
      if !findStores(inUsesOf: rea, index: rea.fieldIndex, stores: &fieldStores) {
        return false
      }
    case let rta as RefTailAddrInst:
      if !findStores(toTailAddress: rta, tailElementIndex: 0, stores: &tailStores) {
        return false
      }
    default:
      if !isValidUseOfObject(use) {
        return false
      }
    }
  }
  return true
}

private func findStores(toTailAddress tailAddr: Value, tailElementIndex: Int, stores: inout [StoreInst?]) -> Bool {
  for use in tailAddr.uses {
    switch use.instruction {
    case let indexAddr as IndexAddrInst:
      guard let indexLiteral = indexAddr.index as? IntegerLiteralInst,
            let tailIdx = indexLiteral.value else
      {
        return false
      }
      if !findStores(toTailAddress: indexAddr, tailElementIndex: tailElementIndex + tailIdx, stores: &stores) {
        return false
      }
    case let tea as TupleElementAddrInst:
      // The tail elements are tuples. There is a separate store for each tuple element.
      let numTupleElements = tea.tuple.type.tupleElements.count
      let tupleIdx = tea.fieldIndex
      if !findStores(inUsesOf: tea, index: tailElementIndex * numTupleElements + tupleIdx, stores: &stores) {
        return false
      }
    case let atp as AddressToPointerInst:
      if !findStores(toTailAddress: atp, tailElementIndex: tailElementIndex, stores: &stores) {
        return false
      }
    case let mdi as MarkDependenceInst:
      if !findStores(toTailAddress: mdi, tailElementIndex: tailElementIndex, stores: &stores) {
        return false
      }
    case let pta as PointerToAddressInst:
      if !findStores(toTailAddress: pta, tailElementIndex: tailElementIndex, stores: &stores) {
        return false
      }
    case let store as StoreInst:
      if store.source.type.isTuple {
        // This kind of SIL is never generated because tuples are stored with separated stores to tuple_element_addr.
        // Just to be on the safe side..
        return false
      }
      if !handleStore(store, index: tailElementIndex, stores: &stores) {
        return false
      }
    default:
      if !isValidUseOfObject(use) {
        return false
      }
    }
  }
  return true
}

private func findStores(inUsesOf address: Value, index: Int, stores: inout [StoreInst?]) -> Bool {
  for use in address.uses {
    if let store = use.instruction as? StoreInst {
      if !handleStore(store, index: index, stores: &stores) {
        return false
      }
    } else if !isValidUseOfObject(use) {
      return false
    }
  }
  return true
}

private func handleStore(_ store: StoreInst, index: Int, stores: inout [StoreInst?]) -> Bool {
  if index >= 0 && index < stores.count,
     store.source.isValidGlobalInitValue,
     stores[index] == nil {
    stores[index] = store
    return true
  }
  return false
}

private func isValidUseOfObject(_ use: Operand) -> Bool {
  let inst = use.instruction
  switch inst {
  case is DebugValueInst,
       is LoadInst,
       is DeallocRefInst,
       is DeallocStackRefInst,
       is StrongRetainInst,
       is StrongReleaseInst,
       is FixLifetimeInst,
       is EndCOWMutationInst:
    return true

  case let mdi as MarkDependenceInst:
    if (use == mdi.baseOperand) {
      return true;
    }
    for mdiUse in mdi.uses {
      if !isValidUseOfObject(mdiUse) {
        return false
      }
    }
    return true

  case is StructElementAddrInst,
       is AddressToPointerInst,
       is StructInst,
       is TupleInst,
       is TupleExtractInst,
       is EnumInst,
       is StructExtractInst,
       is UncheckedRefCastInst,
       is UpcastInst,
       is BeginDeallocRefInst,
       is RefTailAddrInst,
       is RefElementAddrInst,
       is StructInst,
       is PointerToAddressInst,
       is IndexAddrInst:
    for instUse in (inst as! SingleValueInstruction).uses {
      if !isValidUseOfObject(instUse) {
        return false
      }
    }
    return true

  case let bi as BuiltinInst:
    switch bi.id {
    case .ICMP_EQ, .ICMP_NE:
      // Handle the case for comparing addresses. This occurs when the Array
      // comparison function is inlined.
      return true
    case .DestroyArray:
      // We must not try to delete the tail allocated values. Although this would be a no-op
      // (because we only handle trivial types), it would be semantically wrong to apply this
      // builtin on the outlined object.
      return true
    default:
      return false
    }

  default:
    return false
  }
}

private func constructObject(of allocRef: AllocRefInstBase,
                             inInitializerOf global: GlobalVariable,
                             _ storesToClassFields: [StoreInst], _ storesToTailElements: [StoreInst],
                             _ context: FunctionPassContext) {
  var cloner = StaticInitCloner(cloneTo: global, context)
  defer { cloner.deinitialize() }

  // Create the initializers for the fields
  var objectArgs = [Value]()
  for store in storesToClassFields {
    objectArgs.append(cloner.clone(store.source as! SingleValueInstruction))
  }
  let globalBuilder = Builder(staticInitializerOf: global, context)

  // Create the initializers for the tail elements.
  let numTailTupleElems = allocRef.numStoresPerTailElement
  if numTailTupleElems > 1 {
    // The elements are tuples: combine numTailTupleElems elements to a single tuple instruction.
    for elementIdx in 0..<allocRef.numTailElements! {
      var tupleElems = [Value]()
      for tupleIdx in 0..<numTailTupleElems {
        let store = storesToTailElements[elementIdx * numTailTupleElems + tupleIdx]
        tupleElems.append(cloner.clone(store.source as! SingleValueInstruction))
      }
      let tuple = globalBuilder.createTuple(type: allocRef.tailAllocatedTypes[0], elements: tupleElems)
      objectArgs.append(tuple)
    }
  } else {
    // The non-tuple element case.
    for store in storesToTailElements {
      objectArgs.append(cloner.clone(store.source as! SingleValueInstruction))
    }
  }
  globalBuilder.createObject(type: allocRef.type, arguments: objectArgs, numBaseElements: storesToClassFields.count)

  // The initial value can contain a `begin_access` if it references another global variable by address, e.g.
  //   var p = Point(x: 10, y: 20)
  //   let a = [UnsafePointer(&p)]
  //
  global.stripAccessInstructionFromInitializer(context)
}

private func replace(object allocRef: AllocRefInstBase,
                     with global: GlobalVariable,
                     _ context: FunctionPassContext) -> GlobalValueInst {

  // Replace the alloc_ref by global_value + strong_retain instructions.
  let builder = Builder(before: allocRef, context)
  let globalValue = builder.createGlobalValue(global: global, isBare: false)
  builder.createStrongRetain(operand: globalValue)

  rewriteUses(of: allocRef, context)
  allocRef.uses.replaceAll(with: globalValue, context)
  context.erase(instruction: allocRef)
  return globalValue
}

private func rewriteUses(of startValue: Value, _ context: FunctionPassContext) {
  var worklist = InstructionWorklist(context)
  defer { worklist.deinitialize() }
  worklist.pushIfNotVisited(usersOf: startValue)

  while let inst = worklist.pop() {
    switch inst {
    case let beginDealloc as BeginDeallocRefInst:
      worklist.pushIfNotVisited(usersOf: beginDealloc)
      let builder = Builder(before: beginDealloc, context)
      builder.createStrongRelease(operand: beginDealloc.reference)
      beginDealloc.uses.replaceAll(with: beginDealloc.reference, context)
      context.erase(instruction: beginDealloc)
    case let endMutation as EndCOWMutationInst:
      worklist.pushIfNotVisited(usersOf: endMutation)
      endMutation.uses.replaceAll(with: endMutation.instance, context)
      context.erase(instruction: endMutation)
    case let upCast as UpcastInst:
      worklist.pushIfNotVisited(usersOf: upCast)
    case let refCast as UncheckedRefCastInst:
      worklist.pushIfNotVisited(usersOf: refCast)
    case let moveValue as MoveValueInst:
      worklist.pushIfNotVisited(usersOf: moveValue)
    case is DeallocRefInst, is DeallocStackRefInst:
      context.erase(instruction: inst)
    default:
      break
    }
  }
}

private extension InstructionWorklist {
  mutating func pushIfNotVisited(usersOf value: Value) {
    pushIfNotVisited(contentsOf: value.uses.lazy.map { $0.instruction })
  }
}

private extension AllocRefInstBase {
  var fieldsKnownStatically: Bool {
    if let allocDynamic = self as? AllocRefDynamicInst,
       !allocDynamic.isDynamicTypeDeinitAndSizeKnownEquivalentToBaseType {
      return false
    }
    if isObjC {
      return false
    }
    return true
  }

  var numTailElements: Int? {
    // We only support a single tail allocated array.
    // Stdlib's tail allocated arrays don't have any side-effects in the constructor if the element type is trivial.
    // TODO: also exclude custom tail allocated arrays which might have side-effects in the destructor.
    if tailAllocatedCounts.count != 1 {
      return nil
    }

    // The number of tail allocated elements must be constant.
    if let tailCountLiteral = tailAllocatedCounts[0].value as? IntegerLiteralInst,
       let count = tailCountLiteral.value
    {
      return count
    }
    return nil
  }

  var numClassFields: Int {
    assert(type.isClass)
    return type.getNominalFields(in: parentFunction)!.count
  }

  var numStoresPerTailElement: Int {
    let tailType = tailAllocatedTypes[0]
    if tailType.isTuple {
      return tailType.tupleElements.count
    }
    return 1
  }
}

private extension FunctionPassContext {
  func erase(instructions: [Instruction]) {
    for inst in instructions {
      erase(instruction: inst)
    }
  }
}

private func optimizeFindStringCall(stringArray: GlobalValueInst, _ context: FunctionPassContext) {
  if stringArray.numArrayElements > 16,
     let findStringCall = findFindStringCall(stringArray: stringArray),
     let cachedFindStringFunc = getFindStringSwitchCaseWithCacheFunction(context) {
    replace(findStringCall: findStringCall, with: cachedFindStringFunc, context)
  }
}

/// Finds a call to findStringSwitchCase which takes `stringArray` as parameter.
private func findFindStringCall(stringArray: Value) -> ApplyInst? {
  for use in stringArray.uses {
    switch use.instruction {
    case let apply as ApplyInst:
      // There should only be a single call to findStringSwitchCase. But even
      // if there are multiple calls, it's not problem - we'll just optimize the
      // last one we find.
      if apply.hasSemanticsAttribute("findStringSwitchCase") {
        return apply
      }
    case is StructInst,
         is TupleInst,
         is UncheckedRefCastInst,
         is UpcastInst:
      if let foundCall = findFindStringCall(stringArray: use.instruction as! SingleValueInstruction) {
        return foundCall
      }
    default:
      break
    }
  }
  return nil
}

private func getFindStringSwitchCaseWithCacheFunction(_ context: FunctionPassContext) -> Function? {
  if let f = context.lookupStdlibFunction(name: "_findStringSwitchCaseWithCache"),
     f.argumentTypes.count == 3 {
    return f
  }
  return nil
}

private func replace(findStringCall: ApplyInst,
                     with cachedFindStringFunc: Function,
                     _ context: FunctionPassContext) {
  let cacheType = cachedFindStringFunc.argumentTypes[2].objectType
  let wordTy = cacheType.getNominalFields(in: findStringCall.parentFunction)![0]

  let name = context.mangleOutlinedVariable(from: findStringCall.parentFunction)

  // Create an "opaque" global variable which is passed as inout to
  // _findStringSwitchCaseWithCache and into which the function stores the "cache".
  let cacheVar = context.createGlobalVariable(name: name, type: cacheType, isPrivate: true)

  let varBuilder = Builder(staticInitializerOf: cacheVar, context)
  let zero = varBuilder.createIntegerLiteral(0, type: wordTy)
  _ = varBuilder.createStruct(type: cacheType, elements: [zero, zero])

  let builder = Builder(before: findStringCall, context)
  let cacheAddr = builder.createGlobalAddr(global: cacheVar, dependencyToken: nil)
  let findStringRef = builder.createFunctionRef(cachedFindStringFunc)
  let newCall = builder.createApply(function: findStringRef, SubstitutionMap(),
                                    arguments: [findStringCall.arguments[0],
                                                findStringCall.arguments[1],
                                                cacheAddr])

  findStringCall.uses.replaceAll(with: newCall, context)
  context.erase(instruction: findStringCall)
}

private extension GlobalValueInst {
  /// Assuming the global is an Array, returns the number of elements = tail elements.
  var numArrayElements: Int {
    (global.staticInitValue! as! ObjectInst).tailOperands.count
  }
}