ast: Fix for with node locations

We were not setting locations on `with` expressions. This adds in support for doing so. Signed-off-by: Patrick East <[email protected]>
mj520 · Mar 30, 2020 · 0c593cf · 0c593cf
1 parent 6209e15
commit 0c593cf
Show file tree

Hide file tree

Showing 4 changed files with 122 additions and 4 deletions.
diff --git a/ast/parser.go b/ast/parser.go
@@ -572,10 +572,9 @@ func (p *Parser) parseWith() []*With {
 
 	for {
 
-		// NOTE(tsandall): location is not being set correctly on with
-		// statements. need special test case for this.
-
-		var with With
+		with := With{
+			Location: p.s.Loc(),
+		}
 		p.scan()
 
 		if p.s.tok != tokens.Ident {
@@ -605,6 +604,8 @@ func (p *Parser) parseWith() []*With {
 			return nil
 		}
 
+		with.Location.Text = p.s.Text(with.Location.Offset, p.s.lastEnd)
+
 		withs = append(withs, &with)
 
 		if p.s.tok != tokens.With {

diff --git a/ast/parser_test.go b/ast/parser_test.go
@@ -542,6 +542,100 @@ func TestExprWith(t *testing.T) {
 	})
 }
 
+func TestExprWithLocation(t *testing.T) {
+	cases := []struct {
+		note     string
+		input    string
+		expected []*Location
+	}{
+		{
+			note:  "base",
+			input: "a with b as c",
+			expected: []*Location{
+				{
+					Row:    1,
+					Col:    3,
+					Offset: 2,
+					Text:   []byte("with b as c"),
+				},
+			},
+		},
+		{
+			note:  "with line break",
+			input: "a with b\nas c",
+			expected: []*Location{
+				{
+					Row:    1,
+					Col:    3,
+					Offset: 2,
+					Text:   []byte("with b\nas c"),
+				},
+			},
+		},
+		{
+			note:  "multiple withs on single line",
+			input: "a with b as c with d as e",
+			expected: []*Location{
+				{
+					Row:    1,
+					Col:    3,
+					Offset: 2,
+					Text:   []byte("with b as c"),
+				},
+				{
+					Row:    1,
+					Col:    15,
+					Offset: 14,
+					Text:   []byte("with d as e"),
+				},
+			},
+		},
+		{
+			note:  "multiple withs on multiple line",
+			input: "a with b as c\n\t\twith d as e",
+			expected: []*Location{
+				{
+					Row:    1,
+					Col:    3,
+					Offset: 2,
+					Text:   []byte("with b as c"),
+				},
+				{
+					Row:    2,
+					Col:    3,
+					Offset: 16,
+					Text:   []byte("with d as e"),
+				},
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.note, func(t *testing.T) {
+			parsed, err := ParseStatement(tc.input)
+			if err != nil {
+				t.Errorf("Unexpected error on %s: %s", tc.input, err)
+				return
+			}
+
+			body := parsed.(Body)
+			if len(body) != 1 {
+				t.Errorf("Parser returned multiple expressions: %v", body)
+				return
+			}
+			expr := body[0]
+			if len(expr.With) != len(tc.expected) {
+				t.Fatalf("Expected %d with statements, got %d", len(expr.With), len(tc.expected))
+			}
+			for i, with := range expr.With {
+				if !with.Location.Equal(tc.expected[i]) {
+					t.Errorf("Expected location %+v for '%v' but got %+v ", *(tc.expected[i]), with.String(), *with.Location)
+				}
+			}
+		})
+	}
+}
+
 func TestSomeDeclExpr(t *testing.T) {
 
 	assertParseOneExpr(t, "one", "some x", &Expr{

diff --git a/format/testfiles/test.rego b/format/testfiles/test.rego
@@ -195,6 +195,18 @@ declare1 := 1
 
 declare2 := 2 { false }
 
+multi_line_with {
+    fn(1) with input.a as "a"
+                with input.b as "b"
+            with input.c as {
+                "foo": "bar",
+            }
+                with input.d as [
+                    1,
+                    2,
+                    3]
+}
+
 # more comments!
 # more comments!
 # more comments!

diff --git a/format/testfiles/test.rego.formatted b/format/testfiles/test.rego.formatted
@@ -230,6 +230,17 @@ declare2 := 2 {
 	false
 }
 
+multi_line_with {
+	fn(1) with input.a as "a"
+		 with input.b as "b"
+		 with input.c as {"foo": "bar"}
+		 with input.d as [
+			1,
+			2,
+			3,
+		]
+}
+
 # more comments!
 # more comments!
 # more comments!