Treat types obviously word-sized in SIL as simple nonescaping partial_apply contexts

Author: jckarter

include/swift/SIL/SILType.h

@@ -427,7 +427,8 @@ class SILType {
 
   /// True if the given type has at least the size and alignment of a native
   /// pointer.
-  bool isPointerSizeAndAligned();
+  bool isPointerSizeAndAligned(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);
@@ -613,6 +614,12 @@ class SILType {
   llvm::hash_code getHashCode() const {
     return llvm::hash_combine(*this);
   }
+  
+  /// If a type is visibly a singleton aggregate (a tuple with one element, a
+  /// struct with one field, or an enum with a single payload case), return the
+  /// type of its field, which it is guaranteed to have identical layout to.
+  SILType getSingletonAggregateFieldType(SILModule &M,
+                                         ResilienceExpansion expansion) const;
 
   //
   // Accessors for types used in SIL instructions:

lib/IRGen/GenType.h

@@ -305,6 +305,10 @@ class LoweringModeScope {
 /// If a type is visibly a singleton aggregate (a tuple with one element, a
 /// struct with one field, or an enum with a single payload case), return the
 /// type of its field, which it is guaranteed to have identical layout to.
+///
+/// This can use more concrete type layout information than
+/// SILType::getSingletonAggregateFieldType, because we have full access to the
+/// LLVM-level layout of types in IRGen.
 SILType getSingletonAggregateFieldType(IRGenModule &IGM,
                                        SILType t,
                                        ResilienceExpansion expansion);

lib/SIL/IR/SILType.cpp

@@ -174,7 +174,8 @@ std::string SILType::getAsString() const {
   return OS.str();
 }
 
-bool SILType::isPointerSizeAndAligned() {
+bool SILType::isPointerSizeAndAligned(SILModule &M,
+                                      ResilienceExpansion expansion) const {
   auto &C = getASTContext();
   if (isHeapObjectReferenceType()
       || getASTType()->isEqual(C.TheRawPointerType)) {
@@ -183,6 +184,10 @@ bool SILType::isPointerSizeAndAligned() {
   if (auto intTy = dyn_cast<BuiltinIntegerType>(getASTType()))
     return intTy->getWidth().isPointerWidth();
 
+  if (auto underlyingField = getSingletonAggregateFieldType(M, expansion)) {
+    return underlyingField.isPointerSizeAndAligned(M, expansion);
+  }
+  
   return false;
 }
 
@@ -766,3 +771,79 @@ SILType SILType::getSILBoxFieldType(const SILFunction *f, unsigned field) {
   return ::getSILBoxFieldType(f->getTypeExpansionContext(), boxTy,
                               f->getModule().Types, field);
 }
+
+SILType
+SILType::getSingletonAggregateFieldType(SILModule &M,
+                                        ResilienceExpansion expansion) const {
+  if (auto tuple = getAs<TupleType>()) {
+    if (tuple->getNumElements() == 1) {
+      return getTupleElementType(0);
+    }
+  }
+
+  if (auto structDecl = getStructOrBoundGenericStruct()) {
+    // If the struct has to be accessed resiliently from this resilience domain,
+    // we can't assume anything about its layout.
+    if (structDecl->isResilient(M.getSwiftModule(), expansion)) {
+      return SILType();
+    }
+
+    // C ABI wackiness may cause a single-field struct to have different layout
+    // from its field.
+    if (structDecl->hasUnreferenceableStorage()
+        || structDecl->hasClangNode()) {
+      return SILType();
+    }
+
+    // A single-field struct with custom alignment has different layout from its
+    // field.
+    if (structDecl->getAttrs().hasAttribute<AlignmentAttr>()) {
+      return SILType();
+    }
+
+    // If there's only one stored property, we have the layout of its field.
+    auto allFields = structDecl->getStoredProperties();
+    
+    if (allFields.size() == 1) {
+      auto fieldTy = getFieldType(
+          allFields[0], M,
+          TypeExpansionContext(expansion, M.getSwiftModule(),
+                               M.isWholeModule()));
+      if (!M.isTypeABIAccessible(fieldTy,
+                       TypeExpansionContext::maximalResilienceExpansionOnly())){
+        return SILType();
+      }
+      return fieldTy;
+    }
+
+    return SILType();
+  }
+
+  if (auto enumDecl = getEnumOrBoundGenericEnum()) {
+    // If the enum has to be accessed resiliently from this resilience domain,
+    // we can't assume anything about its layout.
+    if (enumDecl->isResilient(M.getSwiftModule(), expansion)) {
+      return SILType();
+    }
+
+    auto allCases = enumDecl->getAllElements();
+    
+    auto theCase = allCases.begin();
+    if (!allCases.empty() && std::next(theCase) == allCases.end()
+        && (*theCase)->hasAssociatedValues()) {
+      auto enumEltTy = getEnumElementType(
+          *theCase, M,
+          TypeExpansionContext(expansion, M.getSwiftModule(),
+                               M.isWholeModule()));
+      if (!M.isTypeABIAccessible(enumEltTy,
+                       TypeExpansionContext::maximalResilienceExpansionOnly())){
+        return SILType();
+      }
+      return enumEltTy;
+    }
+
+    return SILType();
+  }
+
+  return SILType();
+}

lib/SILOptimizer/Transforms/PartialApplySimplification.cpp

@@ -136,7 +136,9 @@ class PartialApplySimplificationPass : public SILModuleTransform {
 ///     sized strictly less than one word
 ///   - the argument ownership convention matches the callee convention of the
 ///     resulting function
-static bool isSimplePartialApply(CanSILFunctionType calleeTy,
+static bool isSimplePartialApply(SILModule &M,
+                                 CanSILFunctionType calleeTy,
+                                 TypeExpansionContext context,
                                  unsigned numPartiallyAppliedArgs,
                                  bool isOnStack) {
   if (calleeTy->isPolymorphic()) {
@@ -165,9 +167,13 @@ static bool isSimplePartialApply(CanSILFunctionType calleeTy,
       return true;
 
     case ParameterConvention::Direct_Guaranteed:
-    case ParameterConvention::Direct_Unowned:
-      // TODO: Handle word-sized direct arguments.
-      return false;
+    case ParameterConvention::Direct_Unowned: {
+      auto argTy = contextParam.getArgumentType(M, calleeTy, context);
+      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:
@@ -191,7 +197,9 @@ static bool isSimplePartialApply(CanSILFunctionType calleeTy,
 }
 
 static bool isSimplePartialApply(PartialApplyInst *i) {
-  return isSimplePartialApply(i->getCallee()->getType().castTo<SILFunctionType>(),
+  return isSimplePartialApply(i->getModule(),
+                              i->getCallee()->getType().castTo<SILFunctionType>(),
+                              i->getFunction()->getTypeExpansionContext(),
                               i->getNumArguments(),
                               i->isOnStack());
 }
@@ -312,7 +320,9 @@ void PartialApplySimplificationPass::processKnownCallee(SILFunction *callee,
   // Would the partial application become simple with a mere convention change?
   auto calleeTyAsMethod = callee->getLoweredFunctionType()
     ->getWithRepresentation(SILFunctionTypeRepresentation::Method);
-  if (isSimplePartialApply(calleeTyAsMethod,
+  if (isSimplePartialApply(callee->getModule(),
+                           calleeTyAsMethod,
+                           examplePA->getFunction()->getTypeExpansionContext(),
                            examplePA->getNumArguments(),
                            examplePA->isOnStack())) {
     return rewriteKnownCalleeConventionOnly(callee, pa, examplePA,