From a5d3cc14298625aed9d81e7aff7c701204b7377f Mon Sep 17 00:00:00 2001
From: Anmol Sethi <hi@nhooyr.io>
Date: Tue, 6 Jun 2023 16:15:01 -0700
Subject: [PATCH] d2parser: Parse and format @../file imports correctly

---
 d2ast/d2ast.go                                |  9 ++++
 d2format/format.go                            |  6 +++
 d2format/format_test.go                       |  8 +++
 d2ir/import.go                                |  6 +--
 d2parser/parse.go                             | 16 ++++--
 d2parser/parse_test.go                        | 14 ++++-
 .../TestCompile/imports/nested/array.exp.json |  1 +
 .../TestCompile/imports/nested/map.exp.json   |  1 +
 .../imports/nested/scalar.exp.json            |  1 +
 .../d2ir/TestCompile/imports/value.exp.json   |  1 +
 .../d2parser/TestParse/import/#00.exp.json    |  1 +
 .../d2parser/TestParse/import/#01.exp.json    |  1 +
 .../d2parser/TestParse/import/#02.exp.json    |  1 +
 .../d2parser/TestParse/import/#03.exp.json    |  1 +
 .../d2parser/TestParse/import/#04.exp.json    |  1 +
 .../d2parser/TestParse/import/#05.exp.json    |  1 +
 .../d2parser/TestParse/import/#06.exp.json    | 26 +++++----
 .../d2parser/TestParse/import/#08.exp.json    | 54 ++++++-------------
 .../d2parser/TestParse/import/#09.exp.json    | 48 +++++++++++++++++
 19 files changed, 141 insertions(+), 56 deletions(-)
 create mode 100644 testdata/d2parser/TestParse/import/#09.exp.json

diff --git a/d2ast/d2ast.go b/d2ast/d2ast.go
index 76f3375ea..eee17d368 100644
--- a/d2ast/d2ast.go
+++ b/d2ast/d2ast.go
@@ -28,6 +28,7 @@ import (
 	"errors"
 	"fmt"
 	"math/big"
+	"path"
 	"strconv"
 	"strings"
 	"unicode"
@@ -746,6 +747,7 @@ type Import struct {
 	Range Range `json:"range"`
 
 	Spread bool         `json:"spread"`
+	Pre    string       `json:"pre"`
 	Path   []*StringBox `json:"path"`
 }
 
@@ -1105,3 +1107,10 @@ func (i *Import) IDA() (ida []string) {
 	}
 	return ida
 }
+
+func (i *Import) PathWithPre() string {
+	if len(i.Path) == 0 {
+		return ""
+	}
+	return path.Join(i.Pre, i.Path[0].Unbox().ScalarString())
+}
diff --git a/d2format/format.go b/d2format/format.go
index 3572f3f06..84df2e945 100644
--- a/d2format/format.go
+++ b/d2format/format.go
@@ -1,6 +1,7 @@
 package d2format
 
 import (
+	"path"
 	"strconv"
 	"strings"
 
@@ -210,6 +211,11 @@ func (p *printer) _import(i *d2ast.Import) {
 		p.sb.WriteString("...")
 	}
 	p.sb.WriteString("@")
+	pre := path.Clean(i.Pre)
+	if pre != "." {
+		p.sb.WriteString(pre)
+		p.sb.WriteRune('/')
+	}
 	p.path(i.Path)
 }
 
diff --git a/d2format/format_test.go b/d2format/format_test.go
index 52a1a6c16..63e185dcf 100644
--- a/d2format/format_test.go
+++ b/d2format/format_test.go
@@ -641,6 +641,14 @@ x: @file."d2"
 x: @./file
 `,
 			exp: `x: @file
+`,
+		},
+		{
+			name: "import/4",
+			in: `
+x: @../file
+`,
+			exp: `x: @../file
 `,
 		},
 	}
diff --git a/d2ir/import.go b/d2ir/import.go
index 52fce2e7a..7c0a5b582 100644
--- a/d2ir/import.go
+++ b/d2ir/import.go
@@ -11,12 +11,12 @@ import (
 )
 
 func (c *compiler) pushImportStack(imp *d2ast.Import) bool {
-	if len(imp.Path) == 0 {
+	if imp.PathWithPre() == "" && imp.Range.Path != "" {
 		c.errorf(imp, "imports must specify a path to import")
 		return false
 	}
 
-	newPath := imp.Path[0].Unbox().ScalarString()
+	newPath := imp.PathWithPre()
 	for i, p := range c.importStack {
 		if newPath == p {
 			c.errorf(imp, "detected cyclic import chain: %s", formatCyclicChain(c.importStack[i:]))
@@ -61,7 +61,7 @@ func (c *compiler) _import(imp *d2ast.Import) (Node, bool) {
 }
 
 func (c *compiler) __import(imp *d2ast.Import) (*Map, bool) {
-	impPath := imp.Path[0].Unbox().ScalarString()
+	impPath := imp.PathWithPre()
 	if path.IsAbs(impPath) {
 		c.errorf(imp, "import paths must be relative")
 		return nil, false
diff --git a/d2parser/parse.go b/d2parser/parse.go
index c30a65621..7970f1fa6 100644
--- a/d2parser/parse.go
+++ b/d2parser/parse.go
@@ -1720,12 +1720,20 @@ func (p *parser) parseImport(spread bool) *d2ast.Import {
 		imp.Range.Start = imp.Range.Start.SubtractString("...", p.utf16)
 	}
 
-	s, eof := p.peekn(2)
-	if eof || s != "./" {
-		p.rewind()
-	} else {
+	var pre strings.Builder
+	for {
+		r, eof := p.peek()
+		if eof {
+			break
+		}
+		if r != '.' && r != '/' {
+			p.rewind()
+			break
+		}
+		pre.WriteRune(r)
 		p.commit()
 	}
+	imp.Pre = pre.String()
 
 	k := p.parseKey()
 	if k == nil {
diff --git a/d2parser/parse_test.go b/d2parser/parse_test.go
index dc4e14186..2dc195d02 100644
--- a/d2parser/parse_test.go
+++ b/d2parser/parse_test.go
@@ -454,7 +454,9 @@ func testImport(t *testing.T) {
 		{
 			text: "...@../file",
 			assert: func(t testing.TB, ast *d2ast.Map, err error) {
-				assert.ErrorString(t, err, "d2/testdata/d2parser/TestParse/import/#06.d2:1:5: unexpected text after import")
+				assert.Success(t, err)
+				assert.True(t, ast.Nodes[0].Import.Spread)
+				assert.Equal(t, "../file", ast.Nodes[0].Import.PathWithPre())
 			},
 		},
 		{
@@ -463,10 +465,18 @@ func testImport(t *testing.T) {
 				assert.ErrorString(t, err, "d2/testdata/d2parser/TestParse/import/#07.d2:1:1: @file is not a valid import, did you mean ...@file?")
 			},
 		},
+		{
+			text: "...@./../.././file",
+			assert: func(t testing.TB, ast *d2ast.Map, err error) {
+				assert.Success(t, err)
+				assert.True(t, ast.Nodes[0].Import.Spread)
+				assert.Equal(t, "../../file", ast.Nodes[0].Import.PathWithPre())
+			},
+		},
 		{
 			text: "meow: ...@file",
 			assert: func(t testing.TB, ast *d2ast.Map, err error) {
-				assert.ErrorString(t, err, "d2/testdata/d2parser/TestParse/import/#08.d2:1:7: unquoted strings cannot begin with ...@ as that's import spread syntax")
+				assert.ErrorString(t, err, "d2/testdata/d2parser/TestParse/import/#09.d2:1:7: unquoted strings cannot begin with ...@ as that's import spread syntax")
 			},
 		},
 	}
diff --git a/testdata/d2ir/TestCompile/imports/nested/array.exp.json b/testdata/d2ir/TestCompile/imports/nested/array.exp.json
index 02624ba49..9e1737148 100644
--- a/testdata/d2ir/TestCompile/imports/nested/array.exp.json
+++ b/testdata/d2ir/TestCompile/imports/nested/array.exp.json
@@ -69,6 +69,7 @@
                 "import": {
                   "range": "index.d2,0:3:3-0:7:7",
                   "spread": false,
+                  "pre": "",
                   "path": [
                     {
                       "unquoted_string": {
diff --git a/testdata/d2ir/TestCompile/imports/nested/map.exp.json b/testdata/d2ir/TestCompile/imports/nested/map.exp.json
index d6a74f941..b1dedaa35 100644
--- a/testdata/d2ir/TestCompile/imports/nested/map.exp.json
+++ b/testdata/d2ir/TestCompile/imports/nested/map.exp.json
@@ -213,6 +213,7 @@
                 "import": {
                   "range": "index.d2,0:3:3-0:7:7",
                   "spread": false,
+                  "pre": "",
                   "path": [
                     {
                       "unquoted_string": {
diff --git a/testdata/d2ir/TestCompile/imports/nested/scalar.exp.json b/testdata/d2ir/TestCompile/imports/nested/scalar.exp.json
index 5ab8f8f40..ed273c2c0 100644
--- a/testdata/d2ir/TestCompile/imports/nested/scalar.exp.json
+++ b/testdata/d2ir/TestCompile/imports/nested/scalar.exp.json
@@ -65,6 +65,7 @@
                 "import": {
                   "range": "index.d2,0:3:3-0:7:7",
                   "spread": false,
+                  "pre": "",
                   "path": [
                     {
                       "unquoted_string": {
diff --git a/testdata/d2ir/TestCompile/imports/value.exp.json b/testdata/d2ir/TestCompile/imports/value.exp.json
index 7fbd78c5c..8e5c0f021 100644
--- a/testdata/d2ir/TestCompile/imports/value.exp.json
+++ b/testdata/d2ir/TestCompile/imports/value.exp.json
@@ -213,6 +213,7 @@
                 "import": {
                   "range": "index.d2,0:3:3-0:8:8",
                   "spread": false,
+                  "pre": "",
                   "path": [
                     {
                       "unquoted_string": {
diff --git a/testdata/d2parser/TestParse/import/#00.exp.json b/testdata/d2parser/TestParse/import/#00.exp.json
index 470b26c05..77bf3dfd6 100644
--- a/testdata/d2parser/TestParse/import/#00.exp.json
+++ b/testdata/d2parser/TestParse/import/#00.exp.json
@@ -26,6 +26,7 @@
             "import": {
               "range": "d2/testdata/d2parser/TestParse/import/#00.d2,0:3:3-0:8:8",
               "spread": false,
+              "pre": "",
               "path": [
                 {
                   "unquoted_string": {
diff --git a/testdata/d2parser/TestParse/import/#01.exp.json b/testdata/d2parser/TestParse/import/#01.exp.json
index ab516c33c..298304786 100644
--- a/testdata/d2parser/TestParse/import/#01.exp.json
+++ b/testdata/d2parser/TestParse/import/#01.exp.json
@@ -26,6 +26,7 @@
             "import": {
               "range": "d2/testdata/d2parser/TestParse/import/#01.d2,0:3:3-0:11:11",
               "spread": false,
+              "pre": "",
               "path": [
                 {
                   "unquoted_string": {
diff --git a/testdata/d2parser/TestParse/import/#02.exp.json b/testdata/d2parser/TestParse/import/#02.exp.json
index a8c9309bf..caf633d2b 100644
--- a/testdata/d2parser/TestParse/import/#02.exp.json
+++ b/testdata/d2parser/TestParse/import/#02.exp.json
@@ -6,6 +6,7 @@
         "import": {
           "range": "d2/testdata/d2parser/TestParse/import/#02.d2,0:0:0-0:11:11",
           "spread": true,
+          "pre": "",
           "path": [
             {
               "unquoted_string": {
diff --git a/testdata/d2parser/TestParse/import/#03.exp.json b/testdata/d2parser/TestParse/import/#03.exp.json
index 95c1a6f95..6109861ef 100644
--- a/testdata/d2parser/TestParse/import/#03.exp.json
+++ b/testdata/d2parser/TestParse/import/#03.exp.json
@@ -30,6 +30,7 @@
                   "import": {
                     "range": "d2/testdata/d2parser/TestParse/import/#03.d2,0:4:4-0:15:15",
                     "spread": true,
+                    "pre": "",
                     "path": [
                       {
                         "unquoted_string": {
diff --git a/testdata/d2parser/TestParse/import/#04.exp.json b/testdata/d2parser/TestParse/import/#04.exp.json
index 7f6f3edf9..8cc0420e8 100644
--- a/testdata/d2parser/TestParse/import/#04.exp.json
+++ b/testdata/d2parser/TestParse/import/#04.exp.json
@@ -6,6 +6,7 @@
         "import": {
           "range": "d2/testdata/d2parser/TestParse/import/#04.d2,0:0:0-0:13:13",
           "spread": true,
+          "pre": "",
           "path": [
             {
               "double_quoted_string": {
diff --git a/testdata/d2parser/TestParse/import/#05.exp.json b/testdata/d2parser/TestParse/import/#05.exp.json
index 49eec25d3..f9491edd1 100644
--- a/testdata/d2parser/TestParse/import/#05.exp.json
+++ b/testdata/d2parser/TestParse/import/#05.exp.json
@@ -6,6 +6,7 @@
         "import": {
           "range": "d2/testdata/d2parser/TestParse/import/#05.d2,0:0:0-0:13:13",
           "spread": true,
+          "pre": "",
           "path": [
             {
               "unquoted_string": {
diff --git a/testdata/d2parser/TestParse/import/#06.exp.json b/testdata/d2parser/TestParse/import/#06.exp.json
index 8d0a4fa93..1a4eafe51 100644
--- a/testdata/d2parser/TestParse/import/#06.exp.json
+++ b/testdata/d2parser/TestParse/import/#06.exp.json
@@ -4,19 +4,25 @@
     "nodes": [
       {
         "import": {
-          "range": "d2/testdata/d2parser/TestParse/import/#06.d2,0:0:0-0:4:4",
+          "range": "d2/testdata/d2parser/TestParse/import/#06.d2,0:0:0-0:11:11",
           "spread": true,
-          "path": null
+          "pre": "../",
+          "path": [
+            {
+              "unquoted_string": {
+                "range": "d2/testdata/d2parser/TestParse/import/#06.d2,0:7:7-0:11:11",
+                "value": [
+                  {
+                    "string": "file",
+                    "raw_string": "file"
+                  }
+                ]
+              }
+            }
+          ]
         }
       }
     ]
   },
-  "err": {
-    "errs": [
-      {
-        "range": "d2/testdata/d2parser/TestParse/import/#06.d2,0:4:4-0:11:11",
-        "errmsg": "d2/testdata/d2parser/TestParse/import/#06.d2:1:5: unexpected text after import"
-      }
-    ]
-  }
+  "err": null
 }
diff --git a/testdata/d2parser/TestParse/import/#08.exp.json b/testdata/d2parser/TestParse/import/#08.exp.json
index 3a0df6896..73e586aeb 100644
--- a/testdata/d2parser/TestParse/import/#08.exp.json
+++ b/testdata/d2parser/TestParse/import/#08.exp.json
@@ -1,48 +1,28 @@
 {
   "ast": {
-    "range": "d2/testdata/d2parser/TestParse/import/#08.d2,0:0:0-0:14:14",
+    "range": "d2/testdata/d2parser/TestParse/import/#08.d2,0:0:0-0:18:18",
     "nodes": [
       {
-        "map_key": {
-          "range": "d2/testdata/d2parser/TestParse/import/#08.d2,0:0:0-0:14:14",
-          "key": {
-            "range": "d2/testdata/d2parser/TestParse/import/#08.d2,0:0:0-0:4:4",
-            "path": [
-              {
-                "unquoted_string": {
-                  "range": "d2/testdata/d2parser/TestParse/import/#08.d2,0:0:0-0:4:4",
-                  "value": [
-                    {
-                      "string": "meow",
-                      "raw_string": "meow"
-                    }
-                  ]
-                }
+        "import": {
+          "range": "d2/testdata/d2parser/TestParse/import/#08.d2,0:0:0-0:18:18",
+          "spread": true,
+          "pre": "./../.././",
+          "path": [
+            {
+              "unquoted_string": {
+                "range": "d2/testdata/d2parser/TestParse/import/#08.d2,0:14:14-0:18:18",
+                "value": [
+                  {
+                    "string": "file",
+                    "raw_string": "file"
+                  }
+                ]
               }
-            ]
-          },
-          "primary": {},
-          "value": {
-            "unquoted_string": {
-              "range": "d2/testdata/d2parser/TestParse/import/#08.d2,0:6:6-0:14:14",
-              "value": [
-                {
-                  "string": "...@file",
-                  "raw_string": "...@file"
-                }
-              ]
             }
-          }
+          ]
         }
       }
     ]
   },
-  "err": {
-    "errs": [
-      {
-        "range": "d2/testdata/d2parser/TestParse/import/#08.d2,0:6:6-0:10:10",
-        "errmsg": "d2/testdata/d2parser/TestParse/import/#08.d2:1:7: unquoted strings cannot begin with ...@ as that's import spread syntax"
-      }
-    ]
-  }
+  "err": null
 }
diff --git a/testdata/d2parser/TestParse/import/#09.exp.json b/testdata/d2parser/TestParse/import/#09.exp.json
new file mode 100644
index 000000000..690c52489
--- /dev/null
+++ b/testdata/d2parser/TestParse/import/#09.exp.json
@@ -0,0 +1,48 @@
+{
+  "ast": {
+    "range": "d2/testdata/d2parser/TestParse/import/#09.d2,0:0:0-0:14:14",
+    "nodes": [
+      {
+        "map_key": {
+          "range": "d2/testdata/d2parser/TestParse/import/#09.d2,0:0:0-0:14:14",
+          "key": {
+            "range": "d2/testdata/d2parser/TestParse/import/#09.d2,0:0:0-0:4:4",
+            "path": [
+              {
+                "unquoted_string": {
+                  "range": "d2/testdata/d2parser/TestParse/import/#09.d2,0:0:0-0:4:4",
+                  "value": [
+                    {
+                      "string": "meow",
+                      "raw_string": "meow"
+                    }
+                  ]
+                }
+              }
+            ]
+          },
+          "primary": {},
+          "value": {
+            "unquoted_string": {
+              "range": "d2/testdata/d2parser/TestParse/import/#09.d2,0:6:6-0:14:14",
+              "value": [
+                {
+                  "string": "...@file",
+                  "raw_string": "...@file"
+                }
+              ]
+            }
+          }
+        }
+      }
+    ]
+  },
+  "err": {
+    "errs": [
+      {
+        "range": "d2/testdata/d2parser/TestParse/import/#09.d2,0:6:6-0:10:10",
+        "errmsg": "d2/testdata/d2parser/TestParse/import/#09.d2:1:7: unquoted strings cannot begin with ...@ as that's import spread syntax"
+      }
+    ]
+  }
+}