ASTGen: Set brace locations for iterable declaration contexts

Author: AnthonyLatsis

include/swift/AST/CASTBridging.h

@@ -47,6 +47,11 @@ typedef struct BridgedSourceLoc {
   const void *_Nullable raw;
 } BridgedSourceLoc;
 
+typedef struct {
+  BridgedSourceLoc startLoc;
+  BridgedSourceLoc endLoc;
+} BridgedSourceRange;
+
 typedef struct BridgedIdentifier {
   const void *_Nullable raw;
 } BridgedIdentifier;
@@ -344,13 +349,13 @@ BridgedDeclContextAndDecl
 StructDecl_create(BridgedASTContext cContext, BridgedDeclContext cDeclContext,
                   BridgedSourceLoc cStructKeywordLoc, BridgedIdentifier cName,
                   BridgedSourceLoc cNameLoc,
-                  void *_Nullable opaqueGenericParamList);
+                  void *_Nullable opaqueGenericParamList,
+                  BridgedSourceRange cBraceRange);
 
-BridgedDeclContextAndDecl ClassDecl_create(BridgedASTContext cContext,
-                                           BridgedDeclContext cDeclContext,
-                                           BridgedSourceLoc cClassKeywordLoc,
-                                           BridgedIdentifier cName,
-                                           BridgedSourceLoc cNameLoc);
+BridgedDeclContextAndDecl
+ClassDecl_create(BridgedASTContext cContext, BridgedDeclContext cDeclContext,
+                 BridgedSourceLoc cClassKeywordLoc, BridgedIdentifier cName,
+                 BridgedSourceLoc cNameLoc, BridgedSourceRange cBraceRange);
 
 void *GenericParamList_create(BridgedASTContext cContext,
                               BridgedSourceLoc cLAngleLoc,

lib/AST/CASTBridging.cpp

@@ -70,6 +70,11 @@ static inline SourceLoc convertSourceLoc(BridgedSourceLoc cLoc) {
   return SourceLoc(smLoc);
 }
 
+static inline SourceRange convertSourceRange(BridgedSourceRange cRange) {
+  return SourceRange(convertSourceLoc(cRange.startLoc),
+                     convertSourceLoc(cRange.endLoc));
+}
+
 static inline Identifier convertIdentifier(BridgedIdentifier cIdentifier) {
   return Identifier::getFromOpaquePointer(const_cast<void *>(cIdentifier.raw));
 }
@@ -533,29 +538,31 @@ BridgedDeclContextAndDecl
 StructDecl_create(BridgedASTContext cContext, BridgedDeclContext cDeclContext,
                   BridgedSourceLoc cStructKeywordLoc, BridgedIdentifier cName,
                   BridgedSourceLoc cNameLoc,
-                  void *_Nullable opaqueGenericParamList) {
+                  void *_Nullable opaqueGenericParamList,
+                  BridgedSourceRange cBraceRange) {
   ASTContext &context = convertASTContext(cContext);
 
   auto *decl = new (context)
       StructDecl(convertSourceLoc(cStructKeywordLoc), convertIdentifier(cName),
                  convertSourceLoc(cNameLoc), {},
                  static_cast<GenericParamList *>(opaqueGenericParamList),
                  convertDeclContext(cDeclContext));
+  decl->setBraces(convertSourceRange(cBraceRange));
 
   return {bridgeDeclContext(decl), static_cast<Decl *>(decl)};
 }
 
-BridgedDeclContextAndDecl ClassDecl_create(BridgedASTContext cContext,
-                                           BridgedDeclContext cDeclContext,
-                                           BridgedSourceLoc cClassKeywordLoc,
-                                           BridgedIdentifier cName,
-                                           BridgedSourceLoc cNameLoc) {
+BridgedDeclContextAndDecl
+ClassDecl_create(BridgedASTContext cContext, BridgedDeclContext cDeclContext,
+                 BridgedSourceLoc cClassKeywordLoc, BridgedIdentifier cName,
+                 BridgedSourceLoc cNameLoc, BridgedSourceRange cBraceRange) {
   ASTContext &context = convertASTContext(cContext);
 
   auto *decl = new (context)
       ClassDecl(convertSourceLoc(cClassKeywordLoc), convertIdentifier(cName),
                 convertSourceLoc(cNameLoc), {}, nullptr,
                 convertDeclContext(cDeclContext), false);
+  decl->setBraces(convertSourceRange(cBraceRange));
 
   return {bridgeDeclContext(decl), static_cast<Decl *>(decl)};
 }

lib/ASTGen/Sources/ASTGen/Bridge.swift

@@ -41,6 +41,13 @@ extension BridgedSourceLoc {
   }
 }
 
+extension BridgedSourceRange {
+  @inline(__always)
+  init(startToken: TokenSyntax, endToken: TokenSyntax, in astgen: ASTGenVisitor) {
+    self.init(startLoc: astgen.bridgedSourceLoc(for: startToken), endLoc: astgen.bridgedSourceLoc(for: endToken))
+  }
+}
+
 extension String {
   mutating func withBridgedString<R>(_ body: (BridgedString) throws -> R) rethrows -> R {
     try withUTF8 { buffer in

lib/ASTGen/Sources/ASTGen/Decls.swift

@@ -33,7 +33,8 @@ extension ASTGenVisitor {
       self.bridgedSourceLoc(for: node.structKeyword),
       name,
       nameLoc,
-      self.visit(node.genericParameterClause)?.rawValue
+      self.visit(node.genericParameterClause)?.rawValue,
+      BridgedSourceRange(startToken: node.memberBlock.leftBrace, endToken: node.memberBlock.rightBrace, in: self)
     )
 
     self.withDeclContext(decl.asDeclContext) {
@@ -51,7 +52,8 @@ extension ASTGenVisitor {
       self.declContext,
       self.bridgedSourceLoc(for: node.classKeyword),
       name,
-      nameLoc
+      nameLoc,
+      BridgedSourceRange(startToken: node.memberBlock.leftBrace, endToken: node.memberBlock.rightBrace, in: self)
     )
 
     self.withDeclContext(decl.asDeclContext) {

test/ASTGen/verify-parse.swift

@@ -83,3 +83,18 @@ typealias
 Alias<T>
 =
 String
+
+struct
+Struct
+<
+  T1,
+  T2
+>
+{
+  /*static*/ func method(_ b: Bool) {}
+}
+
+class
+Class {
+  func method(_ b: Bool) {}
+}