Improve handling of pre-specializations.

This commit does not result in any visible functionality changes yet.
This is a preparation for some changes on the swift-3-indexing branch.

Author: swiftix

lib/SILOptimizer/IPO/UsePrespecialized.cpp

@@ -75,6 +75,9 @@ bool UsePrespecialized::replaceByPrespecialized(SILFunction &F) {
     if (!ReferencedF)
       continue;
 
+    DEBUG(llvm::dbgs() << "Trying to use specialized function for:\n";
+          AI.getInstruction()->dumpInContext());
+
     // Check if it is a call of a generic function.
     // If this is the case, check if there is a specialization
     // available for it already and use this specialization
@@ -114,12 +117,22 @@ bool UsePrespecialized::replaceByPrespecialized(SILFunction &F) {
     // If we already have this specialization, reuse it.
     auto PrevF = M.lookUpFunction(ClonedName);
     if (PrevF) {
+      DEBUG(llvm::dbgs() << "Found a specialization: " << ClonedName << "\n");
       if (PrevF->getLinkage() != SILLinkage::SharedExternal)
         NewF = PrevF;
-    } else {
+      else {
+        DEBUG(llvm::dbgs() << "Wrong linkage: " << (int)PrevF->getLinkage()
+                           << "\n");
+      }
+    }
+
+    if (!PrevF || !NewF) {
       // Check for the existence of this function in another module without
       // loading the function body.
       PrevF = lookupPrespecializedSymbol(M, ClonedName);
+      DEBUG(llvm::dbgs()
+            << "Checked if there is a specialization in a different module: "
+            << PrevF << "\n");
       if (!PrevF)
         continue;
       NewF = PrevF;
@@ -128,10 +141,12 @@ bool UsePrespecialized::replaceByPrespecialized(SILFunction &F) {
     if (!NewF)
       continue;
 
+    assert(NewF->isExternalDeclaration() &&
+           "Prespecialized function should be an external declaration");
+
     // An existing specialization was found.
-    DEBUG(
-        llvm::dbgs() << "Found a specialization of " << ReferencedF->getName()
-        << " : " << NewF->getName() << "\n");
+    DEBUG(llvm::dbgs() << "Found a specialization of " << ReferencedF->getName()
+                       << " : " << NewF->getName() << "\n");
 
     auto NewAI = replaceWithSpecializedFunction(AI, NewF, ReInfo);
     AI.getInstruction()->replaceAllUsesWith(NewAI.getInstruction());

lib/SILOptimizer/Utils/Generics.cpp

@@ -41,6 +41,9 @@ ReabstractionInfo::ReabstractionInfo(SILFunction *OrigF,
   if (hasUnboundGenericTypes(InterfaceSubs)) {
     DEBUG(llvm::dbgs() <<
           "    Cannot specialize with unbound interface substitutions.\n");
+    DEBUG(for (auto Sub : ParamSubs) {
+            Sub.dump();
+          });
     return;
   }
   if (hasDynamicSelfTypes(InterfaceSubs)) {
@@ -167,8 +170,12 @@ SILFunction *GenericFuncSpecializer::lookupSpecialization() {
     assert(ReInfo.getSpecializedType()
            == SpecializedF->getLoweredFunctionType() &&
            "Previously specialized function does not match expected type.");
+    DEBUG(llvm::dbgs() << "Found an existing specialization for: " << ClonedName
+                       << "\n");
     return SpecializedF;
   }
+  DEBUG(llvm::dbgs() << "Could not find an existing specialization for: "
+                     << ClonedName << "\n");
   return nullptr;
 }
 
@@ -400,7 +407,8 @@ void swift::trySpecializeApplyOfGeneric(
   auto *F = Apply.getInstruction()->getFunction();
   auto *RefF = cast<FunctionRefInst>(Apply.getCallee())->getReferencedFunction();
 
-  DEBUG(llvm::dbgs() << "  ApplyInst: " << *Apply.getInstruction());
+  DEBUG(llvm::dbgs() << "  ApplyInst:\n";
+        Apply.getInstruction()->dumpInContext());
 
   // If the caller is fragile but the callee is not, bail out.
   // Specializations have shared linkage, which means they do
@@ -470,6 +478,9 @@ void swift::trySpecializeApplyOfGeneric(
            == SpecializedF->isFragile() &&
            "Previously specialized function does not match expected "
            "resilience level.");
+    // Even if the pre-specialization exists already, try to preserve it
+    // if it is whitelisted.
+    linkSpecialization(M, SpecializedF);
   } else {
     SpecializedF = FuncSpecializer.tryCreateSpecialization();
     if (!SpecializedF)
@@ -534,8 +545,29 @@ void swift::trySpecializeApplyOfGeneric(
 // This uses the SIL linker to checks for the does not load the body of the pres
 // =============================================================================
 
+static void keepSpecializationAsPublic(SILFunction *F) {
+  DEBUG(auto DemangledNameString =
+            swift::Demangle::demangleSymbolAsString(F->getName());
+        StringRef DemangledName = DemangledNameString;
+        llvm::dbgs() << "Keep specialization public: " << DemangledName << " : "
+                     << F->getName() << "\n");
+  // Make it public, so that others can refer to it.
+  //
+  // NOTE: This function may refer to non-public symbols, which may lead to
+  // problems, if you ever try to inline this function. Therefore, these
+  // specializations should only be used to refer to them, but should never
+  // be inlined!  The general rule could be: Never inline specializations
+  // from stdlib!
+  //
+  // NOTE: Making these specializations public at this point breaks
+  // some optimizations. Therefore, just mark the function.
+  // DeadFunctionElimination pass will check if the function is marked
+  // and preserve it if required.
+  F->setKeepAsPublic(true);
+}
+
 /// Link a specialization for generating prespecialized code.
-/// 
+///
 /// For now, it is performed only for specializations in the
 /// standard library. But in the future, one could think of
 /// maintaining a cache of optimized specializations.
@@ -545,34 +577,14 @@ void swift::trySpecializeApplyOfGeneric(
 /// the library, but only used only by client code compiled at -Onone. They
 /// should be never inlined.
 static bool linkSpecialization(SILModule &M, SILFunction *F) {
+  if (F->isKeepAsPublic())
+    return true;
   // Do not remove functions from the white-list. Keep them around.
   // Change their linkage to public, so that other applications can refer to it.
-
   if (M.getOptions().Optimization >= SILOptions::SILOptMode::Optimize &&
-      F->getLinkage() != SILLinkage::Public &&
       F->getModule().getSwiftModule()->getName().str() == SWIFT_ONONE_SUPPORT) {
-    if (F->getLinkage() != SILLinkage::Public &&
-        isWhitelistedSpecialization(F->getName())) {
-
-      DEBUG(
-        auto DemangledNameString =
-          swift::Demangle::demangleSymbolAsString(F->getName());
-        StringRef DemangledName = DemangledNameString;
-        llvm::dbgs() << "Keep specialization: " << DemangledName << " : "
-                     << F->getName() << "\n");
-      // Make it public, so that others can refer to it.
-      //
-      // NOTE: This function may refer to non-public symbols, which may lead to
-      // problems, if you ever try to inline this function. Therefore, these
-      // specializations should only be used to refer to them, but should never
-      // be inlined!  The general rule could be: Never inline specializations
-      // from stdlib!
-      //
-      // NOTE: Making these specializations public at this point breaks
-      // some optimizations. Therefore, just mark the function.
-      // DeadFunctionElimination pass will check if the function is marked
-      // and preserve it if required.
-      F->setKeepAsPublic(true);
+    if (isWhitelistedSpecialization(F->getName())) {
+      keepSpecializationAsPublic(F);
       return true;
     }
   }
@@ -607,7 +619,10 @@ bool swift::isWhitelistedSpecialization(StringRef SpecName) {
 
   StringRef DemangledName = DemangledNameString;
 
+  DEBUG(llvm::dbgs() << "Check if whitelisted: " << DemangledName << "\n");
+
   auto pos = DemangledName.find("generic ", 0);
+  auto oldpos = pos;
   if (pos == StringRef::npos)
     return false;
 
@@ -618,11 +633,22 @@ bool swift::isWhitelistedSpecialization(StringRef SpecName) {
   buffer << STDLIB_NAME <<'.';
 
   StringRef OfStr = buffer.str();
-
-  pos = DemangledName.find(OfStr, pos);
-
-  if (pos == StringRef::npos)
-    return false;
+  DEBUG(llvm::dbgs() << "Check substring: " << OfStr << "\n");
+
+  pos = DemangledName.find(OfStr, oldpos);
+
+  if (pos == StringRef::npos) {
+    // Create "of (extension in Swift).Swift"
+    llvm::SmallString<64> OfString;
+    llvm::raw_svector_ostream buffer(OfString);
+    buffer << "of (extension in " << STDLIB_NAME << "):";
+    buffer << STDLIB_NAME << '.';
+    OfStr = buffer.str();
+    pos = DemangledName.find(OfStr, oldpos);
+    DEBUG(llvm::dbgs() << "Check substring: " << OfStr << "\n");
+    if (pos == StringRef::npos)
+      return false;
+  }
 
   pos += OfStr.size();
 
@@ -659,8 +685,10 @@ SILFunction *swift::lookupPrespecializedSymbol(SILModule &M,
                                                StringRef FunctionName) {
   // First check if the module contains a required specialization already.
   auto *Specialization = M.lookUpFunction(FunctionName);
-  if (Specialization)
-    return Specialization;
+  if (Specialization) {
+    if (Specialization->getLinkage() == SILLinkage::PublicExternal)
+      return Specialization;
+  }
 
   // Then check if the required specialization can be found elsewhere.
   Specialization = lookupExistingSpecialization(M, FunctionName);

stdlib/public/SwiftOnoneSupport/SwiftOnoneSupport.swift

@@ -152,6 +152,10 @@ struct _Prespecialize {
     for i in 0..<a.count {
       count += a[i]
     }
+    // Specialize Range for integers
+    for j in 0...a.count - 1{
+      count += a[j]
+    }
     return count
   }
 }