d2ir: Add filtering on class arrays

nhooyr · Jul 30, 2023 · e61ca0b · e61ca0b
1 parent 571fc62
commit e61ca0b
Show file tree

Hide file tree

Showing 5 changed files with 1,279 additions and 141 deletions.
diff --git a/d2ir/compile.go b/d2ir/compile.go
@@ -410,6 +410,10 @@ func (c *compiler) compileKey(refctx *RefContext) {
 }
 
 func (c *compiler) compileField(dst *Map, kp *d2ast.KeyPath, refctx *RefContext) {
+	if refctx.Key.Ampersand {
+		return
+	}
+
 	fa, err := dst.EnsureField(kp, refctx, true)
 	if err != nil {
 		c.err.Errors = append(c.err.Errors, err.(d2ast.Error))
@@ -451,23 +455,26 @@ func (c *compiler) _ampersandFilter(f *Field, refctx *RefContext) bool {
 	if f2 == nil {
 		return false
 	}
-	if refctx.Key.Primary.Unbox() != nil {
-		if f2.Primary_ == nil {
-			return false
-		}
-		if refctx.Key.Primary.Unbox().ScalarString() != f2.Primary_.Value.ScalarString() {
-			return false
-		}
+	if refctx.Key.Value.ScalarBox().Unbox() == nil {
+		c.errorf(refctx.Key, "ampersand filters cannot be composites")
+		return false
 	}
-	if refctx.Key.Value.ScalarBox().Unbox() != nil {
-		if f2.Primary_ == nil {
-			return false
-		}
-		if refctx.Key.Value.ScalarBox().Unbox().ScalarString() != f2.Primary_.Value.ScalarString() {
-			return false
+
+	if a, ok := f2.Composite.(*Array); ok {
+		for _, v := range a.Values {
+			if s, ok := v.(*Scalar); ok {
+				if refctx.Key.Value.ScalarBox().Unbox().ScalarString() == s.Value.ScalarString() {
+					return true
+				}
+			}
 		}
-	} else if refctx.Key.Value.Unbox() != nil {
-		c.errorf(refctx.Key, "ampersand filters cannot be composites")
+	}
+
+	if f2.Primary_ == nil {
+		return false
+	}
+
+	if refctx.Key.Value.ScalarBox().Unbox().ScalarString() != f2.Primary_.Value.ScalarString() {
 		return false
 	}
 

diff --git a/d2ir/filter_test.go b/d2ir/filter_test.go
@@ -43,11 +43,37 @@ jeremy: {
 	&shape: rectangle
 }`)
 				assert.Success(t, err)
+				assertQuery(t, m, 5, 0, nil, "")
 				assertQuery(t, m, 1, 0, nil, "jacob")
 				assertQuery(t, m, 2, 0, nil, "jeremy")
 				assertQuery(t, m, 0, 0, "I'm a rectangle", "jeremy.label")
 			},
 		},
+		{
+			name: "array",
+			run: func(t testing.TB) {
+				m, err := compile(t, `the-little-cannon: {
+	class: [server; deployed]
+}
+dino: {
+	class: [internal; deployed]
+}
+catapult: {
+	class: [jacob; server]
+}
+
+*: {
+	&class: server
+	style.multiple: true
+}
+`)
+				assert.Success(t, err)
+				assertQuery(t, m, 10, 0, nil, "")
+				assertQuery(t, m, 3, 0, nil, "the-little-cannon")
+				assertQuery(t, m, 1, 0, nil, "dino")
+				assertQuery(t, m, 3, 0, nil, "catapult")
+			},
+		},
 	}
 
 	runa(t, tca)