Add support for ir.Loop in compiler

These changes add basic support for ir.Loop in the compiler. To handle
the loop condition, the WASM library has been extended with a function
to set the value of a local boolean. This allows the compiler to
terminate the loop when a match is found.

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

internal/compiler/wasm/externs.go

@@ -20,8 +20,10 @@ const (
 	opaBoolean
 	opaStringTerminated
 	opaNumberInt
+	opaValueBooleanSet
 	opaValueNotEqual
 	opaValueGet
+	opaValueIter
 )
 
 var externs = [...]module.Import{
@@ -49,6 +51,12 @@ var externs = [...]module.Import{
 			Func: funcInt64retInt32,
 		},
 	},
+	{
+		Name: "opa_value_boolean_set",
+		Descriptor: module.FunctionImport{
+			Func: funcInt32Int32retVoid,
+		},
+	},
 	{
 		Name: "opa_value_not_equal",
 		Descriptor: module.FunctionImport{
@@ -61,4 +69,10 @@ var externs = [...]module.Import{
 			Func: funcInt32Int32retInt32,
 		},
 	},
+	{
+		Name: "opa_value_iter",
+		Descriptor: module.FunctionImport{
+			Func: funcInt32Int32retInt32,
+		},
+	},
 }

internal/compiler/wasm/functypes.go

@@ -10,12 +10,16 @@ import (
 )
 
 const (
-	funcInt32Int32retInt32 uint32 = iota
+	funcInt32Int32retVoid  uint32 = iota
+	funcInt32Int32retInt32        = iota
 	funcInt32retInt32             = iota
 	funcInt64retInt32             = iota
 )
 
 var functypes = [...]module.FunctionType{
+	{
+		Params: []types.ValueType{types.I32, types.I32},
+	},
 	{
 		Params:  []types.ValueType{types.I32, types.I32},
 		Results: []types.ValueType{types.I32},

internal/compiler/wasm/wasm.go

@@ -14,6 +14,7 @@ import (
 	"github.com/open-policy-agent/opa/internal/wasm/instruction"
 	"github.com/open-policy-agent/opa/internal/wasm/module"
 	"github.com/open-policy-agent/opa/internal/wasm/types"
+	"github.com/pkg/errors"
 )
 
 // Compiler implements an IR->WASM compiler backend.
@@ -103,7 +104,7 @@ func (c *Compiler) compilePlan() error {
 
 		instrs, err := c.compileBlock(c.policy.Plan.Blocks[i])
 		if err != nil {
-			return err
+			return errors.Wrapf(err, "block %d", i)
 		}
 
 		if i < len(c.policy.Plan.Blocks)-1 {
@@ -125,6 +126,25 @@ func (c *Compiler) compileBlock(block ir.Block) ([]instruction.Instruction, erro
 		case ir.ReturnStmt:
 			instrs = append(instrs, instruction.I32Const{Value: int32(stmt.Code)})
 			instrs = append(instrs, instruction.Return{})
+		case ir.AssignStmt:
+			switch value := stmt.Value.(type) {
+			case ir.BooleanConst:
+				instrs = append(instrs, instruction.GetLocal{Index: c.local(stmt.Target)})
+				if value.Value {
+					instrs = append(instrs, instruction.I32Const{Value: 1})
+				} else {
+					instrs = append(instrs, instruction.I32Const{Value: 0})
+				}
+				instrs = append(instrs, instruction.Call{Index: opaValueBooleanSet})
+			default:
+				var buf bytes.Buffer
+				ir.Pretty(&buf, stmt)
+				return nil, fmt.Errorf("illegal assignment: %v", buf.String())
+			}
+		case ir.LoopStmt:
+			if err := c.compileLoop(stmt, &instrs); err != nil {
+				return nil, err
+			}
 		case ir.DotStmt:
 			instrs = append(instrs, instruction.GetLocal{Index: c.local(stmt.Source)})
 			instrs = append(instrs, instruction.GetLocal{Index: c.local(stmt.Key)})
@@ -154,14 +174,60 @@ func (c *Compiler) compileBlock(block ir.Block) ([]instruction.Instruction, erro
 			instrs = append(instrs, instruction.Call{Index: opaStringTerminated})
 			instrs = append(instrs, instruction.SetLocal{Index: c.local(stmt.Target)})
 		default:
-			return instrs, fmt.Errorf("unsupported IR statement %v", stmt)
+			var buf bytes.Buffer
+			ir.Pretty(&buf, stmt)
+			return instrs, fmt.Errorf("illegal statement: %v", buf.String())
 		}
 
 	}
 
 	return instrs, nil
 }
 
+func (c *Compiler) compileLoop(loop ir.LoopStmt, result *[]instruction.Instruction) error {
+	var instrs = *result
+	instrs = append(instrs, instruction.I32Const{Value: 0})
+	instrs = append(instrs, instruction.SetLocal{Index: c.local(loop.Key)})
+	body, err := c.compileLoopBody(loop)
+	if err != nil {
+		return err
+	}
+	instrs = append(instrs, instruction.Loop{Instrs: body})
+	*result = instrs
+	return nil
+}
+
+func (c *Compiler) compileLoopBody(loop ir.LoopStmt) ([]instruction.Instruction, error) {
+	var instrs []instruction.Instruction
+
+	// Execute iterator.
+	instrs = append(instrs, instruction.GetLocal{Index: c.local(loop.Source)})
+	instrs = append(instrs, instruction.GetLocal{Index: c.local(loop.Key)})
+	instrs = append(instrs, instruction.Call{Index: opaValueIter})
+
+	// Check for emptiness.
+	instrs = append(instrs, instruction.SetLocal{Index: c.local(loop.Key)})
+	instrs = append(instrs, instruction.GetLocal{Index: c.local(loop.Key)})
+	instrs = append(instrs, instruction.I32Eqz{})
+	instrs = append(instrs, instruction.BrIf{Index: 1})
+
+	// Load value.
+	instrs = append(instrs, instruction.GetLocal{Index: c.local(loop.Source)})
+	instrs = append(instrs, instruction.GetLocal{Index: c.local(loop.Key)})
+	instrs = append(instrs, instruction.Call{Index: opaValueGet})
+	instrs = append(instrs, instruction.SetLocal{Index: c.local(loop.Value)})
+
+	// Loop body.
+	nested, err := c.compileBlock(loop.Block)
+	if err != nil {
+		return nil, err
+	}
+
+	instrs = append(instrs, nested...)
+
+	return instrs, nil
+}
+
 func (c *Compiler) emitLocals() error {
 	c.code.Func.Locals = []module.LocalDeclaration{
 		{

internal/planner/planner.go

@@ -263,6 +263,7 @@ func (p *Planner) planLoop(ref ast.Ref, index int, iter planiter) error {
 
 		prev := p.curr
 		p.curr = &loop.Block
+		p.ltarget = value
 
 		if err := iter(); err != nil {
 			return err

test/wasm/assets/002_iteration.yaml

@@ -0,0 +1,17 @@
+cases:
+  - note: iteration
+    query: input.x[i] = 1
+    input: {"x": [3,2,1]}
+    return_code: 1
+  - note: iteration (negative)
+    query: input.x[i] = 1
+    input: {"x": [3,2,0]}
+    return_code: 0
+  - note: iteration nested
+    query: input.x[i] = 1; input.y[j] = "world"
+    input: {"x": [3,2,1], "y": ["hello", "world"]}
+    return_code: 1
+  - note: iteration nested (negative)
+    query: input.x[i] = 1; input.y[j] = "world"
+    input: {"x": [3,2,0], "y": ["hello", "universe"]}
+    return_code: 0

wasm/src/value.c

@@ -549,6 +549,12 @@ opa_value *opa_object()
     return &ret->hdr;
 }
 
+void opa_value_boolean_set(opa_value *v, int b)
+{
+    opa_boolean_t *ret = opa_cast_boolean(v);
+    ret->v = b;
+}
+
 void opa_array_free(opa_array_t *arr)
 {
     if (arr->elems != NULL)

wasm/src/value.h

@@ -102,6 +102,8 @@ opa_value *opa_array();
 opa_value *opa_array_with_cap(size_t cap);
 opa_value *opa_object();
 
+void opa_value_boolean_set(opa_value *v, int b);
+
 void opa_array_free(opa_array_t *arr);
 void opa_array_append(opa_array_t *arr, opa_value *v);
 void opa_array_sort(opa_array_t *arr, opa_compare_fn cmp_fn);