Merge pull request #32544 from eeckstein/fold-kp-offset

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

Author: eeckstein

docs/SIL.rst

@@ -3573,6 +3573,20 @@ the encoding is ``objc_selector``, the string literal produces a
 reference to a UTF-8-encoded Objective-C selector in the Objective-C
 method name segment.
 
+base_addr_for_offset
+````````````````````
+::
+
+  sil-instruction ::= 'base_addr_for_offset' sil-type
+
+  %1 = base_addr_for_offset $*S
+  // %1 has type $*S
+
+Creates a base address for offset calculations. The result can be used by
+address projections, like ``struct_element_addr``, which themselves return the
+offset of the projected fields.
+IR generation simply creates a null pointer for ``base_addr_for_offset``.
+
 Dynamic Dispatch
 ~~~~~~~~~~~~~~~~
 

include/swift/SIL/SILBuilder.h

@@ -631,6 +631,10 @@ class SILBuilder {
     return insert(new (getModule()) GlobalValueInst(getSILDebugLocation(Loc), g,
                                                     getTypeExpansionContext()));
   }
+  BaseAddrForOffsetInst *createBaseAddrForOffset(SILLocation Loc, SILType Ty) {
+    return insert(new (F->getModule())
+                  BaseAddrForOffsetInst(getSILDebugLocation(Loc), Ty));
+  }
   IntegerLiteralInst *createIntegerLiteral(IntegerLiteralExpr *E);
 
   IntegerLiteralInst *createIntegerLiteral(SILLocation Loc, SILType Ty,

include/swift/SIL/SILCloner.h

@@ -1030,6 +1030,15 @@ SILCloner<ImplClass>::visitGlobalValueInst(GlobalValueInst *Inst) {
                                            Inst->getReferencedGlobal()));
 }
 
+template<typename ImplClass>
+void
+SILCloner<ImplClass>::visitBaseAddrForOffsetInst(BaseAddrForOffsetInst *Inst) {
+  getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
+  recordClonedInstruction(
+      Inst, getBuilder().createBaseAddrForOffset(getOpLocation(Inst->getLoc()),
+                                                  getOpType(Inst->getType())));
+}
+
 template<typename ImplClass>
 void
 SILCloner<ImplClass>::visitIntegerLiteralInst(IntegerLiteralInst *Inst) {

include/swift/SIL/SILInstruction.h

@@ -3213,6 +3213,20 @@ class GlobalAddrInst
       : InstructionBase(DebugLoc, Ty, nullptr) {}
 };
 
+/// Creates a base address for offset calculations.
+class BaseAddrForOffsetInst
+    : public InstructionBase<SILInstructionKind::BaseAddrForOffsetInst,
+                             LiteralInst> {
+  friend SILBuilder;
+
+  BaseAddrForOffsetInst(SILDebugLocation DebugLoc, SILType Ty)
+      : InstructionBase(DebugLoc, Ty) {}
+
+public:
+  ArrayRef<Operand> getAllOperands() const { return {}; }
+  MutableArrayRef<Operand> getAllOperands() { return {}; }
+};
+
 /// Gives the value of a global variable.
 ///
 /// The referenced global variable must be a statically initialized object.

include/swift/SIL/SILNodes.def

@@ -458,6 +458,8 @@ ABSTRACT_VALUE_AND_INST(SingleValueInstruction, ValueBase, SILInstruction)
                       LiteralInst, None, DoesNotRelease)
     SINGLE_VALUE_INST(GlobalAddrInst, global_addr,
                       LiteralInst, None, DoesNotRelease)
+    SINGLE_VALUE_INST(BaseAddrForOffsetInst, base_addr_for_offset,
+                      LiteralInst, None, DoesNotRelease)
     SINGLE_VALUE_INST(GlobalValueInst, global_value,
                       LiteralInst, None, DoesNotRelease)
     SINGLE_VALUE_INST(IntegerLiteralInst, integer_literal,

lib/IRGen/IRGenSIL.cpp

@@ -868,6 +868,7 @@ class IRGenSILFunction :
   void visitAllocGlobalInst(AllocGlobalInst *i);
   void visitGlobalAddrInst(GlobalAddrInst *i);
   void visitGlobalValueInst(GlobalValueInst *i);
+  void visitBaseAddrForOffsetInst(BaseAddrForOffsetInst *i);
 
   void visitIntegerLiteralInst(IntegerLiteralInst *i);
   void visitFloatLiteralInst(FloatLiteralInst *i);
@@ -2074,6 +2075,12 @@ void IRGenSILFunction::visitGlobalValueInst(GlobalValueInst *i) {
   setLoweredExplosion(i, e);
 }
 
+void IRGenSILFunction::visitBaseAddrForOffsetInst(BaseAddrForOffsetInst *i) {
+  auto storagePtrTy = IGM.getStoragePointerType(i->getType());
+  llvm::Value *addr = llvm::ConstantPointerNull::get(storagePtrTy);
+  setLoweredAddress(i, Address(addr, Alignment()));
+}
+
 void IRGenSILFunction::visitMetatypeInst(swift::MetatypeInst *i) {
   auto metaTy = i->getType().castTo<MetatypeType>();
   Explosion e;

lib/SIL/IR/OperandOwnership.cpp

@@ -123,6 +123,7 @@ SHOULD_NEVER_VISIT_INST(DynamicFunctionRef)
 SHOULD_NEVER_VISIT_INST(PreviousDynamicFunctionRef)
 SHOULD_NEVER_VISIT_INST(GlobalAddr)
 SHOULD_NEVER_VISIT_INST(GlobalValue)
+SHOULD_NEVER_VISIT_INST(BaseAddrForOffset)
 SHOULD_NEVER_VISIT_INST(IntegerLiteral)
 SHOULD_NEVER_VISIT_INST(Metatype)
 SHOULD_NEVER_VISIT_INST(ObjCProtocol)

lib/SIL/IR/SILPrinter.cpp

@@ -1319,6 +1319,10 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
     *this << " : " << GVI->getType();
   }
 
+  void visitBaseAddrForOffsetInst(BaseAddrForOffsetInst *BAI) {
+    *this << BAI->getType();
+  }
+
   void visitIntegerLiteralInst(IntegerLiteralInst *ILI) {
     const auto &lit = ILI->getValue();
     *this << ILI->getType() << ", " << lit;

lib/SIL/IR/ValueOwnership.cpp

@@ -103,6 +103,7 @@ CONSTANT_OWNERSHIP_INST(None, FunctionRef)
 CONSTANT_OWNERSHIP_INST(None, DynamicFunctionRef)
 CONSTANT_OWNERSHIP_INST(None, PreviousDynamicFunctionRef)
 CONSTANT_OWNERSHIP_INST(None, GlobalAddr)
+CONSTANT_OWNERSHIP_INST(None, BaseAddrForOffset)
 CONSTANT_OWNERSHIP_INST(None, IndexAddr)
 CONSTANT_OWNERSHIP_INST(None, IndexRawPointer)
 CONSTANT_OWNERSHIP_INST(None, InitEnumDataAddr)

lib/SIL/Parser/ParseSIL.cpp

@@ -4583,6 +4583,15 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
       }
       break;
     }
+    case SILInstructionKind::BaseAddrForOffsetInst: {
+      SILType Ty;
+      if (parseSILType(Ty))
+        return true;
+      if (parseSILDebugLocation(InstLoc, B))
+        return true;
+      ResultVal = B.createBaseAddrForOffset(InstLoc, Ty);
+      break;
+    }
     case SILInstructionKind::SelectEnumInst:
     case SILInstructionKind::SelectEnumAddrInst: {
       if (parseTypedValueRef(Val, B))

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;
 }
 

lib/SILOptimizer/UtilityPasses/SerializeSILPass.cpp

@@ -161,6 +161,7 @@ static bool hasOpaqueArchetype(TypeExpansionContext context,
   case SILInstructionKind::PreviousDynamicFunctionRefInst:
   case SILInstructionKind::GlobalAddrInst:
   case SILInstructionKind::GlobalValueInst:
+  case SILInstructionKind::BaseAddrForOffsetInst:
   case SILInstructionKind::IntegerLiteralInst:
   case SILInstructionKind::FloatLiteralInst:
   case SILInstructionKind::StringLiteralInst:

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -685,6 +685,7 @@ InlineCost swift::instructionInlineCost(SILInstruction &I) {
   case SILInstructionKind::FunctionRefInst:
   case SILInstructionKind::AllocGlobalInst:
   case SILInstructionKind::GlobalAddrInst:
+  case SILInstructionKind::BaseAddrForOffsetInst:
   case SILInstructionKind::EndLifetimeInst:
   case SILInstructionKind::UncheckedOwnershipConversionInst:
     return InlineCost::Free;

lib/Serialization/DeserializeSIL.cpp

@@ -1592,6 +1592,11 @@ bool SILDeserializer::readSILInstruction(SILFunction *Fn,
     }
     break;
   }
+  case SILInstructionKind::BaseAddrForOffsetInst:
+    assert(RecordKind == SIL_ONE_TYPE && "Layout should be OneType.");
+    ResultVal = Builder.createBaseAddrForOffset(Loc,
+              getSILType(MF->getType(TyID), (SILValueCategory)TyCategory, Fn));
+    break;
   case SILInstructionKind::DeallocStackInst: {
     auto Ty = MF->getType(TyID);
     ResultVal = Builder.createDeallocStack(

lib/Serialization/SerializeSIL.cpp

@@ -1091,6 +1091,12 @@ void SILSerializer::writeSILInstruction(const SILInstruction &SI) {
         S.addUniquedStringRef(G->getName()));
     break;
   }
+  case SILInstructionKind::BaseAddrForOffsetInst: {
+    const BaseAddrForOffsetInst *BAI = cast<BaseAddrForOffsetInst>(&SI);
+    writeOneTypeLayout(BAI->getKind(), /*attrs*/ 0, BAI->getType());
+    break;
+  }
+
   case SILInstructionKind::BranchInst: {
     // Format: destination basic block ID, a list of arguments. Use
     // SILOneTypeValuesLayout.