From 9c19637fff715225f253aeb65b9b2c2b28dd3d38 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Wed, 15 Jan 2025 16:30:17 -0700
Subject: [PATCH] migrate sketch

---
 d2js/js/examples/customizable.html | 47 +++++++++++++++++--------
 d2js/js/src/worker.js              |  9 +++++
 d2js/js/test/unit/basic.test.js    | 10 ++++++
 d2renderers/d2sketch/setup.js      |  4 +++
 d2renderers/d2sketch/sketch.go     | 56 ++++++++++++------------------
 d2renderers/d2svg/d2svg.go         | 55 +++++++++++++++--------------
 6 files changed, 107 insertions(+), 74 deletions(-)
diff --git a/d2js/js/examples/customizable.html b/d2js/js/examples/customizable.html
index aa6fac6cf..13b7a2138 100644
--- a/d2js/js/examples/customizable.html
+++ b/d2js/js/examples/customizable.html
@@ -24,7 +24,16 @@
         border-radius: 4px;
         font-family: monospace;
       }
-      .layout-toggle {
+      .options-group {
+        display: flex;
+        flex-direction: column;
+        gap: 12px;
+        padding: 12px;
+        border: 1px solid #eee;
+        border-radius: 4px;
+      }
+      .layout-toggle,
+      .sketch-toggle {
         display: flex;
         gap: 16px;
         align-items: center;
@@ -33,7 +42,8 @@
         display: flex;
         gap: 12px;
       }
-      .radio-label {
+      .radio-label,
+      .checkbox-label {
         display: flex;
         gap: 4px;
         align-items: center;
@@ -66,16 +76,24 @@
   <body>
     <div class="controls">
       <textarea id="input">x -> y</textarea>
-      <div class="layout-toggle">
-        <span>Layout:</span>
-        <div class="radio-group">
-          <label class="radio-label">
-            <input type="radio" name="layout" value="dagre" checked />
-            Dagre
-          </label>
-          <label class="radio-label">
-            <input type="radio" name="layout" value="elk" />
-            ELK
+      <div class="options-group">
+        <div class="layout-toggle">
+          <span>Layout:</span>
+          <div class="radio-group">
+            <label class="radio-label">
+              <input type="radio" name="layout" value="dagre" checked />
+              Dagre
+            </label>
+            <label class="radio-label">
+              <input type="radio" name="layout" value="elk" />
+              ELK
+            </label>
+          </div>
+        </div>
+        <div class="sketch-toggle">
+          <label class="checkbox-label">
+            <input type="checkbox" id="sketch" />
+            Sketch mode
           </label>
         </div>
       </div>
@@ -88,9 +106,10 @@
       window.compile = async () => {
         const input = document.getElementById("input").value;
         const layout = document.querySelector('input[name="layout"]:checked').value;
+        const sketch = document.getElementById("sketch").checked;
         try {
-          const result = await d2.compile(input, { layout });
-          const svg = await d2.render(result.diagram);
+          const result = await d2.compile(input, { layout, sketch });
+          const svg = await d2.render(result.diagram, { sketch });
           document.getElementById("output").innerHTML = svg;
         } catch (err) {
           console.error(err);
diff --git a/d2js/js/src/worker.js b/d2js/js/src/worker.js
index ff7d9ec08..b615fd569 100644
--- a/d2js/js/src/worker.js
+++ b/d2js/js/src/worker.js
@@ -2,6 +2,7 @@ import { parentPort } from "node:worker_threads";
 
 let currentPort;
 let d2;
+let elk;
 
 export function setupMessageHandler(isNode, port, initWasm) {
   currentPort = port;
@@ -14,8 +15,10 @@ export function setupMessageHandler(isNode, port, initWasm) {
         try {
           if (isNode) {
             eval(data.wasmExecContent);
+            eval(data.elkContent);
           }
           d2 = await initWasm(data.wasm);
+          elk = new ELK();
           currentPort.postMessage({ type: "ready" });
         } catch (err) {
           currentPort.postMessage({ type: "error", error: err.message });
@@ -24,6 +27,12 @@ export function setupMessageHandler(isNode, port, initWasm) {
 
       case "compile":
         try {
+          if (data.options.layout === "elk") {
+            const elkGraph = await d2.getELKGraph(JSON.stringify(data));
+            const elkGraph2 = JSON.parse(elkGraph).data;
+            const layout = await elk.layout(elkGraph2);
+            globalThis.elkResult = layout;
+          }
           const result = await d2.compile(JSON.stringify(data));
           const response = JSON.parse(result);
           if (response.error) throw new Error(response.error.message);
diff --git a/d2js/js/test/unit/basic.test.js b/d2js/js/test/unit/basic.test.js
index 7cc71c951..8b7bcea33 100644
--- a/d2js/js/test/unit/basic.test.js
+++ b/d2js/js/test/unit/basic.test.js
@@ -25,6 +25,16 @@ describe("D2 Unit Tests", () => {
     await d2.worker.terminate();
   }, 20000);
 
+  test("sketch render works", async () => {
+    const d2 = new D2();
+    const result = await d2.compile("x -> y", { sketch: true });
+    const svg = await d2.render(result.diagram, { sketch: true });
+    expect(svg).toContain("<svg");
+    expect(svg).toContain("</svg>");
+    expect(svg).toContain("sketch-overlay");
+    await d2.worker.terminate();
+  }, 20000);
+
   test("handles syntax errors correctly", async () => {
     const d2 = new D2();
     try {
diff --git a/d2renderers/d2sketch/setup.js b/d2renderers/d2sketch/setup.js
index 87683d55f..edc6958ae 100644
--- a/d2renderers/d2sketch/setup.js
+++ b/d2renderers/d2sketch/setup.js
@@ -17,3 +17,7 @@ const root = {
 };
 const rc = rough.svg(root, { seed: 1 });
 let node;
+
+if (typeof globalThis !== 'undefined') {
+  globalThis.rc = rc;
+}
diff --git a/d2renderers/d2sketch/sketch.go b/d2renderers/d2sketch/sketch.go
index f873eda78..0e27bafe3 100644
--- a/d2renderers/d2sketch/sketch.go
+++ b/d2renderers/d2sketch/sketch.go
@@ -9,12 +9,11 @@ import (
 
 	_ "embed"
 
-	"github.com/dop251/goja"
-
 	"oss.terrastruct.com/d2/d2target"
 	"oss.terrastruct.com/d2/d2themes"
 	"oss.terrastruct.com/d2/lib/color"
 	"oss.terrastruct.com/d2/lib/geo"
+	"oss.terrastruct.com/d2/lib/jsrunner"
 	"oss.terrastruct.com/d2/lib/label"
 	"oss.terrastruct.com/d2/lib/svg"
 	"oss.terrastruct.com/util-go/go2"
@@ -29,8 +28,6 @@ var setupJS string
 //go:embed streaks.txt
 var streaks string
 
-type Runner goja.Runtime
-
 var baseRoughProps = `fillWeight: 2.0,
 hachureGap: 16,
 fillStyle: "solid",
@@ -44,21 +41,14 @@ const (
 	FG_COLOR = color.N1
 )
 
-func (r *Runner) run(js string) (goja.Value, error) {
-	vm := (*goja.Runtime)(r)
-	return vm.RunString(js)
-}
-
-func InitSketchVM() (*Runner, error) {
-	vm := goja.New()
-	if _, err := vm.RunString(roughJS); err != nil {
-		return nil, err
+func LoadJS(runner jsrunner.JSRunner) error {
+	if _, err := runner.RunString(roughJS); err != nil {
+		return err
 	}
-	if _, err := vm.RunString(setupJS); err != nil {
-		return nil, err
+	if _, err := runner.RunString(setupJS); err != nil {
+		return err
 	}
-	r := Runner(*vm)
-	return &r, nil
+	return nil
 }
 
 // DefineFillPatterns adds reusable patterns that are overlayed on shapes with
@@ -83,7 +73,7 @@ func defineFillPattern(buf *bytes.Buffer, source string, luminanceCategory, fill
 	}
 }
 
-func Rect(r *Runner, shape d2target.Shape) (string, error) {
+func Rect(r jsrunner.JSRunner, shape d2target.Shape) (string, error) {
 	js := fmt.Sprintf(`node = rc.rectangle(0, 0, %d, %d, {
 		fill: "#000",
 		stroke: "#000",
@@ -119,7 +109,7 @@ func Rect(r *Runner, shape d2target.Shape) (string, error) {
 	return output, nil
 }
 
-func DoubleRect(r *Runner, shape d2target.Shape) (string, error) {
+func DoubleRect(r jsrunner.JSRunner, shape d2target.Shape) (string, error) {
 	jsBigRect := fmt.Sprintf(`node = rc.rectangle(0, 0, %d, %d, {
 		fill: "#000",
 		stroke: "#000",
@@ -179,7 +169,7 @@ func DoubleRect(r *Runner, shape d2target.Shape) (string, error) {
 	return output, nil
 }
 
-func Oval(r *Runner, shape d2target.Shape) (string, error) {
+func Oval(r jsrunner.JSRunner, shape d2target.Shape) (string, error) {
 	js := fmt.Sprintf(`node = rc.ellipse(%d, %d, %d, %d, {
 		fill: "#000",
 		stroke: "#000",
@@ -218,7 +208,7 @@ func Oval(r *Runner, shape d2target.Shape) (string, error) {
 	return output, nil
 }
 
-func DoubleOval(r *Runner, shape d2target.Shape) (string, error) {
+func DoubleOval(r jsrunner.JSRunner, shape d2target.Shape) (string, error) {
 	jsBigCircle := fmt.Sprintf(`node = rc.ellipse(%d, %d, %d, %d, {
 		fill: "#000",
 		stroke: "#000",
@@ -281,7 +271,7 @@ func DoubleOval(r *Runner, shape d2target.Shape) (string, error) {
 }
 
 // TODO need to personalize this per shape like we do in Terrastruct app
-func Paths(r *Runner, shape d2target.Shape, paths []string) (string, error) {
+func Paths(r jsrunner.JSRunner, shape d2target.Shape, paths []string) (string, error) {
 	output := ""
 	for _, path := range paths {
 		js := fmt.Sprintf(`node = rc.path("%s", {
@@ -320,7 +310,7 @@ func Paths(r *Runner, shape d2target.Shape, paths []string) (string, error) {
 	return output, nil
 }
 
-func Connection(r *Runner, connection d2target.Connection, path, attrs string) (string, error) {
+func Connection(r jsrunner.JSRunner, connection d2target.Connection, path, attrs string) (string, error) {
 	animatedClass := ""
 	if connection.Animated {
 		animatedClass = " animated-connection"
@@ -388,7 +378,7 @@ func Connection(r *Runner, connection d2target.Connection, path, attrs string) (
 }
 
 // TODO cleanup
-func Table(r *Runner, shape d2target.Shape) (string, error) {
+func Table(r jsrunner.JSRunner, shape d2target.Shape) (string, error) {
 	output := ""
 	js := fmt.Sprintf(`node = rc.rectangle(0, 0, %d, %d, {
 		fill: "#000",
@@ -530,7 +520,7 @@ func Table(r *Runner, shape d2target.Shape) (string, error) {
 	return output, nil
 }
 
-func Class(r *Runner, shape d2target.Shape) (string, error) {
+func Class(r jsrunner.JSRunner, shape d2target.Shape) (string, error) {
 	output := ""
 	js := fmt.Sprintf(`node = rc.rectangle(0, 0, %d, %d, {
 		fill: "#000",
@@ -681,8 +671,8 @@ func classRow(shape d2target.Shape, box *geo.Box, prefix, nameText, typeText str
 	return output
 }
 
-func computeRoughPathData(r *Runner, js string) ([]string, error) {
-	if _, err := r.run(js); err != nil {
+func computeRoughPathData(r jsrunner.JSRunner, js string) ([]string, error) {
+	if _, err := r.RunString(js); err != nil {
 		return nil, err
 	}
 	roughPaths, err := extractRoughPaths(r)
@@ -692,8 +682,8 @@ func computeRoughPathData(r *Runner, js string) ([]string, error) {
 	return extractPathData(roughPaths)
 }
 
-func computeRoughPaths(r *Runner, js string) ([]roughPath, error) {
-	if _, err := r.run(js); err != nil {
+func computeRoughPaths(r jsrunner.JSRunner, js string) ([]roughPath, error) {
+	if _, err := r.RunString(js); err != nil {
 		return nil, err
 	}
 	return extractRoughPaths(r)
@@ -722,8 +712,8 @@ func (rp roughPath) StyleCSS() string {
 	return style
 }
 
-func extractRoughPaths(r *Runner) ([]roughPath, error) {
-	val, err := r.run("JSON.stringify(node.children, null, '  ')")
+func extractRoughPaths(r jsrunner.JSRunner) ([]roughPath, error) {
+	val, err := r.RunString("JSON.stringify(node.children, null, '  ')")
 	if err != nil {
 		return nil, err
 	}
@@ -756,7 +746,7 @@ func extractPathData(roughPaths []roughPath) ([]string, error) {
 	return paths, nil
 }
 
-func ArrowheadJS(r *Runner, arrowhead d2target.Arrowhead, stroke string, strokeWidth int) (arrowJS, extraJS string) {
+func ArrowheadJS(r jsrunner.JSRunner, arrowhead d2target.Arrowhead, stroke string, strokeWidth int) (arrowJS, extraJS string) {
 	// Note: selected each seed that looks the good for consistent renders
 	switch arrowhead {
 	case d2target.ArrowArrowhead:
@@ -854,7 +844,7 @@ func ArrowheadJS(r *Runner, arrowhead d2target.Arrowhead, stroke string, strokeW
 	return
 }
 
-func Arrowheads(r *Runner, connection d2target.Connection, srcAdj, dstAdj *geo.Point) (string, error) {
+func Arrowheads(r jsrunner.JSRunner, connection d2target.Connection, srcAdj, dstAdj *geo.Point) (string, error) {
 	arrowPaths := []string{}
 
 	if connection.SrcArrow != d2target.NoArrowhead {
diff --git a/d2renderers/d2svg/d2svg.go b/d2renderers/d2svg/d2svg.go
index 12ea2e7c0..3857085fa 100644
--- a/d2renderers/d2svg/d2svg.go
+++ b/d2renderers/d2svg/d2svg.go
@@ -30,6 +30,7 @@ import (
 	"oss.terrastruct.com/d2/d2themes/d2themescatalog"
 	"oss.terrastruct.com/d2/lib/color"
 	"oss.terrastruct.com/d2/lib/geo"
+	"oss.terrastruct.com/d2/lib/jsrunner"
 	"oss.terrastruct.com/d2/lib/label"
 	"oss.terrastruct.com/d2/lib/shape"
 	"oss.terrastruct.com/d2/lib/svg"
@@ -496,7 +497,7 @@ func makeLabelMask(labelTL *geo.Point, width, height int, opacity float64) strin
 	)
 }
 
-func drawConnection(writer io.Writer, labelMaskID string, connection d2target.Connection, markers map[string]struct{}, idToShape map[string]d2target.Shape, sketchRunner *d2sketch.Runner, inlineTheme *d2themes.Theme) (labelMask string, _ error) {
+func drawConnection(writer io.Writer, labelMaskID string, connection d2target.Connection, markers map[string]struct{}, idToShape map[string]d2target.Shape, jsRunner jsrunner.JSRunner, inlineTheme *d2themes.Theme) (labelMask string, _ error) {
 	opacityStyle := ""
 	if connection.Opacity != 1.0 {
 		opacityStyle = fmt.Sprintf(" style='opacity:%f'", connection.Opacity)
@@ -552,15 +553,15 @@ func drawConnection(writer io.Writer, labelMaskID string, connection d2target.Co
 	path := pathData(connection, srcAdj, dstAdj)
 	mask := fmt.Sprintf(`mask="url(#%s)"`, labelMaskID)
 
-	if sketchRunner != nil {
-		out, err := d2sketch.Connection(sketchRunner, connection, path, mask)
+	if jsRunner != nil {
+		out, err := d2sketch.Connection(jsRunner, connection, path, mask)
 		if err != nil {
 			return "", err
 		}
 		fmt.Fprint(writer, out)
 
 		// render sketch arrowheads separately
-		arrowPaths, err := d2sketch.Arrowheads(sketchRunner, connection, srcAdj, dstAdj)
+		arrowPaths, err := d2sketch.Arrowheads(jsRunner, connection, srcAdj, dstAdj)
 		if err != nil {
 			return "", err
 		}
@@ -957,7 +958,7 @@ func render3DHexagon(targetShape d2target.Shape, inlineTheme *d2themes.Theme) st
 	return borderMask + mainShapeRendered + renderedSides + renderedBorder
 }
 
-func drawShape(writer, appendixWriter io.Writer, diagramHash string, targetShape d2target.Shape, sketchRunner *d2sketch.Runner, inlineTheme *d2themes.Theme) (labelMask string, err error) {
+func drawShape(writer, appendixWriter io.Writer, diagramHash string, targetShape d2target.Shape, jsRunner jsrunner.JSRunner, inlineTheme *d2themes.Theme) (labelMask string, err error) {
 	closingTag := "</g>"
 	if targetShape.Link != "" {
 
@@ -1021,8 +1022,8 @@ func drawShape(writer, appendixWriter io.Writer, diagramHash string, targetShape
 
 	switch targetShape.Type {
 	case d2target.ShapeClass:
-		if sketchRunner != nil {
-			out, err := d2sketch.Class(sketchRunner, targetShape)
+		if jsRunner != nil {
+			out, err := d2sketch.Class(jsRunner, targetShape)
 			if err != nil {
 				return "", err
 			}
@@ -1035,8 +1036,8 @@ func drawShape(writer, appendixWriter io.Writer, diagramHash string, targetShape
 		fmt.Fprint(writer, closingTag)
 		return labelMask, nil
 	case d2target.ShapeSQLTable:
-		if sketchRunner != nil {
-			out, err := d2sketch.Table(sketchRunner, targetShape)
+		if jsRunner != nil {
+			out, err := d2sketch.Table(jsRunner, targetShape)
 			if err != nil {
 				return "", err
 			}
@@ -1053,8 +1054,8 @@ func drawShape(writer, appendixWriter io.Writer, diagramHash string, targetShape
 			if targetShape.Multiple {
 				fmt.Fprint(writer, renderDoubleOval(multipleTL, width, height, fill, "", stroke, style, inlineTheme))
 			}
-			if sketchRunner != nil {
-				out, err := d2sketch.DoubleOval(sketchRunner, targetShape)
+			if jsRunner != nil {
+				out, err := d2sketch.DoubleOval(jsRunner, targetShape)
 				if err != nil {
 					return "", err
 				}
@@ -1066,8 +1067,8 @@ func drawShape(writer, appendixWriter io.Writer, diagramHash string, targetShape
 			if targetShape.Multiple {
 				fmt.Fprint(writer, renderOval(multipleTL, width, height, fill, "", stroke, style, inlineTheme))
 			}
-			if sketchRunner != nil {
-				out, err := d2sketch.Oval(sketchRunner, targetShape)
+			if jsRunner != nil {
+				out, err := d2sketch.Oval(jsRunner, targetShape)
 				if err != nil {
 					return "", err
 				}
@@ -1111,8 +1112,8 @@ func drawShape(writer, appendixWriter io.Writer, diagramHash string, targetShape
 					el.Rx = borderRadius
 					fmt.Fprint(writer, el.Render())
 				}
-				if sketchRunner != nil {
-					out, err := d2sketch.Rect(sketchRunner, targetShape)
+				if jsRunner != nil {
+					out, err := d2sketch.Rect(jsRunner, targetShape)
 					if err != nil {
 						return "", err
 					}
@@ -1155,8 +1156,8 @@ func drawShape(writer, appendixWriter io.Writer, diagramHash string, targetShape
 					el.Rx = borderRadius
 					fmt.Fprint(writer, el.Render())
 				}
-				if sketchRunner != nil {
-					out, err := d2sketch.DoubleRect(sketchRunner, targetShape)
+				if jsRunner != nil {
+					out, err := d2sketch.DoubleRect(jsRunner, targetShape)
 					if err != nil {
 						return "", err
 					}
@@ -1203,8 +1204,8 @@ func drawShape(writer, appendixWriter io.Writer, diagramHash string, targetShape
 				}
 			}
 
-			if sketchRunner != nil {
-				out, err := d2sketch.Paths(sketchRunner, targetShape, s.GetSVGPathData())
+			if jsRunner != nil {
+				out, err := d2sketch.Paths(jsRunner, targetShape, s.GetSVGPathData())
 				if err != nil {
 					return "", err
 				}
@@ -1235,8 +1236,8 @@ func drawShape(writer, appendixWriter io.Writer, diagramHash string, targetShape
 			}
 		}
 
-		if sketchRunner != nil {
-			out, err := d2sketch.Paths(sketchRunner, targetShape, s.GetSVGPathData())
+		if jsRunner != nil {
+			out, err := d2sketch.Paths(jsRunner, targetShape, s.GetSVGPathData())
 			if err != nil {
 				return "", err
 			}
@@ -1846,7 +1847,7 @@ func appendOnTrigger(buf *bytes.Buffer, source string, triggers []string, newCon
 var DEFAULT_DARK_THEME *int64 = nil // no theme selected
 
 func Render(diagram *d2target.Diagram, opts *RenderOpts) ([]byte, error) {
-	var sketchRunner *d2sketch.Runner
+	var jsRunner jsrunner.JSRunner
 	pad := DEFAULT_PADDING
 	themeID := d2themescatalog.NeutralDefault.ID
 	darkThemeID := DEFAULT_DARK_THEME
@@ -1856,8 +1857,8 @@ func Render(diagram *d2target.Diagram, opts *RenderOpts) ([]byte, error) {
 			pad = int(*opts.Pad)
 		}
 		if opts.Sketch != nil && *opts.Sketch {
-			var err error
-			sketchRunner, err = d2sketch.InitSketchVM()
+			jsRunner = jsrunner.NewJSRunner()
+			err := d2sketch.LoadJS(jsRunner)
 			if err != nil {
 				return nil, err
 			}
@@ -1941,7 +1942,7 @@ func Render(diagram *d2target.Diagram, opts *RenderOpts) ([]byte, error) {
 	}
 	for _, obj := range allObjects {
 		if c, is := obj.(d2target.Connection); is {
-			labelMask, err := drawConnection(buf, isolatedDiagramHash, c, markers, idToShape, sketchRunner, inlineTheme)
+			labelMask, err := drawConnection(buf, isolatedDiagramHash, c, markers, idToShape, jsRunner, inlineTheme)
 			if err != nil {
 				return nil, err
 			}
@@ -1949,7 +1950,7 @@ func Render(diagram *d2target.Diagram, opts *RenderOpts) ([]byte, error) {
 				labelMasks = append(labelMasks, labelMask)
 			}
 		} else if s, is := obj.(d2target.Shape); is {
-			labelMask, err := drawShape(buf, appendixItemBuf, diagramHash, s, sketchRunner, inlineTheme)
+			labelMask, err := drawShape(buf, appendixItemBuf, diagramHash, s, jsRunner, inlineTheme)
 			if err != nil {
 				return nil, err
 			} else if labelMask != "" {
@@ -2003,7 +2004,7 @@ func Render(diagram *d2target.Diagram, opts *RenderOpts) ([]byte, error) {
 			fmt.Fprintf(upperBuf, `<style type="text/css">%s</style>`, css)
 		}
 
-		if sketchRunner != nil {
+		if jsRunner != nil {
 			d2sketch.DefineFillPatterns(upperBuf)
 		}
 	}
