SILGen: Emit literal closures at the abstraction level of their context.

Literal closures are only ever directly referenced in the context of the expression they're written in,
so it's wasteful to emit them at their fully-substituted calling convention and then reabstract them if
they're passed directly to a generic function. Avoid this by saving the abstraction pattern of the context
before emitting the closure, and then lowering its main entry point's calling convention at that
level of abstraction. Generalize some of the prolog/epilog code to handle converting arguments and returns
to the correct representation for a different abstraction level.

Author: jckarter

include/swift/SIL/AbstractionPattern.h

@@ -432,6 +432,7 @@ class AbstractionPattern {
     const clang::ObjCMethodDecl *ObjCMethod;
     const clang::CXXMethodDecl *CXXMethod;
     const AbstractionPattern *OrigTupleElements;
+    const void *RawTypePtr;
   };
   CanGenericSignature GenericSig;
 
@@ -1274,7 +1275,7 @@ class AbstractionPattern {
   /// pattern?
   bool matchesTuple(CanTupleType substType);
 
-  bool isTuple() {
+  bool isTuple() const {
     switch (getKind()) {
     case Kind::Invalid:
       llvm_unreachable("querying invalid abstraction pattern!");
@@ -1386,8 +1387,32 @@ class AbstractionPattern {
       Lowering::TypeConverter &TC
   ) const;
 
+  /// How values are passed or returned according to this abstraction pattern.
+  enum CallingConventionKind {
+    // Value is passed or returned directly as a unit.
+    Direct,
+    // Value is passed or returned indirectly through memory.
+    Indirect,
+    // Value is a tuple that is destructured, and each element is considered
+    // independently.
+    Destructured,
+  };
+  
+  /// If this abstraction pattern appears in function return position, how is
+  /// the corresponding value returned?
+  CallingConventionKind getResultConvention(TypeConverter &TC) const;
+  
+  /// If this abstraction pattern appears in function parameter position, how
+  /// is the corresponding value passed?
+  CallingConventionKind getParameterConvention(TypeConverter &TC) const;
+  
   void dump() const LLVM_ATTRIBUTE_USED;
   void print(raw_ostream &OS) const;
+  
+  bool operator==(const AbstractionPattern &other) const;
+  bool operator!=(const AbstractionPattern &other) const {
+    return !(*this == other);
+  }
 };
 
 inline llvm::raw_ostream &operator<<(llvm::raw_ostream &out,

include/swift/SIL/TypeLowering.h

@@ -734,6 +734,9 @@ class TypeConverter {
   ///
   /// Second element is a ResilienceExpansion.
   llvm::DenseMap<std::pair<SILType, unsigned>, unsigned> TypeFields;
+  
+  llvm::DenseMap<AbstractClosureExpr *, Optional<AbstractionPattern>>
+    ClosureAbstractionPatterns;
 
   CanAnyFunctionType makeConstantInterfaceType(SILDeclRef constant);
 
@@ -1106,6 +1109,22 @@ class TypeConverter {
                                          SILType enumType,
                                          EnumElementDecl *elt);
 
+  /// Get the preferred abstraction pattern, if any, by which to lower a
+  /// declaration.
+  ///
+  /// This can be set using \c setAbstractionPattern , but only before
+  /// the abstraction pattern is queried using this function. Once the
+  /// abstraction pattern has been asked for, it may not be changed.
+  Optional<AbstractionPattern> getConstantAbstractionPattern(SILDeclRef constant);
+  
+  /// Set the preferred abstraction pattern for a closure.
+  ///
+  /// The abstraction pattern can only be set before any calls to
+  /// \c getConstantAbstractionPattern on the same closure. It may not be
+  /// changed once it has been read.
+  void setAbstractionPattern(AbstractClosureExpr *closure,
+                             AbstractionPattern pattern);
+  
 private:
   CanType computeLoweredRValueType(TypeExpansionContext context,
                                    AbstractionPattern origType,

lib/SIL/IR/AbstractionPattern.cpp

@@ -1233,3 +1233,155 @@ AbstractionPattern AbstractionPattern::getAutoDiffDerivativeFunctionType(
     llvm_unreachable("called on unsupported abstraction pattern kind");
   }
 }
+
+AbstractionPattern::CallingConventionKind
+AbstractionPattern::getResultConvention(TypeConverter &TC) const {
+  // Tuples should be destructured.
+  if (isTuple()) {
+    return Destructured;
+  }
+  switch (getKind()) {
+  case Kind::Opaque:
+    // Maximally abstracted values are always passed indirectly.
+    return Indirect;
+  
+  case Kind::OpaqueFunction:
+  case Kind::OpaqueDerivativeFunction:
+  case Kind::PartialCurriedObjCMethodType:
+  case Kind::CurriedObjCMethodType:
+  case Kind::PartialCurriedCFunctionAsMethodType:
+  case Kind::CurriedCFunctionAsMethodType:
+  case Kind::CFunctionAsMethodType:
+  case Kind::ObjCMethodType:
+  case Kind::CXXMethodType:
+  case Kind::CurriedCXXMethodType:
+  case Kind::PartialCurriedCXXMethodType:
+  case Kind::CXXOperatorMethodType:
+  case Kind::CurriedCXXOperatorMethodType:
+  case Kind::PartialCurriedCXXOperatorMethodType:
+    // Function types are always passed directly
+    return Direct;
+      
+  case Kind::ClangType:
+  case Kind::Type:
+  case Kind::Discard:
+    // Pass according to the formal type.
+    return SILType::isFormallyReturnedIndirectly(getType(),
+                                                 TC,
+                                                 getGenericSignatureOrNull())
+      ? Indirect : Direct;
+  
+  case Kind::Invalid:
+  case Kind::Tuple:
+  case Kind::ObjCCompletionHandlerArgumentsType:
+    llvm_unreachable("should not get here");
+  }
+}
+
+AbstractionPattern::CallingConventionKind
+AbstractionPattern::getParameterConvention(TypeConverter &TC) const {
+  // Tuples should be destructured.
+  if (isTuple()) {
+    return Destructured;
+  }
+  switch (getKind()) {
+  case Kind::Opaque:
+    // Maximally abstracted values are always passed indirectly.
+    return Indirect;
+  
+  case Kind::OpaqueFunction:
+  case Kind::OpaqueDerivativeFunction:
+  case Kind::PartialCurriedObjCMethodType:
+  case Kind::CurriedObjCMethodType:
+  case Kind::PartialCurriedCFunctionAsMethodType:
+  case Kind::CurriedCFunctionAsMethodType:
+  case Kind::CFunctionAsMethodType:
+  case Kind::ObjCMethodType:
+  case Kind::CXXMethodType:
+  case Kind::CurriedCXXMethodType:
+  case Kind::PartialCurriedCXXMethodType:
+  case Kind::CXXOperatorMethodType:
+  case Kind::CurriedCXXOperatorMethodType:
+  case Kind::PartialCurriedCXXOperatorMethodType:
+    // Function types are always passed directly
+    return Direct;
+      
+  case Kind::ClangType:
+  case Kind::Type:
+  case Kind::Discard:
+    // Pass according to the formal type.
+    return SILType::isFormallyPassedIndirectly(getType(),
+                                                 TC,
+                                                 getGenericSignatureOrNull())
+      ? Indirect : Direct;
+  
+  case Kind::Invalid:
+  case Kind::Tuple:
+  case Kind::ObjCCompletionHandlerArgumentsType:
+    llvm_unreachable("should not get here");
+  }
+}
+
+bool
+AbstractionPattern::operator==(const AbstractionPattern &other) const {
+  if (TheKind != other.TheKind)
+    return false;
+  
+  switch (getKind()) {
+  case Kind::Opaque:
+  case Kind::Invalid:
+  case Kind::OpaqueFunction:
+  case Kind::OpaqueDerivativeFunction:
+    // No additional info to compare.
+    return true;
+
+  case Kind::Tuple:
+    if (getNumTupleElements() != other.getNumTupleElements()) {
+      return false;
+    }
+    for (unsigned i = 0; i < getNumTupleElements(); ++i) {
+      if (getTupleElementType(i) != other.getTupleElementType(i)) {
+        return false;
+      }
+    }
+    return true;
+  
+  case Kind::Type:
+  case Kind::Discard:
+    return OrigType == other.OrigType
+      && GenericSig == other.GenericSig;
+      
+  case Kind::ClangType:
+    return OrigType == other.OrigType
+      && GenericSig == other.GenericSig
+      && ClangType == other.ClangType;
+
+  case Kind::ObjCCompletionHandlerArgumentsType:
+  case Kind::CFunctionAsMethodType:
+  case Kind::CurriedCFunctionAsMethodType:
+  case Kind::PartialCurriedCFunctionAsMethodType:
+    return OrigType == other.OrigType
+      && GenericSig == other.GenericSig
+      && ClangType == other.ClangType
+      && OtherData == other.OtherData;
+
+  case Kind::ObjCMethodType:
+  case Kind::CurriedObjCMethodType:
+  case Kind::PartialCurriedObjCMethodType:
+    return OrigType == other.OrigType
+      && GenericSig == other.GenericSig
+      && ObjCMethod == other.ObjCMethod
+      && OtherData == other.OtherData;
+      
+  case Kind::CXXMethodType:
+  case Kind::CXXOperatorMethodType:
+  case Kind::CurriedCXXMethodType:
+  case Kind::CurriedCXXOperatorMethodType:
+  case Kind::PartialCurriedCXXMethodType:
+  case Kind::PartialCurriedCXXOperatorMethodType:
+    return OrigType == other.OrigType
+      && GenericSig == other.GenericSig
+      && CXXMethod == other.CXXMethod
+      && OtherData == other.OtherData;
+  }
+}

lib/SIL/IR/SILFunctionType.cpp

@@ -2072,7 +2072,8 @@ static CanSILFunctionType getSILFunctionType(
   // for thick or polymorphic functions.  We don't need to worry about
   // non-opaque patterns because the type-checker forbids non-thick
   // function types from having generic parameters or results.
-  if (origType.isTypeParameter() &&
+  if (!constant &&
+      origType.isTypeParameter() &&
       substFnInterfaceType->getExtInfo().getSILRepresentation()
         != SILFunctionType::Representation::Thick &&
       isa<FunctionType>(substFnInterfaceType)) {
@@ -3183,9 +3184,18 @@ static CanSILFunctionType getUncachedSILFunctionTypeForConstant(
       auto proto = constant.getDecl()->getDeclContext()->getSelfProtocolDecl();
       witnessMethodConformance = ProtocolConformanceRef(proto);
     }
+    
+    // Does this constant have a preferred abstraction pattern set?
+    AbstractionPattern origType = [&]{
+      if (auto abstraction = TC.getConstantAbstractionPattern(constant)) {
+        return *abstraction;
+      } else {
+        return AbstractionPattern(origLoweredInterfaceType);
+      }
+    }();
 
     return ::getNativeSILFunctionType(
-        TC, context, AbstractionPattern(origLoweredInterfaceType),
+        TC, context, origType,
         origLoweredInterfaceType, extInfoBuilder, constant, constant, None,
         witnessMethodConformance);
   }

lib/SIL/IR/TypeLowering.cpp

@@ -3491,6 +3491,28 @@ CanSILBoxType TypeConverter::getBoxTypeForEnumElement(
   return boxTy;
 }
 
+Optional<AbstractionPattern>
+TypeConverter::getConstantAbstractionPattern(SILDeclRef constant) {
+  if (auto closure = constant.getAbstractClosureExpr()) {
+    // Using operator[] here creates an entry in the map if one doesn't exist
+    // yet, marking the fact that the lack of abstraction pattern has been
+    // established and cannot be overridden by `setAbstractionPattern` later.
+    return ClosureAbstractionPatterns[closure];
+  }
+  return None;
+}
+
+void TypeConverter::setAbstractionPattern(AbstractClosureExpr *closure,
+                                          AbstractionPattern pattern) {
+  auto existing = ClosureAbstractionPatterns.find(closure);
+  if (existing != ClosureAbstractionPatterns.end()) {
+    assert(*existing->second == pattern
+     && "closure shouldn't be emitted at different abstraction level contexts");
+  } else {
+    ClosureAbstractionPatterns[closure] = pattern;
+  }
+}
+
 static void countNumberOfInnerFields(unsigned &fieldsCount, TypeConverter &TC,
                                      SILType Ty,
                                      TypeExpansionContext expansion) {

lib/SILGen/Conversion.h

@@ -59,6 +59,11 @@ class Conversion {
   static bool isBridgingKind(KindTy kind) {
     return kind <= LastBridgingKind;
   }
+  
+  static bool isReabstractionKind(KindTy kind) {
+    // Update if we end up with more kinds!
+    return !isBridgingKind(kind);
+  }
 
 private:
   KindTy Kind;
@@ -139,6 +144,10 @@ class Conversion {
   bool isBridging() const {
     return isBridgingKind(getKind());
   }
+  
+  bool isReabstraction() const {
+    return isReabstractionKind(getKind());
+  }
 
   AbstractionPattern getReabstractionOrigType() const {
     return Types.get<ReabstractionTypes>(Kind).OrigType;
@@ -264,12 +273,21 @@ class ConvertingInitialization final : public Initialization {
   StateTy getState() const {
     return State;
   }
+  
+  InitializationPtr OwnedSubInitialization;
 
 public:
   ConvertingInitialization(Conversion conversion, SGFContext finalContext)
     : State(Uninitialized), TheConversion(conversion),
       FinalContext(finalContext) {}
 
+  ConvertingInitialization(Conversion conversion,
+                           InitializationPtr subInitialization)
+    : State(Uninitialized), TheConversion(conversion),
+      FinalContext(SGFContext(subInitialization.get())) {
+    OwnedSubInitialization = std::move(subInitialization);
+  }
+  
   /// Return the conversion to apply to the unconverted value.
   const Conversion &getConversion() const {
     return TheConversion;
@@ -328,11 +346,16 @@ class ConvertingInitialization final : public Initialization {
   ConvertingInitialization *getAsConversion() override {
     return this;
   }
+  
+  // Get the abstraction pattern, if any, the value is converted to.
+  Optional<AbstractionPattern> getAbstractionPattern() const override;
 
   // Bookkeeping.
   void finishInitialization(SILGenFunction &SGF) override {
     assert(getState() == Initialized);
     State = Finished;
+    if (OwnedSubInitialization)
+      OwnedSubInitialization->finishInitialization(SGF);
   }
 };
 

lib/SILGen/Initialization.h

@@ -161,6 +161,16 @@ class Initialization {
                      "uninitialized");
   }
 
+  /// The preferred abstraction pattern to initialize with.
+  ///
+  /// Returning something other than None here gives expression emission the
+  /// opportunity to generate the initial value directly at the proper
+  /// abstraction level, avoiding the need for a conversion in some
+  /// circumstances.
+  virtual Optional<AbstractionPattern> getAbstractionPattern() const {
+    return None;
+  }
+
 protected:
   bool EmitDebugValueOnInit = true;
 

lib/SILGen/SGFContext.h

@@ -167,6 +167,14 @@ class SGFContext {
 
     return SGFContext();
   }
+  
+  /// Return the abstraction pattern of the context we're emitting into.
+  Optional<AbstractionPattern> getAbstractionPattern() const {
+    if (auto *init = getEmitInto()) {
+      return init->getAbstractionPattern();
+    }
+    return None;
+  }
 };
 
 using ValueProducerRef =