Add incremental syntax tree deserialization to SwiftSyntax

swiftlang · ahoppen · Sep 12, 2018 · Apr 10, 2018 · Apr 24, 2018 · Apr 24, 2018
commit 08fb8a0e966f70509c91e980f4dc498f24f56b1a
diff --git a/Sources/SwiftSyntax/RawSyntax.swift b/Sources/SwiftSyntax/RawSyntax.swift
@@ -12,45 +12,84 @@
 
 import Foundation
 
+extension CodingUserInfoKey {
+  /// Callback that will be called whenever a `RawSyntax` node is decoded
+  /// Value must have signature `(RawSyntax) -> Void`
+  static let rawSyntaxDecodedCallback =
+    CodingUserInfoKey(rawValue: "SwiftSyntax.RawSyntax.DecodedCallback")!
+  /// Function that shall be used to look up nodes that were omitted in the
+  /// syntax tree transfer.
+  /// Value must have signature `(SyntaxNodeId) -> RawSyntax`
+  static let omittedNodeLookupFunction =
+    CodingUserInfoKey(rawValue: "SwiftSyntax.RawSyntax.OmittedNodeLookup")!
+}
+
+/// A ID that uniquely identifies a syntax node and stays stable across multiple
+/// incremental parses
+struct SyntaxNodeId: Hashable, Codable {
+  private let rawValue: UInt
+
+  fileprivate init(rawValue: UInt) {
+    self.rawValue = rawValue
+  }
+
+  init(from decoder: Decoder) throws {
+    self.rawValue = try decoder.singleValueContainer().decode(UInt.self)
+  }
+
+  func encode(to encoder: Encoder) throws {
+    var container = encoder.singleValueContainer()
+    try container.encode(rawValue)
+  }
+}
+
 /// Represents the raw tree structure underlying the syntax tree. These nodes
 /// have no notion of identity and only provide structure to the tree. They
 /// are immutable and can be freely shared between syntax nodes.
 indirect enum RawSyntax: Codable {
   /// A tree node with a kind, an array of children, and a source presence.
-  case node(SyntaxKind, [RawSyntax?], SourcePresence)
+  case node(SyntaxKind, [RawSyntax?], SourcePresence, SyntaxNodeId?)
 
   /// A token with a token kind, leading trivia, trailing trivia, and a source
   /// presence.
-  case token(TokenKind, Trivia, Trivia, SourcePresence)
+  case token(TokenKind, Trivia, Trivia, SourcePresence, SyntaxNodeId?)
 
   /// The syntax kind of this raw syntax.
   var kind: SyntaxKind {
     switch self {
-    case .node(let kind, _, _): return kind
-    case .token(_, _, _, _): return .token
+    case .node(let kind, _, _, _): return kind
+    case .token(_, _, _, _, _): return .token
     }
   }
 
   var tokenKind: TokenKind? {
     switch self {
-    case .node(_, _, _): return nil
-    case .token(let kind, _, _, _): return kind
+    case .node(_, _, _, _): return nil
+    case .token(let kind, _, _, _, _): return kind
     }
   }
 
   /// The layout of the children of this Raw syntax node.
   var layout: [RawSyntax?] {
     switch self {
-    case .node(_, let layout, _): return layout
-    case .token(_, _, _, _): return []
+    case .node(_, let layout, _, _): return layout
+    case .token(_, _, _, _, _): return []
     }
   }
 
   /// The source presence of this node.
   var presence: SourcePresence {
     switch self {
-    case .node(_, _, let presence): return presence
-    case .token(_, _, _, let presence): return presence
+    case .node(_, _, let presence, _): return presence
+    case .token(_, _, _, let presence, _): return presence
+    }
+  }
+
+  /// The ID of this node
+  var id: SyntaxNodeId? {
+    switch self {
+    case .node(_, _, _, let id): return id
+    case .token(_, _, _, _, let id): return id
     }
   }
 
@@ -66,37 +105,81 @@ indirect enum RawSyntax: Codable {
 
   /// Keys for serializing RawSyntax nodes.
   enum CodingKeys: String, CodingKey {
+    // Shared keys
+    case id, omitted
+
     // Keys for the `node` case
     case kind, layout, presence
 
     // Keys for the `token` case
     case tokenKind, leadingTrivia, trailingTrivia
   }
 
+  public enum IncrementalDecodingError: Error {
+    /// The JSON to decode is invalid because a node had the `omitted` flag set
+    /// but included no id
+    case omittedNodeHasNoId
+    /// An omitted node was encountered but no node lookup function was
+    /// in the decoder's `userInfo` for the `.omittedNodeLookupFunction` key
+    /// or the lookup function had the wrong type
+    case noLookupFunctionPassed
+    /// A node with the given ID was omitted from the tranferred syntax tree
+    /// but the lookup function was unable to look it up
+    case nodeLookupFailed(SyntaxNodeId)
+  }
+
   /// Creates a RawSyntax from the provided Foundation Decoder.
   init(from decoder: Decoder) throws {
     let container = try decoder.container(keyedBy: CodingKeys.self)
+    let id = try container.decodeIfPresent(SyntaxNodeId.self, forKey: .id)
+    let omitted = try container.decodeIfPresent(Bool.self, forKey: .omitted) ?? false
+
+    if omitted {
+      guard let id = id else {
+        throw IncrementalDecodingError.omittedNodeHasNoId
+      }
+      guard let lookupFunc = decoder.userInfo[.omittedNodeLookupFunction] as?
+                               (SyntaxNodeId) -> RawSyntax? else {
+        throw IncrementalDecodingError.noLookupFunctionPassed
+      }
+      guard let lookupNode = lookupFunc(id) else {
+        throw IncrementalDecodingError.nodeLookupFailed(id)
+      }
+      self = lookupNode
+      return
+    }
+
     let presence = try container.decode(SourcePresence.self, forKey: .presence)
     if let kind = try container.decodeIfPresent(SyntaxKind.self, forKey: .kind) {
       let layout = try container.decode([RawSyntax?].self, forKey: .layout)
-      self = .node(kind, layout, presence)
+      self = .node(kind, layout, presence, id)
     } else {
       let kind = try container.decode(TokenKind.self, forKey: .tokenKind)
       let leadingTrivia = try container.decode(Trivia.self, forKey: .leadingTrivia)
       let trailingTrivia = try container.decode(Trivia.self, forKey: .trailingTrivia)
-      self = .token(kind, leadingTrivia, trailingTrivia, presence)
+      self = .token(kind, leadingTrivia, trailingTrivia, presence, id)
+    }
+    if let callback = decoder.userInfo[.rawSyntaxDecodedCallback] as?
+                        (RawSyntax) -> Void {
+      callback(self)
     }
   }
 
   /// Encodes the RawSyntax to the provided Foundation Encoder.
   func encode(to encoder: Encoder) throws {
     var container = encoder.container(keyedBy: CodingKeys.self)
     switch self {
-    case let .node(kind, layout, presence):
+    case let .node(kind, layout, presence, id):
+      if let id = id {
+        try container.encode(id, forKey: .id)
+      }
       try container.encode(kind, forKey: .kind)
       try container.encode(layout, forKey: .layout)
       try container.encode(presence, forKey: .presence)
-    case let .token(kind, leadingTrivia, trailingTrivia, presence):
+    case let .token(kind, leadingTrivia, trailingTrivia, presence, id):
+      if let id = id {
+        try container.encode(id, forKey: .id)
+      }
       try container.encode(kind, forKey: .tokenKind)
       try container.encode(leadingTrivia, forKey: .leadingTrivia)
       try container.encode(trailingTrivia, forKey: .trailingTrivia)
@@ -112,7 +195,7 @@ indirect enum RawSyntax: Codable {
   /// - Returns: A new RawSyntax `.node` with the provided kind and layout, with
   ///            `.missing` source presence.
   static func missing(_ kind: SyntaxKind) -> RawSyntax {
-    return .node(kind, [], .missing)
+    return .node(kind, [], .missing, nil)
   }
 
   /// Creates a RawSyntax token that's marked missing in the source with the
@@ -121,7 +204,7 @@ indirect enum RawSyntax: Codable {
   /// - Returns: A new RawSyntax `.token` with the provided kind, no
   ///            leading/trailing trivia, and `.missing` source presence.
   static func missingToken(_ kind: TokenKind) -> RawSyntax {
-    return .token(kind, [], [], .missing)
+    return .token(kind, [], [], .missing, nil)
   }
 
   /// Returns a new RawSyntax node with the provided layout instead of the
@@ -131,8 +214,9 @@ indirect enum RawSyntax: Codable {
   /// - Parameter newLayout: The children of the new node you're creating.
   func replacingLayout(_ newLayout: [RawSyntax?]) -> RawSyntax {
     switch self {
-    case let .node(kind, _, presence): return .node(kind, newLayout, presence)
-    case .token(_, _, _, _): return self
+    case let .node(kind, _, presence, _):
+      return .node(kind, newLayout, presence, nil)
+    case .token(_, _, _, _, _): return self
     }
   }
 
@@ -176,12 +260,12 @@ extension RawSyntax: TextOutputStreamable {
   func write<Target>(to target: inout Target)
     where Target: TextOutputStream {
     switch self {
-    case .node(_, let layout, _):
+    case .node(_, let layout, _, _):
       for child in layout {
         guard let child = child else { continue }
         child.write(to: &target)
       }
-    case let .token(kind, leadingTrivia, trailingTrivia, presence):
+    case let .token(kind, leadingTrivia, trailingTrivia, presence, _):
       guard case .present = presence else { return }
       for piece in leadingTrivia {
         piece.write(to: &target)
@@ -197,12 +281,12 @@ extension RawSyntax: TextOutputStreamable {
 extension RawSyntax {
   func accumulateAbsolutePosition(_ pos: AbsolutePosition) {
     switch self {
-    case .node(_, let layout, _):
+    case .node(_, let layout, _, _):
       for child in layout {
         guard let child = child else { continue }
         child.accumulateAbsolutePosition(pos)
       }
-    case let .token(kind, leadingTrivia, trailingTrivia, presence):
+    case let .token(kind, leadingTrivia, trailingTrivia, presence, _):
       guard case .present = presence else { return }
       for piece in leadingTrivia {
         piece.accumulateAbsolutePosition(pos)
@@ -216,29 +300,29 @@ extension RawSyntax {
 
   var leadingTrivia: Trivia? {
     switch self {
-    case .node(_, let layout, _):
+    case .node(_, let layout, _, _):
       for child in layout {
         guard let child = child else { continue }
         guard let result = child.leadingTrivia else { continue }
         return result
       }
       return nil
-    case let .token(_, leadingTrivia, _, presence):
+    case let .token(_, leadingTrivia, _, presence, _):
       guard case .present = presence else { return nil }
       return leadingTrivia
     }
   }
 
   var trailingTrivia: Trivia? {
     switch self {
-    case .node(_, let layout, _):
+    case .node(_, let layout, _, _):
       for child in layout.reversed() {
         guard let child = child else { continue }
         guard let result = child.trailingTrivia else { continue }
         return result
       }
       return nil
-    case let .token(_, _, trailingTrivia, presence):
+    case let .token(_, _, trailingTrivia, presence, _):
       guard case .present = presence else { return nil }
       return trailingTrivia
     }

diff --git a/Sources/SwiftSyntax/SwiftSyntax.swift b/Sources/SwiftSyntax/SwiftSyntax.swift
@@ -41,51 +41,66 @@ public enum ParserError: Error, CustomStringConvertible {
   }
 }
 
-extension Syntax {
-  fileprivate static func encodeSourceFileSyntaxInternal(_ url: URL) throws -> Data {
-    let swiftcRunner = try SwiftcRunner(sourceFile: url)
-    let result = try swiftcRunner.invoke()
-    return result.stdoutData
-  }
+/// Deserializes the syntax tree from its serialized form to an object tree in
+/// Swift. To deserialize incrementally transferred syntax trees, the same
+/// instance of the deserializer must be used for all subsequent
+/// deserializations.
+public final class SyntaxTreeDeserializer {
+  // FIXME: This lookup table just accumulates nodes, we should invalidate nodes
+  // that are no longer used at some point and remove them from the table
+  /// Syntax nodes that have already been parsed and are able to be reused if
+  /// they were omitted in an incremental syntax tree transfer
+  private var nodeLookupTable: [SyntaxNodeId: RawSyntax] = [:]
 
-  /// Parses the Swift file at the provided URL into a `Syntax` tree in Json
-  /// serialization format.
-  /// - Parameter url: The URL you wish to parse.
-  /// - Returns: The syntax tree in Json format string.
-  public static func encodeSourceFileSyntax(_ url: URL) throws -> String {
-    return String(data: try encodeSourceFileSyntaxInternal(url), encoding: .utf8)!
-  }
-
-  /// Parses the Swift file at the provided URL into a full-fidelity `Syntax`
-  /// tree.
-  /// - Parameter url: The URL you wish to parse.
-  /// - Returns: A top-level Syntax node representing the contents of the tree,
-  ///            if the parse was successful.
-  /// - Throws: `ParseError.couldNotFindSwiftc` if `swiftc` could not be
-  ///           located, `ParseError.invalidFile` if the file is invalid.
-  ///           FIXME: Fill this out with all error cases.
-  public static func parse(_ url: URL) throws -> SourceFileSyntax {
-    return try decodeSourceFileSyntax(encodeSourceFileSyntaxInternal(url))
+  public init() {
   }
 
-  /// Decode a serialized form of SourceFileSyntax to a syntax node.
-  /// - Parameter content: The data of the serialized SourceFileSyntax.
+  /// Decode a serialized form of SourceFileSyntax to a syntax tree.
+  /// - Parameter data: The UTF-8 represenation of the serialized syntax tree
   /// - Returns: A top-level Syntax node representing the contents of the tree,
   ///            if the parse was successful.
-  fileprivate static func decodeSourceFileSyntax(_ content: Data) throws -> SourceFileSyntax {
+  public func deserialize(_ data: Data) throws -> SourceFileSyntax {
     let decoder = JSONDecoder()
-    let raw = try decoder.decode(RawSyntax.self, from: content)
+    decoder.userInfo[.rawSyntaxDecodedCallback] = self.addToLookupTable
+    decoder.userInfo[.omittedNodeLookupFunction] = self.lookupNode
+    let raw = try decoder.decode(RawSyntax.self, from: data)
     guard let file = makeSyntax(raw) as? SourceFileSyntax else {
       throw ParserError.invalidFile
     }
     return file
   }
 
-  /// Decode a serialized form of SourceFileSyntax to a syntax node.
-  /// - Parameter content: The string content of the serialized SourceFileSyntax.
+  // MARK: Incremental deserialization helper functions
+  private func lookupNode(id: SyntaxNodeId) -> RawSyntax? {
+    return nodeLookupTable[id]
+  }
+
+  private func addToLookupTable(_ node: RawSyntax) {
+    guard let id = node.id else {
+      return
+    }
+    nodeLookupTable[id] = node
+  }
+}
+
+/// Namespace for functions to retrieve a syntax tree from the swift compiler
+/// and deserializing it.
+public enum SyntaxTreeParser {
+  /// Parses the Swift file at the provided URL into a full-fidelity Syntax tree
+  /// - Parameter url: The URL you wish to parse.
   /// - Returns: A top-level Syntax node representing the contents of the tree,
   ///            if the parse was successful.
-  public static func decodeSourceFileSyntax(_ content: String) throws -> SourceFileSyntax {
-    return try decodeSourceFileSyntax(content.data(using: .utf8)!)
+  /// - Throws: `ParseError.couldNotFindSwiftc` if `swiftc` could not be
+  ///           located, `ParseError.invalidFile` if the file is invalid.
+  ///           FIXME: Fill this out with all error cases.
+  public static func parse(_ url: URL) throws -> SourceFileSyntax {
+    let swiftcRunner = try SwiftcRunner(sourceFile: url)
+    let result = try swiftcRunner.invoke()
+    guard result.wasSuccessful else {
+      throw ParserError.swiftcFailed(result.exitCode, result.stderr)
+    }
+    let syntaxTreeData = result.stdoutData
+    let deserializer = SyntaxTreeDeserializer()
+    return try deserializer.deserialize(syntaxTreeData)
   }
 }