SIL: Replace uses of getReferencedFunction() by getReferencedFunctionOrNull() and getInitialReferencedFunction()

With the advent of dynamic_function_ref the actual callee of such a ref
my vary. Optimizations should not assume to know the content of a
function referenced by dynamic_function_ref. Introduce
getReferencedFunctionOrNull which will return null for such function
refs. And getInitialReferencedFunction to return the referenced
function.
Use as appropriate.

rdar://50959798

Author: aschwaighofer

include/swift/SIL/ApplySite.h

@@ -143,8 +143,17 @@ class ApplySite {
 
   /// Return the referenced function if the callee is a function_ref
   /// instruction.
-  SILFunction *getReferencedFunction() const {
-    FOREACH_IMPL_RETURN(getReferencedFunction());
+  SILFunction *getReferencedFunctionOrNull() const {
+    FOREACH_IMPL_RETURN(getReferencedFunctionOrNull());
+  }
+
+  /// Return the referenced function if the callee is a function_ref like
+  /// instruction.
+  /// WARNING: This not necessarily the function that will be called at runtime.
+  /// If the callee is a (prev_)dynamic_function_ref the actuall function called
+  /// might be different.
+  SILFunction *getInitiallyReferencedFunction() const {
+    FOREACH_IMPL_RETURN(getInitiallyReferencedFunction());
   }
 
   /// Should we optimize this call.

include/swift/SIL/PatternMatch.h

@@ -518,7 +518,7 @@ struct Callee_match<SILFunction &> {
     if (!AI)
       return false;
 
-    return AI->getReferencedFunction() == &Fun;
+    return AI->getReferencedFunctionOrNull() == &Fun;
   }
 };
 

include/swift/SIL/SILCloner.h

@@ -950,7 +950,8 @@ SILCloner<ImplClass>::visitEndApplyInst(EndApplyInst *Inst) {
 template<typename ImplClass>
 void
 SILCloner<ImplClass>::visitFunctionRefInst(FunctionRefInst *Inst) {
-  SILFunction *OpFunction = getOpFunction(Inst->getReferencedFunction());
+  SILFunction *OpFunction =
+      getOpFunction(Inst->getInitiallyReferencedFunction());
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
   recordClonedInstruction(Inst, getBuilder().createFunctionRef(
                                     getOpLocation(Inst->getLoc()), OpFunction));
@@ -959,7 +960,8 @@ SILCloner<ImplClass>::visitFunctionRefInst(FunctionRefInst *Inst) {
 template<typename ImplClass>
 void
 SILCloner<ImplClass>::visitDynamicFunctionRefInst(DynamicFunctionRefInst *Inst) {
-  SILFunction *OpFunction = getOpFunction(Inst->getReferencedFunction());
+  SILFunction *OpFunction =
+      getOpFunction(Inst->getInitiallyReferencedFunction());
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
   recordClonedInstruction(Inst, getBuilder().createDynamicFunctionRef(
                                     getOpLocation(Inst->getLoc()), OpFunction));
@@ -968,7 +970,8 @@ SILCloner<ImplClass>::visitDynamicFunctionRefInst(DynamicFunctionRefInst *Inst)
 template <typename ImplClass>
 void SILCloner<ImplClass>::visitPreviousDynamicFunctionRefInst(
     PreviousDynamicFunctionRefInst *Inst) {
-  SILFunction *OpFunction = getOpFunction(Inst->getReferencedFunction());
+  SILFunction *OpFunction =
+      getOpFunction(Inst->getInitiallyReferencedFunction());
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
   recordClonedInstruction(Inst, getBuilder().createPreviousDynamicFunctionRef(
                                     getOpLocation(Inst->getLoc()), OpFunction));

include/swift/SIL/SILInstruction.h

@@ -1844,9 +1844,22 @@ class ApplyInstBase<Impl, Base, false> : public Base {
   bool isCalleeDynamicallyReplaceable() const;
 
   /// Gets the referenced function if the callee is a function_ref instruction.
-  SILFunction *getReferencedFunction() const {
+  /// Returns null if the callee is dynamic or a (prev_)dynamic_function_ref
+  /// instruction.
+  SILFunction *getReferencedFunctionOrNull() const {
     if (auto *FRI = dyn_cast<FunctionRefBaseInst>(getCallee()))
-      return FRI->getReferencedFunction();
+      return FRI->getReferencedFunctionOrNull();
+    return nullptr;
+  }
+
+  /// Return the referenced function if the callee is a function_ref like
+  /// instruction.
+  /// WARNING: This not necessarily the function that will be called at runtime.
+  /// If this is a (prev_)dynamic_function_ref the actuall function called might
+  /// be different.
+  SILFunction *getInitiallyReferencedFunction() const {
+    if (auto *FRI = dyn_cast<FunctionRefBaseInst>(getCallee()))
+      return FRI->getInitiallyReferencedFunction();
     return nullptr;
   }
 
@@ -2294,8 +2307,23 @@ class FunctionRefBaseInst : public LiteralInst {
 public:
   ~FunctionRefBaseInst();
 
+  /// Return the referenced function if this is a function_ref instruction and
+  /// therefore a client can rely on the dynamically called function being equal
+  /// to the returned value and null otherwise.
+  SILFunction *getReferencedFunctionOrNull() const {
+    auto kind = getKind();
+    if (kind == SILInstructionKind::FunctionRefInst)
+      return f;
+    assert(kind == SILInstructionKind::DynamicFunctionRefInst ||
+           kind == SILInstructionKind::PreviousDynamicFunctionRefInst);
+    return nullptr;
+  }
+
   /// Return the referenced function.
-  SILFunction *getReferencedFunction() const { return f; }
+  /// WARNING: This not necessarily the function that will be called at runtime.
+  /// If this is a (prev_)dynamic_function_ref the actuall function called might
+  /// be different.
+  SILFunction *getInitiallyReferencedFunction() const { return f; }
 
   void dropReferencedFunction();
 
@@ -7805,7 +7833,7 @@ SILFunction *ApplyInstBase<Impl, Base, false>::getCalleeFunction() const {
     // previous_dynamic_function_ref as the target of those functions is not
     // statically known.
     if (auto *FRI = dyn_cast<FunctionRefInst>(Callee))
-      return FRI->getReferencedFunction();
+      return FRI->getReferencedFunctionOrNull();
 
     if (auto *PAI = dyn_cast<PartialApplyInst>(Callee)) {
       Callee = PAI->getCalleeOrigin();

include/swift/SIL/TypeSubstCloner.h

@@ -69,7 +69,7 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
 
       if (!Cloner.Inlining) {
         FunctionRefInst *FRI = dyn_cast<FunctionRefInst>(AI.getCallee());
-        if (FRI && FRI->getReferencedFunction() == AI.getFunction() &&
+        if (FRI && FRI->getInitiallyReferencedFunction() == AI.getFunction() &&
             Subs == Cloner.SubsMap) {
           // Handle recursions by replacing the apply to the callee with an
           // apply to the newly specialized function, but only if substitutions

lib/IRGen/AllocStackHoisting.cpp

@@ -352,7 +352,7 @@ bool indicatesDynamicAvailabilityCheckUse(SILInstruction *I) {
     return false;
   if (Apply->hasSemantics("availability.osversion"))
     return true;
-  auto *FunRef = Apply->getReferencedFunction();
+  auto *FunRef = Apply->getReferencedFunctionOrNull();
   if (!FunRef)
     return false;
   if (FunRef->getName().equals("_swift_stdlib_operatingSystemVersion"))

lib/IRGen/IRGenSIL.cpp

@@ -1842,7 +1842,7 @@ void IRGenSILFunction::visitSILBasicBlock(SILBasicBlock *BB) {
 }
 
 void IRGenSILFunction::visitFunctionRefBaseInst(FunctionRefBaseInst *i) {
-  auto fn = i->getReferencedFunction();
+  auto fn = i->getInitiallyReferencedFunction();
 
   llvm::Constant *fnPtr = IGM.getAddrOfSILFunction(
       fn, NotForDefinition, false /*isDynamicallyReplaceableImplementation*/,

lib/IRGen/LoadableByAddress.cpp

@@ -2790,7 +2790,7 @@ void LoadableByAddress::run() {
     for (SILBasicBlock &BB : CurrF) {
       for (SILInstruction &I : BB) {
         if (auto *FRI = dyn_cast<FunctionRefBaseInst>(&I)) {
-          SILFunction *RefF = FRI->getReferencedFunction();
+          SILFunction *RefF = FRI->getInitiallyReferencedFunction();
           if (modFuncs.count(RefF) != 0) {
             // Go over the uses and add them to lists to modify
             //
@@ -2895,7 +2895,7 @@ void LoadableByAddress::run() {
   // They just contain a pointer to the function
   // The pointer does not change
   for (auto *instr : funcRefs) {
-    SILFunction *F = instr->getReferencedFunction();
+    SILFunction *F = instr->getInitiallyReferencedFunction();
     SILBuilderWithScope refBuilder(instr);
     SingleValueInstruction *newInstr =
         refBuilder.createFunctionRef(instr->getLoc(), F, instr->getKind());

lib/SIL/InstructionUtils.cpp

@@ -358,7 +358,7 @@ SILValue swift::isPartialApplyOfReabstractionThunk(PartialApplyInst *PAI) {
       PAI->getNumArguments() != 2)
     return SILValue();
 
-  auto *Fun = PAI->getReferencedFunction();
+  auto *Fun = PAI->getReferencedFunctionOrNull();
   if (!Fun)
     return SILValue();
 

lib/SIL/Linker.cpp

@@ -170,17 +170,17 @@ void SILLinkerVisitor::visitPartialApplyInst(PartialApplyInst *PAI) {
 }
 
 void SILLinkerVisitor::visitFunctionRefInst(FunctionRefInst *FRI) {
-  maybeAddFunctionToWorklist(FRI->getReferencedFunction());
+  maybeAddFunctionToWorklist(FRI->getInitiallyReferencedFunction());
 }
 
 void SILLinkerVisitor::visitDynamicFunctionRefInst(
     DynamicFunctionRefInst *FRI) {
-  maybeAddFunctionToWorklist(FRI->getReferencedFunction());
+  maybeAddFunctionToWorklist(FRI->getInitiallyReferencedFunction());
 }
 
 void SILLinkerVisitor::visitPreviousDynamicFunctionRefInst(
     PreviousDynamicFunctionRefInst *FRI) {
-  maybeAddFunctionToWorklist(FRI->getReferencedFunction());
+  maybeAddFunctionToWorklist(FRI->getInitiallyReferencedFunction());
 }
 
 // Eagerly visiting all used conformances leads to a large blowup