Merge pull request swiftlang#91 from hartbit/syntax-trivia

Helper methods to modify leading and trailing trivia on nodes

Author: nkcsgexi

Sources/SwiftSyntax/RawSyntax.swift

@@ -1099,6 +1099,44 @@ final class RawSyntax: ManagedBuffer<RawSyntaxBase, RawSyntaxDataElement> {
     }
   }
 
+  func withLeadingTrivia(_ leadingTrivia: Trivia) -> RawSyntax {
+    if isToken {
+      return RawSyntax.createAndCalcLength(
+        kind: formTokenKind()!,
+        leadingTrivia: leadingTrivia,
+        trailingTrivia: formTrailingTrivia()!,
+        presence: presence)
+    } else {
+      var layout = formLayoutArray()
+      for (index, raw) in layout.enumerated() {
+        if let raw = raw {
+          layout[index] = raw.withLeadingTrivia(leadingTrivia)
+          return replacingLayout(layout)
+        }
+      }
+      return self
+    }
+  }
+
+  func withTrailingTrivia(_ trailingTrivia: Trivia) -> RawSyntax {
+    if isToken {
+      return RawSyntax.createAndCalcLength(
+        kind: formTokenKind()!,
+        leadingTrivia: formLeadingTrivia()!,
+        trailingTrivia: trailingTrivia,
+        presence: presence)
+    } else {
+      var layout = formLayoutArray()
+      for (index, raw) in layout.enumerated().reversed() {
+        if let raw = raw {
+          layout[index] = raw.withTrailingTrivia(trailingTrivia)
+          return replacingLayout(layout)
+        }
+      }
+      return self
+    }
+  }
+
   /// Creates a RawSyntax node that's marked missing in the source with the
   /// provided kind and layout.
   /// - Parameters:

Sources/SwiftSyntax/Syntax.swift

@@ -630,11 +630,7 @@ public struct TokenSyntax: _SyntaxBase, Hashable {
     guard raw.kind == .token else {
       fatalError("TokenSyntax must have token as its raw")
     }
-    let newRaw = RawSyntax.createAndCalcLength(kind: raw.formTokenKind()!,
-      leadingTrivia: leadingTrivia, trailingTrivia: raw.formTrailingTrivia()!,
-      presence: raw.presence)
-    let newData = data.replacingSelf(newRaw)
-    return TokenSyntax(newData)
+    return TokenSyntax(data.withLeadingTrivia(leadingTrivia))
   }
 
   /// Returns a new TokenSyntax with its trailing trivia replaced
@@ -643,12 +639,7 @@ public struct TokenSyntax: _SyntaxBase, Hashable {
     guard raw.kind == .token else {
       fatalError("TokenSyntax must have token as its raw")
     }
-    let newRaw = RawSyntax.createAndCalcLength(kind: raw.formTokenKind()!,
-                           leadingTrivia: raw.formLeadingTrivia()!,
-                           trailingTrivia: trailingTrivia,
-                           presence: raw.presence)
-    let newData = data.replacingSelf(newRaw)
-    return TokenSyntax(newData)
+    return TokenSyntax(data.withTrailingTrivia(trailingTrivia))
   }
 
   /// Returns a new TokenSyntax with its leading trivia removed.

Sources/SwiftSyntax/SyntaxCollections.swift.gyb

@@ -144,6 +144,53 @@ public struct ${node.name}: _SyntaxBase, Hashable, SyntaxCollection {
     return replacingLayout(newLayout)
   }
 
+  /// Returns a new `${node.name}` with its leading trivia replaced
+  /// by the provided trivia.
+  public func withLeadingTrivia(_ leadingTrivia: Trivia) -> ${node.name} {
+    return ${node.name}(data.withLeadingTrivia(leadingTrivia))
+  }
+
+  /// Returns a new `${node.name}` with its trailing trivia replaced
+  /// by the provided trivia.
+  public func withTrailingTrivia(_ trailingTrivia: Trivia) -> ${node.name} {
+    return ${node.name}(data.withTrailingTrivia(trailingTrivia))
+  }
+
+  /// Returns a new `${node.name}` with its leading trivia removed.
+  public func withoutLeadingTrivia() -> ${node.name} {
+    return withLeadingTrivia([])
+  }
+
+  /// Returns a new `${node.name}` with its trailing trivia removed.
+  public func withoutTrailingTrivia() -> ${node.name} {
+    return withTrailingTrivia([])
+  }
+
+  /// Returns a new `${node.name}` with all trivia removed.
+  public func withoutTrivia() -> ${node.name} {
+    return withoutLeadingTrivia().withoutTrailingTrivia()
+  }
+
+  /// The leading trivia (spaces, newlines, etc.) associated with this `${node.name}`.
+  public var leadingTrivia: Trivia? {
+    get {
+      return raw.formLeadingTrivia()
+    }
+    set {
+      self = withLeadingTrivia(newValue ?? [])
+    }
+  }
+
+  /// The trailing trivia (spaces, newlines, etc.) associated with this `${node.name}`.
+  public var trailingTrivia: Trivia? {
+    get {
+      return raw.formTrailingTrivia()
+    }
+    set {
+      self = withTrailingTrivia(newValue ?? [])
+    }
+  }
+
   /// Determines if two `${node.name}` nodes are equal to each other.
   public static func ==(lhs: ${node.name}, rhs: ${node.name}) -> Bool {
     return lhs.data.nodeId == rhs.data.nodeId

Sources/SwiftSyntax/SyntaxData.swift

@@ -291,4 +291,12 @@ struct SyntaxData {
     where CursorType.RawValue == Int {
     return replacingChild(child, at: cursor.rawValue)
   }
+
+  func withLeadingTrivia(_ leadingTrivia: Trivia) -> SyntaxData {
+    return replacingSelf(raw.withLeadingTrivia(leadingTrivia))
+  }
+
+  func withTrailingTrivia(_ trailingTrivia: Trivia) -> SyntaxData {
+    return replacingSelf(raw.withTrailingTrivia(trailingTrivia))
+  }
 }

Sources/SwiftSyntax/SyntaxNodes.swift.gyb

@@ -165,6 +165,53 @@ public struct ${node.name}: ${base_type}, _SyntaxBase, Hashable {
   }
 %     end
 
+  /// Returns a new `${node.name}` with its leading trivia replaced
+  /// by the provided trivia.
+  public func withLeadingTrivia(_ leadingTrivia: Trivia) -> ${node.name} {
+    return ${node.name}(data.withLeadingTrivia(leadingTrivia))
+  }
+
+  /// Returns a new `${node.name}` with its trailing trivia replaced
+  /// by the provided trivia.
+  public func withTrailingTrivia(_ trailingTrivia: Trivia) -> ${node.name} {
+    return ${node.name}(data.withTrailingTrivia(trailingTrivia))
+  }
+
+  /// Returns a new `${node.name}` with its leading trivia removed.
+  public func withoutLeadingTrivia() -> ${node.name} {
+    return withLeadingTrivia([])
+  }
+
+  /// Returns a new `${node.name}` with its trailing trivia removed.
+  public func withoutTrailingTrivia() -> ${node.name} {
+    return withTrailingTrivia([])
+  }
+
+  /// Returns a new `${node.name}` with all trivia removed.
+  public func withoutTrivia() -> ${node.name} {
+    return withoutLeadingTrivia().withoutTrailingTrivia()
+  }
+
+  /// The leading trivia (spaces, newlines, etc.) associated with this `${node.name}`.
+  public var leadingTrivia: Trivia? {
+    get {
+      return raw.formLeadingTrivia()
+    }
+    set {
+      self = withLeadingTrivia(newValue ?? [])
+    }
+  }
+
+  /// The trailing trivia (spaces, newlines, etc.) associated with this `${node.name}`.
+  public var trailingTrivia: Trivia? {
+    get {
+      return raw.formTrailingTrivia()
+    }
+    set {
+      self = withTrailingTrivia(newValue ?? [])
+    }
+  }
+
   /// Determines if two `${node.name}` nodes are equal to each other.
   public static func ==(lhs: ${node.name}, rhs: ${node.name}) -> Bool {
     return lhs.data.nodeId == rhs.data.nodeId

Tests/SwiftSyntaxTest/AbsolutePosition.swift

@@ -88,24 +88,26 @@ public class AbsolutePositionTestCase: XCTestCase {
     _ = root.statements[idx].positionAfterSkippingLeadingTrivia
   }
 
+  static let leadingTrivia = Trivia(pieces: [
+    .newlines(1),
+    .backticks(1),
+    .docLineComment("/// some comment"),
+    .carriageReturns(1),
+  ])
+
+  static let trailingTrivia = Trivia(pieces: [
+    .blockComment("/* This is comment \r\r\n */"),
+    .carriageReturnLineFeeds(1),
+  ])
+
   func createSourceFile(_ count: Int) -> SourceFileSyntax {
-    let leading = Trivia(pieces: [
-        .newlines(1),
-        .backticks(1),
-        .docLineComment("/// some comment"),
-        .carriageReturns(1),
-        ])
-    let trailing = Trivia(pieces: [
-        .blockComment("/* This is comment \r\r\n */"),
-        .carriageReturnLineFeeds(1),
-        ])
     let items : [CodeBlockItemSyntax] =
       [CodeBlockItemSyntax](repeating: CodeBlockItemSyntax {
         $0.useItem(ReturnStmtSyntax {
           $0.useReturnKeyword(
             SyntaxFactory.makeReturnKeyword(
-              leadingTrivia: leading,
-              trailingTrivia: trailing))
+              leadingTrivia: AbsolutePositionTestCase.leadingTrivia,
+              trailingTrivia: AbsolutePositionTestCase.trailingTrivia))
         })}, count: count)
     return SyntaxFactory.makeSourceFile(
       statements: SyntaxFactory.makeCodeBlockItemList(items),
@@ -124,6 +126,40 @@ public class AbsolutePositionTestCase: XCTestCase {
       state.leadingTrivia!.byteSize + state.trailingTrivia!.byteSize
         + state.byteSizeAfterTrimmingTrivia)
     XCTAssertFalse(root.statements.isImplicit)
+
+    // Test Node trivia setters and getters
+
+    XCTAssertEqual(AbsolutePositionTestCase.leadingTrivia, root.leadingTrivia)
+    XCTAssertEqual([], root.trailingTrivia)
+
+    var modifiedRoot1 = root.withLeadingTrivia([.spaces(6), .tabs(1)])
+    XCTAssertEqual([.spaces(6), .tabs(1)], modifiedRoot1.leadingTrivia)
+    XCTAssertEqual(AbsolutePositionTestCase.leadingTrivia, root.leadingTrivia)
+    modifiedRoot1.leadingTrivia = [.blockComment("/* this is a comment */")]
+    XCTAssertEqual([.blockComment("/* this is a comment */")], modifiedRoot1.leadingTrivia)
+
+    var modifiedRoot2 = root.withTrailingTrivia([.backticks(2)])
+    XCTAssertEqual([.backticks(2)], modifiedRoot2.trailingTrivia)
+    XCTAssertEqual([], root.trailingTrivia)
+    modifiedRoot2.trailingTrivia = [.carriageReturns(1), .newlines(2)]
+    XCTAssertEqual([.carriageReturns(1), .newlines(2)], modifiedRoot2.trailingTrivia)
+
+    // Test Collection trivia setters and getters
+
+    XCTAssertEqual(AbsolutePositionTestCase.leadingTrivia, root.statements.leadingTrivia)
+    XCTAssertEqual(AbsolutePositionTestCase.trailingTrivia, root.statements.trailingTrivia)
+
+    var modifiedStatements1 = root.withLeadingTrivia([.carriageReturnLineFeeds(3)])
+    XCTAssertEqual([.carriageReturnLineFeeds(3)], modifiedStatements1.leadingTrivia)
+    XCTAssertEqual(AbsolutePositionTestCase.leadingTrivia, root.statements.leadingTrivia)
+    modifiedStatements1.leadingTrivia = [.garbageText("GARBAGE")]
+    XCTAssertEqual([.garbageText("GARBAGE")], modifiedStatements1.leadingTrivia)
+
+    var modifiedStatements2 = root.withTrailingTrivia([.formfeeds(1), .carriageReturns(3)])
+    XCTAssertEqual([.formfeeds(1), .carriageReturns(3)], modifiedStatements2.trailingTrivia)
+    XCTAssertEqual(AbsolutePositionTestCase.trailingTrivia, root.statements.trailingTrivia)
+    modifiedStatements2.trailingTrivia = [.verticalTabs(4)]
+    XCTAssertEqual([.verticalTabs(4)], modifiedStatements2.trailingTrivia)
   }
 
   public func testImplicit() {