From 4e93f70db384fc34e62de708ae42acda59d79285 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Wed, 11 Sep 2024 10:34:39 -0600
Subject: [PATCH] d2ir/imports: fix removing parent scope of nested import
 fields/edges

---
 d2compiler/compile_test.go                    |   4 +-
 d2ir/import.go                                |  11 +-
 .../TestCompile/vars-in-imports.exp.json      | 113 +++++++++++-------
 3 files changed, 74 insertions(+), 54 deletions(-)

diff --git a/d2compiler/compile_test.go b/d2compiler/compile_test.go
index f077bccd4..662fe1062 100644
--- a/d2compiler/compile_test.go
+++ b/d2compiler/compile_test.go
@@ -2979,8 +2979,8 @@ qa: {
 			assertions: func(t *testing.T, g *d2graph.Graph) {
 				tassert.Equal(t, "dev.env", g.Objects[1].AbsID())
 				tassert.Equal(t, "Dev Environment", g.Objects[1].Label.Value)
-				tassert.Equal(t, "qa.env", g.Objects[2].AbsID())
-				tassert.Equal(t, "Qa Environment", g.Objects[2].Label.Value)
+				tassert.Equal(t, "qa.env", g.Objects[4].AbsID())
+				tassert.Equal(t, "Qa Environment", g.Objects[4].Label.Value)
 			},
 		},
 		{
diff --git a/d2ir/import.go b/d2ir/import.go
index c415cf138..ef80d358c 100644
--- a/d2ir/import.go
+++ b/d2ir/import.go
@@ -61,15 +61,18 @@ func (c *compiler) _import(imp *d2ast.Import) (Node, bool) {
 	if !ok {
 		return nil, false
 	}
-	nilScopeMap(ir)
 	if len(imp.IDA()) > 0 {
 		f := ir.GetField(imp.IDA()...)
 		if f == nil {
 			c.errorf(imp, "import key %q doesn't exist inside import", imp.IDA())
 			return nil, false
 		}
+		if f.Map() != nil {
+			nilScopeMap(f.Map())
+		}
 		return f, true
 	}
+	nilScopeMap(ir)
 	return ir, true
 }
 
@@ -132,15 +135,9 @@ func nilScopeMap(n Node) {
 		for _, r := range n.References {
 			r.Context_.ScopeMap = nil
 		}
-		if n.Map() != nil {
-			nilScopeMap(n.Map())
-		}
 	case *Field:
 		for _, r := range n.References {
 			r.Context_.ScopeMap = nil
 		}
-		if n.Map() != nil {
-			nilScopeMap(n.Map())
-		}
 	}
 }
diff --git a/testdata/d2compiler/TestCompile/vars-in-imports.exp.json b/testdata/d2compiler/TestCompile/vars-in-imports.exp.json
index 05c63f5e3..fad97a840 100644
--- a/testdata/d2compiler/TestCompile/vars-in-imports.exp.json
+++ b/testdata/d2compiler/TestCompile/vars-in-imports.exp.json
@@ -344,51 +344,6 @@
         },
         "zIndex": 0
       },
-      {
-        "id": "env",
-        "id_val": "env",
-        "references": [
-          {
-            "key": {
-              "range": "d2/testdata/d2compiler/TestCompile/template.d2,0:0:0-0:3:3",
-              "path": [
-                {
-                  "unquoted_string": {
-                    "range": "d2/testdata/d2compiler/TestCompile/template.d2,0:0:0-0:3:3",
-                    "value": [
-                      {
-                        "string": "env",
-                        "raw_string": "env"
-                      }
-                    ]
-                  }
-                }
-              ]
-            },
-            "key_path_index": 0,
-            "map_key_edge_index": -1
-          }
-        ],
-        "attributes": {
-          "label": {
-            "value": "Qa Environment"
-          },
-          "labelDimensions": {
-            "width": 0,
-            "height": 0
-          },
-          "style": {},
-          "near_key": null,
-          "shape": {
-            "value": "rectangle"
-          },
-          "direction": {
-            "value": ""
-          },
-          "constraint": null
-        },
-        "zIndex": 0
-      },
       {
         "id": "vm",
         "id_val": "vm",
@@ -434,6 +389,74 @@
         },
         "zIndex": 0
       },
+      {
+        "id": "env",
+        "id_val": "env",
+        "attributes": {
+          "label": {
+            "value": "env"
+          },
+          "labelDimensions": {
+            "width": 0,
+            "height": 0
+          },
+          "style": {},
+          "near_key": null,
+          "shape": {
+            "value": "rectangle"
+          },
+          "direction": {
+            "value": ""
+          },
+          "constraint": null
+        },
+        "zIndex": 0
+      },
+      {
+        "id": "env",
+        "id_val": "env",
+        "references": [
+          {
+            "key": {
+              "range": "d2/testdata/d2compiler/TestCompile/template.d2,0:0:0-0:3:3",
+              "path": [
+                {
+                  "unquoted_string": {
+                    "range": "d2/testdata/d2compiler/TestCompile/template.d2,0:0:0-0:3:3",
+                    "value": [
+                      {
+                        "string": "env",
+                        "raw_string": "env"
+                      }
+                    ]
+                  }
+                }
+              ]
+            },
+            "key_path_index": 0,
+            "map_key_edge_index": -1
+          }
+        ],
+        "attributes": {
+          "label": {
+            "value": "Qa Environment"
+          },
+          "labelDimensions": {
+            "width": 0,
+            "height": 0
+          },
+          "style": {},
+          "near_key": null,
+          "shape": {
+            "value": "rectangle"
+          },
+          "direction": {
+            "value": ""
+          },
+          "constraint": null
+        },
+        "zIndex": 0
+      },
       {
         "id": "vm",
         "id_val": "vm",