[WIP][IDE]: surface inferred isolation information via sourcekit

include/swift/Sema/IDETypeChecking.h

@@ -211,6 +211,37 @@ namespace swift {
       std::vector<ExpressionTypeInfo> &scratch, bool FullyQualified,
       bool CanonicalType, llvm::raw_ostream &OS);
 
+  /// An inferred actor isolation suitable for surfacing in an IDE inlay hint.
+  /// Reported as an offset+length anchor (currently the closure's `{`), with
+  /// the pretty-printed isolation and a kind string written into a shared
+  /// string buffer.
+  /// Inferred actor isolation info for
+  struct InferredIsolationInfo {
+    /// Offset of the entity to which the isolation corresponds.
+    uint32_t Offset;
+
+    /// Length of the entity to which the isolation corresponds.
+    uint32_t Length;
+
+    /// Offsets into the shared string buffer for the source-style isolation
+    /// (e.g. "@MainActor", "nonisolated", "@SomeGlobalActor").
+    uint32_t IsolationOffset;
+    uint32_t IsolationLength;
+
+    /// Offsets into the shared string buffer for the kind of entity this
+    /// isolation is attached to. Currently always "closure".
+    uint32_t KindOffset;
+    uint32_t KindLength;
+  };
+
+  /// Collect *inferred* actor isolation for every explicit \c ClosureExpr in
+  /// \c SF. Closures with their isolation written explicitly in the signature
+  /// are skipped. Pretty-printed isolation strings are written to \c OS.
+  void
+  collectInferredIsolations(SourceFile &SF, SourceRange Range,
+                            std::vector<InferredIsolationInfo> &IsolationInfos,
+                            llvm::raw_ostream &OS);
+
   /// Resolve a list of mangled names to accessible protocol decls from
   /// the decl context.
   ProtocolDecl *resolveProtocolName(DeclContext *dc, StringRef Name);

lib/IDE/IDETypeChecking.cpp

@@ -13,6 +13,8 @@
 #include "swift/Sema/IDETypeChecking.h"
 #include "ReadyForTypeCheckingCallback.h"
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/ASTWalker.h"
+#include "swift/AST/ActorIsolation.h"
 #include "swift/AST/ASTDemangler.h"
 #include "swift/AST/ASTPrinter.h"
 #include "swift/AST/Attr.h"
@@ -1032,6 +1034,209 @@ swift::getShorthandShadows(LabeledConditionalStmt *CondStmt, DeclContext *DC) {
   return Result;
 }
 
+//===--------------------------------------------------------------------------===//
+// Isolation collection
+//===--------------------------------------------------------------------------===//
+
+/// Walks a SourceFile and records actor isolation info. Currently only supports
+/// explicit ClosureExprs. Closures whose isolation is written in the signature
+/// are skipped.
+class InferredIsolationCollector : public SourceEntityWalker {
+  SourceManager &SM;
+  unsigned BufferId;
+
+  /// The range in which isolation info is to be collected. An invalid range
+  /// means the whole file.
+  SourceRange TotalRange;
+
+  /// The output vector for InferredIsolationInfo emitted during traversal.
+  std::vector<InferredIsolationInfo> &Results;
+
+  /// Buffer in which all printed isolation strings are stored, null-terminated
+  /// and deduplicated.
+  llvm::raw_ostream &OS;
+
+  /// Map from a printed string to its offset in \c OS. Used for deduplicating
+  /// both isolation strings and kind labels.
+  llvm::StringMap<uint32_t> StringOffsets;
+
+  /// Returns the (offset, length) of \p S within \c OS, printing it first if
+  /// it isn't already present.
+  std::pair<uint32_t, uint32_t> getBufferRangeForString(StringRef S) {
+    auto It = StringOffsets.find(S);
+    if (It == StringOffsets.end()) {
+      StringOffsets[S] = OS.tell();
+      OS << S << '\0';
+    }
+    return {StringOffsets[S], S.size()};
+  }
+
+  /// Whether the given range overlaps the total range in which we collect
+  /// isolation info.
+  bool overlapsTotalRange(SourceRange Range) {
+    return TotalRange.isInvalid() || Range.overlaps(TotalRange);
+  }
+
+  /// Render an isolation in the spelling you'd write in source, suitable for
+  /// an inline IDE badge. Examples: "@MainActor", "@SomeGlobalActor",
+  /// "nonisolated", "nonisolated(unsafe)", "nonisolated(nonsending)",
+  /// "actor-isolated", "@isolated(any)".
+  // TODO: what should be responsible for this formatting?
+  static void printSourceStyle(const ActorIsolation &isolation,
+                               llvm::raw_ostream &Out) {
+    switch (isolation.getKind()) {
+    case ActorIsolation::Unspecified:
+      // TODO: how to handle 'unspecified'?
+      // We currently skip it during the walk so should not hit it here, but
+      // what should we return? Empty string? Pass it through and let clients
+      // decide?
+      Out << "";
+      return;
+    case ActorIsolation::ActorInstance:
+      Out << "actor-isolated";
+      return;
+    case ActorIsolation::Nonisolated:
+      Out << "nonisolated";
+      return;
+    case ActorIsolation::NonisolatedConcurrent:
+      Out << "@concurrent";
+      return;
+    case ActorIsolation::NonisolatedNonsending:
+      Out << "nonisolated(nonsending)";
+      return;
+    case ActorIsolation::NonisolatedUnsafe:
+      Out << "nonisolated(unsafe)";
+      return;
+    case ActorIsolation::GlobalActor: {
+      if (isolation.isMainActor()) {
+        Out << "@MainActor";
+      } else {
+        Out << "@" << isolation.getGlobalActor().getString();
+      }
+      return;
+    case ActorIsolation::Erased:
+      Out << "@isolated(any)";
+      return;
+    }
+    }
+    llvm_unreachable("covered switch");
+  }
+
+  std::pair<uint32_t, uint32_t> printIsolation(const ActorIsolation &iso) {
+    SmallString<64> Buf;
+    {
+      llvm::raw_svector_ostream Out(Buf);
+      printSourceStyle(iso, Out);
+    }
+    return getBufferRangeForString(Buf.str());
+  }
+
+  /// Record one inferred-isolation entry. \p Anchor must be valid and in this
+  /// source file's buffer.
+  void record(SourceLoc Anchor, unsigned Length,
+              const ActorIsolation &Isolation, StringRef Kind) {
+    if (Anchor.isInvalid())
+      return;
+    if (SM.findBufferContainingLoc(Anchor) != BufferId)
+      return;
+
+    unsigned offset = SM.getLocOffsetInBuffer(Anchor, BufferId);
+    auto isolationBufferRange = printIsolation(Isolation);
+    auto kindBufferRange = getBufferRangeForString(Kind);
+
+    InferredIsolationInfo Info;
+    Info.Offset = offset;
+    Info.Length = Length;
+    Info.IsolationOffset = isolationBufferRange.first;
+    Info.IsolationLength = isolationBufferRange.second;
+    Info.KindOffset = kindBufferRange.first;
+    Info.KindLength = kindBufferRange.second;
+    // TODO: the var type collector analog constructs this inline...
+    // is one way better than another?
+    Results.push_back(Info);
+  }
+
+public:
+  InferredIsolationCollector(SourceFile &SF, SourceRange Range,
+                             std::vector<InferredIsolationInfo> &Results,
+                             llvm::raw_ostream &OS)
+      : SM(SF.getASTContext().SourceMgr), BufferId(SF.getBufferID()),
+        TotalRange(Range), Results(Results), OS(OS) {}
+
+  bool walkToExprPre(Expr *E) override {
+    // Skip the subtree if it's outside the requested range.
+    if (!overlapsTotalRange(E->getSourceRange()))
+      return false;
+
+    auto *Closure = dyn_cast<ClosureExpr>(E);
+    if (!Closure)
+      return true;
+
+    if (Closure->isImplicit())
+      return true;
+
+    auto Isolation = Closure->getActorIsolation();
+    if (Isolation.isUnspecified())
+      return true;
+
+    // If the closure signature has an explicit isolation attribute in
+    // its signature, exclude it.
+    // TODO: what is the best way to robustly detect this case?
+    // Should we pass it through regardless of whether it's "implicit"?
+    if (Closure->getInLoc().isValid()) {
+      const auto &Attrs = Closure->getAttrs();
+
+      // Found an explicit @concurrent attribute, so skip this closure.
+      if (Attrs.hasAttribute<ConcurrentAttr>())
+        return true;
+
+      // Check for an explicit global actor attribute.
+      if (auto result = evaluateOrDefault(Closure->getASTContext().evaluator,
+                                          GlobalActorAttributeRequest{Closure},
+                                          std::nullopt)) {
+        // If we found a nominal, then this closure has an explicit global actor
+        // annotation, so we can skip it.
+        if (result->second)
+          return true;
+      }
+    }
+
+    // TODO: should we report the full source range here or infer an "anchor"?
+    // Clients would need to figure out where the right placement is for
+    // closures.
+    const auto closureCharSourceRange =
+        Lexer::getCharSourceRangeFromSourceRange(SM, Closure->getSourceRange());
+
+    record(closureCharSourceRange.getStart(),
+           closureCharSourceRange.getByteLength(), Isolation, "closure");
+    return true;
+  }
+
+  bool walkToDeclPre(Decl *D, CharSourceRange DeclNameRange) override {
+    // Skip this node and subtree if outside the range.
+    return overlapsTotalRange(D->getSourceRange());
+  }
+
+  bool walkToStmtPre(Stmt *S) override {
+    // Skip this node and subtree if outside the range.
+    return overlapsTotalRange(S->getSourceRange());
+  }
+
+  bool walkToPatternPre(Pattern *P) override {
+    // Skip this node and subtree if outside the range.
+    return overlapsTotalRange(P->getSourceRange());
+  }
+};
+
+void swift::collectInferredIsolations(
+    SourceFile &SF, SourceRange Range,
+    std::vector<InferredIsolationInfo> &IsolationInfos, llvm::raw_ostream &OS) {
+  InferredIsolationCollector Walker(SF, Range, IsolationInfos, OS);
+  Walker.walk(SF);
+}
+
+//===--------------------------------------------------------------------------===//
+
 void ReadyForTypeCheckingCallback::doneParsing(SourceFile *SrcFile) {
   // Import resolution will have already been done by IDEInspectionInstance,
   // we need to bind extensions here though since IDEInspectionSecondPassRequest

test/IDE/inferred_isolation.swift

@@ -0,0 +1,56 @@
+// RUN: %target-swift-ide-test -print-inferred-isolation -source-filename %s | %FileCheck %s
+
+@globalActor
+actor MyActor {
+  static let shared = MyActor()
+}
+
+@MainActor
+final class C {
+  var n = 0
+
+  func work() async {
+    // Closure that inherits @MainActor from the enclosing context.
+    // CHECK:      let inheritedMain = <inferred-isolation kind="closure" isolation="@MainActor">{</inferred-isolation>
+    let inheritedMain = {
+      self.n += 1
+    }
+    inheritedMain()
+
+    // Detached closure: nonisolated regardless of enclosing context.
+    // CHECK:      Task.detached <inferred-isolation kind="closure" isolation="nonisolated">{</inferred-isolation>
+    Task.detached {
+      print("detached")
+    }
+
+    // Task closure inherits the actor context.
+    // CHECK:      Task <inferred-isolation kind="closure" isolation="@MainActor">{</inferred-isolation>
+    Task {
+      self.n += 1
+    }
+  }
+}
+
+func freeFunc() {
+  // Closure outside any isolated context -- nonisolated.
+  // CHECK:      let f = <inferred-isolation kind="closure" isolation="nonisolated">{</inferred-isolation> 1 }
+  let f = { 1 }
+  _ = f()
+}
+
+// Closures whose isolation is *written* in the signature must not be tagged,
+// since the information is already on screen.
+
+func explicitMain() {
+  // CHECK:      let g = { @MainActor in 2 }
+  // CHECK-NOT:  <inferred-isolation {{[^>]*}}>{</inferred-isolation> @MainActor
+  let g = { @MainActor in 2 }
+  _ = g
+}
+
+func explicitCustom() {
+  // CHECK:      let h = { @MyActor in 3 }
+  // CHECK-NOT:  <inferred-isolation {{[^>]*}}>{</inferred-isolation> @MyActor
+  let h = { @MyActor in 3 }
+  _ = h
+}

test/SourceKit/InferredIsolation/basic.swift

@@ -0,0 +1,32 @@
+@MainActor
+final class C {
+  var n = 0
+
+  func work() async {
+    let inheritedMain = {
+      self.n += 1
+    }
+    inheritedMain()
+
+    Task.detached {
+      print("detached")
+    }
+  }
+}
+
+// RUN: %sourcekitd-test -req=collect-inferred-isolation %s -- %s | %FileCheck %s
+
+// CHECK:      key.results: [
+// CHECK-NEXT:   {
+// CHECK-NEXT:     key.kind: "closure",
+// CHECK-NEXT:     key.offset: {{[0-9]+}},
+// CHECK-NEXT:     key.length: {{[0-9]+}},
+// CHECK-NEXT:     key.actor_isolation: "@MainActor"
+// CHECK-NEXT:   },
+// CHECK-NEXT:   {
+// CHECK-NEXT:     key.kind: "closure",
+// CHECK-NEXT:     key.offset: {{[0-9]+}},
+// CHECK-NEXT:     key.length: {{[0-9]+}},
+// CHECK-NEXT:     key.actor_isolation: "@concurrent"
+// CHECK-NEXT:   }
+// CHECK-NEXT: ]

test/SourceKit/InferredIsolation/range.swift

@@ -0,0 +1,43 @@
+@MainActor
+final class C {
+  var n = 0
+
+  func work() async {
+    let inheritedMain = {
+      self.n += 1
+    }
+    _ = inheritedMain
+
+    Task.detached {
+      print("detached")
+    }
+  }
+}
+
+// Range covering only the first closure.
+// We should see only the @MainActor entry.
+//
+// RUN: %sourcekitd-test -req=collect-inferred-isolation -pos=6:1 -length=60 %s -- %s | %FileCheck %s --check-prefix=INHERITED
+
+// INHERITED:      key.results: [
+// INHERITED-NEXT:   {
+// INHERITED-NEXT:     key.kind: "closure",
+// INHERITED-NEXT:     key.offset: {{[0-9]+}},
+// INHERITED-NEXT:     key.length: {{[0-9]+}},
+// INHERITED-NEXT:     key.actor_isolation: "@MainActor"
+// INHERITED-NEXT:   }
+// INHERITED-NEXT: ]
+
+// Range covering only the second closure.
+// see only the nonisolated entry; the first closure must be filtered out.
+//
+// RUN: %sourcekitd-test -req=collect-inferred-isolation -pos=9:1 -length=100 %s -- %s | %FileCheck %s --check-prefix=DETACHED
+
+// DETACHED:      key.results: [
+// DETACHED-NEXT:   {
+// DETACHED-NEXT:     key.kind: "closure",
+// DETACHED-NEXT:     key.offset: {{[0-9]+}},
+// DETACHED-NEXT:     key.length: {{[0-9]+}},
+// DETACHED-NEXT:     key.actor_isolation: "@concurrent"
+// DETACHED-NEXT:   }
+// DETACHED-NEXT: ]