Merge pull request swiftlang#72324 from hborla/extract-function-isolation-expr

hborla · web-flow · commit 0a627bc44cd9 · 2024-03-14T06:56:43.000-07:00
[Concurrency] Add a `.isolation` member on dynamically isolated function values.
diff --git a/include/swift/AST/Expr.h b/include/swift/AST/Expr.h
@@ -3518,6 +3518,43 @@ class ActorIsolationErasureExpr : public ImplicitConversionExpr {
   }
 };
 
+/// Extracts the isolation of a dynamically isolated function value.
+class ExtractFunctionIsolationExpr : public Expr {
+  /// The function value expression from which to extract the
+  /// isolation. The type of `fnExpr` must be an ``@isolated(any)`
+  /// funciton.
+  Expr *fnExpr;
+
+  /// The source location of `.isolation`
+  SourceLoc isolationLoc;
+
+public:
+  ExtractFunctionIsolationExpr(Expr *fnExpr, SourceLoc isolationLoc,
+                               Type type, bool implicit = false)
+      : Expr(ExprKind::ExtractFunctionIsolation, implicit, type),
+        fnExpr(fnExpr), isolationLoc(isolationLoc) {}
+
+  Expr *getFunctionExpr() const {
+    return fnExpr;
+  }
+
+  void setFunctionExpr(Expr *fnExpr) {
+    this->fnExpr = fnExpr;
+  }
+
+  SourceLoc getStartLoc() const {
+    return fnExpr->getStartLoc();
+  }
+
+  SourceLoc getEndLoc() const {
+    return isolationLoc;
+  }
+
+  static bool classof(const Expr *E) {
+    return E->getKind() == ExprKind::ExtractFunctionIsolation;
+  }
+};
+
 /// UnresolvedSpecializeExpr - Represents an explicit specialization using
 /// a type parameter list (e.g. "Vector<Int>") that has not been resolved.
 class UnresolvedSpecializeExpr final : public Expr,
diff --git a/include/swift/AST/ExprNodes.def b/include/swift/AST/ExprNodes.def
@@ -134,6 +134,7 @@ EXPR(VarargExpansion, Expr)
 EXPR(PackExpansion, Expr)
 EXPR(PackElement, Expr)
 EXPR(MaterializePack, Expr)
+EXPR(ExtractFunctionIsolation, Expr)
 EXPR(DynamicType, Expr)
 EXPR(RebindSelfInConstructor, Expr)
 EXPR(OpaqueValue, Expr)
diff --git a/include/swift/Sema/OverloadChoice.h b/include/swift/Sema/OverloadChoice.h
@@ -56,6 +56,9 @@ enum class OverloadChoiceKind : int {
   DeclViaUnwrappedOptional,
   /// The overload choice materializes a pack from a tuple.
   MaterializePack,
+  /// The overload choice extracts the isolation of a dynamically isolated
+  /// function value.
+  ExtractFunctionIsolation,
   /// The overload choice indexes into a tuple. Index zero will
   /// have the value of this enumerator, index one will have the value of this
   /// enumerator + 1, and so on. Thus, this enumerator must always be last.
diff --git a/lib/AST/ASTDumper.cpp b/lib/AST/ASTDumper.cpp
@@ -2699,6 +2699,13 @@ class PrintExpr : public ExprVisitor<PrintExpr, void, StringRef>,
     printFoot();
   }
 
+  void visitExtractFunctionIsolationExpr(ExtractFunctionIsolationExpr *E,
+                                         StringRef label) {
+    printCommon(E, "extract_function_isolation", label);
+    printRec(E->getFunctionExpr());
+    printFoot();
+  }
+
   void visitInOutExpr(InOutExpr *E, StringRef label) {
     printCommon(E, "inout_expr", label);
     printRec(E->getSubExpr());
diff --git a/lib/AST/ASTPrinter.cpp b/lib/AST/ASTPrinter.cpp
@@ -5290,6 +5290,11 @@ void PrintAST::visitLinearToDifferentiableFunctionExpr(swift::LinearToDifferenti
 void PrintAST::visitActorIsolationErasureExpr(ActorIsolationErasureExpr *expr) {
 }
 
+void PrintAST::visitExtractFunctionIsolationExpr(ExtractFunctionIsolationExpr *expr) {
+  visit(expr->getFunctionExpr());
+  Printer << ".isolation";
+}
+
 void PrintAST::visitPropertyWrapperValuePlaceholderExpr(swift::PropertyWrapperValuePlaceholderExpr *expr) {
 }
 
diff --git a/lib/AST/ASTWalker.cpp b/lib/AST/ASTWalker.cpp
@@ -1326,6 +1326,17 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     return E;
   }
 
+  Expr *visitExtractFunctionIsolationExpr(ExtractFunctionIsolationExpr *E) {
+    if (auto fn = E->getFunctionExpr()) {
+      if (auto newFn = doIt(fn))
+        E->setFunctionExpr(newFn);
+      else
+        return nullptr;
+    }
+
+    return E;
+  }
+
   Expr *visitKeyPathDotExpr(KeyPathDotExpr *E) { return E; }
 
   Expr *visitSingleValueStmtExpr(SingleValueStmtExpr *E) {
diff --git a/lib/AST/Expr.cpp b/lib/AST/Expr.cpp
@@ -378,6 +378,7 @@ ConcreteDeclRef Expr::getReferencedDecl(bool stopAtParenExpr) const {
   PASS_THROUGH_REFERENCE(Try, getSubExpr);
   PASS_THROUGH_REFERENCE(ForceTry, getSubExpr);
   PASS_THROUGH_REFERENCE(OptionalTry, getSubExpr);
+  PASS_THROUGH_REFERENCE(ExtractFunctionIsolation, getFunctionExpr);
 
   NO_REFERENCE(Tuple);
   SIMPLE_REFERENCE(Array, getInitializer);
@@ -723,6 +724,7 @@ bool Expr::canAppendPostfixExpression(bool appendingPostfixOperator) const {
   case ExprKind::UnresolvedSpecialize:
   case ExprKind::UnresolvedMember:
   case ExprKind::UnresolvedDot:
+  case ExprKind::ExtractFunctionIsolation:
     return true;
 
   case ExprKind::Sequence:
@@ -1032,6 +1034,7 @@ bool Expr::isValidParentOfTypeExpr(Expr *typeExpr) const {
   case ExprKind::MacroExpansion:
   case ExprKind::CurrentContextIsolation:
   case ExprKind::ActorIsolationErasure:
+  case ExprKind::ExtractFunctionIsolation:
     return false;
   }
 
diff --git a/lib/IDE/SelectedOverloadInfo.cpp b/lib/IDE/SelectedOverloadInfo.cpp
@@ -89,6 +89,7 @@ swift::ide::getSelectedOverloadInfo(const Solution &S,
   case OverloadChoiceKind::DynamicMemberLookup:
   case OverloadChoiceKind::TupleIndex:
   case OverloadChoiceKind::MaterializePack:
+  case OverloadChoiceKind::ExtractFunctionIsolation:
     // If it's DynamicMemberLookup, we don't know which function is being
     // called, so we can't extract any information from it.
     // TupleIndex isn't a function call and is not relevant for argument
diff --git a/lib/SILGen/SILGenExpr.cpp b/lib/SILGen/SILGenExpr.cpp
@@ -486,6 +486,8 @@ namespace {
                                        SGFContext C);
     RValue visitActorIsolationErasureExpr(ActorIsolationErasureExpr *E,
                                           SGFContext C);
+    RValue visitExtractFunctionIsolationExpr(ExtractFunctionIsolationExpr *E,
+                                             SGFContext C);
     RValue visitCovariantFunctionConversionExpr(
              CovariantFunctionConversionExpr *E,
              SGFContext C);
@@ -2107,6 +2109,14 @@ RValue RValueEmitter::visitActorIsolationErasureExpr(ActorIsolationErasureExpr *
                                                    nonIsolatedType));
 }
 
+RValue RValueEmitter::visitExtractFunctionIsolationExpr(
+    ExtractFunctionIsolationExpr *E, SGFContext C) {
+  auto arg = SGF.emitRValue(E->getFunctionExpr());
+  auto result = SGF.emitExtractFunctionIsolation(
+      E, ArgumentSource(E, std::move(arg)), C);
+  return RValue(SGF, E, result);
+}
+
 RValue RValueEmitter::visitErasureExpr(ErasureExpr *E, SGFContext C) {
   if (auto result = tryEmitAsBridgingConversion(SGF, E, false, C)) {
     return RValue(SGF, E, *result);
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -3423,6 +3423,15 @@ namespace {
                                         packType));
       }
 
+      case OverloadChoiceKind::ExtractFunctionIsolation: {
+        auto isolationType = solution.getResolvedType(expr);
+        auto *extractExpr = new (cs.getASTContext())
+          ExtractFunctionIsolationExpr(base,
+                                       expr->getEndLoc(),
+                                       isolationType);
+        return cs.cacheType(extractExpr);
+      }
+
       case OverloadChoiceKind::KeyPathApplication:
         llvm_unreachable("should only happen in a subscript");
 
@@ -5372,6 +5381,10 @@ namespace {
       return E;
     }
 
+    Expr *visitExtractFunctionIsolationExpr(ExtractFunctionIsolationExpr *E) {
+      llvm_unreachable("found ExtractFunctionIsolationExpr in CSApply");
+    }
+
     Expr *visitKeyPathDotExpr(KeyPathDotExpr *E) {
       llvm_unreachable("found KeyPathDotExpr in CSApply");
     }
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -3935,6 +3935,10 @@ namespace {
       return OptionalType::get(actorProto->getDeclaredExistentialType());
     }
 
+    Type visitExtractFunctionIsolationExpr(ExtractFunctionIsolationExpr *E) {
+      llvm_unreachable("found ExtractFunctionIsolationExpr in CSGen");
+    }
+
     Type visitKeyPathDotExpr(KeyPathDotExpr *E) {
       llvm_unreachable("found KeyPathDotExpr in CSGen");
     }
diff --git a/lib/Sema/CSRanking.cpp b/lib/Sema/CSRanking.cpp
@@ -205,6 +205,7 @@ static bool sameOverloadChoice(const OverloadChoice &x,
     return x.getTupleIndex() == y.getTupleIndex();
 
   case OverloadChoiceKind::MaterializePack:
+  case OverloadChoiceKind::ExtractFunctionIsolation:
     return true;
   }
 
@@ -1121,6 +1122,7 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
     switch (choice1.getKind()) {
     case OverloadChoiceKind::TupleIndex:
     case OverloadChoiceKind::MaterializePack:
+    case OverloadChoiceKind::ExtractFunctionIsolation:
       continue;
 
     case OverloadChoiceKind::KeyPathApplication:
diff --git a/lib/Sema/CSSimplify.cpp b/lib/Sema/CSSimplify.cpp
@@ -205,6 +205,7 @@ static bool areConservativelyCompatibleArgumentLabels(
   case OverloadChoiceKind::KeyPathDynamicMemberLookup:
   case OverloadChoiceKind::TupleIndex:
   case OverloadChoiceKind::MaterializePack:
+  case OverloadChoiceKind::ExtractFunctionIsolation:
     return true;
   }
 
@@ -9763,6 +9764,16 @@ performMemberLookup(ConstraintKind constraintKind, DeclNameRef memberName,
   if (auto *selfTy = instanceTy->getAs<DynamicSelfType>())
     instanceTy = selfTy->getSelfType();
 
+  // Dynamically isolated function types have a magic '.isolation'
+  // member that extracts the isolation value.
+  if (auto *fn = dyn_cast<FunctionType>(instanceTy)) {
+    if (fn->getIsolation().isErased() &&
+        memberName.getBaseIdentifier().str() == "isolation") {
+      result.ViableCandidates.push_back(
+        OverloadChoice(baseTy, OverloadChoiceKind::ExtractFunctionIsolation));
+    }
+  }
+
   if (!instanceTy->mayHaveMembers())
     return result;
 
diff --git a/lib/Sema/Constraint.cpp b/lib/Sema/Constraint.cpp
@@ -534,6 +534,9 @@ void Constraint::print(llvm::raw_ostream &Out, SourceManager *sm,
     case OverloadChoiceKind::MaterializePack:
       Out << "materialize pack";
       break;
+    case OverloadChoiceKind::ExtractFunctionIsolation:
+      Out << "extract function islation";
+      break;
     case OverloadChoiceKind::KeyPathApplication:
       Out << "key path application";
       break;
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
@@ -2890,6 +2890,7 @@ Type ConstraintSystem::getEffectiveOverloadType(ConstraintLocator *locator,
   case OverloadChoiceKind::KeyPathApplication:
   case OverloadChoiceKind::TupleIndex:
   case OverloadChoiceKind::MaterializePack:
+  case OverloadChoiceKind::ExtractFunctionIsolation:
     return Type();
   }
 
@@ -3569,6 +3570,7 @@ void ConstraintSystem::bindOverloadType(
   case OverloadChoiceKind::DeclViaUnwrappedOptional:
   case OverloadChoiceKind::TupleIndex:
   case OverloadChoiceKind::MaterializePack:
+  case OverloadChoiceKind::ExtractFunctionIsolation:
   case OverloadChoiceKind::KeyPathApplication:
     bindTypeOrIUO(openedType);
     return;
@@ -3829,6 +3831,15 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
     break;
   }
 
+  case OverloadChoiceKind::ExtractFunctionIsolation: {
+    // The type of `.isolation` is `(any Actor)?`
+    auto actor = getASTContext().getProtocol(KnownProtocolKind::Actor);
+    adjustedRefType =
+        OptionalType::get(actor->getDeclaredExistentialType());
+    refType = adjustedRefType;
+    break;
+  }
+
   case OverloadChoiceKind::KeyPathApplication: {
     // Key path application looks like a subscript(keyPath: KeyPath<Base, T>).
     // The element type is T or @lvalue T based on the key path subtype and
@@ -4466,6 +4477,7 @@ DeclName OverloadChoice::getName() const {
 
     case OverloadChoiceKind::MaterializePack:
     case OverloadChoiceKind::TupleIndex:
+    case OverloadChoiceKind::ExtractFunctionIsolation:
       llvm_unreachable("no name!");
   }
 
@@ -5830,6 +5842,7 @@ bool ConstraintSystem::diagnoseAmbiguity(ArrayRef<Solution> solutions) {
 
       case OverloadChoiceKind::TupleIndex:
       case OverloadChoiceKind::MaterializePack:
+      case OverloadChoiceKind::ExtractFunctionIsolation:
         // FIXME: Actually diagnose something here.
         break;
       }
diff --git a/lib/Sema/TypeCheckConstraints.cpp b/lib/Sema/TypeCheckConstraints.cpp
@@ -1304,6 +1304,10 @@ void OverloadChoice::dump(Type adjustedOpenedType, SourceManager *sm,
   case OverloadChoiceKind::MaterializePack:
     out << "materialize pack from tuple " << getBaseType()->getString(PO);
     break;
+
+  case OverloadChoiceKind::ExtractFunctionIsolation:
+    out << "extract isolation from " << getBaseType()->getString(PO);
+    break;
   }
 }
 
@@ -1610,6 +1614,11 @@ void ConstraintSystem::print(raw_ostream &out) const {
         out << "materialize pack from tuple "
             << choice.getBaseType()->getString(PO);
         break;
+
+      case OverloadChoiceKind::ExtractFunctionIsolation:
+        out << "extract isolation from "
+            << choice.getBaseType()->getString(PO);
+        break;
       }
       out << " for ";
       elt.first->dump(&getASTContext().SourceMgr, out);
diff --git a/test/Concurrency/isolated_any.swift b/test/Concurrency/isolated_any.swift
@@ -83,3 +83,16 @@ func extractFunctionIsolation(_ fn: @isolated(any) @escaping () async -> Void) {
   let _: (any Actor)? = extractIsolation(myActor.asyncThrowsActorFunction)
   let _: (any Actor)? = extractIsolation(myActor.actorFunctionWithArgs(value:))
 }
+
+func extractFunctionIsolationExpr(
+  _ fn1: @isolated(any) @escaping () async -> Void,
+  _ fn2: @isolated(any) @escaping (Int, String) -> Bool
+) {
+  let _: (any Actor)? = fn1.isolation
+  let _: (any Actor)? = fn2.isolation
+
+  // Only `@isolated(any)` functions have `.isolation`
+  let myActor = A()
+  let _: (any Actor)? = myActor.asyncActorFunction.isolation // expected-error {{value of type '@Sendable () async -> ()' has no member 'isolation'}}
+  let _: (any Actor)? = globalNonisolatedFunction.isolation // expected-error {{value of type '@Sendable () -> ()' has no member 'isolation'}}
+}
diff --git a/test/SILGen/isolated_any.swift b/test/SILGen/isolated_any.swift
@@ -452,3 +452,17 @@ extension MyActor {
 func testEraseAsyncActorIsolatedPartialApplication(a: MyActor) {
   takeAsyncIsolatedAny(fn: a.asyncAction)
 }
+
+/*-- Isolation extraction --*/
+
+// CHECK-LABEL: sil hidden [ossa] @$s4test16extractIsolation2fnScA_pSgyyYAc_tF
+// CHECK: [[FN:%.*]] = copy_value %0 : $@isolated(any) @callee_guaranteed () -> ()
+// CHECK-NEXT: [[FN_BORROW:%.*]] = begin_borrow [[FN]] : $@isolated(any) @callee_guaranteed () -> ()
+// CHECK-NEXT: [[ISOLATION:%.*]] = function_extract_isolation [[FN_BORROW]] : $@isolated(any) @callee_guaranteed () -> ()
+// CHECK-NEXT: [[RESULT:%.*]] = copy_value [[ISOLATION]] : $Optional<any Actor>
+// CHECK-NEXT: end_borrow [[FN_BORROW]] : $@isolated(any) @callee_guaranteed () -> ()
+// CHECK-NEXT: destroy_value [[FN]] : $@isolated(any) @callee_guaranteed () -> ()
+// CHECK-NEXT: return [[RESULT]] : $Optional<any Actor>
+func extractIsolation(fn: @escaping @isolated(any) () -> Void) -> (any Actor)? {
+  fn.isolation
+}

Original file line number	Diff line number	Diff line change
`@@ -5290,6 +5290,11 @@ void PrintAST::visitLinearToDifferentiableFunctionExpr(swift::LinearToDifferenti`
`5290`	`5290`	`void PrintAST::visitActorIsolationErasureExpr(ActorIsolationErasureExpr *expr) {`
`5291`	`5291`	`}`
`5292`	`5292`
	`5293`	`+void PrintAST::visitExtractFunctionIsolationExpr(ExtractFunctionIsolationExpr *expr) {`
	`5294`	`+ visit(expr->getFunctionExpr());`
	`5295`	`+ Printer << ".isolation";`
	`5296`	`+}`
	`5297`	`+`
`5293`	`5298`	`void PrintAST::visitPropertyWrapperValuePlaceholderExpr(swift::PropertyWrapperValuePlaceholderExpr *expr) {`
`5294`	`5299`	`}`
`5295`	`5300`