Sketch out a SILType::isSingleSwiftRefcounted function.

Author: jckarter

include/swift/SIL/SILType.h

@@ -430,6 +430,11 @@ class SILType {
   bool isPointerSizeAndAligned(SILModule &M,
                                ResilienceExpansion expansion) const;
 
+  /// True if the layout of the given type consists of a single native Swift-
+  /// refcounted object reference, possibly nullable.
+  bool isSingleSwiftRefcounted(SILModule &M,
+                               ResilienceExpansion expansion) const;
+
   /// True if `operTy` can be cast by single-reference value into `resultTy`.
   static bool canRefCast(SILType operTy, SILType resultTy, SILModule &M);
 

lib/SIL/IR/SILType.cpp

@@ -191,6 +191,79 @@ bool SILType::isPointerSizeAndAligned(SILModule &M,
   return false;
 }
 
+static bool isSingleSwiftRefcounted(SILModule &M,
+                                    SILType SILTy,
+                                    ResilienceExpansion expansion,
+                                    bool didUnwrapOptional) {
+  auto &C = M.getASTContext();
+  
+  // Unwrap one layer of optionality.
+  // TODO: Or more generally, any fragile enum with a single payload and single
+  // no-payload case.
+  if (!didUnwrapOptional) {
+    if (auto objectTy = SILTy.getOptionalObjectType()) {
+      return ::isSingleSwiftRefcounted(M, objectTy, expansion, true);
+    }
+  }
+
+  // Unwrap singleton aggregates.
+  if (auto underlyingField = SILTy.getSingletonAggregateFieldType(M, expansion)) {
+    return ::isSingleSwiftRefcounted(M, underlyingField, expansion,
+                                     didUnwrapOptional);
+  }
+  
+  auto Ty = SILTy.getASTType();
+  
+  // Easy cases: Builtin.NativeObject and boxes are always Swift-refcounted.
+  if (Ty == C.TheNativeObjectType)
+    return true;
+  if (isa<SILBoxType>(Ty))
+    return true;
+  
+  // Is the type a Swift-refcounted class?
+  // For a generic type, consider its superclass constraint, if any.
+  auto ClassTy = Ty;
+  if (auto archety = dyn_cast<ArchetypeType>(Ty)) {
+    if (auto superclass = Ty->getSuperclass()) {
+      ClassTy = superclass->getCanonicalType();
+    }
+  }
+  // For an existential type, consider its superclass constraint, if it carries
+  // no witness tables.
+  if (Ty->isAnyExistentialType()) {
+    auto layout = Ty->getExistentialLayout();
+    // Must be no protocol constraints that aren't @objc or @_marker.
+    if (layout.containsNonObjCProtocol) {
+      for (auto proto : layout.getProtocols()) {
+        if (!proto->isObjC() && !proto->isMarkerProtocol()) {
+          return false;
+        }
+      }
+    }
+    
+    // The Error existential has its own special layout.
+    if (layout.isErrorExistential()) {
+      return false;
+    }
+    
+    // We can look at the superclass constraint, if any, to see if it's
+    // Swift-refcounted.
+    if (!layout.getSuperclass()) {
+      return false;
+    }
+    ClassTy = layout.getSuperclass()->getCanonicalType();
+  }
+  
+  // TODO: Does the base class we found have fully native Swift ancestry,
+  // so we can use Swift native refcounting on it?
+  return false;
+}
+
+bool SILType::isSingleSwiftRefcounted(SILModule &M,
+                                      ResilienceExpansion expansion) const {
+  return ::isSingleSwiftRefcounted(M, *this, expansion, false);
+}
+
 // Reference cast from representations with single pointer low bits.
 // Only reference cast to simple single pointer representations.
 //

lib/SILOptimizer/Transforms/PartialApplySimplification.cpp

@@ -139,6 +139,7 @@ class PartialApplySimplificationPass : public SILModuleTransform {
 static bool isSimplePartialApply(SILModule &M,
                                  CanSILFunctionType calleeTy,
                                  TypeExpansionContext context,
+                                 ParameterConvention calleeConvention,
                                  unsigned numPartiallyAppliedArgs,
                                  bool isOnStack) {
   if (calleeTy->isPolymorphic()) {
@@ -157,6 +158,7 @@ static bool isSimplePartialApply(SILModule &M,
 
   auto contextParam = calleeTy->getSelfParameter();
 
+  auto argTy = contextParam.getArgumentType(M, calleeTy, context);
   if (isOnStack) {
     switch (contextParam.getConvention()) {
     case ParameterConvention::Indirect_Inout:
@@ -167,13 +169,11 @@ static bool isSimplePartialApply(SILModule &M,
       return true;
 
     case ParameterConvention::Direct_Guaranteed:
-    case ParameterConvention::Direct_Unowned: {
-      auto argTy = contextParam.getArgumentType(M, calleeTy, context);
+    case ParameterConvention::Direct_Unowned:
       return SILType::getPrimitiveObjectType(argTy)
         .isPointerSizeAndAligned(M, context.getResilienceExpansion());
       // TODO: If we're running as an IRGen pass, use IRGen's version of
       // `isPointerSizeAndAligned` as a more accurate check.
-    }
 
     // +1 arguments need a thunk to stage a copy for the callee to consume.
     case ParameterConvention::Direct_Owned:
@@ -185,12 +185,16 @@ static bool isSimplePartialApply(SILModule &M,
       return false;
     }
 
-    // TODO: Handle native-refcounted classes, single-refcounted aggregates,
-    // and bit-packable trivial types using knowledge from IRGen if this becomes
-    // an IRGenPrepare pass
-    if (!isa<SILBoxType>(contextParam.getInterfaceType())) {
+    // The context parameter's convention must match the callee convention of
+    // the resulting closure.
+    if (contextParam.getConvention() != calleeConvention) {
       return false;
     }
+    
+    // The context type must consist of only a swift-refcounted object
+    // reference.
+    return SILType::getPrimitiveObjectType(argTy)
+      .isSingleSwiftRefcounted(M, context.getResilienceExpansion());
   }
 
   return true;
@@ -200,6 +204,7 @@ static bool isSimplePartialApply(PartialApplyInst *i) {
   return isSimplePartialApply(i->getModule(),
                               i->getCallee()->getType().castTo<SILFunctionType>(),
                               i->getFunction()->getTypeExpansionContext(),
+                              i->getFunctionType()->getCalleeConvention(),
                               i->getNumArguments(),
                               i->isOnStack());
 }
@@ -323,6 +328,7 @@ void PartialApplySimplificationPass::processKnownCallee(SILFunction *callee,
   if (isSimplePartialApply(callee->getModule(),
                            calleeTyAsMethod,
                            examplePA->getFunction()->getTypeExpansionContext(),
+                           examplePA->getFunctionType()->getCalleeConvention(),
                            examplePA->getNumArguments(),
                            examplePA->isOnStack())) {
     return rewriteKnownCalleeConventionOnly(callee, pa, examplePA,