From 387c33660fcd4c185edfbfdffe13631a030a357c Mon Sep 17 00:00:00 2001
From: Anmol Sethi <hi@nhooyr.io>
Date: Tue, 29 Aug 2023 21:21:26 -0700
Subject: [PATCH] d2parser: Add not ampersand support

See #1567
---
 d2ast/d2ast.go                                |  3 +
 d2parser/parse.go                             | 19 ++++-
 d2parser/parse_test.go                        | 12 +++
 .../d2parser/TestParse/not-amper.exp.json     | 76 +++++++++++++++++++
 4 files changed, 106 insertions(+), 4 deletions(-)
 create mode 100644 testdata/d2parser/TestParse/not-amper.exp.json

diff --git a/d2ast/d2ast.go b/d2ast/d2ast.go
index fdbe9c9c32..8434c4b960 100644
--- a/d2ast/d2ast.go
+++ b/d2ast/d2ast.go
@@ -613,6 +613,9 @@ type Key struct {
 	// Indicates this MapKey is a filter selector.
 	Ampersand bool `json:"ampersand,omitempty"`
 
+	// Indicates this MapKey is a not filter selector.
+	NotAmpersand bool `json:"not_ampersand,omitempty"`
+
 	// At least one of Key and Edges will be set but all four can also be set.
 	// The following are all valid MapKeys:
 	// Key:
diff --git a/d2parser/parse.go b/d2parser/parse.go
index 7320384807..84d032a586 100644
--- a/d2parser/parse.go
+++ b/d2parser/parse.go
@@ -660,16 +660,27 @@ func (p *parser) parseMapKey() (mk *d2ast.Key) {
 		}
 	}()
 
-	// Check for ampersand/@.
+	// Check for not ampersand/@.
 	r, eof := p.peek()
 	if eof {
 		return mk
 	}
-	if r != '&' {
-		p.rewind()
-	} else {
+	if r == '!' {
+		r, eof := p.peek()
+		if eof {
+			return mk
+		}
+		if r == '&' {
+			p.commit()
+			mk.NotAmpersand = true
+		} else {
+			p.rewind()
+		}
+	} else if r == '&' {
 		p.commit()
 		mk.Ampersand = true
+	} else {
+		p.rewind()
 	}
 
 	r, eof = p.peek()
diff --git a/d2parser/parse_test.go b/d2parser/parse_test.go
index 605be68f7f..f4be7de0f7 100644
--- a/d2parser/parse_test.go
+++ b/d2parser/parse_test.go
@@ -380,6 +380,18 @@ b-
 c-
 `,
 		},
+		{
+			name: "not-amper",
+			text: `
+&k: amper
+!&k: not amper
+`,
+			assert: func(t testing.TB, ast *d2ast.Map, err error) {
+				assert.Success(t, err)
+				assert.True(t, ast.Nodes[0].MapKey.Ampersand)
+				assert.True(t, ast.Nodes[1].MapKey.NotAmpersand)
+			},
+		},
 		{
 			name: "whitespace_range",
 			text: ` a -> b -> c `,
diff --git a/testdata/d2parser/TestParse/not-amper.exp.json b/testdata/d2parser/TestParse/not-amper.exp.json
new file mode 100644
index 0000000000..974490d9e2
--- /dev/null
+++ b/testdata/d2parser/TestParse/not-amper.exp.json
@@ -0,0 +1,76 @@
+{
+  "ast": {
+    "range": "d2/testdata/d2parser/TestParse/not-amper.d2,0:0:0-3:0:26",
+    "nodes": [
+      {
+        "map_key": {
+          "range": "d2/testdata/d2parser/TestParse/not-amper.d2,1:0:1-1:9:10",
+          "ampersand": true,
+          "key": {
+            "range": "d2/testdata/d2parser/TestParse/not-amper.d2,1:1:2-1:2:3",
+            "path": [
+              {
+                "unquoted_string": {
+                  "range": "d2/testdata/d2parser/TestParse/not-amper.d2,1:1:2-1:2:3",
+                  "value": [
+                    {
+                      "string": "k",
+                      "raw_string": "k"
+                    }
+                  ]
+                }
+              }
+            ]
+          },
+          "primary": {},
+          "value": {
+            "unquoted_string": {
+              "range": "d2/testdata/d2parser/TestParse/not-amper.d2,1:4:5-1:9:10",
+              "value": [
+                {
+                  "string": "amper",
+                  "raw_string": "amper"
+                }
+              ]
+            }
+          }
+        }
+      },
+      {
+        "map_key": {
+          "range": "d2/testdata/d2parser/TestParse/not-amper.d2,2:0:11-2:14:25",
+          "not_ampersand": true,
+          "key": {
+            "range": "d2/testdata/d2parser/TestParse/not-amper.d2,2:2:13-2:3:14",
+            "path": [
+              {
+                "unquoted_string": {
+                  "range": "d2/testdata/d2parser/TestParse/not-amper.d2,2:2:13-2:3:14",
+                  "value": [
+                    {
+                      "string": "k",
+                      "raw_string": "k"
+                    }
+                  ]
+                }
+              }
+            ]
+          },
+          "primary": {},
+          "value": {
+            "unquoted_string": {
+              "range": "d2/testdata/d2parser/TestParse/not-amper.d2,2:5:16-2:14:25",
+              "value": [
+                {
+                  "string": "not amper",
+                  "raw_string": "not amper"
+                }
+              ]
+            }
+          }
+        }
+      }
+    ]
+  },
+  "err": null
+}