Validate a node's layout after it has been rewritten using the syntax rewriter

Otherwise it is possible to generate invalid syntax trees using the syntax
rewriter.

Author: ahoppen

Sources/SwiftSyntax/Syntax.swift

@@ -24,6 +24,11 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
     self.data = data
   }
 
+  public func _validateLayout() {
+    // Check the layout of the concrete type
+    return self.as(SyntaxProtocol.self)._validateLayout()
+  }
+
   /// Create a `Syntax` node from a specialized syntax node.
   public init<S: SyntaxProtocol>(_ syntax: S) {
     self = syntax._syntaxNode
@@ -85,6 +90,11 @@ public protocol SyntaxProtocol: CustomStringConvertible,
   /// Converts the given `Syntax` node to this type. Returns `nil` if the
   /// conversion is not possible.
   init?(_ syntaxNode: Syntax)
+
+  /// Check that the raw layout of this node is valid. Used to verify a node's
+  /// integrity after it has been rewritten by the syntax rewriter.
+  /// Results in an assertion failure if the layout is invalid.
+  func _validateLayout()
 }
 
 internal extension SyntaxProtocol {

Sources/SwiftSyntax/SyntaxBaseNodes.swift.gyb

@@ -90,6 +90,11 @@ public struct ${node.name}: ${node.name}Protocol, SyntaxHashable {
     self._syntaxNode = Syntax(data)
   }
 
+  public func _validateLayout() {
+    // Check the layout of the concrete type
+    return Syntax(self)._validateLayout()
+  }
+
   public func `is`<S: ${node.name}Protocol>(_ syntaxType: S.Type) -> Bool {
     return self.as(syntaxType) != nil
   }

Sources/SwiftSyntax/SyntaxChildren.swift

@@ -27,7 +27,7 @@ struct RawSyntaxChildren: Sequence {
       self.nextChildInfo = startFrom
     }
 
-    mutating func next() -> (RawSyntax?, AbsoluteSyntaxInfo)? {
+    mutating func next() -> (raw: RawSyntax?, syntaxInfo: AbsoluteSyntaxInfo)? {
       let idx = Int(nextChildInfo.indexInParent)
       guard idx < parent.numberOfChildren else {
         return nil

Sources/SwiftSyntax/SyntaxCollections.swift.gyb

@@ -200,6 +200,13 @@ public struct ${node.name}: SyntaxCollection, SyntaxHashable {
       self = withTrailingTrivia(newValue ?? [])
     }
   }
+
+  public func _validateLayout() {
+    // Check that all children match the expected element type
+    assert(self.allSatisfy { node in
+      return Syntax(node).is(${node.collection_element_type}.self)
+    })
+  }
 }
 
 /// Conformance for `${node.name}`` to the Sequence protocol.

Sources/SwiftSyntax/SyntaxNodes.swift.gyb.template

@@ -155,6 +155,33 @@ public struct ${node.name}: ${base_type}Protocol, SyntaxHashable {
     return ${node.name}(newData)
   }
 %     end
+
+%     # ===================
+%     # Validate layout
+%     # ===================
+
+  public func _validateLayout() {
+%     if node.is_unknown():
+    // We are verifying an unknown node. Since we don’t know anything about it
+    // we need to assume it’s valid.
+%     else:
+    let rawChildren = Array(RawSyntaxChildren(Syntax(self)))
+    assert(rawChildren.count == ${len(node.children)})
+%       for i, child in enumerate(node.children):
+    // Check ${i}-th child is ${child.type_name} ${"or missing" if child.is_optional else ""}
+%         if not child.is_optional:
+    assert(rawChildren[${i}].raw != nil)
+%         end
+    if let raw = rawChildren[${i}].raw {
+      let info = rawChildren[${i}].syntaxInfo
+      let absoluteRaw = AbsoluteRawSyntax(raw: raw, info: info)
+      let syntaxData = SyntaxData(absoluteRaw, parent: Syntax(self))
+      let syntaxChild = Syntax(syntaxData)
+      assert(syntaxChild.is(${child.type_name}.self))
+    }
+%       end
+%     end
+  }
 }
 
 extension ${node.name}: CustomReflectable {

Sources/SwiftSyntax/SyntaxOtherNodes.swift

@@ -25,6 +25,11 @@ public struct UnknownSyntax: SyntaxProtocol, SyntaxHashable {
     self._syntaxNode = syntax
   }
 
+  public func _validateLayout() {
+    // We are verifying an unknown node. Since we don’t know anything about it
+    // we need to assume it’s valid.
+  }
+
   /// Creates an `UnknownSyntax` node from the given `SyntaxData`. This assumes
   /// that the `SyntaxData` is of the correct kind. If it is not, the behaviour
   /// is undefined.
@@ -61,6 +66,10 @@ public struct TokenSyntax: SyntaxProtocol, SyntaxHashable {
     self._syntaxNode = Syntax(data)
   }
 
+  public func _validateLayout() {
+    /// A token is always valid as it has no children. Nothing to do here.
+  }
+
   public var presence: SourcePresence {
     return raw.presence
   }

Sources/SwiftSyntax/SyntaxRewriter.swift.gyb

@@ -203,7 +203,13 @@ open class SyntaxRewriter {
       assert(newLayout.count == node.raw.numberOfChildren)
 
       let newRaw = node.raw.replacingLayout(Array(newLayout))
-      return SyntaxType(Syntax(SyntaxData.forRoot(newRaw)))!
+      let newNode = SyntaxType(Syntax(SyntaxData.forRoot(newRaw)))!
+      assert({
+        // In assertion builds inovoke the _validateLayout method
+        newNode._validateLayout()
+        return true
+      }())
+      return newNode
     } else {
       // No child node was rewritten. So no need to change this node as well.
       return node

Sources/SwiftSyntax/gyb_generated/SyntaxBaseNodes.swift

@@ -71,6 +71,11 @@ public struct DeclSyntax: DeclSyntaxProtocol, SyntaxHashable {
     self._syntaxNode = Syntax(data)
   }
 
+  public func _validateLayout() {
+    // Check the layout of the concrete type
+    return Syntax(self)._validateLayout()
+  }
+
   public func `is`<S: DeclSyntaxProtocol>(_ syntaxType: S.Type) -> Bool {
     return self.as(syntaxType) != nil
   }
@@ -159,6 +164,11 @@ public struct ExprSyntax: ExprSyntaxProtocol, SyntaxHashable {
     self._syntaxNode = Syntax(data)
   }
 
+  public func _validateLayout() {
+    // Check the layout of the concrete type
+    return Syntax(self)._validateLayout()
+  }
+
   public func `is`<S: ExprSyntaxProtocol>(_ syntaxType: S.Type) -> Bool {
     return self.as(syntaxType) != nil
   }
@@ -247,6 +257,11 @@ public struct StmtSyntax: StmtSyntaxProtocol, SyntaxHashable {
     self._syntaxNode = Syntax(data)
   }
 
+  public func _validateLayout() {
+    // Check the layout of the concrete type
+    return Syntax(self)._validateLayout()
+  }
+
   public func `is`<S: StmtSyntaxProtocol>(_ syntaxType: S.Type) -> Bool {
     return self.as(syntaxType) != nil
   }
@@ -335,6 +350,11 @@ public struct TypeSyntax: TypeSyntaxProtocol, SyntaxHashable {
     self._syntaxNode = Syntax(data)
   }
 
+  public func _validateLayout() {
+    // Check the layout of the concrete type
+    return Syntax(self)._validateLayout()
+  }
+
   public func `is`<S: TypeSyntaxProtocol>(_ syntaxType: S.Type) -> Bool {
     return self.as(syntaxType) != nil
   }
@@ -423,6 +443,11 @@ public struct PatternSyntax: PatternSyntaxProtocol, SyntaxHashable {
     self._syntaxNode = Syntax(data)
   }
 
+  public func _validateLayout() {
+    // Check the layout of the concrete type
+    return Syntax(self)._validateLayout()
+  }
+
   public func `is`<S: PatternSyntaxProtocol>(_ syntaxType: S.Type) -> Bool {
     return self.as(syntaxType) != nil
   }