Merge pull request swiftlang#31633 from jckarter/keypath-kvc-string-fold

SILOptimizer: Constant fold the _kvcKeyPathString of literal key paths.

Author: jckarter

include/swift/AST/ASTContext.h

@@ -1027,6 +1027,10 @@ class ASTContext final {
 
   /// Retrieve the IRGen specific SIL passes.
   SILTransformCtors getIRGenSILTransforms() const;
+  
+  /// Check whether a given string would be considered "pure ASCII" by the
+  /// standard library's String implementation.
+  bool isASCIIString(StringRef s) const;
 
 private:
   friend Decl;

include/swift/AST/Attr.h

@@ -2209,6 +2209,19 @@ class DeclAttributes {
         make_range(begin(), end()), ToAttributeKind<ATTR, AllowInvalid>());
   }
 
+  /// Return the range of semantics attributes attached to this attribute set.
+  auto getSemanticsAttrs() const
+      -> decltype(getAttributes<SemanticsAttr>()) {
+    return getAttributes<SemanticsAttr>();
+  }
+
+  /// Return whether this attribute set includes the given semantics attribute.
+  bool hasSemanticsAttr(StringRef attrValue) const {
+    return llvm::any_of(getSemanticsAttrs(), [&](const SemanticsAttr *attr) {
+      return attrValue.equals(attr->Value);
+    });
+  }
+
   // Remove the given attribute from the list of attributes. Used when
   // the attribute was semantically invalid.
   void removeAttribute(const DeclAttribute *attr) {

include/swift/AST/Decl.h

@@ -3524,14 +3524,12 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
 
   /// Return the range of semantics attributes attached to this NominalTypeDecl.
   auto getSemanticsAttrs() const
-      -> decltype(getAttrs().getAttributes<SemanticsAttr>()) {
-    return getAttrs().getAttributes<SemanticsAttr>();
+      -> decltype(getAttrs().getSemanticsAttrs()) {
+    return getAttrs().getSemanticsAttrs();
   }
 
   bool hasSemanticsAttr(StringRef attrValue) const {
-    return llvm::any_of(getSemanticsAttrs(), [&](const SemanticsAttr *attr) {
-      return attrValue.equals(attr->Value);
-    });
+    return getAttrs().hasSemanticsAttr(attrValue);
   }
 
   /// Whether this declaration has a synthesized memberwise initializer.

include/swift/AST/SemanticAttrs.def

@@ -87,5 +87,7 @@ SEMANTICS_ATTR(SLOWPATH, "slowpath")
 SEMANTICS_ATTR(PROGRAMTERMINATION_POINT, "programtermination_point")
 SEMANTICS_ATTR(CONVERT_TO_OBJECTIVE_C, "convertToObjectiveC")
 
+SEMANTICS_ATTR(KEYPATH_KVC_KEY_PATH_STRING, "keypath.kvcKeyPathString")
+
 #undef SEMANTICS_ATTR
 

include/swift/SILOptimizer/Utils/KeyPathProjector.h

@@ -24,6 +24,8 @@
 
 namespace swift {
 
+class KeyPathInst;
+
 /// Projects a statically known key path expression to
 /// a direct property access.
 class KeyPathProjector {
@@ -53,6 +55,10 @@ class KeyPathProjector {
   static std::unique_ptr<KeyPathProjector>
   create(SILValue keyPath, SILValue root, SILLocation loc, SILBuilder &builder);
 
+  /// Extract the literal KeyPathInst underlying a value, or return null if there is none.
+  static KeyPathInst *
+  getLiteralKeyPath(SILValue keyPath);
+  
   /// Projects the key path to an address. Sets up the projection,
   /// invokes the callback, then tears down the projection.
   /// \param accessType The access type of the projected address.

lib/AST/ASTContext.cpp

@@ -4827,3 +4827,11 @@ llvm::LLVMContext &ASTContext::getIntrinsicScratchContext() const {
 #endif
 }
 
+bool ASTContext::isASCIIString(StringRef s) const {
+  for (unsigned char c : s) {
+    if (c > 127) {
+      return false;
+    }
+  }
+  return true;
+}

lib/SILGen/SILGenApply.cpp

@@ -1614,13 +1614,7 @@ static PreparedArguments emitStringLiteral(SILGenFunction &SGF, Expr *E,
                                            StringRef Str, SGFContext C,
                                         StringLiteralExpr::Encoding encoding) {
   uint64_t Length;
-  bool isASCII = true;
-  for (unsigned char c : Str) {
-    if (c > 127) {
-      isASCII = false;
-      break;
-    }
-  }
+  bool isASCII = SGF.getASTContext().isASCIIString(Str);
   StringLiteralInst::Encoding instEncoding;
   switch (encoding) {
   case StringLiteralExpr::UTF8:

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -241,7 +241,9 @@ class SILCombiner :
 
   bool tryOptimizeKeypath(ApplyInst *AI);
   bool tryOptimizeInoutKeypath(BeginApplyInst *AI);
-
+  bool tryOptimizeKeypathApplication(ApplyInst *AI, SILFunction *callee);
+  bool tryOptimizeKeypathKVCString(ApplyInst *AI, SILDeclRef callee);
+      
   // Optimize concatenation of string literals.
   // Constant-fold concatenation of string literals known at compile-time.
   SILInstruction *optimizeConcatenationOfStringLiterals(ApplyInst *AI);

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -226,11 +226,8 @@ SILCombiner::optimizeApplyOfConvertFunctionInst(FullApplySite AI,
 ///   %addr = struct_element_addr/ref_element_addr %root_object
 ///   ...
 ///   load/store %addr
-bool SILCombiner::tryOptimizeKeypath(ApplyInst *AI) {
-  SILFunction *callee = AI->getReferencedFunctionOrNull();
-  if (!callee)
-    return false;
-
+bool SILCombiner::tryOptimizeKeypathApplication(ApplyInst *AI,
+                                          SILFunction *callee) {
   if (AI->getNumArguments() != 3)
     return false;
 
@@ -274,6 +271,99 @@ bool SILCombiner::tryOptimizeKeypath(ApplyInst *AI) {
   return true;
 }
 
+/// Try to optimize a keypath KVC string access on a literal key path.
+///
+/// Replace:
+///   %kp = keypath (objc "blah", ...)
+///   %string = apply %keypath_kvcString_method(%kp)
+/// With:
+///   %string = string_literal "blah"
+bool SILCombiner::tryOptimizeKeypathKVCString(ApplyInst *AI,
+                                              SILDeclRef callee) {
+  if (AI->getNumArguments() != 1) {
+    return false;
+  }
+  if (!callee.hasDecl()) {
+    return false;
+  }
+  auto calleeFn = dyn_cast<FuncDecl>(callee.getDecl());
+  if (!calleeFn)
+    return false;
+  
+  if (!calleeFn->getAttrs()
+        .hasSemanticsAttr(semantics::KEYPATH_KVC_KEY_PATH_STRING))
+    return false;
+  
+  // Method should return `String?`
+  auto &C = calleeFn->getASTContext();
+  auto objTy = AI->getType().getOptionalObjectType();
+  if (!objTy || objTy.getStructOrBoundGenericStruct() != C.getStringDecl())
+    return false;
+  
+  KeyPathInst *kp
+    = KeyPathProjector::getLiteralKeyPath(AI->getArgument(0));
+  if (!kp || !kp->hasPattern())
+    return false;
+  
+  auto objcString = kp->getPattern()->getObjCString();
+  
+  SILValue literalValue;
+  if (objcString.empty()) {
+    // Replace with a nil String value.
+    literalValue = Builder.createEnum(AI->getLoc(), SILValue(),
+                                      C.getOptionalNoneDecl(),
+                                      AI->getType());
+  } else {
+    // Construct a literal String from the ObjC string.
+    auto init = C.getStringBuiltinInitDecl(C.getStringDecl());
+    if (!init)
+      return false;
+    auto initRef = SILDeclRef(init.getDecl(), SILDeclRef::Kind::Allocator);
+    auto initFn = AI->getModule().findFunction(initRef.mangle(),
+                                               SILLinkage::PublicExternal);
+    if (!initFn)
+      return false;
+
+    auto stringValue = Builder.createStringLiteral(AI->getLoc(), objcString,
+                                             StringLiteralInst::Encoding::UTF8);
+    auto stringLen = Builder.createIntegerLiteral(AI->getLoc(),
+                                                SILType::getBuiltinWordType(C),
+                                                objcString.size());
+    auto isAscii = Builder.createIntegerLiteral(AI->getLoc(),
+                                          SILType::getBuiltinIntegerType(1, C),
+                                          C.isASCIIString(objcString));
+    auto metaTy =
+      CanMetatypeType::get(objTy.getASTType(), MetatypeRepresentation::Thin);
+    auto self = Builder.createMetatype(AI->getLoc(),
+                                     SILType::getPrimitiveObjectType(metaTy));
+    
+    auto initFnRef = Builder.createFunctionRef(AI->getLoc(), initFn);
+    auto string = Builder.createApply(AI->getLoc(),
+                                      initFnRef, {},
+                                      {stringValue, stringLen, isAscii, self});
+    
+    literalValue = Builder.createEnum(AI->getLoc(), string,
+                                      C.getOptionalSomeDecl(), AI->getType());
+  }
+
+  AI->replaceAllUsesWith(literalValue);
+  eraseInstFromFunction(*AI);
+  ++NumOptimizedKeypaths;
+  return true;
+}
+
+bool SILCombiner::tryOptimizeKeypath(ApplyInst *AI) {
+  if (SILFunction *callee = AI->getReferencedFunctionOrNull()) {
+    return tryOptimizeKeypathApplication(AI, callee);
+  }
+  
+  if (auto method = dyn_cast<ClassMethodInst>(AI->getCallee())) {
+    return tryOptimizeKeypathKVCString(AI, method->getMember());
+  }
+  
+  return false;
+}
+
 /// Try to optimize a keypath application with an apply instruction.
 ///
 /// Replaces (simplified SIL):

lib/SILOptimizer/Utils/KeyPathProjector.cpp

@@ -656,14 +656,19 @@ class CompleteKeyPathProjector : public KeyPathProjector {
   }
 };
 
+KeyPathInst *
+KeyPathProjector::getLiteralKeyPath(SILValue keyPath) {
+  if (auto *upCast = dyn_cast<UpcastInst>(keyPath))
+    keyPath = upCast->getOperand();
+  // TODO: Look through other conversions, copies, etc.?
+  return dyn_cast<KeyPathInst>(keyPath);
+}
+
 std::unique_ptr<KeyPathProjector>
 KeyPathProjector::create(SILValue keyPath, SILValue root,
                          SILLocation loc, SILBuilder &builder) {
-  if (auto *upCast = dyn_cast<UpcastInst>(keyPath))
-    keyPath = upCast->getOperand();
-  
   // Is it a keypath instruction at all?
-  auto *kpInst = dyn_cast<KeyPathInst>(keyPath);
+  auto *kpInst = getLiteralKeyPath(keyPath);
   if (!kpInst || !kpInst->hasPattern())
     return nullptr;