[SwiftSyntax] Add accessors for source locations and test diagnostic emission (#16141)

* [SwiftSyntax] Add accessors for source locations and test diagnostic emission

* Add tests for endLocation

* Pre-emptively copy AbsolutePosition to avoid mutating it twice

Author: harlanhaskins

test/SwiftSyntax/DiagnosticTest.swift

@@ -12,6 +12,14 @@ func loc(_ file: String = #file, line: Int = #line,
   return SourceLocation(line: line, column: column, offset: 0, file: file)
 }
 
+func getInput(_ file: String) -> URL {
+  var result = URL(fileURLWithPath: #file)
+  result.deleteLastPathComponent()
+  result.appendPathComponent("Inputs")
+  result.appendPathComponent(file)
+  return result
+}
+
 /// Adds static constants to Diagnostic.Message.
 extension Diagnostic.Message {
   /// Error thrown when a conversion between two types is impossible.
@@ -24,6 +32,13 @@ extension Diagnostic.Message {
   /// Suggestion for the user to explicitly check a value does not equal zero.
   static let checkEqualToZero =
     Diagnostic.Message(.note, "check for explicit equality to '0'")
+
+  static func badFunction(_ name: TokenSyntax) -> Diagnostic.Message {
+    return .init(.error, "bad function '\(name.text)'")
+  }
+  static func endOfFunction(_ name: TokenSyntax) -> Diagnostic.Message {
+    return .init(.warning, "end of function '\(name.text)'")
+  }
 }
 
 var Diagnostics = TestSuite("Diagnostics")
@@ -58,4 +73,39 @@ Diagnostics.test("DiagnosticEmission") {
   expectEqual(fixIt.text, " != 0")
 }
 
+Diagnostics.test("SourceLocations") {
+  let engine = DiagnosticEngine()
+  engine.addConsumer(PrintingDiagnosticConsumer())
+  let url = getInput("diagnostics.swift")
+
+  class Visitor: SyntaxVisitor {
+    let url: URL
+    let engine: DiagnosticEngine
+    init(url: URL, engine: DiagnosticEngine) {
+      self.url = url
+      self.engine = engine
+    }
+    override func visit(_ function: FunctionDeclSyntax) {
+      let startLoc = function.identifier.startLocation(in: url)
+      let endLoc = function.endLocation(in: url)
+      print("\(function.identifier.text): startLoc: \(startLoc), endLoc: \(endLoc)")
+      engine.diagnose(.badFunction(function.identifier), location: startLoc) {
+        $0.highlight(function.identifier.sourceRange(in: self.url))
+      }
+      engine.diagnose(.endOfFunction(function.identifier), location: endLoc)
+    }
+  }
+
+  expectDoesNotThrow({
+    let file = try SourceFileSyntax.parse(url)
+    Visitor(url: url, engine: engine).visit(file)
+  })
+
+  expectEqual(6, engine.diagnostics.count)
+  let lines = Set(engine.diagnostics.compactMap { $0.location?.line })
+  expectEqual([1, 3, 5, 7, 9, 11], lines)
+  let columns = Set(engine.diagnostics.compactMap { $0.location?.column })
+  expectEqual([6, 2], columns)
+}
+
 runAllTests()

test/SwiftSyntax/Inputs/diagnostics.swift

@@ -0,0 +1,11 @@
+func foo() {
+
+}
+
+func bar() {
+
+}
+
+func baz() {
+
+}

tools/SwiftSyntax/RawSyntax.swift

@@ -251,6 +251,13 @@ extension RawSyntax {
     }
   }
 
+  func accumulateTrailingTrivia(_ pos: AbsolutePosition) {
+    guard let trivia = trailingTrivia else { return }
+    for piece in trivia {
+      piece.accumulateAbsolutePosition(pos)
+    }
+  }
+
   var isSourceFile: Bool {
     switch self {
     case .node(let kind, _, _):

tools/SwiftSyntax/Syntax.swift

@@ -202,6 +202,56 @@ extension Syntax {
     where Target: TextOutputStream {
     data.raw.write(to: &target)
   }
+
+  /// The starting location, in the provided file, of this Syntax node.
+  /// - Parameters:
+  ///   - file: The file URL this node resides in.
+  ///   - afterLeadingTrivia: Whether to skip leading trivia when getting
+  ///                         the node's location. Defaults to `true`.
+  public func startLocation(
+    in file: URL,
+    afterLeadingTrivia: Bool = true
+  ) -> SourceLocation {
+    let pos = afterLeadingTrivia ? 
+      data.position.copy() :
+      data.positionAfterSkippingLeadingTrivia.copy()
+    return SourceLocation(file: file.path, position: pos)
+  }
+
+
+  /// The ending location, in the provided file, of this Syntax node.
+  /// - Parameters:
+  ///   - file: The file URL this node resides in.
+  ///   - afterTrailingTrivia: Whether to skip trailing trivia when getting
+  ///                          the node's location. Defaults to `false`.
+  public func endLocation(
+    in file: URL,
+    afterTrailingTrivia: Bool = false
+  ) -> SourceLocation {
+    let pos = data.position.copy()
+    raw.accumulateAbsolutePosition(pos)
+    if afterTrailingTrivia {
+      raw.accumulateTrailingTrivia(pos)
+    }
+    return SourceLocation(file: file.path, position: pos)
+  }
+
+  /// The source range, in the provided file, of this Syntax node.
+  /// - Parameters:
+  ///   - file: The file URL this node resides in.
+  ///   - afterLeadingTrivia: Whether to skip leading trivia when getting
+  ///                          the node's start location. Defaults to `true`.
+  ///   - afterTrailingTrivia: Whether to skip trailing trivia when getting
+  ///                          the node's end location. Defaults to `false`.
+  public func sourceRange(
+    in file: URL,
+    afterLeadingTrivia: Bool = true,
+    afterTrailingTrivia: Bool = false
+  ) -> SourceRange {
+    let start = startLocation(in: file, afterLeadingTrivia: afterLeadingTrivia)
+    let end = endLocation(in: file, afterTrailingTrivia: afterTrailingTrivia)
+    return SourceRange(start: start, end: end)
+  }
 }
 
 /// Determines if two nodes are equal to each other.