Remove unused nodes from SyntaxTreeDeserializer.nodeLookupTable

Sources/SwiftSyntax/RawSyntax.swift

@@ -85,7 +85,7 @@ fileprivate enum RawSyntaxData {
 /// 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.
-struct RawSyntax {
+final class RawSyntax {
   fileprivate let data: RawSyntaxData
   let presence: SourcePresence
 
@@ -124,6 +124,13 @@ struct RawSyntax {
       }
     }).value
   }
+
+  /// Creates a copy of `other`.
+  init(_ other: RawSyntax) {
+    self.data = other.data
+    self.presence = other.presence
+    self.id = other.id
+  }
 
   init(kind: SyntaxKind, layout: [RawSyntax?], presence: SourcePresence,
        id: SyntaxNodeId? = nil) {
@@ -336,7 +343,7 @@ extension RawSyntax {
 
 extension RawSyntax: Codable {
   /// Creates a RawSyntax from the provided Foundation Decoder.
-  init(from decoder: Decoder) throws {
+  convenience 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
@@ -352,7 +359,7 @@ extension RawSyntax: Codable {
       guard let lookupNode = lookupFunc(id) else {
         throw IncrementalDecodingError.nodeLookupFailed(id)
       }
-      self = lookupNode
+      self.init(lookupNode)
       return
     }
 

Sources/SwiftSyntax/SwiftSyntax.swift

@@ -46,6 +46,14 @@ public enum SerializationFormat {
   case byteTree
 }
 
+fileprivate struct WeakReference<T: AnyObject> {
+  weak private(set) var value: T?
+
+  init(_ value: T) {
+    self.value = value
+  }
+}
+
 /// 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
@@ -56,7 +64,14 @@ public final class SyntaxTreeDeserializer {
 
   /// 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] = [:]
+  private var nodeLookupTable: [SyntaxNodeId: WeakReference<RawSyntax>] = [:]
+
+  /// Keep a strong reference to the syntax tree that contains the nodes in the
+  /// `nodeLookupTable`. Because `nodeLookupTable` only holds a weak reference 
+  /// to the RawSyntax nodes, all retired `RawSyntax` nodes will be deallocated
+  /// once we set a new tree. The weak references in `nodeLookupTable` will then
+  /// become `nil` but will also never be accessed again.
+  private var nodeLookupTree: RawSyntax? = nil
 
   /// The IDs of the nodes that were reused as part of incremental syntax
   /// parsing during the last deserialization
@@ -70,19 +85,23 @@ public final class SyntaxTreeDeserializer {
     let decoder = JSONDecoder()
     decoder.userInfo[.rawSyntaxDecodedCallback] = self.addToLookupTable
     decoder.userInfo[.omittedNodeLookupFunction] = self.lookupNode
-    return try decoder.decode(RawSyntax.self, from: data)
+    let tree = try decoder.decode(RawSyntax.self, from: data)
+    self.nodeLookupTree = tree
+    return tree
   }
 
   /// Deserialize the given data as a ByteTree encoded syntax tree
   private func deserializeByteTree(_ data: Data) throws -> RawSyntax {
     var userInfo: [ByteTreeUserInfoKey: Any] = [:]
     userInfo[.rawSyntaxDecodedCallback] = self.addToLookupTable
     userInfo[.omittedNodeLookupFunction] = self.lookupNode
-    return try ByteTreeReader.read(RawSyntax.self, from: data,
+    let tree = try ByteTreeReader.read(RawSyntax.self, from: data,
                                             userInfo: &userInfo) {
       (version: ByteTreeProtocolVersion) in
       return version.major == 1
     }
+    self.nodeLookupTree = tree
+    return tree
   }
 
   /// Decode a serialized form of SourceFileSyntax to a syntax tree.
@@ -112,11 +131,17 @@ public final class SyntaxTreeDeserializer {
 
   private func lookupNode(id: SyntaxNodeId) -> RawSyntax? {
     reusedNodeIds.insert(id)
-    return nodeLookupTable[id]
+    guard let weakRef = nodeLookupTable[id] else {
+      return nil
+    }
+    guard let value = weakRef.value else {
+      fatalError("Trying to retrieve a node that has since been deallocated")
+    }
+    return value
   }
 
   private func addToLookupTable(_ node: RawSyntax) {
-    nodeLookupTable[node.id] = node
+    nodeLookupTable[node.id] = WeakReference(node)
   }
 }