[ASTGen] Generate LifetimeAttr

rintaro · rintaro · commit 94402b6c9dfa · 2025-02-17T08:56:29.000-08:00
diff --git a/include/swift/AST/ASTBridging.h b/include/swift/AST/ASTBridging.h
@@ -48,6 +48,7 @@ class Fingerprint;
 class Identifier;
 class IfConfigClauseRangeInfo;
 struct LabeledStmtInfo;
+struct LifetimeDescriptor;
 enum class MacroRole : uint32_t;
 class MacroIntroducedDeclName;
 enum class MacroIntroducedDeclNameKind;
@@ -964,6 +965,79 @@ BridgedInlineAttr BridgedInlineAttr_createParsed(BridgedASTContext cContext,
                                                  BridgedSourceRange cRange,
                                                  BridgedInlineKind cKind);
 
+enum ENUM_EXTENSIBILITY_ATTR(closed) BridgedParsedLifetimeDependenceKind {
+  BridgedParsedLifetimeDependenceKindDefault,
+  BridgedParsedLifetimeDependenceKindScope,
+  BridgedParsedLifetimeDependenceKindInherit,
+};
+
+class BridgedLifetimeDescriptor {
+  union Value {
+    BridgedIdentifier name;
+    unsigned index;
+
+    Value(BridgedIdentifier name) : name(name) {}
+    Value(unsigned index) : index(index) {}
+    Value() : name() {}
+  } value;
+
+  enum DescriptorKind {
+    Named,
+    Ordered,
+    Self,
+  } kind;
+
+  BridgedParsedLifetimeDependenceKind dependenceKind;
+  BridgedSourceLoc loc;
+
+  BridgedLifetimeDescriptor(Value value, DescriptorKind kind,
+                            BridgedParsedLifetimeDependenceKind dependenceKind,
+                            BridgedSourceLoc loc)
+      : value(value), kind(kind), dependenceKind(dependenceKind), loc(loc) {}
+
+public:
+  SWIFT_NAME("forNamed(_:dependenceKind:loc:)")
+  static BridgedLifetimeDescriptor
+  forNamed(BridgedIdentifier name,
+           BridgedParsedLifetimeDependenceKind dependenceKind,
+           BridgedSourceLoc loc) {
+    return BridgedLifetimeDescriptor(name, DescriptorKind::Named,
+                                     dependenceKind, loc);
+  }
+  SWIFT_NAME("forOrdered(_:dependenceKind:loc:)")
+  static BridgedLifetimeDescriptor
+  forOrdered(size_t index, BridgedParsedLifetimeDependenceKind dependenceKind,
+             BridgedSourceLoc loc) {
+    return BridgedLifetimeDescriptor(index, DescriptorKind::Ordered,
+                                     dependenceKind, loc);
+  }
+  SWIFT_NAME("forSelf(dependenceKind:loc:)")
+  static BridgedLifetimeDescriptor
+  forSelf(BridgedParsedLifetimeDependenceKind dependenceKind,
+          BridgedSourceLoc loc) {
+    return BridgedLifetimeDescriptor({}, DescriptorKind::Self, dependenceKind,
+                                     loc);
+  }
+
+  swift::LifetimeDescriptor unbridged();
+};
+
+SWIFT_NAME("BridgedLifetimeEntry.createParsed(_:range:sources:)")
+BridgedLifetimeEntry
+BridgedLifetimeEntry_createParsed(BridgedASTContext cContext,
+                                  BridgedSourceRange cRange,
+                                  BridgedArrayRef cSources);
+
+SWIFT_NAME("BridgedLifetimeEntry.createParsed(_:range:sources:target:)")
+BridgedLifetimeEntry BridgedLifetimeEntry_createParsed(
+    BridgedASTContext cContext, BridgedSourceRange cRange,
+    BridgedArrayRef cSources, BridgedLifetimeDescriptor cTarget);
+
+SWIFT_NAME("BridgedLifetimeAttr.createParsed(_:atLoc:range:entry:)")
+BridgedLifetimeAttr BridgedLifetimeAttr_createParsed(
+    BridgedASTContext cContext, BridgedSourceLoc cAtLoc,
+    BridgedSourceRange cRange, BridgedLifetimeEntry cEntry);
+
 enum ENUM_EXTENSIBILITY_ATTR(closed) BridgedMacroSyntax {
   BridgedMacroSyntaxFreestanding,
   BridgedMacroSyntaxAttached,
diff --git a/include/swift/AST/ASTBridgingWrappers.def b/include/swift/AST/ASTBridgingWrappers.def
@@ -115,6 +115,7 @@ AST_BRIDGING_WRAPPER_NULLABLE(ArgumentList)
 AST_BRIDGING_WRAPPER_NULLABLE(AvailabilitySpec)
 AST_BRIDGING_WRAPPER_CONST_NONNULL(AvailabilityMacroMap)
 AST_BRIDGING_WRAPPER_NONNULL(PoundAvailableInfo)
+AST_BRIDGING_WRAPPER_NONNULL(LifetimeEntry)
 
 // Non-AST types to generate wrappers for.
 AST_BRIDGING_WRAPPER_NULLABLE(DiagnosticEngine)
diff --git a/lib/AST/ASTDumper.cpp b/lib/AST/ASTDumper.cpp
@@ -5120,7 +5120,12 @@ class PrintAttribute : public AttributeVisitor<PrintAttribute, void, Label>,
   }
   void visitLifetimeAttr(LifetimeAttr *Attr, Label label) {
     printCommon(Attr, "lifetime_attr", label);
-    // TODO: Implement.
+    // FIXME: Improve, more detailed info.
+    printFieldRaw(
+        [&](raw_ostream &out) {
+          out << " " << Attr->getLifetimeEntry()->getString() << " ";
+        },
+        Label::optional("lifetime_entry"));
     printFoot();
   }
   void visitMacroRoleAttr(MacroRoleAttr *Attr, Label label) {
diff --git a/lib/AST/Bridging/DeclAttributeBridging.cpp b/lib/AST/Bridging/DeclAttributeBridging.cpp
@@ -360,6 +360,73 @@ BridgedInlineAttr BridgedInlineAttr_createParsed(BridgedASTContext cContext,
       InlineAttr(cAtLoc.unbridged(), cRange.unbridged(), unbridged(cKind));
 }
 
+static swift::ParsedLifetimeDependenceKind
+unbridged(BridgedParsedLifetimeDependenceKind kind) {
+  switch (kind) {
+  case BridgedParsedLifetimeDependenceKindDefault:
+    return swift::ParsedLifetimeDependenceKind::Default;
+  case BridgedParsedLifetimeDependenceKindScope:
+    return swift::ParsedLifetimeDependenceKind::Scope;
+  case BridgedParsedLifetimeDependenceKindInherit:
+    return swift::ParsedLifetimeDependenceKind::Inherit;
+  }
+  llvm_unreachable("unhandled enum value");
+}
+
+swift::LifetimeDescriptor BridgedLifetimeDescriptor::unbridged() {
+  switch (kind) {
+  case DescriptorKind::Named:
+    return LifetimeDescriptor::forNamed(value.name.unbridged().str(),
+                                        ::unbridged(dependenceKind),
+                                        loc.unbridged());
+  case DescriptorKind::Ordered:
+    return LifetimeDescriptor::forOrdered(
+        value.index, ::unbridged(dependenceKind), loc.unbridged());
+  case DescriptorKind::Self:
+    return LifetimeDescriptor::forSelf(::unbridged(dependenceKind),
+                                       loc.unbridged());
+  }
+  llvm_unreachable("unhandled enum value");
+}
+
+static BridgedLifetimeEntry BridgedLifetimeEntry_createParsedImpl(
+    BridgedASTContext cContext, BridgedSourceRange cRange,
+    BridgedArrayRef cSources,
+    std::optional<BridgedLifetimeDescriptor> cTarget) {
+  SmallVector<LifetimeDescriptor> sources;
+  for (auto cSource : cSources.unbridged<BridgedLifetimeDescriptor>())
+    sources.push_back(cSource.unbridged());
+  std::optional<LifetimeDescriptor> target;
+  if (cTarget)
+    target = cTarget->unbridged();
+
+  return LifetimeEntry::create(cContext.unbridged(), cRange.Start.unbridged(),
+                               cRange.End.unbridged(), sources, target);
+}
+
+BridgedLifetimeEntry
+BridgedLifetimeEntry_createParsed(BridgedASTContext cContext,
+                                  BridgedSourceRange cRange,
+                                  BridgedArrayRef cSources) {
+  return BridgedLifetimeEntry_createParsedImpl(cContext, cRange, cSources,
+                                               std::nullopt);
+}
+
+BridgedLifetimeEntry BridgedLifetimeEntry_createParsed(
+    BridgedASTContext cContext, BridgedSourceRange cRange,
+    BridgedArrayRef cSources, BridgedLifetimeDescriptor cTarget) {
+  return BridgedLifetimeEntry_createParsedImpl(cContext, cRange, cSources,
+                                               cTarget);
+}
+
+BridgedLifetimeAttr BridgedLifetimeAttr_createParsed(
+    BridgedASTContext cContext, BridgedSourceLoc cAtLoc,
+    BridgedSourceRange cRange, BridgedLifetimeEntry cEntry) {
+  return LifetimeAttr::create(cContext.unbridged(), cAtLoc.unbridged(),
+                              cRange.unbridged(), /*implicit=*/false,
+                              cEntry.unbridged());
+}
+
 BridgedMacroRole BridgedMacroRole_fromString(BridgedStringRef str) {
   // Match the role string to the known set of roles.
   auto role =
diff --git a/lib/ASTGen/Sources/ASTGen/DeclAttrs.swift b/lib/ASTGen/Sources/ASTGen/DeclAttrs.swift
@@ -146,7 +146,7 @@ extension ASTGenVisitor {
       case .inline:
         return handle(self.generateInlineAttr(attribute: node)?.asDeclAttribute)
       case .lifetime:
-        fatalError("unimplemented")
+        return handle(self.generateLifetimeAttr(attribute: node)?.asDeclAttribute)
       case .macroRole:
         return handle(self.generateMacroRoleAttr(attribute: node, attrName: attrName)?.asDeclAttribute)
       case .nonSendable:
@@ -880,6 +880,129 @@ extension ASTGenVisitor {
     )
   }
 
+  func generateLifetimeDescriptor(nameToken node: TokenSyntax, lifetimeDependenceKind: BridgedParsedLifetimeDependenceKind = .default) -> BridgedLifetimeDescriptor {
+    let ident = self.generateIdentifier(node)
+    let loc = self.generateSourceLoc(node)
+    if ident == ctx.id_self {
+      return .forSelf(
+        dependenceKind: lifetimeDependenceKind,
+        loc: loc
+      )
+    } else {
+      return .forNamed(
+        ident,
+        dependenceKind: lifetimeDependenceKind,
+        loc: loc
+      );
+    }
+  }
+
+  func generateLifetimeDescriptor(expr node: ExprSyntax) -> BridgedLifetimeDescriptor? {
+    let lifetimeDependenceKind: BridgedParsedLifetimeDependenceKind
+    let descriptorExpr: ExprSyntax
+    if let copyExpr = node.as(CopyExprSyntax.self) {
+      lifetimeDependenceKind = .inherit
+      descriptorExpr = copyExpr.expression
+    } else if let borrowExpr = node.as(BorrowExprSyntax.self) {
+      lifetimeDependenceKind = .scope
+      descriptorExpr = borrowExpr.expression
+    } else {
+      lifetimeDependenceKind = .default
+      descriptorExpr = node
+    }
+
+    let loc = self.generateSourceLoc(descriptorExpr)
+    if
+      let declRefExpr = descriptorExpr.as(DeclReferenceExprSyntax.self),
+      declRefExpr.argumentNames == nil
+    {
+      return generateLifetimeDescriptor(
+        nameToken: declRefExpr.baseName,
+        lifetimeDependenceKind: lifetimeDependenceKind
+      )
+    }
+
+    if let index = descriptorExpr.as(IntegerLiteralExprSyntax.self)?.representedLiteralValue {
+      return .forOrdered(
+        index,
+        dependenceKind: lifetimeDependenceKind,
+        loc: loc
+      )
+    }
+
+    // TODO: Diangose
+    fatalError("expected identifier, 'self', or integer in @lifetime")
+  }
+
+  func generateLifetimeEntry(attribute node: AttributeSyntax) -> BridgedLifetimeEntry? {
+    self.generateWithLabeledExprListArguments(attribute: node) { args in
+      guard !args.isEmpty else {
+        // TODO: Diagnose
+        fatalError("expected arguments in @lifetime attribute")
+      }
+
+      var target: BridgedLifetimeDescriptor? = nil
+      var sources: [BridgedLifetimeDescriptor] = []
+      var first = true
+      while let arg = args.popFirst() {
+        if first {
+          if let targetToken = arg.label {
+            target = self.generateLifetimeDescriptor(nameToken: targetToken)
+          }
+          first = false
+        } else {
+          if arg.label != nil {
+            // TODO: Diagnose.
+            fatalError("invalid argument label in @lifetime attribute")
+          }
+        }
+
+        if let src = self.generateLifetimeDescriptor(expr: arg.expression) {
+          sources.append(src)
+        }
+      }
+
+      if let target {
+        return .createParsed(
+          self.ctx,
+          range: self.generateAttrSourceRange(node),
+          sources: sources.lazy.bridgedArray(in: self),
+          target: target
+        )
+      } else {
+        return .createParsed(
+          self.ctx,
+          range: self.generateAttrSourceRange(node),
+          sources: sources.lazy.bridgedArray(in: self)
+        )
+      }
+    }
+  }
+
+  /// E.g.
+  ///   ```
+  ///   @lifetime(src1, src2)
+  ///   @lifetime(target: borrow src1, copy src2)
+  ///   @lifetime(2)
+  ///   @lifetime(self)
+  ///   ```
+  func generateLifetimeAttr(attribute node: AttributeSyntax) -> BridgedLifetimeAttr? {
+    guard self.ctx.langOptsHasFeature(.LifetimeDependence) else {
+      // TODO: Diagnose
+      fatalError("@lifetime attribute requires 'LifetimeDependence' feature")
+    }
+    guard let entry = self.generateLifetimeEntry(attribute: node) else {
+      // TODO: Diagnose?
+      return nil
+    }
+    return .createParsed(
+      self.ctx,
+      atLoc: self.generateSourceLoc(node.atSign),
+      range: self.generateAttrSourceRange(node),
+      entry: entry
+    )
+  }
+
   func generateMacroIntroducedDeclNameKind(declReferenceExpr node: DeclReferenceExprSyntax) -> BridgedMacroIntroducedDeclNameKind? {
     if node.argumentNames != nil {
       // TODO: Diagnose
diff --git a/test/ASTGen/attrs.swift b/test/ASTGen/attrs.swift
@@ -1,19 +1,21 @@
 // RUN: %empty-directory(%t)
 
-// RUN: %target-swift-frontend-dump-parse -disable-availability-checking \
-// RUN:   -enable-experimental-feature SymbolLinkageMarkers \
+// RUN: %target-swift-frontend-dump-parse \
 // RUN:   -enable-experimental-feature ABIAttribute \
 // RUN:   -enable-experimental-feature Extern \
+// RUN:   -enable-experimental-feature LifetimeDependence \
 // RUN:   -enable-experimental-feature NonIsolatedAsyncInheritsIsolationFromContext \
+// RUN:   -enable-experimental-feature SymbolLinkageMarkers \
 // RUN:   -enable-experimental-move-only \
 // RUN:   -enable-experimental-feature ParserASTGen \
 // RUN:   | %sanitize-address > %t/astgen.ast
 
-// RUN: %target-swift-frontend-dump-parse -disable-availability-checking \
-// RUN:   -enable-experimental-feature SymbolLinkageMarkers \
+// RUN: %target-swift-frontend-dump-parse \
 // RUN:   -enable-experimental-feature ABIAttribute \
 // RUN:   -enable-experimental-feature Extern \
+// RUN:   -enable-experimental-feature LifetimeDependence \
 // RUN:   -enable-experimental-feature NonIsolatedAsyncInheritsIsolationFromContext \
+// RUN:   -enable-experimental-feature SymbolLinkageMarkers \
 // RUN:   -enable-experimental-move-only \
 // RUN:   | %sanitize-address > %t/cpp-parser.ast
 
@@ -22,19 +24,21 @@
 // RUN: %target-typecheck-verify-swift \
 // RUN:   -module-abi-name ASTGen \
 // RUN:   -enable-experimental-feature ParserASTGen \
-// RUN:   -enable-experimental-feature SymbolLinkageMarkers \
 // RUN:   -enable-experimental-feature ABIAttribute \
 // RUN:   -enable-experimental-feature Extern \
-// RUN:   -enable-experimental-move-only \
-// RUN:   -enable-experimental-feature NonIsolatedAsyncInheritsIsolationFromContext
+// RUN:   -enable-experimental-feature LifetimeDependence \
+// RUN:   -enable-experimental-feature NonIsolatedAsyncInheritsIsolationFromContext \
+// RUN:   -enable-experimental-feature SymbolLinkageMarkers \
+// RUN:   -enable-experimental-move-only
 
 // REQUIRES: executable_test
 // REQUIRES: swift_swift_parser
-// REQUIRES: swift_feature_SymbolLinkageMarkers
-// REQUIRES: swift_feature_Extern
 // REQUIRES: swift_feature_ParserASTGen
 // REQUIRES: swift_feature_ABIAttribute
+// REQUIRES: swift_feature_Extern
+// REQUIRES: swift_feature_LifetimeDependence
 // REQUIRES: swift_feature_NonIsolatedAsyncInheritsIsolationFromContext
+// REQUIRES: swift_feature_SymbolLinkageMarkers
 
 // rdar://116686158
 // UNSUPPORTED: asan
@@ -201,3 +205,13 @@ struct OpTest {
   func opResult() -> some OpProto { OpStruct() }
   typealias Result = @_opaqueReturnTypeOf("$s6ASTGen6OpTestV8opResultQryF", 0) __
 }
+
+struct E {}
+struct NE : ~Escapable {}
+@lifetime(ne) func derive(_ ne: NE) -> NE { ne }
+@lifetime(borrow ne1, ne2) func derive(_ ne1: NE, _ ne2: NE) -> NE {
+  if (Int.random(in: 1..<100) < 50) { return ne1 }
+  return ne2
+}
+@lifetime(borrow borrow) func testNameConflict(_ borrow: E) -> NE { NE() }
+@lifetime(result: source) func testTarget(_ result: inout NE, _ source: consuming NE) { result = source }

Original file line number	Diff line number	Diff line change
`@@ -5120,7 +5120,12 @@ class PrintAttribute : public AttributeVisitor<PrintAttribute, void, Label>,`
`5120`	`5120`	`}`
`5121`	`5121`	`void visitLifetimeAttr(LifetimeAttr *Attr, Label label) {`
`5122`	`5122`	`printCommon(Attr, "lifetime_attr", label);`
`5123`		`- // TODO: Implement.`
	`5123`	`+ // FIXME: Improve, more detailed info.`
	`5124`	`+ printFieldRaw(`
	`5125`	`+ [&](raw_ostream &out) {`
	`5126`	`+ out << " " << Attr->getLifetimeEntry()->getString() << " ";`
	`5127`	`+ },`
	`5128`	`+ Label::optional("lifetime_entry"));`
`5124`	`5129`	`printFoot();`
`5125`	`5130`	`}`
`5126`	`5131`	`void visitMacroRoleAttr(MacroRoleAttr *Attr, Label label) {`