Merge pull request swiftlang#28043 from zoecarver/fix/substitution-map-composition

Devirtualize calls to protocol composition type methods

Author: atrick

benchmark/CMakeLists.txt

@@ -62,6 +62,7 @@ set(SWIFT_BENCH_MODULES
     single-source/Combos
     single-source/DataBenchmarks
     single-source/DeadArray
+    single-source/DevirtualizeProtocolComposition
     single-source/DictOfArraysToArrayOfDicts
     single-source/DictTest
     single-source/DictTest2

benchmark/single-source/DevirtualizeProtocolComposition.swift

@@ -0,0 +1,45 @@
+//===--- DevirtualizeProtocolComposition.swift -------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2019 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
+//
+//===----------------------------------------------------------------------===//
+
+import TestsUtils
+
+public let DevirtualizeProtocolComposition = [
+  BenchmarkInfo(name: "DevirtualizeProtocolComposition", runFunction: run_DevirtualizeProtocolComposition, tags: [.validation, .api]),
+]
+
+protocol Pingable { func ping() -> Int; func pong() -> Int}
+
+public class Game<T> {
+  func length() -> Int { return 10 }
+}
+
+public class PingPong: Game<String> { }
+
+extension PingPong : Pingable {
+  func ping() -> Int { return 1 }
+  func pong() -> Int { return 2 }
+}
+
+func playGame<T>(_ x: Game<T> & Pingable) -> Int {
+  var sum = 0
+  for _ in 0..<x.length() {
+    sum += x.ping() + x.pong()
+  }
+  return sum
+}
+
+@inline(never)
+public func run_DevirtualizeProtocolComposition(N: Int) {
+  for _ in 0..<N * 20_000 {
+    blackHole(playGame(PingPong()))
+  }
+}

benchmark/utils/main.swift

@@ -50,6 +50,7 @@ import Codable
 import Combos
 import DataBenchmarks
 import DeadArray
+import DevirtualizeProtocolComposition
 import DictOfArraysToArrayOfDicts
 import DictTest
 import DictTest2
@@ -231,6 +232,7 @@ registerBenchmark(Combos)
 registerBenchmark(ClassArrayGetter)
 registerBenchmark(DataBenchmarks)
 registerBenchmark(DeadArray)
+registerBenchmark(DevirtualizeProtocolComposition)
 registerBenchmark(DictOfArraysToArrayOfDicts)
 registerBenchmark(Dictionary)
 registerBenchmark(Dictionary2)

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -783,10 +783,10 @@ bool SILCombiner::canReplaceArg(FullApplySite Apply,
 /// (@in or @owned convention).
 struct ConcreteArgumentCopy {
   SILValue origArg;
-  CopyAddrInst *tempArgCopy;
+  AllocStackInst *tempArg;
 
-  ConcreteArgumentCopy(SILValue origArg, CopyAddrInst *tempArgCopy)
-      : origArg(origArg), tempArgCopy(tempArgCopy) {
+  ConcreteArgumentCopy(SILValue origArg, AllocStackInst *tempArg)
+      : origArg(origArg), tempArg(tempArg) {
     assert(origArg->getType().isAddress());
   }
 
@@ -825,9 +825,18 @@ struct ConcreteArgumentCopy {
     SILBuilderWithScope B(apply.getInstruction(), BuilderCtx);
     auto loc = apply.getLoc();
     auto *ASI = B.createAllocStack(loc, CEI.ConcreteValue->getType());
-    auto *CAI = B.createCopyAddr(loc, CEI.ConcreteValue, ASI, IsNotTake,
-                                 IsInitialization_t::IsInitialization);
-    return ConcreteArgumentCopy(origArg, CAI);
+    // If the type is an address, simple copy it.
+    if (CEI.ConcreteValue->getType().isAddress()) {
+      B.createCopyAddr(loc, CEI.ConcreteValue, ASI, IsNotTake,
+                       IsInitialization_t::IsInitialization);
+    } else {
+      // Otherwise, we probably got the value from the source of a store
+      // instruction so, create a store into the temporary argument.
+      B.createStrongRetain(loc, CEI.ConcreteValue, B.getDefaultAtomicity());
+      B.createStore(loc, CEI.ConcreteValue, ASI,
+                    StoreOwnershipQualifier::Unqualified);
+    }
+    return ConcreteArgumentCopy(origArg, ASI);
   }
 };
 
@@ -858,20 +867,19 @@ SILInstruction *SILCombiner::createApplyWithConcreteType(
     FullApplySite Apply,
     const llvm::SmallDenseMap<unsigned, ConcreteOpenedExistentialInfo> &COAIs,
     SILBuilderContext &BuilderCtx) {
-
   // Ensure that the callee is polymorphic.
   assert(Apply.getOrigCalleeType()->isPolymorphic());
 
   // Create the new set of arguments to apply including their substitutions.
   SubstitutionMap NewCallSubs = Apply.getSubstitutionMap();
   SmallVector<SILValue, 8> NewArgs;
-  bool UpdatedArgs = false;
   unsigned ArgIdx = 0;
   // Push the indirect result arguments.
   for (unsigned EndIdx = Apply.getSubstCalleeConv().getSILArgIndexOfFirstParam();
        ArgIdx < EndIdx; ++ArgIdx) {
       NewArgs.push_back(Apply.getArgument(ArgIdx));
   }
+
   // Transform the parameter arguments.
   SmallVector<ConcreteArgumentCopy, 4> concreteArgCopies;
   for (unsigned EndIdx = Apply.getNumArguments(); ArgIdx < EndIdx; ++ArgIdx) {
@@ -890,16 +898,18 @@ SILInstruction *SILCombiner::createApplyWithConcreteType(
       NewArgs.push_back(Apply.getArgument(ArgIdx));
       continue;
     }
-    UpdatedArgs = true;
+
     // Ensure that we have a concrete value to propagate.
     assert(CEI.ConcreteValue);
+
     auto argSub =
         ConcreteArgumentCopy::generate(CEI, Apply, ArgIdx, BuilderCtx);
     if (argSub) {
       concreteArgCopies.push_back(*argSub);
-      NewArgs.push_back(argSub->tempArgCopy->getDest());
-    } else
+      NewArgs.push_back(argSub->tempArg);
+    } else {
       NewArgs.push_back(CEI.ConcreteValue);
+    }
 
     // Form a new set of substitutions where the argument is
     // replaced with a concrete type.
@@ -922,7 +932,33 @@ SILInstruction *SILCombiner::createApplyWithConcreteType(
         });
   }
 
-  if (!UpdatedArgs) {
+  // We need to make sure that we can a) update Apply to use the new args and b)
+  // at least one argument has changed. If no arguments have changed, we need
+  // to return nullptr. Otherwise, we will have an infinite loop.
+  auto substTy =
+      Apply.getCallee()
+          ->getType()
+          .substGenericArgs(Apply.getModule(), NewCallSubs,
+                            Apply.getFunction()->getTypeExpansionContext())
+          .getAs<SILFunctionType>();
+  SILFunctionConventions conv(substTy,
+                              SILModuleConventions(Apply.getModule()));
+  bool canUpdateArgs = true;
+  bool madeUpdate = false;
+  for (unsigned index = 0; index < conv.getNumSILArguments(); ++index) {
+    // Make sure that *all* the arguments in both the new substitution function
+    // and our vector of new arguments have the same type.
+    canUpdateArgs &=
+        conv.getSILArgumentType(index) == NewArgs[index]->getType();
+    // Make sure that we have changed at least one argument.
+    madeUpdate |=
+        NewArgs[index]->getType() != Apply.getArgument(index)->getType();
+  }
+
+  // If we can't update the args (because of a type mismatch) or the args don't
+  // change, bail out by removing the instructions we've added and returning
+  // nullptr.
+  if (!canUpdateArgs || !madeUpdate) {
     // Remove any new instructions created while attempting to optimize this
     // apply. Since the apply was never rewritten, if they aren't removed here,
     // they will be removed later as dead when visited by SILCombine, causing
@@ -969,8 +1005,7 @@ SILInstruction *SILCombiner::createApplyWithConcreteType(
     auto cleanupLoc = RegularLocation::getAutoGeneratedLocation();
     for (ConcreteArgumentCopy &argCopy : llvm::reverse(concreteArgCopies)) {
       cleanupBuilder.createDestroyAddr(cleanupLoc, argCopy.origArg);
-      cleanupBuilder.createDeallocStack(cleanupLoc,
-                                        argCopy.tempArgCopy->getDest());
+      cleanupBuilder.createDeallocStack(cleanupLoc, argCopy.tempArg);
     }
   }
   return NewApply.getInstruction();

lib/SILOptimizer/Utils/Existential.cpp

@@ -69,9 +69,9 @@ findInitExistentialFromGlobalAddr(GlobalAddrInst *GAI, SILInstruction *Insn) {
 
 /// Returns the instruction that initializes the given stack address. This is
 /// currently either a init_existential_addr, unconditional_checked_cast_addr,
-/// or copy_addr (if the instruction initializing the source of the copy cannot
-/// be determined). Returns nullptr if the initializer does not dominate the
-/// alloc_stack user \p ASIUser.  If the value is copied from another stack
+/// store, or copy_addr (if the instruction initializing the source of the copy
+/// cannot be determined). Returns nullptr if the initializer does not dominate
+/// the alloc_stack user \p ASIUser.  If the value is copied from another stack
 /// location, \p isCopied is set to true.
 ///
 /// allocStackAddr may either itself be an AllocStackInst or an
@@ -111,6 +111,19 @@ static SILInstruction *getStackInitInst(SILValue allocStackAddr,
       }
       continue;
     }
+    if (auto *store = dyn_cast<StoreInst>(User)) {
+      if (store->getDest() == allocStackAddr) {
+        if (SingleWrite)
+          return nullptr;
+        SingleWrite = store;
+        // When we support OSSA here, we need to insert a new copy of the value
+        // before `store` (and make sure that the copy is destroyed when
+        // replacing the apply operand).
+        assert(store->getOwnershipQualifier() ==
+               StoreOwnershipQualifier::Unqualified);
+      }
+      continue;
+    }
     if (isa<InitExistentialAddrInst>(User)) {
       if (SingleWrite)
         return nullptr;
@@ -144,6 +157,9 @@ static SILInstruction *getStackInitInst(SILValue allocStackAddr,
   if (BB != allocStackAddr->getParentBlock() && BB != ASIUser->getParent())
     return nullptr;
 
+  if (auto *store = dyn_cast<StoreInst>(SingleWrite))
+    return store;
+
   if (auto *IE = dyn_cast<InitExistentialAddrInst>(SingleWrite))
     return IE;
 
@@ -180,24 +196,6 @@ static SILInstruction *getStackInitInst(SILValue allocStackAddr,
   return CAI;
 }
 
-/// Return the address of the value used to initialize the given stack location.
-/// If the value originates from init_existential_addr, then it will be a
-/// different type than \p allocStackAddr.
-static SILValue getAddressOfStackInit(SILValue allocStackAddr,
-                                      SILInstruction *ASIUser, bool &isCopied) {
-  SILInstruction *initI = getStackInitInst(allocStackAddr, ASIUser, isCopied);
-  if (!initI)
-    return SILValue();
-
-  if (auto *IEA = dyn_cast<InitExistentialAddrInst>(initI))
-    return IEA;
-
-  if (auto *CAI = dyn_cast<CopyAddrInst>(initI))
-    return CAI->getSrc();
-
-  return SILValue();
-}
-
 /// Check if the given operand originates from a recognized OpenArchetype
 /// instruction. If so, return the Opened, otherwise return nullptr.
 OpenedArchetypeInfo::OpenedArchetypeInfo(Operand &use) {
@@ -207,11 +205,17 @@ OpenedArchetypeInfo::OpenedArchetypeInfo(Operand &use) {
     // Handle:
     //   %opened = open_existential_addr
     //   %instance = alloc $opened
-    //   copy_addr %opened to %stack
+    //   <copy|store> %opened to %stack
     //   <opened_use> %instance
-    if (auto stackInitVal =
-            getAddressOfStackInit(instance, user, isOpenedValueCopied)) {
-      openedVal = stackInitVal;
+    if (auto *initI = getStackInitInst(instance, user, isOpenedValueCopied)) {
+      // init_existential_addr isn't handled here because it isn't considered an
+      // "opened" archtype. init_existential_addr should be handled by
+      // ConcreteExistentialInfo.
+
+      if (auto *CAI = dyn_cast<CopyAddrInst>(initI))
+        openedVal = CAI->getSrc();
+      if (auto *store = dyn_cast<StoreInst>(initI))
+        openedVal = store->getSrc();
     }
   }
   if (auto *Open = dyn_cast<OpenExistentialAddrInst>(openedVal)) {

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -827,9 +827,6 @@ SILValue swift::castValueToABICompatibleType(SILBuilder *builder,
   if (srcTy == destTy)
     return value;
 
-  assert(srcTy.isAddress() == destTy.isAddress()
-         && "Addresses aren't compatible with values");
-
   if (srcTy.isAddress() && destTy.isAddress()) {
     // Cast between two addresses and that's it.
     return builder->createUncheckedAddrCast(loc, value, destTy);