Optimize LSP discriminated union decoding

tsc/internal/lsp/lsproto/_generate/generate.mts

@@ -2196,6 +2196,44 @@
         return exhaustive;
     }
 
+    /**
+     * Generate streaming discriminator dispatch for unions with at most one
+     * unmapped fallback arm. Fields after the discriminator decode directly
+     * from the active decoder; fields before it are replayed by the helper.
+     */
+    function generateStreamingDiscriminatorDispatch(
+        name: string,
+        disc: NonNullable<ReturnType<typeof findDiscriminatorField>>,
+        indent: string,
+    ) {
+        writeLine(`${indent}state, err := scanDiscriminatedStruct(dec, "${name}", ${JSON.stringify(disc.fieldName)})`);
+        writeLine(`${indent}if err != nil {`);
+        writeLine(`${indent}\treturn err`);
+        writeLine(`${indent}}`);
+        writeLine(`${indent}switch string(state.discriminatorValue) {`);
+        for (const [value, entry] of disc.mapping) {
+            writeLine(`${indent}case \`"${value}"\`:`);
+            writeLine(`${indent}\treturn unmarshalDiscriminatedArm(state, &o.${entry.fieldName})`);
+        }
+        writeLine(`${indent}default:`);
+        if (disc.unmapped.length === 1) {
+            writeLine(`${indent}\treturn unmarshalDiscriminatedFallbackArm(state, &o.${disc.unmapped[0].fieldName})`);
+        }
+        else {
+            writeLine(`${indent}\treturn state.invalidDiscriminator()`);
+        }
+        writeLine(`${indent}}`);
+    }
+
+    function canStreamDiscriminator(
+        disc: NonNullable<ReturnType<typeof findDiscriminatorField>>,
+    ): boolean {
+        return disc.unmapped.length <= 1 && disc.unmapped.every(entry =>
+            entry.originalType.kind === "reference"
+            && model.structures.some(structure => structure.name === entry.originalType.name)
+        );
+    }
+
     /**
      * Generate try-each fallback code for unmapped entries, chaining into
      * presence dispatch if possible before falling back to raw try-each.
@@ -3381,17 +3419,24 @@
                     }
                 }
                 else {
-                    // Ambiguous: buffer and dispatch
-                    writeLine(`\t\tdata, err := dec.ReadValue()`);
-                    writeLine(`\t\tif err != nil {`);
-                    writeLine(`\t\t\treturn err`);
-                    writeLine(`\t\t}`);
                     let exhaustive = false;
                     const disc = findDiscriminatorField(entries);
-                    if (disc) {
-                        exhaustive = generateDiscriminatorDispatch(disc, "\t\t");
+                    if (disc && canStreamDiscriminator(disc)) {
+                        generateStreamingDiscriminatorDispatch(name, disc, "\t\t");
+                        exhaustive = true;
                     }
                     else {
+                        // Ambiguous non-discriminated objects need the complete
+                        // value for presence checks or speculative decoding.
+                        writeLine(`\t\tdata, err := dec.ReadValue()`);
+                        writeLine(`\t\tif err != nil {`);
+                        writeLine(`\t\t\treturn err`);
+                        writeLine(`\t\t}`);
+                    }
+                    if (disc && !canStreamDiscriminator(disc)) {
+                        exhaustive = generateDiscriminatorDispatch(disc, "\t\t");
+                    }
+                    else if (!disc) {
                         const pres = findPresenceDiscriminator(entries);
                         if (pres) {
                             exhaustive = generatePresenceDispatch(pres, "\t\t");
@@ -3417,25 +3462,30 @@
             writeLine(`\t}`);
         }
         else {
-            // Fallback: unknown kinds present (e.g. `any`), use ReadValue + try-each.
-            writeLine("\tdata, err := dec.ReadValue()");
-            writeLine("\tif err != nil {");
-            writeLine("\t\treturn err");
-            writeLine("\t}");
-
-            if (unionContainedNull) {
-                writeLine(`\tif string(data) == "null" {`);
-                writeLine(`\t\treturn nil`);
-                writeLine(`\t}`);
-                writeLine("");
-            }
-
+            // Fallback for unknown kinds (e.g. `any`). Discriminated object
+            // unions can still stream; other unions use ReadValue + try-each.
             let exhaustive = false;
             const disc = findDiscriminatorField(fieldEntries);
-            if (disc) {
-                exhaustive = generateDiscriminatorDispatch(disc, "\t");
+            if (disc && canStreamDiscriminator(disc)) {
+                generateStreamingDiscriminatorDispatch(name, disc, "\t");
+                exhaustive = true;
             }
             else {
+                writeLine("\tdata, err := dec.ReadValue()");
+                writeLine("\tif err != nil {");
+                writeLine("\t\treturn err");
+                writeLine("\t}");
+                if (unionContainedNull) {
+                    writeLine(`\tif string(data) == "null" {`);
+                    writeLine(`\t\treturn nil`);
+                    writeLine(`\t}`);
+                    writeLine("");
+                }
+            }
+            if (disc && !canStreamDiscriminator(disc)) {
+                exhaustive = generateDiscriminatorDispatch(disc, "\t");
+            }
+            else if (!disc) {
                 const pres = findPresenceDiscriminator(fieldEntries);
                 if (pres) {
                     exhaustive = generatePresenceDispatch(pres, "\t");

tsc/internal/lsp/lsproto/lsp.go

@@ -125,36 +125,6 @@ func jsonKeyCheck(name []byte, key string) bool {
 	return len(name) == len(key)+2 && name[0] == '"' && string(name[1:len(name)-1]) == key
 }
 
-// jsonObjectRawField scans the top-level keys of a JSON object looking for the
-// given field name, and returns its raw JSON value (e.g. `"full"` with quotes).
-// Returns nil if the field is not found.
-func jsonObjectRawField(data []byte, field string) json.Value {
-	dec := json.NewDecoder(bytes.NewBuffer(data))
-	if dec.PeekKind() != '{' {
-		return nil
-	}
-	if _, err := dec.ReadToken(); err != nil {
-		return nil
-	}
-	for dec.PeekKind() != '}' {
-		name, err := dec.ReadValue()
-		if err != nil {
-			return nil
-		}
-		if jsonKeyCheck(name, field) {
-			val, err := dec.ReadValue()
-			if err != nil {
-				return nil
-			}
-			return val
-		}
-		if err := dec.SkipValue(); err != nil {
-			return nil
-		}
-	}
-	return nil
-}
-
 // jsonObjectHasKey scans the top-level keys of a JSON object looking for any of the
 // given keys. Returns the index of the first key found, or -1 if none match.
 // Bails early on first match without decoding any values.