Concrete type propagation using ProtocolConformanceAnalysis

Author: Raj Barik

include/swift/SILOptimizer/Analysis/ProtocolConformanceAnalysis.h

@@ -35,6 +35,8 @@ class ProtocolConformanceAnalysis : public SILAnalysis {
   typedef SmallVector<NominalTypeDecl *, 8> NominalTypeList;
   typedef llvm::DenseMap<ProtocolDecl *, NominalTypeList>
       ProtocolConformanceMap;
+  typedef llvm::DenseMap<ProtocolDecl *, NominalTypeDecl *>
+      SoleConformingTypeMap;
 
   ProtocolConformanceAnalysis(SILModule *Mod)
       : SILAnalysis(SILAnalysisKind::ProtocolConformance), M(Mod) {
@@ -71,6 +73,10 @@ class ProtocolConformanceAnalysis : public SILAnalysis {
                                              ConformsListIt->second.end())
                : ArrayRef<NominalTypeDecl *>();
   }
+  
+  /// Traverse ProtocolConformanceMapCache recursively to determine sole
+  /// conforming concrete type. 
+  NominalTypeDecl *findSoleConformingType(ProtocolDecl *Protocol);
 
 private:
   /// Compute inheritance properties.
@@ -79,8 +85,11 @@ class ProtocolConformanceAnalysis : public SILAnalysis {
   /// The module.
   SILModule *M;
 
-  /// A cache that maps a protocol to its conformances
+  /// A cache that maps a protocol to its conformances.
   ProtocolConformanceMap ProtocolConformanceCache;
+
+  /// A cache that holds SoleConformingType for protocols.
+  SoleConformingTypeMap SoleConformingTypeCache;
 };
 
 } // namespace swift

include/swift/SILOptimizer/Utils/Existential.h

@@ -13,7 +13,10 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_EXISTENTIAL_H
 #define SWIFT_SILOPTIMIZER_UTILS_EXISTENTIAL_H
 
+#include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILInstruction.h"
+#include "swift/SILOptimizer/Analysis/ClassHierarchyAnalysis.h"
+#include "swift/SILOptimizer/Analysis/ProtocolConformanceAnalysis.h"
 
 namespace swift {
 
@@ -86,13 +89,24 @@ struct ConcreteExistentialInfo {
   // Constructs a valid ConcreteExistentialInfo object if successfull.
   ConcreteExistentialInfo(Operand &openedUse);
 
+  // This constructor initializes a ConcreteExistentialInfo based on already
+  // known ConcreteType and ProtocolDecl pair. It determines the
+  // OpenedArchetypeDef for the ArgOperand that will be used by unchecked_cast
+  // instructions to cast OpenedArchetypeDef to ConcreteType.
+  ConcreteExistentialInfo(Operand &ArgOperand, CanType ConcreteType,
+                          ProtocolDecl *Protocol);
+
   ConcreteExistentialInfo(ConcreteExistentialInfo &) = delete;
 
-  /// We do not not need to check for InitExistential not being not null
-  /// since if that were the case, we would have everything else null too.
+  /// For scenerios where ConcreteExistentialInfo is created using a known
+  /// ConcreteType and ProtocolDecl, both of InitExistential
+  /// and ConcreteValue can be null. So there is no need for explicit check for
+  /// not null for them instead we assert on (!InitExistential ||
+  /// ConcreteValue). 
   bool isValid() const {
+    assert(!InitExistential || ConcreteValue);
     return OpenedArchetype && OpenedArchetypeDef && ConcreteType &&
-           !ExistentialSubs.empty() && ConcreteValue;
+           !ExistentialSubs.empty();
   }
 
   // Do a conformance lookup on ConcreteType with the given requirement, P. If P

lib/SILOptimizer/Analysis/ProtocolConformanceAnalysis.cpp

@@ -67,4 +67,60 @@ void ProtocolConformanceAnalysis::init() {
   }
 }
 
+/// Recursively traverse the conformance lists to determine sole conforming
+/// class, struct or enum type.
+NominalTypeDecl *
+ProtocolConformanceAnalysis::findSoleConformingType(ProtocolDecl *Protocol) {
+
+  /// First check in the SoleConformingTypeCache.
+  auto SoleConformingTypeIt = SoleConformingTypeCache.find(Protocol);
+  if (SoleConformingTypeIt != SoleConformingTypeCache.end())
+    return SoleConformingTypeIt->second;
+
+  SmallVector<ProtocolDecl *, 8> PDWorkList;
+  SmallPtrSet<ProtocolDecl *, 8> VisitedPDs;
+  NominalTypeDecl *SoleConformingNTD = nullptr;
+  PDWorkList.push_back(Protocol);
+  while (!PDWorkList.empty()) {
+    auto *PD = PDWorkList.pop_back_val();
+    // Protocols must have internal or lower access.
+    if (PD->getEffectiveAccess() > AccessLevel::Internal) {
+      return nullptr;
+    }
+    VisitedPDs.insert(PD);
+    auto NTDList = getConformances(PD);
+    for (auto *ConformingNTD : NTDList) {
+      // Recurse on protocol types.
+      if (auto *Proto = dyn_cast<ProtocolDecl>(ConformingNTD)) {
+        // Ignore visited protocol decls.
+        if (!VisitedPDs.count(Proto))
+          PDWorkList.push_back(Proto);
+      } else { // Classes, Structs and Enums are added here.
+        // Bail if more than one conforming types were found.
+        if (SoleConformingNTD && ConformingNTD != SoleConformingNTD) {
+          return nullptr;
+        } else {
+          SoleConformingNTD = ConformingNTD;
+        }
+      }
+    }
+  }
+
+  // Bail if we did not find a sole conforming type.
+  if (!SoleConformingNTD)
+    return nullptr;
+
+  // Generic declarations are ignored.
+  if (SoleConformingNTD->isGenericContext()) {
+    return nullptr;
+  }
+
+  // Populate SoleConformingTypeCache.
+  SoleConformingTypeCache.insert(std::pair<ProtocolDecl *, NominalTypeDecl *>(
+      Protocol, SoleConformingNTD));
+
+  // Return SoleConformingNTD.
+  return SoleConformingNTD;
+}
+
 ProtocolConformanceAnalysis::~ProtocolConformanceAnalysis() {}

lib/SILOptimizer/SILCombiner/SILCombine.cpp

@@ -347,12 +347,15 @@ class SILCombine : public SILFunctionTransform {
   void run() override {
     auto *AA = PM->getAnalysis<AliasAnalysis>();
     auto *DA = PM->getAnalysis<DominanceAnalysis>();
+    auto *PCA = PM->getAnalysis<ProtocolConformanceAnalysis>();
+    auto *CHA = PM->getAnalysis<ClassHierarchyAnalysis>();
 
     SILOptFunctionBuilder FuncBuilder(*this);
     // Create a SILBuilder with a tracking list for newly added
     // instructions, which we will periodically move to our worklist.
     SILBuilder B(*getFunction(), &TrackingList);
-    SILCombiner Combiner(FuncBuilder, B, AA, DA, getOptions().RemoveRuntimeAsserts);
+    SILCombiner Combiner(FuncBuilder, B, AA, DA, PCA, CHA,
+                         getOptions().RemoveRuntimeAsserts);
     bool Changed = Combiner.runOnFunction(*getFunction());
     assert(TrackingList.empty() &&
            "TrackingList should be fully processed by SILCombiner");

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -25,6 +25,8 @@
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILValue.h"
 #include "swift/SIL/SILVisitor.h"
+#include "swift/SILOptimizer/Analysis/ClassHierarchyAnalysis.h"
+#include "swift/SILOptimizer/Analysis/ProtocolConformanceAnalysis.h"
 #include "swift/SILOptimizer/Utils/CastOptimizer.h"
 #include "swift/SILOptimizer/Utils/Existential.h"
 #include "swift/SILOptimizer/Utils/Local.h"
@@ -118,6 +120,14 @@ class SILCombiner :
 
   DominanceAnalysis *DA;
 
+  // Determine the set of types a protocol conforms to in whole-module
+  // compilation mode.
+  ProtocolConformanceAnalysis *PCA;
+
+  // Class hierarchy analysis needed to confirm no derived classes of a sole
+  // conforming class.
+  ClassHierarchyAnalysis *CHA;
+
   /// Worklist containing all of the instructions primed for simplification.
   SILCombineWorklist Worklist;
 
@@ -139,8 +149,9 @@ class SILCombiner :
 public:
   SILCombiner(SILOptFunctionBuilder &FuncBuilder,
               SILBuilder &B, AliasAnalysis *AA, DominanceAnalysis *DA,
+              ProtocolConformanceAnalysis *PCA, ClassHierarchyAnalysis *CHA,
               bool removeCondFails)
-      : AA(AA), DA(DA), Worklist(), MadeChange(false),
+      : AA(AA), DA(DA), PCA(PCA), CHA(CHA), Worklist(), MadeChange(false),
         RemoveCondFails(removeCondFails), Iteration(0), Builder(B),
         CastOpt(FuncBuilder,
                 /* ReplaceInstUsesAction */
@@ -288,17 +299,21 @@ class SILCombiner :
                                               const ConcreteExistentialInfo &CEI,
                                               SILBuilderContext &BuilderCtx);
 
+  // Common utility function to replace the WitnessMethodInst using a
+  // BuilderCtx.
+  void replaceWitnessMethodInst(WitnessMethodInst *WMI,
+                                SILBuilderContext &BuilderCtx,
+                                CanType ConcreteType,
+                                const ProtocolConformanceRef ConformanceRef);
+
   SILInstruction *
   propagateConcreteTypeOfInitExistential(FullApplySite Apply,
                                          WitnessMethodInst *WMI);
   SILInstruction *propagateConcreteTypeOfInitExistential(FullApplySite Apply);
 
-  // Common utility function to replace the WitnessMethodInst using a
-  // BuilderCtx.
-  void replaceWitnessMethodInst(FullApplySite Apply, WitnessMethodInst *WMI,
-                                SILBuilderContext &BuilderCtx,
-                                const CanType &ConcreteType,
-                                const ProtocolConformanceRef ConformanceRef);
+  /// Propagate concrete types from ProtocolConformanceAnalysis.
+  SILInstruction *propagateSoleConformingType(FullApplySite Apply,
+                                              WitnessMethodInst *WMI);
 
   /// Perform one SILCombine iteration.
   bool doOneIteration(SILFunction &F, unsigned Iteration);

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -605,21 +605,100 @@ SILCombiner::optimizeConcatenationOfStringLiterals(ApplyInst *AI) {
 
 /// This routine replaces the old witness method inst with a new one.
 void SILCombiner::replaceWitnessMethodInst(
-    FullApplySite Apply, WitnessMethodInst *WMI, SILBuilderContext &BuilderCtx,
-    const CanType &ConcreteType, const ProtocolConformanceRef ConformanceRef) {
+    WitnessMethodInst *WMI, SILBuilderContext &BuilderCtx, CanType ConcreteType,
+    const ProtocolConformanceRef ConformanceRef) {
   SILBuilderWithScope WMIBuilder(WMI, BuilderCtx);
   auto *NewWMI = WMIBuilder.createWitnessMethod(
       WMI->getLoc(), ConcreteType, ConformanceRef, WMI->getMember(),
       WMI->getType());
-  MutableArrayRef<Operand> Operands = Apply.getInstruction()->getAllOperands();
-  for (auto &Op : Operands) {
-    if (Op.get() == WMI)
-      Op.set(NewWMI);
-  }
+  WMI->replaceAllUsesWith(NewWMI);
   if (WMI->use_empty())
     eraseInstFromFunction(*WMI);
 }
 
+//  This function propagates concrete type of existential self argument using
+//  ProtocolConformanceAnalysis. The concrete type of self can be a class,
+//  struct, or an enum. It replaces the witness_method instruction
+//  with one that has a concrete type, allowing other optimizations to
+//  devirtualize it later.
+SILInstruction *
+SILCombiner::propagateSoleConformingType(FullApplySite Apply,
+                                         WitnessMethodInst *WMI) {
+  // If WMI has concrete conformance, it can be optimized.
+  if (WMI->getConformance().isConcrete())
+    return nullptr;
+
+  // If the lookup type is an opened existential type,
+  // it cannot be made concrete.
+  if (!WMI->getLookupType()->isOpenedExistential())
+    return nullptr;
+
+  // If the witness method mutates self, we cannot replace self.
+  if (Apply.getOrigCalleeType()->getSelfParameter().isIndirectMutating())
+    return nullptr;
+
+  // Only applicable in whole-module compilation.
+  if (!Apply.getModule().isWholeModule())
+    return nullptr;
+
+  auto *PD = WMI->getLookupProtocol();
+
+  // Determine the sole conforming type.
+  auto *NTD = PCA->findSoleConformingType(PD);
+  if (!NTD)
+    return nullptr;
+
+  // Sole conforming class should not be open access or have any derived class.
+  ClassDecl *CD;
+  if ((CD = dyn_cast<ClassDecl>(NTD)) &&
+      (CD->getEffectiveAccess() == AccessLevel::Open ||
+       CHA->hasKnownDirectSubclasses(CD))) {
+    return nullptr;
+  }
+
+  // Create SIL type for the concrete type.
+  auto ElementType = NTD->getDeclaredType();
+  auto ConcreteType = ElementType->getCanonicalType();
+  auto &M = Builder.getModule();
+
+  /// Determine OpenedArchetypeDef and SubstituionMap.
+  ConcreteExistentialInfo CEI(Apply.getSelfArgumentOperand(), ConcreteType, PD);
+  if (!CEI.isValid())
+    return nullptr;
+
+  if (!CEI.InitExistential) {
+    // Create SIL type for the concrete type.
+    SILType ConcreteSILType = M.Types.getLoweredType(ConcreteType);
+
+    // Prepare the code by adding UncheckedCast instructions that cast opened
+    // existentials to concrete types. Set the ConcreteValue of CEI.
+    if (auto *OER = dyn_cast<OpenExistentialRefInst>(CEI.OpenedArchetypeDef)) {
+      auto *URCI =
+          Builder.createUncheckedRefCast(OER->getLoc(), OER, ConcreteSILType);
+      CEI.ConcreteValue = URCI;
+    } else if (auto *OEA =
+                   dyn_cast<OpenExistentialAddrInst>(CEI.OpenedArchetypeDef)) {
+      auto *UACI = Builder.createUncheckedAddrCast(
+          OEA->getLoc(), OEA, ConcreteSILType.getAddressType());
+      CEI.ConcreteValue = UACI;
+    } else {
+      llvm_unreachable(
+          "Unhandled Argument Type in propagateSoleConformingType");
+    }
+  }
+
+  assert(CEI.ConcreteValue);
+
+  /// Replace the old WitnessMethod with a new one that has concrete type and
+  /// conformance.
+  SILBuilderContext BuilderCtx(M, Builder.getTrackingList());
+  replaceWitnessMethodInst(WMI, BuilderCtx, ConcreteType,
+                           *(CEI.ExistentialSubs.getConformances().begin()));
+  /// Create the new apply instruction using the concrete type.
+  auto *NewAI = createApplyWithConcreteType(Apply, CEI, BuilderCtx);
+  return NewAI;
+}
+
 /// Given an Apply and an argument value produced by InitExistentialAddrInst,
 /// return true if the argument can be replaced by a copy of its value.
 ///
@@ -799,7 +878,6 @@ SILCombiner::propagateConcreteTypeOfInitExistential(FullApplySite Apply,
   if (WMI->getConformance().isConcrete())
     return nullptr;
 
-
   // If the lookup type is not an opened existential type,
   // it cannot be made more concrete.
   if (!WMI->getLookupType()->isOpenedExistential())
@@ -840,7 +918,7 @@ SILCombiner::propagateConcreteTypeOfInitExistential(FullApplySite Apply,
   // with it.
   if (CEI.ConcreteType != WMI->getLookupType() ||
       SelfConformance != WMI->getConformance()) {
-    replaceWitnessMethodInst(Apply, WMI, BuilderCtx, CEI.ConcreteType,
+    replaceWitnessMethodInst(WMI, BuilderCtx, CEI.ConcreteType,
                              SelfConformance);
   }
   // Try to rewrite the apply.
@@ -1101,7 +1179,12 @@ SILInstruction *SILCombiner::visitApplyInst(ApplyInst *AI) {
   // (apply (witness_method)) -> propagate information about
   // a concrete type from init_existential_addr or init_existential_ref.
   if (auto *WMI = dyn_cast<WitnessMethodInst>(AI->getCallee())) {
-    propagateConcreteTypeOfInitExistential(AI, WMI);
+    if (!propagateConcreteTypeOfInitExistential(AI, WMI)) {
+      if (auto *WitnessMethod = dyn_cast<WitnessMethodInst>(AI->getCallee())) {
+        // Propagate concrete type from ProtocolConformanceAnalysis.
+        propagateSoleConformingType(AI, WitnessMethod);
+      }
+    }
     return nullptr;
   }
 
@@ -1219,10 +1302,16 @@ SILInstruction *SILCombiner::visitTryApplyInst(TryApplyInst *AI) {
     if (!AI->getOrigCalleeType()->isCalleeConsumed())
       return rewriteApplyCallee(AI, TTTFI->getOperand()).getInstruction();
   }
+
   // (apply (witness_method)) -> propagate information about
   // a concrete type from init_existential_addr or init_existential_ref.
   if (auto *WMI = dyn_cast<WitnessMethodInst>(AI->getCallee())) {
-    propagateConcreteTypeOfInitExistential(AI, WMI);
+    if (!propagateConcreteTypeOfInitExistential(AI, WMI)) {
+      if (auto *WitnessMethod = dyn_cast<WitnessMethodInst>(AI->getCallee())) {
+        // Propagate concrete type from ProtocolConformanceAnalysis.
+        propagateSoleConformingType(AI, WitnessMethod);
+      }
+    }
     return nullptr;
   }