Skip to content

Commit 644bde2

Browse files
jakebaileyjakebailey
authored
Optimize LSP marshalling via jsonv2 methods (#1502)
1 parent ec1a434 commit 644bde2

File tree

6 files changed

+12099
-5537
lines changed

6 files changed

+12099
-5537
lines changed

‎internal/api/proto.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,7 @@ type GetSourceFileParams struct {
200200
FileName string `json:"fileName"`
201201
}
202202

203-
func unmarshalPayload(method string, payload json.RawMessage) (any, error) {
203+
func unmarshalPayload(method string, payload json.Value) (any, error) {
204204
unmarshaler, ok := unmarshalers[Method(method)]
205205
if !ok {
206206
return nil, fmt.Errorf("unknown API method %q", method)

‎internal/incremental/buildInfo.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,7 @@ func (b *BuildInfoDiagnosticsOfFile) MarshalJSON() ([]byte, error) {
213213
}
214214

215215
func (b *BuildInfoDiagnosticsOfFile) UnmarshalJSON(data []byte) error {
216-
var fileIdAndDiagnostics []json.RawMessage
216+
var fileIdAndDiagnostics []json.Value
217217
if err := json.Unmarshal(data, &fileIdAndDiagnostics); err != nil {
218218
return fmt.Errorf("invalid BuildInfoDiagnosticsOfFile: %s", data)
219219
}

‎internal/json/json.go

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55

66
jsonv2 "github.com/go-json-experiment/json"
77
"github.com/go-json-experiment/json/jsontext"
8-
jsonv1 "github.com/go-json-experiment/json/v1"
98
)
109

1110
func Marshal(in any, opts ...jsonv2.Options) (out []byte, err error) {
@@ -47,7 +46,7 @@ func UnmarshalRead(in io.Reader, out any, opts ...jsonv2.Options) (err error) {
4746
}
4847

4948
type (
50-
RawMessage = jsonv1.RawMessage
49+
Value = jsontext.Value
5150
UnmarshalerFrom = jsonv2.UnmarshalerFrom
5251
MarshalerTo = jsonv2.MarshalerTo
5352
)

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

Lines changed: 83 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -479,7 +479,8 @@ function generateCode() {
479479
writeLine(`import (`);
480480
writeLine(`\t"fmt"`);
481481
writeLine("");
482-
writeLine(`\t"github.com/microsoft/typescript-go/internal/json"`);
482+
writeLine(`\t"github.com/go-json-experiment/json"`);
483+
writeLine(`\t"github.com/go-json-experiment/json/jsontext"`);
483484
writeLine(`)`);
484485
writeLine("");
485486
writeLine("// Meta model version " + model.metaData.version);
@@ -519,35 +520,64 @@ function generateCode() {
519520
generateStructFields(structure.name, true);
520521
writeLine("");
521522

522-
// Generate UnmarshalJSON method for structure validation
523+
// Generate UnmarshalJSONFrom method for structure validation
523524
const requiredProps = structure.properties?.filter(p => !p.optional) || [];
524525
if (requiredProps.length > 0) {
525-
writeLine(`func (s *${structure.name}) UnmarshalJSON(data []byte) error {`);
526-
writeLine(`\t// Check required props`);
527-
writeLine(`\ttype requiredProps struct {`);
526+
writeLine(`\tvar _ json.UnmarshalerFrom = (*${structure.name})(nil)`);
527+
writeLine("");
528+
529+
writeLine(`func (s *${structure.name}) UnmarshalJSONFrom(dec *jsontext.Decoder) error {`);
530+
writeLine(`\tvar (`);
528531
for (const prop of requiredProps) {
529-
writeLine(`\t\t${titleCase(prop.name)} requiredProp \`json:"${prop.name}"\``);
532+
writeLine(`\t\tseen${titleCase(prop.name)} bool`);
530533
}
531-
writeLine(`}`);
534+
writeLine(`\t)`);
532535
writeLine("");
533536

534-
writeLine(`\tvar keys requiredProps`);
535-
writeLine(`\tif err := json.Unmarshal(data, &keys); err != nil {`);
537+
writeLine(`\tif k := dec.PeekKind(); k != '{' {`);
538+
writeLine(`\t\treturn fmt.Errorf("expected object start, but encountered %v", k)`);
539+
writeLine(`\t}`);
540+
writeLine(`\tif _, err := dec.ReadToken(); err != nil {`);
541+
writeLine(`\t\treturn err`);
542+
writeLine(`\t}`);
543+
writeLine("");
544+
545+
writeLine(`\tfor dec.PeekKind() != '}' {`);
546+
writeLine("name, err := dec.ReadValue()");
547+
writeLine(`\t\tif err != nil {`);
548+
writeLine(`\t\t\treturn err`);
549+
writeLine(`\t\t}`);
550+
writeLine(`\t\tswitch string(name) {`);
551+
552+
for (const prop of structure.properties) {
553+
writeLine(`\t\tcase \`"${prop.name}"\`:`);
554+
if (!prop.optional) {
555+
writeLine(`\t\t\tseen${titleCase(prop.name)} = true`);
556+
}
557+
writeLine(`\t\t\tif err := json.UnmarshalDecode(dec, &s.${titleCase(prop.name)}); err != nil {`);
558+
writeLine(`\t\t\t\treturn err`);
559+
writeLine(`\t\t\t}`);
560+
}
561+
562+
writeLine(`\t\tdefault:`);
563+
writeLine(`\t\t// Ignore unknown properties.`);
564+
writeLine(`\t\t}`);
565+
writeLine(`\t}`);
566+
writeLine("");
567+
568+
writeLine(`\tif _, err := dec.ReadToken(); err != nil {`);
536569
writeLine(`\t\treturn err`);
537570
writeLine(`\t}`);
538571
writeLine("");
539572

540-
// writeLine(`\t// Check for missing required keys`);
541573
for (const prop of requiredProps) {
542-
writeLine(`if !keys.${titleCase(prop.name)} {`);
543-
writeLine(`\t\treturn fmt.Errorf("required key '${prop.name}' is missing")`);
544-
writeLine(`}`);
574+
writeLine(`\tif !seen${titleCase(prop.name)} {`);
575+
writeLine(`\t\treturn fmt.Errorf("required property '${prop.name}' is missing")`);
576+
writeLine(`\t}`);
545577
}
546578

547-
writeLine(``);
548-
writeLine(`\t// Redeclare the struct to prevent infinite recursion`);
549-
generateStructFields("temp", false);
550-
writeLine(`\treturn json.Unmarshal(data, (*temp)(s))`);
579+
writeLine("");
580+
writeLine(`\treturn nil`);
551581
writeLine(`}`);
552582
writeLine("");
553583
}
@@ -606,17 +636,6 @@ function generateCode() {
606636

607637
writeLine(")");
608638
writeLine("");
609-
610-
// Add custom JSON unmarshaling
611-
writeLine(`func (e *${enumeration.name}) UnmarshalJSON(data []byte) error {`);
612-
writeLine(`\tvar v ${baseType}`);
613-
writeLine(`\tif err := json.Unmarshal(data, &v); err != nil {`);
614-
writeLine(`\t\treturn err`);
615-
writeLine(`\t}`);
616-
writeLine(`\t*e = ${enumeration.name}(v)`);
617-
writeLine(`\treturn nil`);
618-
writeLine(`}`);
619-
writeLine("");
620639
}
621640

622641
const requestsAndNotifications: (Request | Notification)[] = [...model.requests, ...model.notifications];
@@ -736,14 +755,21 @@ function generateCode() {
736755
const fieldEntries = Array.from(uniqueTypeFields.entries()).map(([typeName, fieldName]) => ({ fieldName, typeName }));
737756

738757
// Marshal method
739-
writeLine(`func (o ${name}) MarshalJSON() ([]byte, error) {`);
758+
writeLine(`var _ json.MarshalerTo = (*${name})(nil)`);
759+
writeLine("");
760+
761+
writeLine(`func (o *${name}) MarshalJSONTo(enc *jsontext.Encoder) error {`);
740762

741763
// Determine if this union contained null (check if any member has containedNull = true)
742764
const unionContainedNull = members.some(member => member.containedNull);
743-
const assertionFunc = unionContainedNull ? "assertAtMostOne" : "assertOnlyOne";
765+
if (unionContainedNull) {
766+
write(`\tassertAtMostOne("more than one element of ${name} is set", `);
767+
}
768+
else {
769+
write(`\tassertOnlyOne("exactly one element of ${name} should be set", `);
770+
}
744771

745772
// Create assertion to ensure at most one field is set at a time
746-
write(`\t${assertionFunc}("more than one element of ${name} is set", `);
747773

748774
// Write the assertion conditions
749775
for (let i = 0; i < fieldEntries.length; i++) {
@@ -755,14 +781,13 @@ function generateCode() {
755781

756782
for (const entry of fieldEntries) {
757783
writeLine(`\tif o.${entry.fieldName} != nil {`);
758-
writeLine(`\t\treturn json.Marshal(*o.${entry.fieldName})`);
784+
writeLine(`\t\treturn json.MarshalEncode(enc, o.${entry.fieldName})`);
759785
writeLine(`\t}`);
760786
}
761787

762788
// If all fields are nil, marshal as null (only for unions that can contain null)
763789
if (unionContainedNull) {
764-
writeLine(`\t// All fields are nil, represent as null`);
765-
writeLine(`\treturn []byte("null"), nil`);
790+
writeLine(`\treturn enc.WriteToken(jsontext.Null)`);
766791
}
767792
else {
768793
writeLine(`\tpanic("unreachable")`);
@@ -771,13 +796,19 @@ function generateCode() {
771796
writeLine("");
772797

773798
// Unmarshal method
774-
writeLine(`func (o *${name}) UnmarshalJSON(data []byte) error {`);
799+
writeLine(`var _ json.UnmarshalerFrom = (*${name})(nil)`);
800+
writeLine("");
801+
802+
writeLine(`func (o *${name}) UnmarshalJSONFrom(dec *jsontext.Decoder) error {`);
775803
writeLine(`\t*o = ${name}{}`);
776804
writeLine("");
777805

778-
// Handle null case only for unions that can contain null
806+
writeLine("\tdata, err := dec.ReadValue()");
807+
writeLine("\tif err != nil {");
808+
writeLine("\t\treturn err");
809+
writeLine("\t}");
810+
779811
if (unionContainedNull) {
780-
writeLine(`\t// Handle null case`);
781812
writeLine(`\tif string(data) == "null" {`);
782813
writeLine(`\t\treturn nil`);
783814
writeLine(`\t}`);
@@ -808,14 +839,24 @@ function generateCode() {
808839
writeLine(`type ${name} struct{}`);
809840
writeLine("");
810841

811-
writeLine(`func (o ${name}) MarshalJSON() ([]byte, error) {`);
812-
writeLine(`\treturn []byte(\`${jsonValue}\`), nil`);
842+
writeLine(`var _ json.MarshalerTo = ${name}{}`);
843+
writeLine("");
844+
845+
writeLine(`func (o ${name}) MarshalJSONTo(enc *jsontext.Encoder) error {`);
846+
writeLine(`\treturn enc.WriteValue(jsontext.Value(\`${jsonValue}\`))`);
813847
writeLine(`}`);
814848
writeLine("");
815849

816-
writeLine(`func (o *${name}) UnmarshalJSON(data []byte) error {`);
817-
writeLine(`\tif string(data) != \`${jsonValue}\` {`);
818-
writeLine(`\t\treturn fmt.Errorf("invalid ${name}: %s", data)`);
850+
writeLine(`var _ json.UnmarshalerFrom = &${name}{}`);
851+
writeLine("");
852+
853+
writeLine(`func (o *${name}) UnmarshalJSONFrom(dec *jsontext.Decoder) error {`);
854+
writeLine(`\tv, err := dec.ReadValue();`);
855+
writeLine(`\tif err != nil {`);
856+
writeLine(`\t\treturn err`);
857+
writeLine(`\t}`);
858+
writeLine(`\tif string(v) != \`${jsonValue}\` {`);
859+
writeLine(`\t\treturn fmt.Errorf("expected ${name} value %s, got %s", \`${jsonValue}\`, v)`);
819860
writeLine(`\t}`);
820861
writeLine(`\treturn nil`);
821862
writeLine(`}`);

‎internal/lsp/lsproto/jsonrpc.go

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -100,12 +100,12 @@ func (m *Message) AsResponse() *ResponseMessage {
100100

101101
func (m *Message) UnmarshalJSON(data []byte) error {
102102
var raw struct {
103-
JSONRPC JSONRPCVersion `json:"jsonrpc"`
104-
Method Method `json:"method"`
105-
ID *ID `json:"id,omitzero"`
106-
Params json.RawMessage `json:"params"`
107-
Result any `json:"result,omitzero"`
108-
Error *ResponseError `json:"error,omitzero"`
103+
JSONRPC JSONRPCVersion `json:"jsonrpc"`
104+
Method Method `json:"method"`
105+
ID *ID `json:"id,omitzero"`
106+
Params json.Value `json:"params"`
107+
Result any `json:"result,omitzero"`
108+
Error *ResponseError `json:"error,omitzero"`
109109
}
110110
if err := json.Unmarshal(data, &raw); err != nil {
111111
return fmt.Errorf("%w: %w", ErrInvalidRequest, err)
@@ -182,10 +182,10 @@ func (r *RequestMessage) Message() *Message {
182182

183183
func (r *RequestMessage) UnmarshalJSON(data []byte) error {
184184
var raw struct {
185-
JSONRPC JSONRPCVersion `json:"jsonrpc"`
186-
ID *ID `json:"id"`
187-
Method Method `json:"method"`
188-
Params json.RawMessage `json:"params"`
185+
JSONRPC JSONRPCVersion `json:"jsonrpc"`
186+
ID *ID `json:"id"`
187+
Method Method `json:"method"`
188+
Params json.Value `json:"params"`
189189
}
190190
if err := json.Unmarshal(data, &raw); err != nil {
191191
return fmt.Errorf("%w: %w", ErrInvalidRequest, err)

0 commit comments

Comments
 (0)