wasm: Improve calling convention of eval() function

This change updates the calling convention for the eval()
function. Instead of accepting input and data values and returning the
result set directly, the eval() function now accepts a pointer to the
opa_eval_ctx_t structure defined in the wasm library. The caller will
be responsible for setting the input and data addresses in the struct
before invoking eval() and reading the result address out of the
struct once eval() returns. The wasm library exposes getters and
setters for the caller.

This change will make it easier to extend the API in the future
without impacting the caller. For the time being the eval() function
always returns zero however in the future this could be changed.

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

internal/compiler/wasm/wasm.go

@@ -73,6 +73,8 @@ type Compiler struct {
 
 	nextLocal uint32
 	locals    map[ir.Local]uint32
+	lctx      uint32 // local pointing to eval context
+	lrs       uint32 // local pointing to result set
 }
 
 const (
@@ -171,7 +173,7 @@ func (c *Compiler) initModule() error {
 	}
 
 	c.emitFunctionDecl("eval", module.FunctionType{
-		Params:  []types.ValueType{types.I32, types.I32},
+		Params:  []types.ValueType{types.I32},
 		Results: []types.ValueType{types.I32},
 	}, true)
 
@@ -233,13 +235,27 @@ func (c *Compiler) compileFuncs() error {
 // the module.
 func (c *Compiler) compilePlan() error {
 
-	// reset local variables and declare raw ptr, len, and input ptr.
+	c.code = &module.CodeEntry{}
 	c.nextLocal = 0
 	c.locals = map[ir.Local]uint32{}
-	_ = c.local(ir.Input)
-	_ = c.local(ir.Data)
+	c.lctx = c.genLocal()
+	c.lrs = c.genLocal()
 
-	c.code = &module.CodeEntry{}
+	// Initialize the input and data locals.
+	c.appendInstr(instruction.GetLocal{Index: c.lctx})
+	c.appendInstr(instruction.I32Load{Offset: 0, Align: 2})
+	c.appendInstr(instruction.SetLocal{Index: c.local(ir.Input)})
+
+	c.appendInstr(instruction.GetLocal{Index: c.lctx})
+	c.appendInstr(instruction.I32Load{Offset: 4, Align: 2})
+	c.appendInstr(instruction.SetLocal{Index: c.local(ir.Data)})
+
+	// Initialize the result set.
+	c.appendInstr(instruction.Call{Index: c.function(opaSet)})
+	c.appendInstr(instruction.SetLocal{Index: c.lrs})
+	c.appendInstr(instruction.GetLocal{Index: c.lctx})
+	c.appendInstr(instruction.GetLocal{Index: c.lrs})
+	c.appendInstr(instruction.I32Store{Offset: 8, Align: 2})
 
 	for i := range c.policy.Plan.Blocks {
 
@@ -248,13 +264,11 @@ func (c *Compiler) compilePlan() error {
 			return errors.Wrapf(err, "block %d", i)
 		}
 
-		if i < len(c.policy.Plan.Blocks)-1 {
-			c.appendInstr(instruction.Block{Instrs: instrs})
-		} else {
-			c.appendInstrs(instrs)
-		}
+		c.appendInstr(instruction.Block{Instrs: instrs})
 	}
 
+	c.appendInstr(instruction.I32Const{Value: int32(0)})
+
 	c.code.Func.Locals = []module.LocalDeclaration{
 		{
 			Count: c.nextLocal,
@@ -316,6 +330,10 @@ func (c *Compiler) compileBlock(block *ir.Block) ([]instruction.Instruction, err
 
 	for _, stmt := range block.Stmts {
 		switch stmt := stmt.(type) {
+		case *ir.ResultSetAdd:
+			instrs = append(instrs, instruction.GetLocal{Index: c.lrs})
+			instrs = append(instrs, instruction.GetLocal{Index: c.local(stmt.Value)})
+			instrs = append(instrs, instruction.Call{Index: c.function(opaSetAdd)})
 		case *ir.ReturnLocalStmt:
 			instrs = append(instrs, instruction.GetLocal{Index: c.local(stmt.Source)})
 			instrs = append(instrs, instruction.Return{})

internal/ir/ir.go

@@ -94,18 +94,6 @@ type (
 	}
 )
 
-const (
-	// Undefined represents an undefined return value. An undefined return value
-	// indicates the policy did not return a definitive answer.
-	Undefined int32 = iota
-
-	// Defined represents a defined return value.
-	Defined
-
-	// Error indicates a runtime error occurred during evaluation.
-	Error
-)
-
 const (
 	// Input is the local variable that refers to the global input document.
 	Input Local = iota
@@ -393,3 +381,8 @@ type WithStmt struct {
 	Value Local
 	Block *Block
 }
+
+// ResultSetAdd adds a value into the result set returned by the query plan.
+type ResultSetAdd struct {
+	Value Local
+}

internal/planner/planner.go

@@ -334,8 +334,6 @@ func (p *Planner) planQueries() error {
 
 	// Initialize the plan with a block that prepares the query result.
 	p.curr = &ir.Block{}
-	lresultset := p.newLocal()
-	p.appendStmt(&ir.MakeSetStmt{Target: lresultset})
 
 	// Build a set of variables appearing in the query and allocate strings for
 	// each one. The strings will be used in the result set objects.
@@ -360,8 +358,11 @@ func (p *Planner) planQueries() error {
 		}
 	}
 
+	if len(p.curr.Stmts) > 0 {
+		p.appendBlock(p.curr)
+	}
+
 	lnext := p.lnext
-	p.appendBlock(p.curr)
 
 	for _, q := range p.queries {
 		p.lnext = lnext
@@ -390,9 +391,8 @@ func (p *Planner) planQueries() error {
 				}
 			}
 
-			p.appendStmt(&ir.SetAddStmt{
+			p.appendStmt(&ir.ResultSetAdd{
 				Value: lr,
-				Set:   lresultset,
 			})
 
 			defined = true
@@ -408,14 +408,6 @@ func (p *Planner) planQueries() error {
 		}
 	}
 
-	p.appendBlock(&ir.Block{
-		Stmts: []ir.Stmt{
-			&ir.ReturnLocalStmt{
-				Source: lresultset,
-			},
-		},
-	})
-
 	return nil
 }
 

internal/wasm/instruction/memory.go

@@ -0,0 +1,39 @@
+// Copyright 2019 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package instruction
+
+import "github.com/open-policy-agent/opa/internal/wasm/opcode"
+
+// I32Load represents the WASM i32.load instruction.
+type I32Load struct {
+	Offset int32
+	Align  int32 // expressed as a power of two
+}
+
+// Op returns the opcode of the instruction.
+func (I32Load) Op() opcode.Opcode {
+	return opcode.I32Load
+}
+
+// ImmediateArgs returns the static offset and alignment operands.
+func (i I32Load) ImmediateArgs() []interface{} {
+	return []interface{}{i.Align, i.Offset}
+}
+
+// I32Store represents the WASM i32.store instruction.
+type I32Store struct {
+	Offset int32
+	Align  int32 // expressed as a power of two
+}
+
+// Op returns the opcode of the instruction.
+func (I32Store) Op() opcode.Opcode {
+	return opcode.I32Store
+}
+
+// ImmediateArgs returns the static offset and alignment operands.
+func (i I32Store) ImmediateArgs() []interface{} {
+	return []interface{}{i.Align, i.Offset}
+}

test/wasm/assets/test.js

@@ -116,10 +116,20 @@ async function instantiate(bytes, memory, data) {
 }
 
 function evaluate(policy, input) {
+
     policy.module.instance.exports.opa_heap_ptr_set(policy.heapPtr);
     policy.module.instance.exports.opa_heap_top_set(policy.heapTop);
+
     const inputAddr = loadJSON(policy.module, policy.memory, input);
-    const resultAddr = policy.module.instance.exports.eval(inputAddr, policy.dataAddr);
+    const ctxAddr = policy.module.instance.exports.opa_eval_ctx_new();
+
+    policy.module.instance.exports.opa_eval_ctx_set_input(ctxAddr, inputAddr);
+    policy.module.instance.exports.opa_eval_ctx_set_data(ctxAddr, policy.dataAddr);
+
+    policy.module.instance.exports.eval(ctxAddr);
+
+    const resultAddr = policy.module.instance.exports.opa_eval_ctx_get_result(ctxAddr);
+
     return { addr: resultAddr };
 }
 

wasm/Makefile

@@ -52,10 +52,14 @@ $(WASM_OBJ_DIR)/json.wasm: src/json.c
 $(WASM_OBJ_DIR)/value.wasm: src/value.c
 	@$(CC) $(CFLAGS) -c $^ -o $@
 
+$(WASM_OBJ_DIR)/context.wasm: src/context.c
+	@$(CC) $(CFLAGS) -c $^ -o $@
+
 $(WASM_OBJ_DIR)/opa.wasm: $(WASM_OBJ_DIR)/malloc.wasm \
 		$(WASM_OBJ_DIR)/value.wasm \
 		$(WASM_OBJ_DIR)/printf.wasm \
 		$(WASM_OBJ_DIR)/string.wasm \
+		$(WASM_OBJ_DIR)/context.wasm \
 		$(WASM_OBJ_DIR)/json.wasm
 	@wasm-ld-8 \
 			--allow-undefined-file=src/undefined.symbols \
@@ -74,6 +78,7 @@ $(WASM_OBJ_DIR)/opa-test.wasm: $(WASM_OBJ_DIR)/test.wasm \
 		$(WASM_OBJ_DIR)/value.wasm \
 		$(WASM_OBJ_DIR)/printf.wasm \
 		$(WASM_OBJ_DIR)/string.wasm \
+		$(WASM_OBJ_DIR)/context.wasm \
 		$(WASM_OBJ_DIR)/json.wasm
 	@wasm-ld-8 \
 			--allow-undefined-file=tests/undefined.symbols \

wasm/src/context.c

@@ -0,0 +1,26 @@
+#include "malloc.h"
+#include "context.h"
+
+opa_eval_ctx_t *opa_eval_ctx_new()
+{
+    opa_eval_ctx_t *ctx = (opa_eval_ctx_t *)opa_malloc(sizeof(opa_eval_ctx_t));
+    ctx->input = NULL;
+    ctx->data = NULL;
+    ctx->result = NULL;
+    return ctx;
+}
+
+void opa_eval_ctx_set_input(opa_eval_ctx_t *ctx, opa_value *v)
+{
+    ctx->input = v;
+}
+
+void opa_eval_ctx_set_data(opa_eval_ctx_t *ctx, opa_value *v)
+{
+    ctx->data = v;
+}
+
+opa_value *opa_eval_ctx_get_result(opa_eval_ctx_t *ctx)
+{
+    return ctx->result;
+}

wasm/src/context.h

@@ -0,0 +1,18 @@
+#ifndef OPA_CONTEXT_H
+#define OPA_CONTEXT_H
+
+#include "value.h"
+
+typedef struct
+{
+    opa_value *input;
+    opa_value *data;
+    opa_value *result;
+} opa_eval_ctx_t;
+
+opa_eval_ctx_t *opa_eval_ctx_new();
+void opa_eval_ctx_set_input(opa_eval_ctx_t *ctx, opa_value *v);
+void opa_eval_ctx_set_data(opa_eval_ctx_t *ctx, opa_value *v);
+opa_value *opa_eval_ctx_get_result(opa_eval_ctx_t *ctx);
+
+#endif