[ASTGen] Lower @abi in ASTGen

Author: beccadax

include/swift/AST/ASTBridging.h

@@ -567,11 +567,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, BridgedNullableDecl 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,
@@ -913,6 +936,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,
@@ -948,6 +975,14 @@ BridgedPatternBindingDecl BridgedPatternBindingDecl_createParsed(
     BridgedSourceLoc cBindingKeywordLoc, BridgedArrayRef cBindingEntries, BridgedDeclAttributes cAttrs,
                                                                  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(
@@ -2105,6 +2140,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);

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)

lib/AST/Bridging/DeclAttributeBridging.cpp

@@ -93,6 +93,42 @@ static std::optional<AccessLevel> unbridge(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,
+                                           BridgedNullableDecl 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,

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,

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());

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 {}
@@ -38,6 +40,7 @@ extension BridgedNullablePatternBindingInitializer: /*@retroactive*/ swiftASTGen
 extension BridgedNullableArgumentList: /*@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 {
@@ -71,6 +74,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
 }
@@ -101,6 +110,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`.
@@ -253,6 +265,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

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
@@ -26,6 +42,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)
+        }
+      }
     }
   }
 
@@ -330,8 +351,55 @@ extension ASTGenVisitor {
   ///   @abi(func fn())
   ///   ```
   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: BridgedDecl?
+    switch arg.provider {
+    case .associatedType(let assocTyDecl):
+      abiDecl = self.generate(associatedTypeDecl: assocTyDecl)?.asDecl
+    case .deinitializer(let deinitDecl):
+      abiDecl = self.generate(deinitializerDecl: deinitDecl).asDecl
+    case .enumCase(let caseDecl):
+      abiDecl = self.generate(enumCaseDecl: caseDecl).asDecl
+    case .function(let funcDecl):
+      abiDecl = self.generate(functionDecl: funcDecl)?.asDecl
+    case .initializer(let initDecl):
+      abiDecl = self.generate(initializerDecl: initDecl).asDecl
+    case .`subscript`(let subscriptDecl):
+      abiDecl = self.generate(subscriptDecl: subscriptDecl).asDecl
+    case .typeAlias(let typealiasDecl):
+      abiDecl = self.generate(typeAliasDecl: typealiasDecl)?.asDecl
+    case .variable(let varDecl):
+      abiDecl = self.generate(variableDecl: varDecl).asDecl
+    case .missing(_):
+      // This error condition will have been diagnosed in SwiftSyntax.
+      abiDecl = nil
+    }
+
+    // Attach empty inverse @abi attributes to each (relevant) decl
+    if let abiDecl {
+      // TODO: Diagnose if `abiDecl` has a body/initial value/etc.--the C++ parser considers it syntactically invalid
+      //       but SwiftSyntax does not.
+
+      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.generateAttrSourceRange(node),
+      abiDecl: abiDecl.asNullable
+    )
   }
 
   /// E.g.:

lib/ASTGen/Sources/ASTGen/SourceFile.swift

@@ -76,6 +76,7 @@ extension Parser.ExperimentalFeatures {
     mapFeature(.TrailingComma, to: .trailingComma)
     mapFeature(.CoroutineAccessors, to: .coroutineAccessors)
     mapFeature(.ValueGenerics, to: .valueGenerics)
+    mapFeature(.ABIAttribute, to: .abiAttribute)
   }
 }