SILCombine: Constant-fold MemoryLayout<T>.offset(of: \.literalKeyPath)

Replace a call of the getter of AnyKeyPath._storedInlineOffset with a "constant" offset, in case of a keypath literal.
"Constant" offset means a series of struct_element_addr and tuple_element_addr instructions with a 0-pointer as base address.
These instructions can then be lowered to "real" constants in IRGen for concrete types, or to metatype offset lookups for generic or resilient types.

Replace:
  %kp = keypath ...
  %offset = apply %_storedInlineOffset_method(%kp)
with:
  %zero = integer_literal $Builtin.Word, 0
  %null_ptr = unchecked_trivial_bit_cast %zero to $Builtin.RawPointer
  %null_addr = pointer_to_address %null_ptr
  %projected_addr = struct_element_addr %null_addr
   ... // other address projections
  %offset_ptr = address_to_pointer %projected_addr
  %offset_builtin_int = unchecked_trivial_bit_cast %offset_ptr
  %offset_int = struct $Int (%offset_builtin_int)
  %offset = enum $Optional<Int>, #Optional.some!enumelt, %offset_int

rdar://problem/53309403

Author: eeckstein

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -246,8 +246,11 @@ class SILCombiner :
   bool tryOptimizeKeypath(ApplyInst *AI);
   bool tryOptimizeInoutKeypath(BeginApplyInst *AI);
   bool tryOptimizeKeypathApplication(ApplyInst *AI, SILFunction *callee);
-  bool tryOptimizeKeypathKVCString(ApplyInst *AI, SILDeclRef callee);
-      
+  bool tryOptimizeKeypathOffsetOf(ApplyInst *AI, FuncDecl *calleeFn,
+                                  KeyPathInst *kp);
+  bool tryOptimizeKeypathKVCString(ApplyInst *AI, FuncDecl *calleeFn,
+                                  KeyPathInst *kp);
+
   // Optimize concatenation of string literals.
   // Constant-fold concatenation of string literals known at compile-time.
   SILInstruction *optimizeConcatenationOfStringLiterals(ApplyInst *AI);

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -271,6 +271,122 @@ bool SILCombiner::tryOptimizeKeypathApplication(ApplyInst *AI,
   return true;
 }
 
+/// Replaces a call of the getter of AnyKeyPath._storedInlineOffset with a
+/// "constant" offset, in case of a keypath literal.
+///
+/// "Constant" offset means a series of struct_element_addr and
+/// tuple_element_addr instructions with a 0-pointer as base address.
+/// These instructions can then be lowered to "real" constants in IRGen for
+/// concrete types, or to metatype offset lookups for generic or resilient types.
+///
+/// Replaces:
+///   %kp = keypath ...
+///   %offset = apply %_storedInlineOffset_method(%kp)
+/// with:
+///   %zero = integer_literal $Builtin.Word, 0
+///   %null_ptr = unchecked_trivial_bit_cast %zero to $Builtin.RawPointer
+///   %null_addr = pointer_to_address %null_ptr
+///   %projected_addr = struct_element_addr %null_addr
+///    ... // other address projections
+///   %offset_ptr = address_to_pointer %projected_addr
+///   %offset_builtin_int = unchecked_trivial_bit_cast %offset_ptr
+///   %offset_int = struct $Int (%offset_builtin_int)
+///   %offset = enum $Optional<Int>, #Optional.some!enumelt, %offset_int
+bool SILCombiner::tryOptimizeKeypathOffsetOf(ApplyInst *AI,
+                                             FuncDecl *calleeFn,
+                                             KeyPathInst *kp) {
+  auto *accessor = dyn_cast<AccessorDecl>(calleeFn);
+  if (!accessor || !accessor->isGetter())
+    return false;
+
+  AbstractStorageDecl *storage = accessor->getStorage();
+  DeclName name = storage->getName();
+  if (!name.isSimpleName() ||
+      (name.getBaseIdentifier().str() != "_storedInlineOffset"))
+    return false;
+
+  KeyPathPattern *pattern = kp->getPattern();
+  SubstitutionMap patternSubs = kp->getSubstitutions();
+  CanType rootTy = pattern->getRootType().subst(patternSubs)->getCanonicalType();
+  CanType parentTy = rootTy;
+  
+  // First check if _storedInlineOffset would return an offset or nil. Basically
+  // only stored struct and tuple elements produce an offset. Everything else
+  // (e.g. computed properties, class properties) result in nil.
+  bool hasOffset = true;
+  for (const KeyPathPatternComponent &component : pattern->getComponents()) {
+    switch (component.getKind()) {
+    case KeyPathPatternComponent::Kind::StoredProperty: {
+      if (!parentTy.getStructOrBoundGenericStruct())
+        hasOffset = false;
+      break;
+    }
+    case KeyPathPatternComponent::Kind::TupleElement:
+      break;
+    case KeyPathPatternComponent::Kind::GettableProperty:
+    case KeyPathPatternComponent::Kind::SettableProperty:
+      // We cannot predict the offset of fields in resilient types, because it's
+      // unknown if a resilient field is a computed or stored property.
+      if (component.getExternalDecl())
+        return false;
+      hasOffset = false;
+      break;
+    case KeyPathPatternComponent::Kind::OptionalChain:
+    case KeyPathPatternComponent::Kind::OptionalForce:
+    case KeyPathPatternComponent::Kind::OptionalWrap:
+      hasOffset = false;
+      break;
+    }
+    parentTy = component.getComponentType();
+  }
+
+  SILLocation loc = AI->getLoc();
+  SILValue result;
+
+  if (hasOffset) {
+    SILType rootAddrTy = SILType::getPrimitiveAddressType(rootTy);
+    SILValue rootAddr = Builder.createBaseAddrForOffset(loc, rootAddrTy);
+
+    auto projector = KeyPathProjector::create(kp, rootAddr, loc, Builder);
+    if (!projector)
+      return false;
+
+    // Create the address projections of the keypath.
+    SILType ptrType = SILType::getRawPointerType(Builder.getASTContext());
+    SILValue offsetPtr;
+    projector->project(KeyPathProjector::AccessType::Get, [&](SILValue addr) {
+      offsetPtr = Builder.createAddressToPointer(loc, addr, ptrType);
+    });
+
+    // The result of the _storedInlineOffset call should be Optional<Int>. If
+    // not, something is wrong with the stdlib. Anyway, if it's not like we
+    // expect, bail.
+    SILType intType = AI->getType().getOptionalObjectType();
+    if (!intType)
+      return false;
+    StructDecl *intDecl = intType.getStructOrBoundGenericStruct();
+    if (!intDecl || intDecl->getStoredProperties().size() != 1)
+      return false;
+    VarDecl *member = intDecl->getStoredProperties()[0];
+    CanType builtinIntTy = member->getType()->getCanonicalType();
+    if (!isa<BuiltinIntegerType>(builtinIntTy))
+      return false;
+
+    // Convert the projected address back to an optional integer.
+    SILValue offset = Builder.createUncheckedBitCast(loc, offsetPtr,
+                                SILType::getPrimitiveObjectType(builtinIntTy));
+    SILValue offsetInt = Builder.createStruct(loc, intType, { offset });
+    result = Builder.createOptionalSome(loc, offsetInt, AI->getType());
+  } else {
+    // The keypath has no offset.
+    result = Builder.createOptionalNone(loc, AI->getType());
+  }
+  AI->replaceAllUsesWith(result);
+  eraseInstFromFunction(*AI);
+  ++NumOptimizedKeypaths;
+  return true;
+}
+
 /// Try to optimize a keypath KVC string access on a literal key path.
 ///
 /// Replace:
@@ -279,17 +395,8 @@ bool SILCombiner::tryOptimizeKeypathApplication(ApplyInst *AI,
 /// 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;
-  
+                                              FuncDecl *calleeFn,
+                                              KeyPathInst *kp) {
   if (!calleeFn->getAttrs()
         .hasSemanticsAttr(semantics::KEYPATH_KVC_KEY_PATH_STRING))
     return false;
@@ -300,11 +407,6 @@ bool SILCombiner::tryOptimizeKeypathKVCString(ApplyInst *AI,
   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;
@@ -357,10 +459,33 @@ bool SILCombiner::tryOptimizeKeypath(ApplyInst *AI) {
     return tryOptimizeKeypathApplication(AI, callee);
   }
 
-  if (auto method = dyn_cast<ClassMethodInst>(AI->getCallee())) {
-    return tryOptimizeKeypathKVCString(AI, method->getMember());
+  // Try optimize keypath method calls.
+  auto *methodInst = dyn_cast<ClassMethodInst>(AI->getCallee());
+  if (!methodInst)
+    return false;
+  
+  if (AI->getNumArguments() != 1) {
+    return false;
+  }
+
+  SILDeclRef callee = methodInst->getMember();
+  if (!callee.hasDecl()) {
+    return false;
   }
+  auto *calleeFn = dyn_cast<FuncDecl>(callee.getDecl());
+  if (!calleeFn)
+    return false;
+
+  KeyPathInst *kp = KeyPathProjector::getLiteralKeyPath(AI->getArgument(0));
+  if (!kp || !kp->hasPattern())
+    return false;
 
+  if (tryOptimizeKeypathOffsetOf(AI, calleeFn, kp))
+    return true;
+
+  if (tryOptimizeKeypathKVCString(AI, calleeFn, kp))
+    return true;
+
   return false;
 }
 

test/SILOptimizer/Inputs/struct_with_fields.swift

@@ -0,0 +1,5 @@
+
+public struct TestStruct {
+  public var x: Int
+  public var y: Int
+}