planner: when planning a ref's extent, plan the keys as values (#5257)

The compiler ensures that all the keys we see there are scalars. For strings,
nothing changes -- they're handled just like before -- but this now also allows
numbers and booleans.

An example policy that exploded with a panic before is

    package p
    a[0] = true

when querying the full extent of `data.p` or `data.p.a`.

Fixes #5252.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

internal/planner/planner.go

@@ -1308,8 +1308,11 @@ func (p *Planner) planUnifyObjectsRec(a, b ast.Object, keys []*ast.Term, index i
 }
 
 func (p *Planner) planTerm(t *ast.Term, iter planiter) error {
+	return p.planValue(t.Value, t.Loc(), iter)
+}
 
-	switch v := t.Value.(type) {
+func (p *Planner) planValue(t ast.Value, loc *ast.Location, iter planiter) error {
+	switch v := t.(type) {
 	case ast.Null:
 		return p.planNull(v, iter)
 	case ast.Boolean:
@@ -1329,13 +1332,13 @@ func (p *Planner) planTerm(t *ast.Term, iter planiter) error {
 	case ast.Set:
 		return p.planSet(v, iter)
 	case *ast.SetComprehension:
-		p.loc = t.Loc()
+		p.loc = loc
 		return p.planSetComprehension(v, iter)
 	case *ast.ArrayComprehension:
-		p.loc = t.Loc()
+		p.loc = loc
 		return p.planArrayComprehension(v, iter)
 	case *ast.ObjectComprehension:
-		p.loc = t.Loc()
+		p.loc = loc
 		return p.planObjectComprehension(v, iter)
 	default:
 		return fmt.Errorf("%v term not implemented", ast.TypeName(v))
@@ -1918,46 +1921,48 @@ func (p *Planner) planRefDataExtent(virtual *ruletrie, base *baseptr, iter plani
 					continue
 				}
 
-				lkey := ir.StringIndex(p.getStringConst(string(key.(ast.String))))
 				rules := child.Rules()
-
-				// Build object hierarchy depth-first.
-				if len(rules) == 0 {
-					err := p.planRefDataExtent(child, nil, func() error {
-						p.appendStmt(&ir.ObjectInsertStmt{
-							Object: vtarget,
-							Key:    op(lkey),
-							Value:  p.ltarget,
+				err := p.planValue(key, nil, func() error {
+					lkey := p.ltarget
+
+					// Build object hierarchy depth-first.
+					if len(rules) == 0 {
+						return p.planRefDataExtent(child, nil, func() error {
+							p.appendStmt(&ir.ObjectInsertStmt{
+								Object: vtarget,
+								Key:    lkey,
+								Value:  p.ltarget,
+							})
+							return nil
 						})
-						return nil
-					})
+					}
+
+					// Generate virtual document for leaf.
+					lvalue := p.newLocal()
+
+					funcName, err := p.planRules(rules, false)
 					if err != nil {
 						return err
 					}
-					continue
-				}
 
-				// Generate virtual document for leaf.
-				lvalue := p.newLocal()
-
-				funcName, err := p.planRules(rules, false)
+					// Add leaf to object if defined.
+					b := &ir.Block{}
+					p.appendStmtToBlock(&ir.CallStmt{
+						Func:   funcName,
+						Args:   p.defaultOperands(),
+						Result: lvalue,
+					}, b)
+					p.appendStmtToBlock(&ir.ObjectInsertStmt{
+						Object: vtarget,
+						Key:    lkey,
+						Value:  op(lvalue),
+					}, b)
+					p.appendStmt(&ir.BlockStmt{Blocks: []*ir.Block{b}})
+					return nil
+				})
 				if err != nil {
 					return err
 				}
-
-				// Add leaf to object if defined.
-				b := &ir.Block{}
-				p.appendStmtToBlock(&ir.CallStmt{
-					Func:   funcName,
-					Args:   p.defaultOperands(),
-					Result: lvalue,
-				}, b)
-				p.appendStmtToBlock(&ir.ObjectInsertStmt{
-					Object: vtarget,
-					Key:    op(lkey),
-					Value:  op(lvalue),
-				}, b)
-				p.appendStmt(&ir.BlockStmt{Blocks: []*ir.Block{b}})
 			}
 		}
 

test/cases/testdata/refheads/test-regressions.yaml

@@ -113,4 +113,42 @@ cases:
     p[{"a": "b"}] = true
   query: data.test.p[{"a":"b"}] = x
   want_result:
-  - x: true
+  - x: true
+- note: regression/full extent with non-string (number) last term
+  modules:
+  - |
+    package test
+    p[0] = true
+  query: data.test = x
+  want_result:
+  - x:
+     p:
+       "0": true # stringified key
+- note: regression/full extent with non-string last term, comparison
+  modules:
+  - |
+    package test
+    p[0] = 1
+    q { p[0] == 1 }
+  query: data.test.q = x
+  want_result:
+  - x: true
+- note: regression/full extent with non-string (boolean) last term
+  modules:
+  - |
+    package test
+    p[true] = true
+  query: data.test = x
+  want_result:
+  - x:
+     p:
+       "true": true # stringified key
+- note: regression/full extent with non-string last term, comparison
+  modules:
+  - |
+    package test
+    p[true] = false
+    q { p[true] == false }
+  query: data.test.q = x
+  want_result:
+  - x: true