Merge pull request #39080 from jckarter/emit-closure-literal-at-context-abstraction-2

SILGen: Emit literal closures at the abstraction level of their context. [take 2]

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 =