SwiftCompilerSources: Replace BlockArgument with Phi and TermResult.

All SILArgument types are "block arguments". There are three kinds:
1. Function arguments
2. Phis
3. Terminator results

In every situation where the source of the block argument matters, we
need to distinguish between these three. Accidentally failing to
handle one of the cases is an perpetual source of compiler
bugs. Attempting to handle both phis and terminator results uniformly
is *always* a bug, especially once OSSA has phi flags. Even when all
cases are handled correctly, the code that deals with data flow across
blocks is incomprehensible without giving each case a type. This
continues to be a massive waste of time literally every time I review
code that involves cross-block control flow.

Unfortunately, we don't have these C++ types yet (nothing big is
blocking that, it just wasn't done). That's manageable because we can
use wrapper types on the Swift side for now. Wrapper types don't
create any more complexity than protocols, but they do sacrifice some
usability in switch cases.

There is no reason for a BlockArgument type. First, a function
argument is a block argument just as much as any other. BlockArgument
provides no useful information beyond Argument. And it is nearly
always a mistake to care about whether a value is a function argument
and not care whether it is a phi or terminator result.

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ObjCBridgingOptimization.swift

@@ -211,11 +211,11 @@ private func optimizeNonOptionalBridging(_ apply: ApplyInst,
 
   // In the some-case just forward the original NSString.
   let objCType = emptyObjCValue.type
-  let forwardedValue = someBlock.addBlockArgument(type: objCType, ownership: .owned, context)
+  let forwardedValue = someBlock.addArgument(type: objCType, ownership: .owned, context)
   let someBuilder = Builder(atEndOf: someBlock, location: bridgeToSwiftCall.location, context)
   someBuilder.createBranch(to: continueBlock, arguments: [forwardedValue])
 
-  let s = continueBlock.addBlockArgument(type: objCType, ownership: .owned, context)
+  let s = continueBlock.addArgument(type: objCType, ownership: .owned, context)
 
   // Now replace the bridged value with the original value in the destination block.
   let replacement = s.makeAvailable(in: bridgeToObjcCall.parentBlock, context)
@@ -272,16 +272,13 @@ private func lookThroughOwnershipInsts(_ value: Value) -> Value {
 /// continue_bb(%value):         // passed value
 /// ```
 private func isOptionalBridging(of value: Value, isBridging: (Value) -> ApplyInst?) -> SwitchEnumInst? {
-  guard let arg = value as? BlockArgument,
-        arg.isPhiArgument else {
-    return nil
-  }
+  guard let phi = Phi(value) else { return nil }
 
   var noneSwitch: SwitchEnumInst?
   var someSwitch: SwitchEnumInst?
 
   // Check if one incoming value is the none-case and the other is the some-case.
-  for incomingVal in arg.incomingPhiValues {
+  for incomingVal in phi.incomingValues {
     // In both branches, the result must be an `enum` which is passed to the
     // continue_bb's phi-argument.
     guard let enumInst = incomingVal as? EnumInst,
@@ -338,10 +335,9 @@ private func isOptionalBridging(of value: Value, isBridging: (Value) -> ApplyIns
 /// Returns the `switch_enum` together with the enum case index, if `value` is
 /// the payload block argument of the `switch_enum`.
 private func isPayloadOfSwitchEnum(_ value: Value) -> (SwitchEnumInst, case: Int)? {
-  if let payloadArg = value as? BlockArgument,
-     let pred = payloadArg.parentBlock.singlePredecessor,
-     let se = pred.terminator as? SwitchEnumInst,
-     let caseIdx = se.getUniqueCase(forSuccessor: payloadArg.parentBlock) {
+  if let payloadArg = TerminatorResult(value),
+     let se = payloadArg.terminator as? SwitchEnumInst,
+     let caseIdx = se.getUniqueCase(forSuccessor: payloadArg.successor) {
     return (se, caseIdx)
   }
   return nil

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/StackPromotion.swift

@@ -267,7 +267,7 @@ private struct ComputeOuterBlockrange : EscapeVisitorWithResult {
     // instructions (for which the `visitUse` closure is not called).
     result.insert(operandsDefinitionBlock)
 
-    // We need to explicitly add predecessor blocks of phi-arguments becaues they
+    // We need to explicitly add predecessor blocks of phis becaues they
     // are not necesesarily visited during the down-walk in `isEscaping()`.
     // This is important for the special case where there is a back-edge from the
     // inner range to the inner rage's begin-block:
@@ -278,8 +278,8 @@ private struct ComputeOuterBlockrange : EscapeVisitorWithResult {
     //     %k = alloc_ref $Klass   // innerInstRange.begin
     //     cond_br bb2, bb1(%k)    // back-edge to bb1 == innerInstRange.blockRange.begin
     //
-    if value is BlockArgument {
-      result.insert(contentsOf: operandsDefinitionBlock.predecessors)
+    if let phi = Phi(value) {
+      result.insert(contentsOf: phi.predecessors)
     }
     return .continueWalk
   }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyRefCasts.swift

@@ -111,7 +111,7 @@ private extension UnaryInstruction {
 private func insertCompensatingInstructions(for inst: Instruction, in failureBlock: BasicBlock, _ context: SimplifyContext) {
   assert(failureBlock.arguments.count == 1)
   let sourceValue = inst.operands[0].value
-  let newArg = failureBlock.addBlockArgument(type: sourceValue.type, ownership: sourceValue.ownership, context)
+  let newArg = failureBlock.addArgument(type: sourceValue.type, ownership: sourceValue.ownership, context)
   let builder = Builder(atBeginOf: failureBlock, context)
   let newInst: SingleValueInstruction
   switch inst {

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -400,9 +400,9 @@ extension Undef {
 }
 
 extension BasicBlock {
-  func addBlockArgument(type: Type, ownership: Ownership, _ context: some MutatingContext) -> BlockArgument {
+  func addArgument(type: Type, ownership: Ownership, _ context: some MutatingContext) -> Argument {
     context.notifyInstructionsChanged()
-    return bridged.addBlockArgument(type.bridged, ownership._bridged).blockArgument
+    return bridged.addBlockArgument(type.bridged, ownership._bridged).argument
   }
 
   func eraseArgument(at index: Int, _ context: some MutatingContext) {

SwiftCompilerSources/Sources/Optimizer/TestPasses/SILPrinter.swift

@@ -29,8 +29,8 @@ func runSILPrinter(function: Function, context: FunctionPassContext) {
       for use in arg.uses {
         print("      user: \(use.instruction)")
       }
-      if let blockArg = arg as? BlockArgument, blockArg.isPhiArgument {
-        for incoming in blockArg.incomingPhiValues {
+      if let phi = Phi(arg) {
+        for incoming in phi.incomingValues {
           print("      incoming: \(incoming)")
         }
       }

SwiftCompilerSources/Sources/Optimizer/Utilities/AccessUtils.swift

@@ -389,9 +389,8 @@ extension PointerToAddressInst {
 
       mutating func walkUp(value: Value, path: SmallProjectionPath) -> WalkResult {
         switch value {
-        case is BlockArgument, is MarkDependenceInst, is CopyValueInst,
-             is StructExtractInst, is TupleExtractInst, is StructInst, is TupleInst,
-             is FunctionArgument, is AddressToPointerInst:
+        case is Argument, is MarkDependenceInst, is CopyValueInst,
+             is StructExtractInst, is TupleExtractInst, is StructInst, is TupleInst, is AddressToPointerInst:
           return walkUpDefault(value: value, path: path)
         default:
           return .abortWalk

SwiftCompilerSources/Sources/Optimizer/Utilities/EscapeUtils.swift

@@ -680,11 +680,11 @@ fileprivate struct EscapeWalker<V: EscapeVisitor> : ValueDefUseWalker,
       return cachedWalkDown(addressOrValue: def, path: path.with(knownType: def.type))
     case is AllocBoxInst:
       return cachedWalkDown(addressOrValue: def, path: path.with(knownType: nil))
-    case let arg as BlockArgument:
-      let block = arg.parentBlock
-      switch block.singlePredecessor!.terminator {
+    case let arg as Argument:
+      guard let termResult = TerminatorResult(arg) else { return isEscaping }
+      switch termResult.terminator {
       case let ta as TryApplyInst:
-        if block != ta.normalBlock { return isEscaping }
+        if termResult.successor != ta.normalBlock { return isEscaping }
         return walkUpApplyResult(apply: ta, path: path.with(knownType: nil))
       default:
         return isEscaping

SwiftCompilerSources/Sources/Optimizer/Utilities/WalkUtils.swift

@@ -673,9 +673,9 @@ extension ValueUseDefWalker {
          is BeginDeallocRefInst,
          is RefToBridgeObjectInst, is BridgeObjectToRefInst, is MarkUnresolvedNonCopyableValueInst:
       return walkUp(value: (def as! Instruction).operands[0].value, path: path)
-    case let arg as BlockArgument:
-      if arg.isPhiArgument {
-        for incoming in arg.incomingPhiValues {
+    case let arg as Argument:
+      if let phi = Phi(arg) {
+        for incoming in phi.incomingValues {
           // Check the cache to avoid cycles in the walk
           if let path = walkUpCache.needWalk(for: incoming, path: path) {
             if walkUp(value: incoming, path: path) == .abortWalk {
@@ -685,14 +685,13 @@ extension ValueUseDefWalker {
         }
         return .continueWalk
       }
-      
-      let block = arg.parentBlock
-      if let pred = block.singlePredecessor,
-         let se = pred.terminator as? SwitchEnumInst,
-         let caseIdx = se.getUniqueCase(forSuccessor: block) {
-        return walkUp(value: se.enumOp, path: path.push(.enumCase, index: caseIdx))
+      if let termResult = TerminatorResult(arg) {
+        let pred = termResult.predecessor
+        if let se = pred.terminator as? SwitchEnumInst,
+           let caseIdx = se.getUniqueCase(forSuccessor: termResult.successor) {
+          return walkUp(value: se.enumOp, path: path.push(.enumCase, index: caseIdx))
+        }
       }
-      
       return rootDef(value: def, path: path)
     default:
       return rootDef(value: def, path: path)

SwiftCompilerSources/Sources/SIL/Argument.swift

@@ -52,37 +52,91 @@ final public class FunctionArgument : Argument {
   }
 }
 
-final public class BlockArgument : Argument {
-  public var isPhiArgument: Bool {
-    parentBlock.predecessors.allSatisfy {
-      let term = $0.terminator
-      return term is BranchInst || term is CondBranchInst
-    }
+public struct Phi {
+  public let value: Argument
+
+  // TODO: Remove the CondBr case. All passes avoid critical edges. It is only included here for compatibility with .sil tests that have not been migrated.
+  public init?(_ value: Value) {
+    guard let argument = value as? Argument else { return nil }
+    var preds = argument.parentBlock.predecessors
+    guard let pred = preds.next() else { return nil }
+    let term = pred.terminator
+    guard term is BranchInst || term is CondBranchInst else { return nil }
+    self.value = argument
+  }
+  
+  public var predecessors: PredecessorList {
+    return value.parentBlock.predecessors
   }
 
-  public var incomingPhiOperands: LazyMapSequence<PredecessorList, Operand> {
-    assert(isPhiArgument)
-    let idx = index
-    return parentBlock.predecessors.lazy.map {
+  public var successor: BasicBlock {
+    return value.parentBlock
+  }
+
+  public var incomingOperands: LazyMapSequence<PredecessorList, Operand> {
+    let blockArgIdx = value.index
+    return predecessors.lazy.map {
       switch $0.terminator {
         case let br as BranchInst:
-          return br.operands[idx]
+          return br.operands[blockArgIdx]
         case let condBr as CondBranchInst:
-          if condBr.trueBlock == self.parentBlock {
-            assert(condBr.falseBlock != self.parentBlock)
-            return condBr.trueOperands[idx]
+          if condBr.trueBlock == successor {
+            assert(condBr.falseBlock != successor)
+            return condBr.trueOperands[blockArgIdx]
           } else {
-            assert(condBr.falseBlock == self.parentBlock)
-            return condBr.falseOperands[idx]
+            assert(condBr.falseBlock == successor)
+            return condBr.falseOperands[blockArgIdx]
           }
         default:
           fatalError("wrong terminator for phi-argument")
       }
     }
   }
 
-  public var incomingPhiValues: LazyMapSequence<LazyMapSequence<PredecessorList, Operand>, Value> {
-    incomingPhiOperands.lazy.map { $0.value }
+  public var incomingValues: LazyMapSequence<LazyMapSequence<PredecessorList, Operand>, Value> {
+    incomingOperands.lazy.map { $0.value }
+  }
+
+  public static func ==(lhs: Phi, rhs: Phi) -> Bool {
+    lhs.value === rhs.value
+  }
+
+  public func hash(into hasher: inout Hasher) {
+    value.hash(into: &hasher)
+  }
+}
+
+public struct TerminatorResult {
+  public let value: Argument
+
+  public init?(_ value: Value) {
+    guard let argument = value as? Argument else { return nil }
+    var preds = argument.parentBlock.predecessors
+    guard let pred = preds.next() else { return nil }
+    let term = pred.terminator
+    if term is BranchInst || term is CondBranchInst { return nil }
+    self.value = argument
+  }
+
+  public var terminator: TermInst {
+    var preds = value.parentBlock.predecessors
+    return preds.next()!.terminator
+  }
+  
+  public var predecessor: BasicBlock {
+    return terminator.parentBlock
+  }
+
+  public var successor: BasicBlock {
+    return value.parentBlock
+  }
+
+  public static func ==(lhs: TerminatorResult, rhs: TerminatorResult) -> Bool {
+    lhs.value == rhs.value
+  }
+
+  public func hash(into hasher: inout Hasher) {
+    value.hash(into: &hasher)
   }
 }
 
@@ -242,7 +296,6 @@ public enum ArgumentConvention {
 
 extension BridgedArgument {
   public var argument: Argument { obj.getAs(Argument.self) }
-  public var blockArgument: BlockArgument { obj.getAs(BlockArgument.self) }
   public var functionArgument: FunctionArgument { obj.getAs(FunctionArgument.self) }
 }
 

SwiftCompilerSources/Sources/SIL/Registration.swift

@@ -36,7 +36,7 @@ public func registerSILClasses() {
   register(PlaceholderValue.self)
 
   register(FunctionArgument.self)
-  register(BlockArgument.self)
+  register(Argument.self)
 
   register(StoreInst.self)
   register(StoreWeakInst.self)