SyntaxTreeDeserializer: Hold on to omitted nodes across multiple updates

RawSyntax.init(from decoder: Decoder) looks up omitted nodes by id from a
WeakLookupTable, initializing itself as a copy. This copy is then used in the
new tree, but wasn't being written back to the lookup table, leaving it with a
weak reference to the original node. This was an issue because the
SyntaxTreeDeserializer only holds a strong reference to the most recent syntax
tree, meaning the orignal node (and table entry) could have been cleaned up,
even though further incremental updates could still refer to that id.

Resolves rdar://problem/44877860

Author: Nathan Hawes

Sources/SwiftSyntax/RawSyntax.swift

@@ -360,19 +360,18 @@ extension RawSyntax: Codable {
         throw IncrementalDecodingError.nodeLookupFailed(id)
       }
       self.init(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.init(kind: kind, layout: layout, presence: presence, id: 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.init(kind: kind, leadingTrivia: leadingTrivia,
-                trailingTrivia: trailingTrivia, presence: presence, id: id)
+      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.init(kind: kind, layout: layout, presence: presence, id: 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.init(kind: kind, leadingTrivia: leadingTrivia,
+                  trailingTrivia: trailingTrivia, presence: presence, id: id)
+      }
     }
     if let callback = decoder.userInfo[.rawSyntaxDecodedCallback] as?
                         (RawSyntax) -> Void {

Sources/SwiftSyntax/WeakLookupTable.swift

@@ -189,10 +189,14 @@ class WeakLookupTable<Element: Identifiable & AnyObject> {
   }
 
   /// Inserts the given object into the table.
+  /// - returns: true if no existing table object was overwritten
   @discardableResult
   func insert(_ obj: Element) -> Bool {
     var (pos, alreadyExists) = _findHole(obj.id)
     if alreadyExists {
+      // The existing and new object have the same identifier, but aren't
+      // necessarily the same object, so we still need to update it
+      buckets[pos].value = obj
       return false
     }
 

Tests/SwiftSyntaxTest/DeserializeFile.swift

@@ -15,4 +15,18 @@ public class DecodeSyntaxTestCase: XCTestCase {
       XCTAssertEqual("\(parsed)", source)
     }())
   }
+
+  public func testIncrementalDeserialize() {
+    XCTAssertNoThrow(try {
+      let deserializer = SyntaxTreeDeserializer()
+      var tree: SourceFileSyntax? = nil
+      let responses = try (1...3).map {
+        try Data(contentsOf: getInput("incremental-deserialize-\($0).json"))
+      }
+      for data in responses {
+        tree = try deserializer.deserialize(data, serializationFormat: .json)
+      }
+      XCTAssertEqual(tree?.description, "x\ny")
+    }())
+  }
 }

Tests/SwiftSyntaxTest/Inputs/incremental-deserialize-1.json

@@ -0,0 +1,51 @@
+{
+  "id":39,
+  "kind":"SourceFile",
+  "layout":[
+    {
+      "id":38,
+      "kind":"CodeBlockItemList",
+      "layout":[
+        {
+          "id":35,
+          "kind":"CodeBlockItem",
+          "layout":[
+            {
+              "id":34,
+              "kind":"IdentifierExpr",
+              "layout":[
+                {
+                  "id":33,
+                  "tokenKind":{
+                    "kind":"identifier",
+                    "text":"x"
+                  },
+                  "leadingTrivia":[],
+                  "trailingTrivia":[],
+                  "presence":"Present"
+                },
+                null
+              ],
+              "presence":"Present"
+            },
+            null,
+            null
+          ],
+          "presence":"Present"
+        }
+      ],
+      "presence":"Present"
+    },
+    {
+      "id":37,
+      "tokenKind":{
+        "kind":"eof",
+        "text":""
+      },
+      "leadingTrivia":[],
+      "trailingTrivia":[],
+      "presence":"Present"
+    }
+  ],
+  "presence":"Present"
+}

Tests/SwiftSyntaxTest/Inputs/incremental-deserialize-2.json

@@ -0,0 +1,33 @@
+{
+  "id":45,
+  "kind":"SourceFile",
+  "layout":[
+    {
+      "id":44,
+      "kind":"CodeBlockItemList",
+      "layout":[
+        {
+          "id":35,
+          "omitted":true
+        }
+      ],
+      "presence":"Present"
+    },
+    {
+      "id":43,
+      "tokenKind":{
+        "kind":"eof",
+        "text":""
+      },
+      "leadingTrivia":[
+        {
+          "kind":"Newline",
+          "value":1
+        }
+      ],
+      "trailingTrivia":[],
+      "presence":"Present"
+    }
+  ],
+  "presence":"Present"
+}

Tests/SwiftSyntaxTest/Inputs/incremental-deserialize-3.json

@@ -0,0 +1,60 @@
+{
+  "id":54,
+  "kind":"SourceFile",
+  "layout":[
+    {
+      "id":53,
+      "kind":"CodeBlockItemList",
+      "layout":[
+        {
+          "id":35,
+          "omitted":true
+        },
+        {
+          "id":50,
+          "kind":"CodeBlockItem",
+          "layout":[
+            {
+              "id":49,
+              "kind":"IdentifierExpr",
+              "layout":[
+                {
+                  "id":48,
+                  "tokenKind":{
+                    "kind":"identifier",
+                    "text":"y"
+                  },
+                  "leadingTrivia":[
+                    {
+                      "kind":"Newline",
+                      "value":1
+                    }
+                  ],
+                  "trailingTrivia":[],
+                  "presence":"Present"
+                },
+                null
+              ],
+              "presence":"Present"
+            },
+            null,
+            null
+          ],
+          "presence":"Present"
+        }
+      ],
+      "presence":"Present"
+    },
+    {
+      "id":52,
+      "tokenKind":{
+        "kind":"eof",
+        "text":""
+      },
+      "leadingTrivia":[],
+      "trailingTrivia":[],
+      "presence":"Present"
+    }
+  ],
+  "presence":"Present"
+}

Tests/SwiftSyntaxTest/WeakLookupTableTest.swift

@@ -108,6 +108,18 @@ public class WeakLookupTableTestCase: XCTestCase {
 
     _ = tree // Prevents 'never read' warning
   }
+
+  func testSameIdentifierDifferentObject() {
+    let table = WeakLookupTable<MinimalNode>()
+
+    let object1 = MinimalNode(id: "1", children: [])
+    let object2 = MinimalNode(id: "1", children: [])
+
+    XCTAssertTrue(table.insert(object1))
+    XCTAssert(table["1"] === object1)
+    XCTAssertFalse(table.insert(object2))
+    XCTAssert(table["1"] === object2)
+  }
 }
 
 #endif // DEBUG