From 9c578b035fb4ae9287e43427fea7dff4916b81d4 Mon Sep 17 00:00:00 2001
From: Gavin Nishizawa <gavin@terrastruct.com>
Date: Tue, 28 Nov 2023 14:06:48 -0800
Subject: [PATCH] lsp mode refactor

---
 d2cli/main.go         | 22 ++++++++++++++++++++++
 d2compiler/compile.go | 16 ----------------
 d2lib/d2.go           | 13 +++++++++++++
 3 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/d2cli/main.go b/d2cli/main.go
index 793c04839..b0b29d43e 100644
--- a/d2cli/main.go
+++ b/d2cli/main.go
@@ -2,6 +2,7 @@ package d2cli
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -21,6 +22,7 @@ import (
 	"oss.terrastruct.com/util-go/go2"
 	"oss.terrastruct.com/util-go/xmain"
 
+	"oss.terrastruct.com/d2/d2ast"
 	"oss.terrastruct.com/d2/d2graph"
 	"oss.terrastruct.com/d2/d2lib"
 	"oss.terrastruct.com/d2/d2parser"
@@ -389,6 +391,26 @@ func compile(ctx context.Context, ms *xmain.State, plugins []d2plugin.Plugin, fs
 		FS:             fs,
 	}
 
+	if os.Getenv("D2_LSP_MODE") == "1" {
+		// only the parse result is needed if running d2 for lsp
+		ast, err := d2lib.Parse(ctx, string(input), opts)
+		if err != nil {
+			return nil, false, err
+		}
+
+		type LspOutputData struct {
+			Ast *d2ast.Map
+			Err error
+		}
+		jsonOutput, err := json.Marshal(LspOutputData{Ast: ast, Err: err})
+		if err != nil {
+			return nil, false, err
+		}
+		fmt.Print(string(jsonOutput))
+		os.Exit(42)
+		return nil, false, nil
+	}
+
 	cancel := background.Repeat(func() {
 		ms.Log.Info.Printf("compiling & running layout algorithms...")
 	}, time.Second*5)
diff --git a/d2compiler/compile.go b/d2compiler/compile.go
index b37f79c9d..4a51d2457 100644
--- a/d2compiler/compile.go
+++ b/d2compiler/compile.go
@@ -1,13 +1,11 @@
 package d2compiler
 
 import (
-	"encoding/json"
 	"encoding/xml"
 	"fmt"
 	"io"
 	"io/fs"
 	"net/url"
-	"os"
 	"strconv"
 	"strings"
 
@@ -29,12 +27,6 @@ type CompileOptions struct {
 	FS fs.FS
 }
 
-// Changes for Language Server 'mode'
-type LspOutputData struct {
-	Ast *d2ast.Map
-	Err error
-}
-
 func Compile(p string, r io.Reader, opts *CompileOptions) (*d2graph.Graph, *d2target.Config, error) {
 	if opts == nil {
 		opts = &CompileOptions{}
@@ -43,14 +35,6 @@ func Compile(p string, r io.Reader, opts *CompileOptions) (*d2graph.Graph, *d2ta
 	ast, err := d2parser.Parse(p, r, &d2parser.ParseOptions{
 		UTF16Pos: opts.UTF16Pos,
 	})
-
-	if os.Getenv("D2_LSP_MODE") == "1" {
-		jsonOutput, _ := json.Marshal(LspOutputData{Ast: ast, Err: err})
-		fmt.Print(string(jsonOutput))
-		os.Exit(42)
-		return nil, nil, err
-	}
-
 	if err != nil {
 		return nil, nil, err
 	}
diff --git a/d2lib/d2.go b/d2lib/d2.go
index 7115a7ffd..f34bbc6e5 100644
--- a/d2lib/d2.go
+++ b/d2lib/d2.go
@@ -7,11 +7,13 @@ import (
 	"os"
 	"strings"
 
+	"oss.terrastruct.com/d2/d2ast"
 	"oss.terrastruct.com/d2/d2compiler"
 	"oss.terrastruct.com/d2/d2exporter"
 	"oss.terrastruct.com/d2/d2graph"
 	"oss.terrastruct.com/d2/d2layouts"
 	"oss.terrastruct.com/d2/d2layouts/d2dagrelayout"
+	"oss.terrastruct.com/d2/d2parser"
 	"oss.terrastruct.com/d2/d2renderers/d2fonts"
 	"oss.terrastruct.com/d2/d2renderers/d2svg"
 	"oss.terrastruct.com/d2/d2target"
@@ -39,6 +41,17 @@ type CompileOptions struct {
 	InputPath string
 }
 
+func Parse(ctx context.Context, input string, compileOpts *CompileOptions) (*d2ast.Map, error) {
+	if compileOpts == nil {
+		compileOpts = &CompileOptions{}
+	}
+
+	ast, err := d2parser.Parse(compileOpts.InputPath, strings.NewReader(input), &d2parser.ParseOptions{
+		UTF16Pos: compileOpts.UTF16Pos,
+	})
+	return ast, err
+}
+
 func Compile(ctx context.Context, input string, compileOpts *CompileOptions, renderOpts *d2svg.RenderOpts) (*d2target.Diagram, *d2graph.Graph, error) {
 	if compileOpts == nil {
 		compileOpts = &CompileOptions{}