[ASTGen] Lower @abi in ASTGen

beccadax · beccadax · commit d8490716fad5 · 2024-10-15T17:38:01.000-07:00
diff --git a/include/swift/AST/ASTBridging.h b/include/swift/AST/ASTBridging.h
@@ -601,11 +601,34 @@ SWIFT_NAME("BridgedDeclAttributes.add(self:_:)")
 void BridgedDeclAttributes_add(BridgedDeclAttributes *_Nonnull attrs,
                                BridgedDeclAttribute add);
 
+/// Retrieve the attribute from \p cAttrs that comes after \p cPriorAttr .
+/// If \c cPriorAttr is null, retrieves the first attribute. If there are no 
+/// more attributes, returns null.
+SWIFT_NAME("BridgedDeclAttributes.attr(self:after:)")
+BridgedNullableDeclAttribute BridgedDeclAttributes_getAttrAfter(
+    BridgedDeclAttributes cAttrs, BridgedNullableDeclAttribute cPriorAttr);
+
 SWIFT_NAME("BridgedDeclAttribute.createSimple(_:kind:atLoc:nameLoc:)")
 BridgedDeclAttribute BridgedDeclAttribute_createSimple(
     BridgedASTContext cContext, BridgedDeclAttrKind cKind,
     BridgedSourceLoc cAtLoc, BridgedSourceLoc cNameLoc);
 
+SWIFT_NAME("getter:BridgedDeclAttribute.asABIAttr(self:)")
+BridgedNullableABIAttr BridgedDeclAttribute_asABIAttr(
+    BridgedDeclAttribute cAttr);
+
+SWIFT_NAME("BridgedABIAttr.createParsed(_:atLoc:range:abiDecl:)")
+BridgedABIAttr BridgedABIAttr_createParsed(
+    BridgedASTContext cContext, BridgedSourceLoc atLoc,
+    BridgedSourceRange range, BridgedDecl abiDecl);
+
+SWIFT_NAME("BridgedABIAttr.createImplicitInverse(_:)")
+BridgedABIAttr BridgedABIAttr_createImplicitInverse(
+    BridgedASTContext cContext);
+
+SWIFT_NAME("BridgedABIAttr.connectToInverse(self:attachedTo:)")
+void BridgedABIAttr_connectToInverse(BridgedABIAttr cAttr, BridgedDecl cOwner);
+
 enum ENUM_EXTENSIBILITY_ATTR(closed) BridgedAccessLevel {
   BridgedAccessLevelPrivate,
   BridgedAccessLevelFilePrivate,
@@ -859,6 +882,10 @@ BridgedDeclAttributes BridgedDecl_getAttrs(BridgedDecl decl);
 SWIFT_NAME("setter:BridgedDecl.attrs(self:newValue:)")
 void BridgedDecl_setAttrs(BridgedDecl decl, BridgedDeclAttributes attrs);
 
+SWIFT_NAME("getter:BridgedDecl.asPatternBindingDecl(self:)")
+BridgedNullablePatternBindingDecl BridgedDecl_getAsPatternBindingDecl(
+    BridgedDecl decl);
+
 enum ENUM_EXTENSIBILITY_ATTR(closed) BridgedStaticSpelling {
   BridgedStaticSpellingNone,
   BridgedStaticSpellingStatic,
@@ -894,6 +921,14 @@ BridgedPatternBindingDecl BridgedPatternBindingDecl_createParsed(
     BridgedSourceLoc cBindingKeywordLoc, BridgedArrayRef cBindingEntries,
     bool isStatic, bool isLet);
 
+SWIFT_NAME("getter:BridgedPatternBindingDecl.patternCount(self:)")
+SwiftInt BridgedPatternBindingDecl_getNumPatternEntries(
+    BridgedPatternBindingDecl cPBD);
+
+SWIFT_NAME("BridgedPatternBindingDecl.pattern(self:at:)")
+BridgedPattern BridgedPatternBindingDecl_getPattern(
+    BridgedPatternBindingDecl cPBD, SwiftInt i);
+
 SWIFT_NAME("BridgedParamDecl.createParsed(_:declContext:specifierLoc:argName:"
            "argNameLoc:paramName:paramNameLoc:type:defaultValue:)")
 BridgedParamDecl BridgedParamDecl_createParsed(
@@ -1921,6 +1956,21 @@ void BridgedTypeRepr_dump(BridgedTypeRepr type);
 SWIFT_NAME("getter:BridgedPattern.singleVar(self:)")
 BridgedNullableVarDecl BridgedPattern_getSingleVar(BridgedPattern cPattern);
 
+/// Iterates over the pattern, collecting and returning all of the BridgedVarDecls
+/// referenced by its subpatterns.
+///
+/// The exact semantics of this are as follows:
+///
+/// \li This method always returns the current count of BridgedVarDecls in the
+///     pattern.
+/// \li If \p capacity is greater than or equal to the current count \em and
+///     \p varDeclArrayOut is non-null, this method also writes the
+///     BridgedVarDecls into the buffer starting at \p varDeclArrayOut .
+SWIFT_NAME("BridgedPattern.unsafeFetchVarDecls(self:into:capacity:)")
+SwiftInt BridgedPattern_unsafeFetchVarDecls(
+    BridgedPattern cPattern, BridgedVarDecl * _Nullable varDeclArrayOut,
+    SwiftInt capacity);
+
 SWIFT_NAME("BridgedAnyPattern.createParsed(_:loc:)")
 BridgedAnyPattern BridgedAnyPattern_createParsed(BridgedASTContext cContext,
                                                  BridgedSourceLoc cLoc);
diff --git a/include/swift/AST/ASTBridgingWrappers.def b/include/swift/AST/ASTBridgingWrappers.def
@@ -62,13 +62,13 @@
 
 #ifndef DECL_ATTR
 #define SIMPLE_DECL_ATTR(...)
-#define DECL_ATTR(_, Id, ...) AST_BRIDGING_WRAPPER_NONNULL(Id##Attr)
+#define DECL_ATTR(_, Id, ...) AST_BRIDGING_WRAPPER_NULLABLE(Id##Attr)
 #endif
 #include "swift/AST/DeclAttr.def"
 
 // Some of the base classes need to be nullable to allow them to be used as
 // optional parameters.
-AST_BRIDGING_WRAPPER_NONNULL(Decl)
+AST_BRIDGING_WRAPPER_NULLABLE(Decl)
 AST_BRIDGING_WRAPPER_NONNULL(DeclContext)
 AST_BRIDGING_WRAPPER_NONNULL(SourceFile)
 AST_BRIDGING_WRAPPER_NULLABLE(Stmt)
diff --git a/lib/AST/Bridging/DeclAttributeBridging.cpp b/lib/AST/Bridging/DeclAttributeBridging.cpp
@@ -90,6 +90,42 @@ static AccessLevel unbridged(BridgedAccessLevel level) {
   llvm_unreachable("unhandled BridgedAccessLevel");
 }
 
+BridgedNullableDeclAttribute BridgedDeclAttributes_getAttrAfter(
+    BridgedDeclAttributes cAttrs, BridgedNullableDeclAttribute cPriorAttr) {
+  auto attrs = cAttrs.unbridged();
+  auto priorAttr = cPriorAttr.unbridged();
+
+  auto i = priorAttr ? std::next(DeclAttributes::iterator(priorAttr))
+                     : attrs.begin();
+  return *i;
+}
+
+BridgedNullableABIAttr BridgedDeclAttribute_asABIAttr(
+    BridgedDeclAttribute cAttr) {
+  return dyn_cast<ABIAttr>(cAttr.unbridged());
+}
+
+BridgedABIAttr BridgedABIAttr_createParsed(BridgedASTContext cContext,
+                                           BridgedSourceLoc atLoc,
+                                           BridgedSourceRange range,
+                                           BridgedDecl abiDecl) {
+  return new (cContext.unbridged()) ABIAttr(abiDecl.unbridged(),
+                                            atLoc.unbridged(),
+                                            range.unbridged(),
+                                            /*isInverse=*/false,
+                                            /*isImplicit=*/false);
+}
+
+BridgedABIAttr BridgedABIAttr_createImplicitInverse(BridgedASTContext cContext){
+  return new (cContext.unbridged()) ABIAttr(nullptr,
+                                            /*isInverse=*/true,
+                                            /*isImplicit=*/true);
+}
+
+void BridgedABIAttr_connectToInverse(BridgedABIAttr cAttr, BridgedDecl cOwner) {
+  return cAttr.unbridged()->connectToInverse(cOwner.unbridged());
+}
+
 BridgedAccessControlAttr
 BridgedAccessControlAttr_createParsed(BridgedASTContext cContext,
                                       BridgedSourceRange cRange,
diff --git a/lib/AST/Bridging/DeclBridging.cpp b/lib/AST/Bridging/DeclBridging.cpp
@@ -123,6 +123,13 @@ void BridgedDecl_setAttrs(BridgedDecl decl, BridgedDeclAttributes attrs) {
   decl.unbridged()->getAttrs() = attrs.unbridged();
 }
 
+// FIXME: We should do this systematically for downcasting BridgedDecl to all
+//        subclasses.
+BridgedNullablePatternBindingDecl BridgedDecl_getAsPatternBindingDecl(
+    BridgedDecl decl) {
+  return dyn_cast<PatternBindingDecl>(decl.unbridged());
+}
+
 BridgedAccessorDecl BridgedAccessorDecl_createParsed(
     BridgedASTContext cContext, BridgedDeclContext cDeclContext,
     BridgedAccessorKind cKind, BridgedAbstractStorageDecl cStorage,
@@ -167,6 +174,17 @@ BridgedPatternBindingDecl BridgedPatternBindingDecl_createParsed(
       cBindingKeywordLoc.unbridged(), entries, declContext);
 }
 
+SwiftInt BridgedPatternBindingDecl_getNumPatternEntries(
+    BridgedPatternBindingDecl cPBD) {
+  return cPBD.unbridged()->getNumPatternEntries();
+}
+
+BridgedPattern BridgedPatternBindingDecl_getPattern(
+    BridgedPatternBindingDecl cPBD, SwiftInt i) {
+  ASSERT(i >= 0 && i < cPBD.unbridged()->getNumPatternEntries());
+  return cPBD.unbridged()->getPattern(i);
+}
+
 BridgedParamDecl BridgedParamDecl_createParsed(
     BridgedASTContext cContext, BridgedDeclContext cDeclContext,
     BridgedSourceLoc cSpecifierLoc, BridgedIdentifier cArgName,
diff --git a/lib/AST/Bridging/PatternBridging.cpp b/lib/AST/Bridging/PatternBridging.cpp
@@ -37,6 +37,21 @@ BridgedNullableVarDecl BridgedPattern_getSingleVar(BridgedPattern cPattern) {
   return cPattern.unbridged()->getSingleVar();
 }
 
+SwiftInt BridgedPattern_unsafeFetchVarDecls(BridgedPattern cPattern,
+                                            BridgedVarDecl *varDeclArrayOut,
+                                            SwiftInt capacity) {
+  SmallVector<VarDecl *, 32> vars;
+  cPattern.unbridged()->collectVariables(vars);
+
+  if (varDeclArrayOut && vars.size() <= (size_t)capacity) {
+    for (auto i : indices(vars)) {
+      varDeclArrayOut[i] = vars[i];
+    }
+  }
+
+  return vars.size();
+}
+
 BridgedAnyPattern BridgedAnyPattern_createParsed(BridgedASTContext cContext,
                                                  BridgedSourceLoc cLoc) {
   return new (cContext.unbridged()) AnyPattern(cLoc.unbridged());
diff --git a/lib/ASTGen/Sources/ASTGen/Bridge.swift b/lib/ASTGen/Sources/ASTGen/Bridge.swift
@@ -27,6 +27,8 @@ extension BridgedNullable {
 
 extension BridgedSourceLoc: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 extension BridgedIdentifier: /*@retroactive*/ swiftASTGen.BridgedNullable {}
+extension BridgedNullableDeclAttribute: /*@retroactive*/ swiftASTGen.BridgedNullable {}
+extension BridgedNullableDecl: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 extension BridgedNullableExpr: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 extension BridgedNullableStmt: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 extension BridgedNullableTypeRepr: /*@retroactive*/ swiftASTGen.BridgedNullable {}
@@ -37,6 +39,7 @@ extension BridgedNullableParameterList: /*@retroactive*/ swiftASTGen.BridgedNull
 extension BridgedNullablePatternBindingInitializer: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 extension BridgedNullablePatternBindingDecl: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 extension BridgedNullableVarDecl: /*@retroactive*/ swiftASTGen.BridgedNullable {}
+extension BridgedNullableABIAttr: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 
 extension BridgedIdentifier: /*@retroactive*/ Swift.Equatable {
   public static func == (lhs: Self, rhs: Self) -> Bool {
@@ -70,6 +73,12 @@ extension BridgedHasNullable {
 extension BridgedStmt: BridgedHasNullable {
   typealias Nullable = BridgedNullableStmt
 }
+extension BridgedDeclAttribute: BridgedHasNullable {
+  typealias Nullable = BridgedNullableDeclAttribute
+}
+extension BridgedDecl: BridgedHasNullable {
+  typealias Nullable = BridgedNullableDecl
+}
 extension BridgedExpr: BridgedHasNullable {
   typealias Nullable = BridgedNullableExpr
 }
@@ -97,6 +106,9 @@ extension BridgedPatternBindingDecl: BridgedHasNullable {
 extension BridgedVarDecl: BridgedHasNullable {
   typealias Nullable = BridgedNullableVarDecl
 }
+extension BridgedABIAttr: BridgedHasNullable {
+  typealias Nullable = BridgedNullableABIAttr
+}
 
 public extension BridgedSourceLoc {
   /// Form a source location at the given absolute position in `buffer`.
@@ -249,6 +261,56 @@ extension BridgedSourceRange {
   }
 }
 
+extension BridgedPatternBindingDecl {
+  var patterns: [BridgedPattern] {
+    return (0 ..< patternCount).map(pattern(at:))
+  }
+}
+
+extension BridgedPattern {
+  private enum FetchVarDeclsError: Error {
+    case insufficientCapacity(Int)
+  }
+
+  private func fetchVarDecls(capacity: Int) throws -> [BridgedVarDecl] {
+    return try Array(unsafeUninitializedCapacity: capacity) { buffer, initializedCount in
+      let fullCount = self.unsafeFetchVarDecls(
+        into: buffer.baseAddress,
+        capacity: buffer.count
+      )
+
+      guard fullCount <= buffer.count else {
+        throw FetchVarDeclsError.insufficientCapacity(fullCount)
+      }
+
+      // `unsafeFetchVarDecls` writes nothing if `fullCount` > `capacity`, so
+      // it's correct to put this after the `throw`.
+      initializedCount = fullCount
+    }
+  }
+
+  var varDecls: [BridgedVarDecl] {
+    do {
+      return try fetchVarDecls(capacity: 8) // "probably big enough" guess
+    }
+    catch FetchVarDeclsError.insufficientCapacity(let neededCapacity) {
+      return try! fetchVarDecls(capacity: neededCapacity)
+    }
+    catch {
+      fatalError("Unknown error \(error)")
+    }
+  }
+}
+
+extension BridgedDeclAttributes {
+  var attrs: UnfoldSequence<BridgedDeclAttribute, BridgedNullableDeclAttribute> {
+    return sequence(state: nil) { prior in
+      prior = self.attr(after: prior)
+      return BridgedDeclAttribute(prior)
+    }
+  }
+}
+
 /// Helper collection type that lazily concatenates two collections.
 struct ConcatCollection<C1: Collection, C2: Collection> where C1.Element == C2.Element {
   let c1: C1
diff --git a/lib/ASTGen/Sources/ASTGen/DeclAttrs.swift b/lib/ASTGen/Sources/ASTGen/DeclAttrs.swift
@@ -17,6 +17,22 @@ import SwiftIfConfig
 
 @_spi(ExperimentalLanguageFeatures) @_spi(RawSyntax) import SwiftSyntax
 
+/// Calls `visitor` on each declaration that an `@abi` attribute `decl` ought to apply to.
+/// This is usually just `decl` itself, but for pattern binding decls, it's all of the
+/// variables bound to it.
+func visitABIDecls(for decl: BridgedDecl, visitor: (BridgedDecl) -> Void) {
+  if let pbd = BridgedPatternBindingDecl(decl.asPatternBindingDecl) {
+    for pattern in pbd.patterns {
+      for varDecl in pattern.varDecls {
+        visitABIDecls(for: varDecl.asDecl, visitor: visitor)
+      }
+    }
+    return
+  }
+
+  visitor(decl)
+}
+
 extension ASTGenVisitor {
   struct DeclAttributesResult {
     var attributes: BridgedDeclAttributes
@@ -25,6 +41,11 @@ extension ASTGenVisitor {
 
     func attach(to decl: BridgedDecl) {
       decl.attrs = attributes
+      if let abiAttr = attributes.attrs.lazy.map(\.asABIAttr).compactMap(BridgedABIAttr.init).first {
+        visitABIDecls(for: decl) { declToConnect in
+          abiAttr.connectToInverse(attachedTo: declToConnect)
+        }
+      }
     }
   }
 
@@ -329,8 +350,29 @@ extension ASTGenVisitor {
   }
 
   func generateABIAttr(attribute node: AttributeSyntax) -> BridgedABIAttr? {
-    #warning("TODO: implement generateABIAttr")
-    fatalError("TODO: implement generateABIAttr")
+    guard
+      let arg = node.arguments?.as(ABIAttributeArgumentsSyntax.self)
+    else {
+      // TODO: diagnose
+      return nil
+    }
+
+    let abiDecl = self.generate(decl: arg.provider)
+
+    // Attach empty inverse @abi attributes to each (relevant) decl
+    visitABIDecls(for: abiDecl) { decl in
+      let attr = BridgedABIAttr.createImplicitInverse(self.ctx)
+      var attrs = decl.attrs
+      attrs.add(attr.asDeclAttribute)
+      decl.attrs = attrs
+    }
+
+    return .createParsed(
+      self.ctx,
+      atLoc: self.generateSourceLoc(node.atSign),
+      range: self.generateSourceRange(node),
+      abiDecl: abiDecl
+    )
   }
 
   func generateAlignmentAttr(attribute node: AttributeSyntax) -> BridgedAlignmentAttr? {
diff --git a/lib/ASTGen/Sources/ASTGen/SourceFile.swift b/lib/ASTGen/Sources/ASTGen/SourceFile.swift
@@ -75,6 +75,7 @@ extension Parser.ExperimentalFeatures {
     mapFeature(.NonescapableTypes, to: .nonescapableTypes)
     mapFeature(.TrailingComma, to: .trailingComma)
     mapFeature(.CoroutineAccessors, to: .coroutineAccessors)
+    mapFeature(.ABIAttribute, to: .abiAttribute)
   }
 }
 
diff --git a/test/ASTGen/attrs.swift b/test/ASTGen/attrs.swift
@@ -1,14 +1,14 @@
 // RUN: %empty-directory(%t)
-// RUN: %target-swift-frontend %s -dump-parse -disable-availability-checking -enable-experimental-move-only -enable-experimental-feature ParserASTGen > %t/astgen.ast.raw
-// RUN: %target-swift-frontend %s -dump-parse -disable-availability-checking -enable-experimental-move-only > %t/cpp-parser.ast.raw
+// RUN: %target-swift-frontend %s -dump-parse -disable-availability-checking -enable-experimental-move-only -enable-experimental-feature ABIAttribute -enable-experimental-feature ParserASTGen > %t/astgen.ast.raw
+// RUN: %target-swift-frontend %s -dump-parse -disable-availability-checking -enable-experimental-move-only -enable-experimental-feature ABIAttribute > %t/cpp-parser.ast.raw
 
 // Filter out any addresses in the dump, since they can differ.
 // RUN: sed -E 's#0x[0-9a-fA-F]+##g' %t/cpp-parser.ast.raw > %t/cpp-parser.ast
 // RUN: sed -E 's#0x[0-9a-fA-F]+##g' %t/astgen.ast.raw > %t/astgen.ast
 
 // RUN: %diff -u %t/astgen.ast %t/cpp-parser.ast
 
-// RUN: %target-typecheck-verify-swift -enable-experimental-move-only -enable-experimental-feature ParserASTGen 
+// RUN: %target-typecheck-verify-swift -enable-experimental-move-only -enable-experimental-feature ABIAttribute -enable-experimental-feature ParserASTGen
 
 // REQUIRES: executable_test
 // REQUIRES: swift_swift_parser
@@ -62,3 +62,8 @@ struct S4 {}
 @implementation(Category) extension ObjCClass1 {} // expected-error {{cannot find type 'ObjCClass1' in scope}}
 @_objcImplementation extension ObjCClass2 {} // expected-error {{cannot find type 'ObjCClass2' in scope}}
 @_objcImplementation(Category) extension ObjCClass2 {} // expected-error {{cannot find type 'ObjCClass2' in scope}}
+
+@abi(func fn_abi()) // expected-error {{cannot give global function 'fn' the ABI of a global function with a different number of low-level parameters}}
+func fn(_: Int) {}
+
+// FIXME: Test @abi on a PatternBindingDecl once we've figured out VarDecl attrs

Original file line number	Diff line number	Diff line change
`@@ -75,6 +75,7 @@ extension Parser.ExperimentalFeatures {`
`75`	`75`	`mapFeature(.NonescapableTypes, to: .nonescapableTypes)`
`76`	`76`	`mapFeature(.TrailingComma, to: .trailingComma)`
`77`	`77`	`mapFeature(.CoroutineAccessors, to: .coroutineAccessors)`
	`78`	`+ mapFeature(.ABIAttribute, to: .abiAttribute)`
`78`	`79`	`}`
`79`	`80`	`}`
`80`	`81`