[concurrency] Implement bit masking for TBI when available or in tagged pointer bits otherwise.

Specifically, when TBI is available we use the bottom two bits of the top nibble
(bits 60,61). On platforms without TBI, we use the bottom two tagged pointer
bits (bits 0, 1).

rdar://156525771
(cherry picked from commit 390afe3)

Author: gottesmm

include/swift/SIL/ConcurrencyUtils.h

@@ -312,7 +312,9 @@ class LoweredValue {
   bool isBoxWithAddress() const {
     return kind == Kind::OwnedAddress;
   }
-  
+
+  bool isExplosionVector() const { return kind == Kind::ExplosionVector; }
+
   const StackAddress &getStackAddress() const {
     return Storage.get<StackAddress>(kind);
   }
@@ -3387,12 +3389,10 @@ void IRGenSILFunction::visitExistentialMetatypeInst(
   setLoweredExplosion(i, result);
 }
 
-static void emitApplyArgument(IRGenSILFunction &IGF,
-                              SILValue arg,
-                              SILType paramType,
-                              Explosion &out,
-                              SILInstruction *apply = nullptr,
-                              unsigned idx = 0) {
+static void emitApplyArgument(IRGenSILFunction &IGF, SILValue arg,
+                              SILType paramType, Explosion &out,
+                              SILInstruction *apply = nullptr, unsigned idx = 0,
+                              bool isImplicitIsolatedParameter = false) {
   bool isSubstituted = (arg->getType() != paramType);
 
   // For indirect arguments, we just need to pass a pointer.
@@ -3434,7 +3434,20 @@ static void emitApplyArgument(IRGenSILFunction &IGF,
       }
       canForwardLoadToIndirect = true;
     }();
-    IGF.getLoweredExplosion(arg, out);
+
+    // If we are emitting a parameter for an implicit isolated parameter, then
+    // we need to clear the implicit isolated actor bits.
+    if (isImplicitIsolatedParameter) {
+      auto &loweredValue = IGF.getLoweredValue(arg);
+      assert(loweredValue.isExplosionVector() &&
+             "Should be an explosion of two pointers");
+      auto explosionVector = loweredValue.getKnownExplosionVector();
+      assert(explosionVector.size() == 2 && "We should have two values");
+      out.add(explosionVector[0]);
+      out.add(clearImplicitIsolatedActorBits(IGF, explosionVector[1]));
+    } else {
+      IGF.getLoweredExplosion(arg, out);
+    }
     if (canForwardLoadToIndirect) {
       IGF.setForwardableArgument(idx);
     }
@@ -3852,6 +3865,20 @@ void IRGenSILFunction::visitFullApplySite(FullApplySite site) {
     }
   }
 
+  // Extract the implicit isolated parameter so that we can mask it as
+  // appropriate.
+  //
+  // NOTE: We cannot just drop_front since we could be between the indirect
+  // results and the parameters.
+  std::optional<unsigned> implicitIsolatedParameterIndex;
+  if (auto actorIsolation = site.getFunction()->getActorIsolation();
+      actorIsolation && actorIsolation->isCallerIsolationInheriting() &&
+      site.isCallerIsolationInheriting()) {
+    auto *iso = site.getIsolatedArgumentOperandOrNullPtr();
+    assert(iso);
+    implicitIsolatedParameterIndex = site.getAppliedArgIndex(*iso);
+  }
+
   // Lower the arguments and return value in the callee's generic context.
   GenericContextScope scope(IGM,
                             origCalleeType->getInvocationGenericSignature());
@@ -3912,8 +3939,11 @@ void IRGenSILFunction::visitFullApplySite(FullApplySite site) {
       emission->setIndirectTypedErrorResultSlot(addr.getAddress());
       continue;
     }
+
     emitApplyArgument(*this, args[index], emission->getParameterType(index),
-                      llArgs, site.getInstruction(), index);
+                      llArgs, site.getInstruction(), index,
+                      implicitIsolatedParameterIndex &&
+                          *implicitIsolatedParameterIndex == index);
   }
 
   // Pass the generic arguments.

lib/IRGen/IRGenSIL.cpp

@@ -2,7 +2,6 @@ target_sources(swiftSIL PRIVATE
   BasicBlockUtils.cpp
   BitDataflow.cpp
   CalleeCache.cpp
-  ConcurrencyUtils.cpp
   DebugUtils.cpp
   Dominance.cpp
   DynamicCasts.cpp

lib/SIL/Utils/CMakeLists.txt

@@ -13,7 +13,6 @@
 #include "ArgumentScope.h"
 #include "ArgumentSource.h"
 #include "Callee.h"
-#include "ConcurrencyUtils.h"
 #include "Conversion.h"
 #include "ExecutorBreadcrumb.h"
 #include "FormalEvaluation.h"

lib/SIL/Utils/ConcurrencyUtils.cpp

@@ -562,6 +562,13 @@ class SILGenBuilder : public SILBuilder {
         convertToImplicitActor(loc, value.borrow(SGF, loc).getValue());
     return ManagedValue::forBorrowedRValue(result);
   }
+
+  using SILBuilder::createImplicitActorToOpaqueIsolationCast;
+  ManagedValue createImplicitActorToOpaqueIsolationCast(SILLocation loc,
+                                                        ManagedValue mv) {
+    return ManagedValue::forBorrowedRValue(
+        createImplicitActorToOpaqueIsolationCast(loc, mv.getUnmanagedValue()));
+  }
 };
 
 } // namespace Lowering

lib/SILGen/ConcurrencyUtils.h

@@ -13,7 +13,6 @@
 #include "ArgumentScope.h"
 #include "ArgumentSource.h"
 #include "Callee.h"
-#include "ConcurrencyUtils.h"
 #include "Condition.h"
 #include "Conversion.h"
 #include "Initialization.h"
@@ -7321,9 +7320,8 @@ RValue RValueEmitter::visitCurrentContextIsolationExpr(
     assert(isolatedArg &&
            "Caller Isolation Inheriting without isolated parameter");
     auto isolatedMV = ManagedValue::forBorrowedRValue(isolatedArg);
-    return clearImplicitActorBits(
-        SGF, E, isolatedMV,
-        SILType::getOpaqueIsolationType(SGF.getASTContext()));
+    return RValue(
+        SGF, E, SGF.B.createImplicitActorToOpaqueIsolationCast(E, isolatedMV));
   }
 
   if (isolation == ActorIsolation::ActorInstance) {

lib/SILGen/SILGenApply.cpp

@@ -15,7 +15,6 @@
 #include "swift/AST/ConformanceLookup.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/FrozenMultiMap.h"
-#include "swift/SIL/ConcurrencyUtils.h"
 #include "swift/SIL/Dominance.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILFunction.h"
@@ -353,8 +352,8 @@ static SILValue getExecutorForImplicitActor(SILOptFunctionBuilder &funcBuilder,
     auto *front = newFunc->createBasicBlock();
     SILBuilder builder(front);
     auto *fArg = front->createFunctionArgument(implicitIsolatedActorType);
-    auto value = clearImplicitActorBits(builder, autoGenLoc, fArg,
-                                        SILType::getOpaqueIsolationType(ctx));
+    auto value = SILValue(
+        builder.createImplicitActorToOpaqueIsolationCast(autoGenLoc, fArg));
     value = getExecutorForOptionalActor(builder, autoGenLoc, value);
     builder.createReturn(autoGenLoc, value);
   }

lib/SILGen/SILGenBuilder.h

@@ -0,0 +1,65 @@
+// RUN: %target-swift-frontend -parse-as-library -emit-ir -disable-llvm-merge-functions-pass %s | %FileCheck --check-prefix=NO-TBI %s
+// RUN: %target-swift-frontend -parse-as-library -Xllvm -aarch64-use-tbi -emit-ir -disable-llvm-merge-functions-pass %s | %FileCheck --check-prefix=TBI %s
+
+// This test makes sure that we can properly fold the mask for the witness table
+// when we have a #isolation.
+
+// REQUIRES: concurrency
+// REQUIRES: CODEGENERATOR=AArch64
+// REQUIRES: PTRSIZE=64
+// REQUIRES: OS=macosx || OS=ios
+// REQUIRES: CPU=arm64
+
+@inline(never)
+func useActor(iso: (any Actor)?) {
+  print(iso!.unownedExecutor)
+}
+
+@inline(never)
+func implicitParam(_ x: (any Actor)? = #isolation) {
+  print(x!.unownedExecutor)
+}
+
+// #isolation via direct usage
+//
+// TBI: define swifttailcc void @"$s18isolation_macro_ir46nonisolatedNonsendingUsePoundIsolationDirectlyyyYaF"(ptr swiftasync %0, i64 %1, i64 [[SECOND_WORD:%.*]])
+// TBI: [[MASKED_SECOND_WORD:%.*]] = and i64 [[SECOND_WORD]], -3458764513820540929
+// TBI: call swiftcc void @"$s18isolation_macro_ir8useActor3isoyScA_pSg_tF"(i64 %1, i64 [[MASKED_SECOND_WORD]])
+
+// NO-TBI: define swifttailcc void @"$s18isolation_macro_ir46nonisolatedNonsendingUsePoundIsolationDirectlyyyYaF"(ptr swiftasync %0, i64 [[FIRST_WORD:%.*]], i64 [[SECOND_WORD:%.*]])
+// NO-TBI: [[MASKED_SECOND_WORD:%.*]] = and i64 [[SECOND_WORD]], -4
+// NO-TBI: call swiftcc void @"$s18isolation_macro_ir8useActor3isoyScA_pSg_tF"(i64 [[FIRST_WORD]], i64 [[MASKED_SECOND_WORD]])
+
+public nonisolated(nonsending) func nonisolatedNonsendingUsePoundIsolationDirectly() async {
+  let iso = #isolation
+  useActor(iso: iso)
+}
+
+// #isolation via default arg
+//
+// TBI: define swifttailcc void @"$s18isolation_macro_ir45nonisolatedNonsendingPoundIsolationDefaultArgyyYaF"(ptr swiftasync {{%.*}}, i64 {{%.*}}, i64 [[WORD_2:%.*]])
+// TBI: [[MASKED_WORD_2:%.*]] = and i64 [[WORD_2]], -3458764513820540929
+// TBI: call swiftcc void @"$s18isolation_macro_ir13implicitParamyyScA_pSgF"(i64 {{%.*}}, i64 [[MASKED_WORD_2]])
+
+// NO-TBI: define swifttailcc void @"$s18isolation_macro_ir45nonisolatedNonsendingPoundIsolationDefaultArgyyYaF"(ptr swiftasync %0, i64 [[WORD_1:%.*]], i64 [[WORD_2:%.*]])
+// NO-TBI: [[MASKED_WORD_2:%.*]] = and i64 [[WORD_2]], -4
+// NO-TBI: call swiftcc void @"$s18isolation_macro_ir13implicitParamyyScA_pSgF"(i64 {{%.*}}, i64 [[MASKED_WORD_2]])
+public nonisolated(nonsending) func nonisolatedNonsendingPoundIsolationDefaultArg() async {
+  implicitParam()
+}
+
+@inline(never)
+public nonisolated(nonsending) func calleeFunction() async {
+}
+
+// TBI: define swifttailcc void @"$s18isolation_macro_ir14callerFunctionyyYaF"(ptr swiftasync %0, i64 %1, i64 [[WORD_2:%.*]])
+// TBI: [[MASKED_WORD_2:%.*]] = and i64 [[WORD_2]], -3458764513820540929
+// TBI: musttail call swifttailcc void @"$s18isolation_macro_ir14calleeFunctionyyYaF"(ptr swiftasync {{%.*}}, i64 {{%.*}}, i64 [[MASKED_WORD_2]])
+
+// NO-TBI: define swifttailcc void @"$s18isolation_macro_ir14callerFunctionyyYaF"(ptr swiftasync %0, i64 %1, i64 [[WORD:%.*]])
+// NO-TBI: [[MASKED_WORD:%.*]] = and i64 [[WORD]], -4
+// NO-TBI: musttail call swifttailcc void @"$s18isolation_macro_ir14calleeFunctionyyYaF"(ptr swiftasync {{%.*}}, i64 {{%.*}}, i64 [[MASKED_WORD]])
+@inline(never)
+public nonisolated(nonsending) func callerFunction() async {
+  await calleeFunction()
+}