Merge pull request #7277 from swiftix/wip-partial-pre-specializations-wip-specialize-attribute-2

Teach EagerSpecializer how to use the new form of the @_specialize attribute

Author: swiftix

include/swift/SIL/Mangle.h

@@ -30,6 +30,8 @@ class AbstractClosureExpr;
 enum class SpecializationKind : uint8_t {
   Generic,
   NotReAbstractedGeneric,
+  Partial,
+  NotReAbstractedPartial,
   FunctionSignature,
 };
 
@@ -73,6 +75,12 @@ class SpecializationManglerBase {
     case SpecializationKind::NotReAbstractedGeneric:
       M.append("r");
       break;
+    case SpecializationKind::Partial:
+      M.append("p");
+      break;
+    case SpecializationKind::NotReAbstractedPartial:
+      M.append("P");
+      break;
     case SpecializationKind::FunctionSignature:
       M.append("f");
       break;
@@ -145,14 +153,44 @@ class GenericSpecializationMangler :
   };
 
   GenericSpecializationMangler(Mangle::Mangler &M, SILFunction *F,
-                               SubstitutionList Subs,
+                               SubstitutionList Subs, IsFragile_t Fragile,
+                               ReAbstractionMode isReAbstracted = ReAbstracted)
+      : SpecializationMangler(isReAbstracted == ReAbstracted
+                                  ? SpecializationKind::Generic
+                                  : SpecializationKind::NotReAbstractedGeneric,
+                              SpecializationPass::GenericSpecializer, M,
+                              Fragile, F),
+        Subs(Subs) {
+  }
+
+private:
+  void mangleSpecialization();
+};
+
+class PartialSpecializationMangler :
+  public SpecializationMangler<PartialSpecializationMangler> {
+
+  friend class SpecializationMangler<PartialSpecializationMangler>;
+
+  CanSILFunctionType SpecializedFnTy;
+public:
+
+  enum ReAbstractionMode {
+    ReAbstracted,
+    NotReabstracted
+  };
+
+  PartialSpecializationMangler(Mangle::Mangler &M,
+                               SILFunction *F,
+                               CanSILFunctionType SpecializedFnTy,
                                IsFragile_t Fragile,
                                ReAbstractionMode isReAbstracted = ReAbstracted)
     : SpecializationMangler(isReAbstracted == ReAbstracted ?
-                              SpecializationKind::Generic :
-                              SpecializationKind::NotReAbstractedGeneric,
+                              SpecializationKind::Partial :
+                              SpecializationKind::NotReAbstractedPartial,
                             SpecializationPass::GenericSpecializer,
-                            M, Fragile, F), Subs(Subs) {}
+                            M, Fragile, F), SpecializedFnTy(SpecializedFnTy) {
+    }
 
 private:
   void mangleSpecialization();

include/swift/SILOptimizer/Utils/Generics.h

@@ -53,20 +53,69 @@ class ReabstractionInfo {
   /// out-parameters.
   unsigned NumFormalIndirectResults;
 
-  /// The function type after applying the substitutions of the original
-  /// apply site.
+  /// The function type after applying the substitutions used to call the
+  /// specialized function.
   CanSILFunctionType SubstitutedType;
 
   /// The function type after applying the re-abstractions on the
   /// SubstitutedType.
   CanSILFunctionType SpecializedType;
 
+  /// The generic environment to be used by the specialization.
+  GenericEnvironment *SpecializedGenericEnv;
+
+  /// The generic signature of the specialization.
+  /// It is nullptr if the specialization is not polymorphic.
+  GenericSignature *SpecializedGenericSig;
+
+  // Set of the substitutions used by the caller's apply instruction before
+  // any transformations performed by the generic specializer.
+  // It uses archetypes.
+  SubstitutionList OriginalParamSubs;
+
+  // Set of substitutions to be used by the caller's apply when it calls a
+  // specialized function.
+  // It uses archetypes.
+  SubstitutionList CallerParamSubs;
+
+  // Set of substitutions to be used by the cloner during cloning.
+  // It maps to concrete types for any types which were replaced by
+  // concrete types in the caller's apply substitution list. All other
+  // types are replaced by their respective archetypes.
+  SubstitutionList ClonerParamSubs;
+
+  // Reference to the original generic non-specialized function.
+  SILFunction *OriginalF;
+
+  // The apply site which invokes the generic function.
+  ApplySite Apply;
+
+  // Set if a specialized function has unbound generic parameters.
+  bool HasUnboundGenericParams;
+
+  // Substitutions to be used for creating a new function type
+  // for the specialized function.
+  // It uses interface types.
+  SubstitutionMap CallerInterfaceSubs;
+
+  // Create a new substituted type with the updated signature.
+  CanSILFunctionType createSubstitutedType(SILFunction *OrigF,
+                                           const SubstitutionMap &SubstMap,
+                                           bool HasUnboundGenericParams);
+
+  void createSubstitutedAndSpecializedTypes();
 public:
   /// Constructs the ReabstractionInfo for generic function \p Orig with
   /// substitutions \p ParamSubs.
   /// If specialization is not possible getSpecializedType() will return an
   /// invalid type.
-  ReabstractionInfo(SILFunction *Orig, SubstitutionList ParamSubs);
+  ReabstractionInfo(ApplySite Apply, SILFunction *Callee,
+                    SubstitutionList ParamSubs);
+
+  /// Constructs the ReabstractionInfo for generic function \p Orig with
+  /// additional requirements. Requirements may contain new layout,
+  /// conformances or same concrete type requirements.
+  ReabstractionInfo(SILFunction *Orig, ArrayRef<Requirement> Requirements);
 
   /// Returns true if the \p ParamIdx'th (non-result) formal parameter is
   /// converted from indirect to direct.
@@ -81,6 +130,16 @@ class ReabstractionInfo {
     return Conversions.test(ResultIdx);
   }
 
+  /// Gets the total number of original function arguments.
+  unsigned getNumArguments() const { return Conversions.size(); }
+
+  /// Returns true if the \p ArgIdx'th argument is converted from an
+  /// indirect
+  /// result or parameter to a direct result or parameter.
+  bool isArgConverted(unsigned ArgIdx) const {
+    return Conversions.test(ArgIdx);
+  }
+
   /// Returns true if there are any conversions from indirect to direct values.
   bool hasConversions() const { return Conversions.any(); }
 
@@ -92,6 +151,15 @@ class ReabstractionInfo {
     Conversions.resize(Conversions.size() - numPartialApplyArgs);
   }
 
+  /// Returns the index of the first argument of an apply site, which may be
+  /// > 0 in case of a partial_apply.
+  unsigned getIndexOfFirstArg(ApplySite Apply) const {
+    unsigned numArgs = Apply.getNumArguments();
+    assert(numArgs == Conversions.size() ||
+           (numArgs < Conversions.size() && isa<PartialApplyInst>(Apply)));
+    return Conversions.size() - numArgs;
+  }
+
   /// Get the function type after applying the substitutions to the original
   /// generic function.
   CanSILFunctionType getSubstitutedType() const { return SubstitutedType; }
@@ -101,10 +169,45 @@ class ReabstractionInfo {
   /// possible.
   CanSILFunctionType getSpecializedType() const { return SpecializedType; }
 
+  GenericEnvironment *getSpecializedGenericEnvironment() const {
+    return SpecializedGenericEnv;
+  }
+
+  SubstitutionList getCallerParamSubstitutions() const {
+    return CallerParamSubs;
+  }
+
+  SubstitutionList getClonerParamSubstitutions() const {
+    return ClonerParamSubs;
+  }
+
+  SubstitutionList getOriginalParamSubstitutions() const {
+    return OriginalParamSubs;
+  }
+
   /// Create a specialized function type for a specific substituted type \p
   /// SubstFTy by applying the re-abstractions.
   CanSILFunctionType createSpecializedType(CanSILFunctionType SubstFTy,
                                            SILModule &M) const;
+
+  SILFunction *getNonSpecializedFunction() const { return OriginalF; }
+
+  /// Map type into a context of the specialized function.
+  Type mapTypeIntoContext(Type type) const;
+
+  /// Map SIL type into a context of the specialized function.
+  SILType mapTypeIntoContext(SILType type) const;
+
+  SILModule &getModule() const { return OriginalF->getModule(); }
+
+  /// Returns true if generic specialization is possible.
+  bool canBeSpecialized() const;
+
+  /// Returns true if it is a full generic specialization.
+  bool isFullSpecialization() const;
+
+  /// Returns true if it is a partial generic specialization.
+  bool isPartialSpecialization() const;
 };
 
 /// Helper class for specializing a generic function given a list of
@@ -118,6 +221,7 @@ class GenericFuncSpecializer {
 
   SubstitutionMap ContextSubs;
   std::string ClonedName;
+
 public:
   GenericFuncSpecializer(SILFunction *GenericFunc,
                          SubstitutionList ParamSubs,
@@ -129,7 +233,7 @@ class GenericFuncSpecializer {
 
   /// Return a newly created specialized function.
   SILFunction *tryCreateSpecialization();
-  
+
   /// Try to specialize GenericFunc given a list of ParamSubs.
   /// Returns either a new or existing specialized function, or nullptr.
   SILFunction *trySpecialization() {
@@ -142,6 +246,10 @@ class GenericFuncSpecializer {
 
     return SpecializedF;
   }
+
+  StringRef getClonedName() {
+    return ClonedName;
+  }
 };
 
 // =============================================================================

lib/IRGen/GenArchetype.cpp

@@ -280,14 +280,39 @@ class ClassArchetypeTypeInfo
   }
 };
 
+class FixedSizeArchetypeTypeInfo
+  : public PODSingleScalarTypeInfo<FixedSizeArchetypeTypeInfo, LoadableTypeInfo>,
+    public ArchetypeTypeInfoBase
+{
+  FixedSizeArchetypeTypeInfo(llvm::Type *type, Size size, Alignment align,
+                             const SpareBitVector &spareBits,
+                             ArrayRef<ProtocolEntry> protocols)
+      : PODSingleScalarTypeInfo(type, size, spareBits, align),
+        ArchetypeTypeInfoBase(this + 1, protocols) {}
+
+public:
+  static const FixedSizeArchetypeTypeInfo *
+  create(llvm::Type *type, Size size, Alignment align,
+         const SpareBitVector &spareBits, ArrayRef<ProtocolEntry> protocols) {
+    void *buffer = operator new(sizeof(FixedSizeArchetypeTypeInfo) +
+                                protocols.size() * sizeof(ProtocolEntry));
+    return ::new (buffer)
+        FixedSizeArchetypeTypeInfo(type, size, align, spareBits, protocols);
+  }
+};
+
 } // end anonymous namespace
 
 /// Return the ArchetypeTypeInfoBase information from the TypeInfo for any
 /// archetype.
 static const ArchetypeTypeInfoBase &
 getArchetypeInfo(IRGenFunction &IGF, CanArchetypeType t, const TypeInfo &ti) {
-  if (t->requiresClass())
+  LayoutConstraint LayoutInfo = t->getLayoutConstraint();
+  if (t->requiresClass() || (LayoutInfo && LayoutInfo->isRefCountedObject()))
     return ti.as<ClassArchetypeTypeInfo>();
+  if (LayoutInfo && LayoutInfo->isFixedSizeTrivial())
+    return ti.as<FixedSizeArchetypeTypeInfo>();
+  // TODO: Handle LayoutConstraintInfo::Trivial
   return ti.as<OpaqueArchetypeTypeInfo>();
 }
 
@@ -374,9 +399,12 @@ const TypeInfo *TypeConverter::convertArchetypeType(ArchetypeType *archetype) {
     protocols.push_back(ProtocolEntry(protocol, impl));
   }
 
+  LayoutConstraint LayoutInfo = archetype->getLayoutConstraint();
+
   // If the archetype is class-constrained, use a class pointer
   // representation.
-  if (archetype->requiresClass()) {
+  if (archetype->requiresClass() ||
+      (LayoutInfo && LayoutInfo->isRefCountedObject())) {
     ReferenceCounting refcount;
     llvm::PointerType *reprTy;
 
@@ -412,6 +440,28 @@ const TypeInfo *TypeConverter::convertArchetypeType(ArchetypeType *archetype) {
                                       protocols, refcount);
   }
 
+  // If the archetype is trivial fixed-size layout-constrained, use a fixed size
+  // representation.
+  if (LayoutInfo && LayoutInfo->isFixedSizeTrivial()) {
+    Size size(LayoutInfo->getTrivialSizeInBytes());
+    Alignment align(LayoutInfo->getTrivialSizeInBytes());
+    auto spareBits =
+      SpareBitVector::getConstant(size.getValueInBits(), false);
+    // Get an integer type of the required size.
+    auto ProperlySizedIntTy = SILType::getBuiltinIntegerType(
+        size.getValueInBits(), IGM.getSwiftModule()->getASTContext());
+    auto storageType = IGM.getStorageType(ProperlySizedIntTy);
+    return FixedSizeArchetypeTypeInfo::create(storageType, size, align,
+                                              spareBits, protocols);
+  }
+
+  // If the archetype is a trivial layout-constrained, use a POD
+  // representation. This type is not loadable, but it is known
+  // to be a POD.
+  if (LayoutInfo && LayoutInfo->isAddressOnlyTrivial()) {
+    // TODO: Create NonFixedSizeArchetypeTypeInfo and return it.
+  }
+
   // Otherwise, for now, always use an opaque indirect type.
   llvm::Type *storageType = IGM.OpaquePtrTy->getElementType();
   return OpaqueArchetypeTypeInfo::create(storageType, protocols);

lib/IRGen/GenOpaque.cpp

@@ -669,6 +669,9 @@ void irgen::emitAssignWithTakeCall(IRGenFunction &IGF,
 void irgen::emitDestroyCall(IRGenFunction &IGF,
                             SILType T,
                             Address object) {
+  // If T is a trivial/POD type, nothing needs to be done.
+  if (T.getObjectType().isTrivial(IGF.getSILModule()))
+    return;
   auto metadata = IGF.emitTypeMetadataRefForLayout(T);
   llvm::Value *fn = IGF.emitValueWitnessForLayout(T,
                                    ValueWitness::Destroy);
@@ -683,6 +686,9 @@ void irgen::emitDestroyArrayCall(IRGenFunction &IGF,
                                  SILType T,
                                  Address object,
                                  llvm::Value *count) {
+  // If T is a trivial/POD type, nothing needs to be done.
+  if (T.getObjectType().isTrivial(IGF.getSILModule()))
+    return;
   auto metadata = IGF.emitTypeMetadataRefForLayout(T);
   llvm::Value *fn = IGF.emitValueWitnessForLayout(T,
                                    ValueWitness::DestroyArray);

lib/SIL/Mangle.cpp

@@ -55,7 +55,7 @@ static void mangleSubstitution(Mangler &M, Substitution Sub) {
 
 void GenericSpecializationMangler::mangleSpecialization() {
   Mangler &M = getMangler();
-
+  // This is a full specialization.
   SILFunctionType *FTy = Function->getLoweredFunctionType();
   CanGenericSignature Sig = FTy->getGenericSignature();
 
@@ -73,6 +73,25 @@ void GenericSpecializationMangler::mangleSpecialization() {
   assert(idx == Subs.size() && "subs not parallel to dependent types");
 }
 
+void PartialSpecializationMangler::mangleSpecialization() {
+  Mangler &M = getMangler();
+  // If the only change to the generic signature during specialization is
+  // addition of new same-type requirements, which happens in case of a
+  // full specialization, it would be enough to mangle only the substitutions.
+  //
+  // If the types of function arguments have not changed, but some new
+  // conformances were added to the generic parameters, e.g. in case of
+  // a pre-specialization, then it would be enough to mangle only the new
+  // generic signature.
+  //
+  // If the types of function arguments have changed as a result of a partial
+  // specialization, we need to mangle the entire new function type.
+
+  // This is a partial specialization.
+  M.mangleType(SpecializedFnTy, 0);
+  M.append("_");
+}
+
 //===----------------------------------------------------------------------===//
 //                      Function Signature Optimizations
 //===----------------------------------------------------------------------===//