Merge pull request #72388 from eeckstein/simplify_addr_casts

Simplify `checked_cast_addr_br` to enable a shortcut check  in embedded swift's `Array.append(contentsOf:)`

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/CMakeLists.txt

@@ -12,6 +12,7 @@ swift_compiler_sources(Optimizer
   SimplifyBeginCOWMutation.swift
   SimplifyBranch.swift
   SimplifyBuiltin.swift
+  SimplifyCheckedCast.swift
   SimplifyCondBranch.swift
   SimplifyCondFail.swift
   SimplifyConvertEscapeToNoEscape.swift

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyBuiltin.swift

@@ -285,61 +285,21 @@ private func typesOfValuesAreEqual(_ lhs: Value, _ rhs: Value, in function: Func
   let typesAreExact = lhsExistential.metatype is MetatypeInst &&
                       rhsExistential.metatype is MetatypeInst
 
-  switch (lhsTy.typeKind, rhsTy.typeKind) {
-  case (_, .unknown), (.unknown, _):
-    return nil
-  case (let leftKind, let rightKind) where leftKind != rightKind:
-    // E.g. a function type is always different than a struct, regardless of what archetypes
-    // the two types may contain.
-    return false
-  case (.struct, .struct), (.enum, .enum):
-    // Two different structs/enums are always not equal, regardless of what archetypes
-    // the two types may contain.
-    if lhsTy.nominal != rhsTy.nominal {
-      return false
+  if typesAreExact {
+    if lhsTy == rhsTy {
+      return true
     }
-  case (.class, .class):
-    // In case of classes this only holds if we know the exact types.
-    // Otherwise one class could be a sub-class of the other class.
-    if typesAreExact && lhsTy.nominal != rhsTy.nominal {
+    // Comparing types of different classes which are in a sub-class relation is not handled by the
+    // cast optimizer (below).
+    if lhsTy.isClass && rhsTy.isClass && lhsTy.nominal != rhsTy.nominal {
       return false
     }
-  default:
-    break
   }
 
-  if !typesAreExact {
-    // Types which e.g. come from type parameters may differ at runtime while the declared AST types are the same.
-    return nil
-  }
-
-  if lhsTy.hasArchetype || rhsTy.hasArchetype {
-    // We don't know anything about archetypes. They may be identical at runtime or not.
-    // We could do something more sophisticated here, e.g. look at conformances. But for simplicity,
-    // we are just conservative.
-    return nil
-  }
-
-  // Generic ObjectiveC class, which are specialized for different NSObject types have different AST types
-  // but the same runtime metatype.
-  if lhsTy.isOrContainsObjectiveCClass || rhsTy.isOrContainsObjectiveCClass {
-    return nil
-  }
-
-  return lhsTy == rhsTy
-}
-
-private extension Type {
-  enum TypeKind {
-    case `struct`, `class`, `enum`, tuple, function, unknown
-  }
-
-  var typeKind: TypeKind {
-    if isStruct  { return .struct }
-    if isClass  { return .class }
-    if isEnum  { return .enum }
-    if isTuple    { return .tuple }
-    if isFunction { return .function }
-    return .unknown
+  // If casting in either direction doesn't work, the types cannot be equal.
+  if !(canDynamicallyCast(from: lhsTy, to: rhsTy, in: function, sourceTypeIsExact: typesAreExact) ?? true) ||
+     !(canDynamicallyCast(from: rhsTy, to: lhsTy, in: function, sourceTypeIsExact: typesAreExact) ?? true) {
+    return false
   }
+  return nil
 }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyCheckedCast.swift

@@ -0,0 +1,52 @@
+//===--- SimplifyCheckedCast.swift ----------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 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
+
+extension CheckedCastAddrBranchInst : OnoneSimplifyable {
+  func simplify(_ context: SimplifyContext) {
+    guard let castWillSucceed = self.dynamicCastResult else {
+      return
+    }
+    if castWillSucceed {
+      replaceSuccess(context)
+    } else {
+      replaceFailure(context)
+    }
+  }
+}
+
+private extension CheckedCastAddrBranchInst {
+  func replaceSuccess(_ context: SimplifyContext) {
+    let builder = Builder(before: self, context)
+    switch consumptionKind {
+    case .TakeAlways, .TakeOnSuccess:
+      builder.createCopyAddr(from: source, to: destination, takeSource: true, initializeDest: true)
+    case .CopyOnSuccess:
+      builder.createCopyAddr(from: source, to: destination, takeSource: false, initializeDest: true)
+    }
+    builder.createBranch(to: successBlock)
+    context.erase(instruction: self)
+  }
+
+  func replaceFailure(_ context: SimplifyContext) {
+    let builder = Builder(before: self, context)
+    switch consumptionKind {
+    case .TakeAlways:
+      builder.createDestroyAddr(address: source)
+    case .CopyOnSuccess, .TakeOnSuccess:
+      break
+    }
+    builder.createBranch(to: failureBlock)
+    context.erase(instruction: self)
+  }
+}

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -695,3 +695,23 @@ func getGlobalInitialization(
   }
   return nil
 }
+
+func canDynamicallyCast(from sourceType: Type, to destType: Type, in function: Function, sourceTypeIsExact: Bool) -> Bool? {
+  switch classifyDynamicCastBridged(sourceType.bridged, destType.bridged, function.bridged, sourceTypeIsExact) {
+    case .willSucceed: return true
+    case .maySucceed:  return nil
+    case .willFail:    return false
+    default: fatalError("unknown result from classifyDynamicCastBridged")
+  }
+}
+
+extension CheckedCastAddrBranchInst {
+  var dynamicCastResult: Bool? {
+    switch classifyDynamicCastBridged(bridged) {
+      case .willSucceed: return true
+      case .maySucceed:  return nil
+      case .willFail:    return false
+      default: fatalError("unknown result from classifyDynamicCastBridged")
+    }
+  }
+}

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -1481,5 +1481,35 @@ final public class CheckedCastBranchInst : TermInst, UnaryInstruction {
   public var failureBlock: BasicBlock { bridged.CheckedCastBranch_getFailureBlock().block }
 }
 
-final public class CheckedCastAddrBranchInst : TermInst, UnaryInstruction {
+final public class CheckedCastAddrBranchInst : TermInst {
+  public var source: Value { operands[0].value }
+  public var destination: Value { operands[1].value }
+
+  public var successBlock: BasicBlock { bridged.CheckedCastAddrBranch_getSuccessBlock().block }
+  public var failureBlock: BasicBlock { bridged.CheckedCastAddrBranch_getFailureBlock().block }
+
+  public enum CastConsumptionKind {
+    /// The source value is always taken, regardless of whether the cast
+    /// succeeds.  That is, if the cast fails, the source value is
+    /// destroyed.
+    case TakeAlways
+
+    /// The source value is taken only on a successful cast; otherwise,
+    /// it is left in place.
+    case TakeOnSuccess
+
+    /// The source value is always left in place, and the destination
+    /// value is copied into on success.
+    case CopyOnSuccess
+  }
+
+  public var consumptionKind: CastConsumptionKind {
+    switch bridged.CheckedCastAddrBranch_getConsumptionKind() {
+    case .TakeAlways:    return .TakeAlways
+    case .TakeOnSuccess: return .TakeOnSuccess
+    case .CopyOnSuccess: return .CopyOnSuccess
+    default:
+      fatalError("invalid cast consumption kind")
+    }
+  }
 }

include/swift/SIL/SILBridging.h

@@ -823,6 +823,12 @@ struct BridgedInstruction {
     SwiftInt numFunctions;
   };
 
+  enum class CastConsumptionKind {
+    TakeAlways,
+    TakeOnSuccess,
+    CopyOnSuccess
+  };
+
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedStringRef CondFailInst_getMessage() const;
   BRIDGED_INLINE SwiftInt LoadInst_getLoadOwnership() const ;
   BRIDGED_INLINE BuiltinValueKind BuiltinInst_getID() const;
@@ -904,6 +910,9 @@ struct BridgedInstruction {
   BRIDGED_INLINE void LoadInst_setOwnership(SwiftInt ownership) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedBasicBlock CheckedCastBranch_getSuccessBlock() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedBasicBlock CheckedCastBranch_getFailureBlock() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedBasicBlock CheckedCastAddrBranch_getSuccessBlock() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedBasicBlock CheckedCastAddrBranch_getFailureBlock() const;
+  BRIDGED_INLINE CastConsumptionKind CheckedCastAddrBranch_getConsumptionKind() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedSubstitutionMap ApplySite_getSubstitutionMap() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedASTType ApplySite_getSubstitutedCalleeType() const;
   BRIDGED_INLINE SwiftInt ApplySite_getNumArguments() const;

include/swift/SIL/SILBridgingImpl.h

@@ -1244,6 +1244,27 @@ BridgedBasicBlock BridgedInstruction::CheckedCastBranch_getFailureBlock() const
   return {getAs<swift::CheckedCastBranchInst>()->getFailureBB()};
 }
 
+BridgedBasicBlock BridgedInstruction::CheckedCastAddrBranch_getSuccessBlock() const {
+  return {getAs<swift::CheckedCastAddrBranchInst>()->getSuccessBB()};
+}
+
+BridgedBasicBlock BridgedInstruction::CheckedCastAddrBranch_getFailureBlock() const {
+  return {getAs<swift::CheckedCastAddrBranchInst>()->getFailureBB()};
+}
+
+BridgedInstruction::CastConsumptionKind BridgedInstruction::CheckedCastAddrBranch_getConsumptionKind() const {
+  static_assert((int)BridgedInstruction::CastConsumptionKind::TakeAlways ==
+                (int)swift::CastConsumptionKind::TakeAlways);
+  static_assert((int)BridgedInstruction::CastConsumptionKind::TakeOnSuccess ==
+                (int)swift::CastConsumptionKind::TakeOnSuccess);
+  static_assert((int)BridgedInstruction::CastConsumptionKind::CopyOnSuccess ==
+                (int)swift::CastConsumptionKind::CopyOnSuccess);
+
+  return static_cast<BridgedInstruction::CastConsumptionKind>(
+           getAs<swift::CheckedCastAddrBranchInst>()->getConsumptionKind());
+}
+
+
 BridgedSubstitutionMap BridgedInstruction::ApplySite_getSubstitutionMap() const {
   auto as = swift::ApplySite(unbridged());
   return as.getSubstitutionMap();

include/swift/SILOptimizer/OptimizerBridging.h

@@ -331,6 +331,18 @@ struct BridgedPassContext {
 
 bool FullApplySite_canInline(BridgedInstruction apply);
 
+enum class BridgedDynamicCastResult {
+  willSucceed,
+  maySucceed,
+  willFail
+};
+
+BridgedDynamicCastResult classifyDynamicCastBridged(BridgedType sourceTy, BridgedType destTy,
+                                                    BridgedFunction function,
+                                                    bool sourceTypeIsExact);
+
+BridgedDynamicCastResult classifyDynamicCastBridged(BridgedInstruction inst);
+
 //===----------------------------------------------------------------------===//
 //                          Pass registration
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -18,6 +18,7 @@
 #include "swift/Demangling/Demangle.h"
 #include "../../IRGen/IRGenModule.h"
 #include "swift/SIL/ApplySite.h"
+#include "swift/SIL/DynamicCasts.h"
 #include "swift/SIL/SILCloner.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILModule.h"
@@ -1838,6 +1839,25 @@ bool FullApplySite_canInline(BridgedInstruction apply) {
       swift::FullApplySite(apply.unbridged()));
 }
 
+BridgedDynamicCastResult classifyDynamicCastBridged(BridgedType sourceTy, BridgedType destTy,
+                                                    BridgedFunction function,
+                                                    bool sourceTypeIsExact) {
+  static_assert((int)DynamicCastFeasibility::WillSucceed == (int)BridgedDynamicCastResult::willSucceed);
+  static_assert((int)DynamicCastFeasibility::MaySucceed  == (int)BridgedDynamicCastResult::maySucceed);
+  static_assert((int)DynamicCastFeasibility::WillFail    == (int)BridgedDynamicCastResult::willFail);
+
+  return static_cast<BridgedDynamicCastResult>(
+    classifyDynamicCast(function.getFunction()->getModule().getSwiftModule(),
+                        sourceTy.unbridged().getASTType(),
+                        destTy.unbridged().getASTType(),
+                        sourceTypeIsExact));
+}
+
+BridgedDynamicCastResult classifyDynamicCastBridged(BridgedInstruction inst) {
+  SILDynamicCastInst castInst(inst.unbridged());
+  return static_cast<BridgedDynamicCastResult>(castInst.classifyFeasibility(/*allowWholeModule=*/ false));
+}
+
 // TODO: can't be inlined to work around https://github.com/apple/swift/issues/64502
 BridgedCalleeAnalysis::CalleeList BridgedCalleeAnalysis::getCallees(BridgedValue callee) const {
   return ca->getCalleeListOfValue(callee.getSILValue());

stdlib/public/core/Array.swift

@@ -1254,7 +1254,6 @@ extension Array: RangeReplaceableCollection {
 
     if _slowPath(writtenUpTo == buf.endIndex) {
 
-#if !$Embedded
       // A shortcut for appending an Array: If newElements is an Array then it's
       // guaranteed that buf.initialize(from: newElements) already appended all
       // elements. It reduces code size, because the following code
@@ -1264,7 +1263,6 @@ extension Array: RangeReplaceableCollection {
         _internalInvariant(remainder.next() == nil)
         return
       }
-#endif
 
       // there may be elements that didn't fit in the existing buffer,
       // append them in slow sequence-only mode