wasm: Add support for comprehensions

This commit adds suppor for set/array/object comprehensions. No
changes are required in the wasm backend because we already have the
built-ins for constructing collections dynamically.

Fixes open-policy-agent#1120

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

internal/planner/planner.go

@@ -538,7 +538,7 @@ func (p *Planner) planCallArgs(terms []*ast.Term, idx int, args []ir.Local, iter
 func (p *Planner) planUnify(a, b *ast.Term, iter planiter) error {
 
 	switch va := a.Value.(type) {
-	case ast.Null, ast.Boolean, ast.Number, ast.String, ast.Ref, ast.Set:
+	case ast.Null, ast.Boolean, ast.Number, ast.String, ast.Ref, ast.Set, *ast.SetComprehension, *ast.ArrayComprehension, *ast.ObjectComprehension:
 		return p.planTerm(a, func() error {
 			return p.planUnifyLocal(p.ltarget, b, iter)
 		})
@@ -596,7 +596,7 @@ func (p *Planner) planUnifyVar(a ast.Var, b *ast.Term, iter planiter) error {
 
 func (p *Planner) planUnifyLocal(a ir.Local, b *ast.Term, iter planiter) error {
 	switch vb := b.Value.(type) {
-	case ast.Null, ast.Boolean, ast.Number, ast.String, ast.Ref, ast.Set:
+	case ast.Null, ast.Boolean, ast.Number, ast.String, ast.Ref, ast.Set, *ast.SetComprehension, *ast.ArrayComprehension, *ast.ObjectComprehension:
 		return p.planTerm(b, func() error {
 			p.appendStmt(&ir.EqualStmt{
 				A: a,
@@ -791,6 +791,12 @@ func (p *Planner) planTerm(t *ast.Term, iter planiter) error {
 		return p.planObject(v, iter)
 	case ast.Set:
 		return p.planSet(v, iter)
+	case *ast.SetComprehension:
+		return p.planSetComprehension(v, iter)
+	case *ast.ArrayComprehension:
+		return p.planArrayComprehension(v, iter)
+	case *ast.ObjectComprehension:
+		return p.planObjectComprehension(v, iter)
 	default:
 		return fmt.Errorf("%v term not implemented", ast.TypeName(v))
 	}
@@ -953,6 +959,91 @@ func (p *Planner) planSetRec(set ast.Set, index int, elems []*ast.Term, lset ir.
 	})
 }
 
+func (p *Planner) planSetComprehension(sc *ast.SetComprehension, iter planiter) error {
+
+	lset := p.newLocal()
+
+	p.appendStmt(&ir.MakeSetStmt{
+		Target: lset,
+	})
+
+	return p.planComprehension(sc.Body, func() error {
+		return p.planTerm(sc.Term, func() error {
+			p.appendStmt(&ir.SetAddStmt{
+				Value: p.ltarget,
+				Set:   lset,
+			})
+			return nil
+		})
+	}, lset, iter)
+}
+
+func (p *Planner) planArrayComprehension(ac *ast.ArrayComprehension, iter planiter) error {
+
+	larr := p.newLocal()
+
+	p.appendStmt(&ir.MakeArrayStmt{
+		Target: larr,
+	})
+
+	return p.planComprehension(ac.Body, func() error {
+		return p.planTerm(ac.Term, func() error {
+			p.appendStmt(&ir.ArrayAppendStmt{
+				Value: p.ltarget,
+				Array: larr,
+			})
+			return nil
+		})
+	}, larr, iter)
+}
+
+func (p *Planner) planObjectComprehension(oc *ast.ObjectComprehension, iter planiter) error {
+
+	lobj := p.newLocal()
+
+	p.appendStmt(&ir.MakeObjectStmt{
+		Target: lobj,
+	})
+
+	return p.planComprehension(oc.Body, func() error {
+		return p.planTerm(oc.Key, func() error {
+			lkey := p.ltarget
+			return p.planTerm(oc.Value, func() error {
+				p.appendStmt(&ir.ObjectInsertOnceStmt{
+					Key:    lkey,
+					Value:  p.ltarget,
+					Object: lobj,
+				})
+				return nil
+			})
+		})
+	}, lobj, iter)
+}
+
+func (p *Planner) planComprehension(body ast.Body, closureIter planiter, target ir.Local, iter planiter) error {
+
+	prev := p.curr
+	p.curr = &ir.Block{}
+
+	if err := p.planQuery(body, 0, func() error {
+		return closureIter()
+	}); err != nil {
+		return err
+	}
+
+	block := p.curr
+	p.curr = prev
+
+	p.appendStmt(&ir.BlockStmt{
+		Blocks: []*ir.Block{
+			block,
+		},
+	})
+
+	p.ltarget = target
+	return iter()
+}
+
 func (p *Planner) planRef(ref ast.Ref, iter planiter) error {
 
 	head, ok := ref[0].Value.(ast.Var)

internal/planner/planner_test.go

@@ -146,6 +146,14 @@ func TestPlannerHelloWorld(t *testing.T) {
 				q = 2 { false }
 			`},
 		},
+		{
+			note:    "comprehension",
+			queries: []string{`{x | input[_] = x}`},
+		},
+		{
+			note:    "closure",
+			queries: []string{`a = [1]; {x | a[_] = x}`},
+		},
 	}
 
 	for _, tc := range tests {

test/wasm/assets/014_comprehensions.yaml

@@ -0,0 +1,45 @@
+cases:
+- note: set comprehension
+  query: |
+    {x | input[_] = x} == {1,2,3}
+  input: [1,2,3]
+  return_code: 1
+- note: set comprehension (negative)
+  query: |
+    {x | input[_] = x; x > 1} == {1,2,3}
+  input: [1,2,3]
+  return_code: 0
+- note: array comprehension
+  query: |
+    [x | input[_] = x] == [1,2,3]
+  input: [1,2,3]
+  return_code: 1
+- note: array comprehension (negative)
+  query: |
+    [x | input[_] = x; x > 1] == [1,2,3]
+  input: [1,2,3]
+  return_code: 0
+- note: array comprehension unify
+  query: |
+    [x | input[_] = x] = [1,y,3]
+  input: [1,2,3]
+  return_code: 1
+- note: object comprehension
+  query: |
+    {k: v | input[k] = v} == {"a": 1, "b": 2}
+  input: {"a": 1, "b": 2}
+  return_code: 1
+- note: object comprehension (negative)
+  query: |
+    {k: v | input[k] = v; v > 1} == {"a": 1, "b": 2}
+  input: {"a": 1, "b": 2}
+  return_code: 0
+- note: object comprehension unify
+  query: |
+    {k: v | input[k] = v} = {"a": y, "b": z}
+  input: {"a": 1, "b": 2}
+  return_code: 1
+- note: closure
+  query: |
+    a = [1,2,3]; b = 1; {x | a[_] = x; x > b} == {2,3}
+  return_code: 1