Merge pull request #41120 from eeckstein/capture-propagate-keypaths

CapturePropagation: specialize closures which capture a constant keypath

Author: eeckstein

docs/ABI/Mangling.rst

@@ -1139,6 +1139,7 @@ Some kinds need arguments, which precede ``Tf``.
   CONST-PROP ::= 'i' NATURAL_ZERO            // 64-bit-integer
   CONST-PROP ::= 'd' NATURAL_ZERO            // float-as-64-bit-integer
   CONST-PROP ::= 's' ENCODING                // string literal. Consumes one identifier argument.
+  CONST-PROP ::= 'k'                         // keypath. Consumes one identifier - the SHA1 of the keypath and two types (root and value).
 
   ENCODING ::= 'b'                           // utf8
   ENCODING ::= 'w'                           // utf16

include/swift/Basic/MD5Stream.h

@@ -0,0 +1,45 @@
+//===--- MD5Stream.h - raw_ostream that compute MD5  ------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_MD5STREAM_H
+#define SWIFT_MD5STREAM_H
+
+#include "llvm/Support/MD5.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace swift {
+
+/// An output stream which calculates the MD5 hash of the streamed data.
+class MD5Stream : public llvm::raw_ostream {
+private:
+
+  uint64_t Pos = 0;
+  llvm::MD5 Hash;
+
+  void write_impl(const char *Ptr, size_t Size) override {
+    Hash.update(ArrayRef<uint8_t>(reinterpret_cast<const uint8_t *>(Ptr), Size));
+    Pos += Size;
+  }
+
+  uint64_t current_pos() const override { return Pos; }
+
+public:
+
+  void final(llvm::MD5::MD5Result &Result) {
+    flush();
+    Hash.final(Result);
+  }
+};
+
+} // namespace swift
+
+#endif

include/swift/Demangling/Demangle.h

@@ -109,6 +109,7 @@ enum class FunctionSigSpecializationParamKind : unsigned {
   BoxToValue = 6,
   BoxToStack = 7,
   InOutToOut = 8,
+  ConstantPropKeyPath = 9,
 
   // Option Set Flags use bits 6-31. This gives us 26 bits to use for option
   // flags.

include/swift/SIL/SILInstruction.h

@@ -78,6 +78,7 @@ class StringLiteralExpr;
 class ValueDecl;
 class VarDecl;
 class FunctionRefBaseInst;
+class SILPrintContext;
 
 template <typename ImplClass> class SILClonerWithScopes;
 
@@ -3489,7 +3490,9 @@ class KeyPathPatternComponent {
 
   void incrementRefCounts() const;
   void decrementRefCounts() const;
-  
+
+  void print(SILPrintContext &ctxt) const;
+
   void Profile(llvm::FoldingSetNodeID &ID);
 };
 

include/swift/SILOptimizer/Utils/SpecializationMangler.h

@@ -95,7 +95,9 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
   FunctionSignatureSpecializationMangler(SpecializationPass Pass,
                                          IsSerialized_t Serialized,
                                          SILFunction *F);
-  void setArgumentConstantProp(unsigned OrigArgIdx, LiteralInst *LI);
+  void setArgumentConstantProp(unsigned OrigArgIdx, SILInstruction *constInst);
+  void appendStringAsIdentifier(StringRef str);
+
   void setArgumentClosureProp(unsigned OrigArgIdx, PartialApplyInst *PAI);
   void setArgumentClosureProp(unsigned OrigArgIdx,
                               ThinToThickFunctionInst *TTTFI);
@@ -112,7 +114,7 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
   std::string mangle();
 
 private:
-  void mangleConstantProp(LiteralInst *LI);
+  void mangleConstantProp(SILInstruction *constInst);
   void mangleClosureProp(SILInstruction *Inst);
   void mangleArgument(ArgumentModifierIntBase ArgMod,
                       NullablePtr<SILInstruction> Inst);

lib/Demangling/Demangler.cpp

@@ -2816,10 +2816,12 @@ NodePointer Demangler::demangleFunctionSpecialization() {
       case FunctionSigSpecializationParamKind::ConstantPropFunction:
       case FunctionSigSpecializationParamKind::ConstantPropGlobal:
       case FunctionSigSpecializationParamKind::ConstantPropString:
+      case FunctionSigSpecializationParamKind::ConstantPropKeyPath:
       case FunctionSigSpecializationParamKind::ClosureProp: {
         size_t FixedChildren = Param->getNumChildren();
         while (NodePointer Ty = popNode(Node::Kind::Type)) {
-          if (ParamKind != FunctionSigSpecializationParamKind::ClosureProp)
+          if (ParamKind != FunctionSigSpecializationParamKind::ClosureProp &&
+              ParamKind != FunctionSigSpecializationParamKind::ConstantPropKeyPath)
             return nullptr;
           Param = addChild(Param, Ty);
         }
@@ -2901,6 +2903,14 @@ NodePointer Demangler::demangleFuncSpecParam(Node::Kind Kind) {
                   Node::Kind::FunctionSignatureSpecializationParamPayload,
                   Encoding));
         }
+        case 'k': {
+          // Consumes two types and a SHA1 identifier.
+          return addChild(
+              Param,
+              createNode(Node::Kind::FunctionSignatureSpecializationParamKind,
+                         Node::IndexType(FunctionSigSpecializationParamKind::
+                                             ConstantPropKeyPath)));
+        }
         default:
           return nullptr;
       }

lib/Demangling/NodePrinter.cpp

@@ -1054,6 +1054,17 @@ void NodePrinter::printFunctionSigSpecializationParams(NodePointer Node,
       Printer << "'";
       Printer << "]";
       break;
+    case FunctionSigSpecializationParamKind::ConstantPropKeyPath:
+      Printer << "[";
+      print(Node->getChild(Idx++), depth + 1);
+      Printer << " : ";
+      print(Node->getChild(Idx++), depth + 1);
+      Printer << "<";
+      print(Node->getChild(Idx++), depth + 1);
+      Printer << ",";
+      print(Node->getChild(Idx++), depth + 1);
+      Printer << ">]";
+      break;
     case FunctionSigSpecializationParamKind::ClosureProp:
       Printer << "[";
       print(Node->getChild(Idx++), depth + 1);
@@ -1611,6 +1622,9 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
     case FunctionSigSpecializationParamKind::ConstantPropString:
       Printer << "Constant Propagated String";
       return nullptr;
+    case FunctionSigSpecializationParamKind::ConstantPropKeyPath:
+      Printer << "Constant Propagated KeyPath";
+      return nullptr;
     case FunctionSigSpecializationParamKind::ClosureProp:
       Printer << "Closure Propagated";
       return nullptr;

lib/Demangling/Remangler.cpp

@@ -1322,6 +1322,7 @@ ManglingError Remangler::mangleFunctionSignatureSpecialization(Node *node,
           break;
         }
         case FunctionSigSpecializationParamKind::ClosureProp:
+        case FunctionSigSpecializationParamKind::ConstantPropKeyPath:
           RETURN_IF_ERROR(mangleIdentifier(Param->getChild(1), depth + 1));
           for (unsigned i = 2, e = Param->getNumChildren(); i != e; ++i) {
             RETURN_IF_ERROR(mangleType(Param->getChild(i), depth + 1));
@@ -1399,6 +1400,9 @@ Remangler::mangleFunctionSignatureSpecializationParam(Node *node,
       }
       break;
     }
+    case FunctionSigSpecializationParamKind::ConstantPropKeyPath:
+      Buffer << "pk";
+      break;
     case FunctionSigSpecializationParamKind::ClosureProp:
       Buffer << 'c';
       break;

lib/IRGen/IRGen.cpp

@@ -27,6 +27,7 @@
 #include "swift/AST/TBDGenRequests.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/Dwarf.h"
+#include "swift/Basic/MD5Stream.h"
 #include "swift/Basic/Platform.h"
 #include "swift/Basic/Statistic.h"
 #include "swift/Basic/Version.h"
@@ -69,7 +70,6 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FormattedStream.h"
-#include "llvm/Support/MD5.h"
 #include "llvm/Support/Mutex.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Target/TargetMachine.h"
@@ -376,30 +376,6 @@ void swift::performLLVMOptimizations(const IRGenOptions &Opts,
   }
 }
 
-namespace {
-/// An output stream which calculates the MD5 hash of the streamed data.
-class MD5Stream : public llvm::raw_ostream {
-private:
-
-  uint64_t Pos = 0;
-  llvm::MD5 Hash;
-
-  void write_impl(const char *Ptr, size_t Size) override {
-    Hash.update(ArrayRef<uint8_t>(reinterpret_cast<const uint8_t *>(Ptr), Size));
-    Pos += Size;
-  }
-
-  uint64_t current_pos() const override { return Pos; }
-
-public:
-
-  void final(MD5::MD5Result &Result) {
-    flush();
-    Hash.final(Result);
-  }
-};
-} // end anonymous namespace
-
 /// Computes the MD5 hash of the llvm \p Module including the compiler version
 /// and options which influence the compilation.
 static MD5::MD5Result getHashOfModule(const IRGenOptions &Opts,

lib/SIL/IR/SILPrinter.cpp

@@ -3861,6 +3861,11 @@ void SILSpecializeAttr::print(llvm::raw_ostream &OS) const {
   }
 }
 
+void KeyPathPatternComponent::print(SILPrintContext &ctxt) const {
+  SILPrinter printer(ctxt);
+  printer.printKeyPathPatternComponent(*this);
+}
+
 //===----------------------------------------------------------------------===//
 // SILPrintContext members
 //===----------------------------------------------------------------------===//