wasm: Add WithStmt to IR and implement on backend

This change updates the compiler to support a WithStmt instruction in
the IR. The WithStmt has the same upsert behaviour as in Rego however
the statement returns the modified value as opposed to modifying it in
place. The WithStmt will allow the planner to easily insert values
into the input and data documents. In theory, the WithStmt could also
be used on locals.

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

internal/compiler/wasm/wasm.go

@@ -52,6 +52,7 @@ const (
 	opaValueIter         = "opa_value_iter"
 	opaValueLength       = "opa_value_length"
 	opaValueMerge        = "opa_value_merge"
+	opaValueShallowCopy  = "opa_value_shallow_copy"
 	opaValueType         = "opa_value_type"
 )
 
@@ -334,6 +335,10 @@ func (c *Compiler) compileBlock(block *ir.Block) ([]instruction.Instruction, err
 			instrs = append(instrs, instruction.GetLocal{Index: c.local(stmt.Result)})
 			instrs = append(instrs, instruction.I32Eqz{})
 			instrs = append(instrs, instruction.BrIf{Index: 0})
+		case *ir.WithStmt:
+			if err := c.compileWithStmt(stmt, &instrs); err != nil {
+				return instrs, err
+			}
 		case *ir.AssignVarStmt:
 			instrs = append(instrs, instruction.GetLocal{Index: c.local(stmt.Source)})
 			instrs = append(instrs, instruction.SetLocal{Index: c.local(stmt.Target)})
@@ -630,6 +635,118 @@ func (c *Compiler) compileNot(not *ir.NotStmt, result *[]instruction.Instruction
 	return nil
 }
 
+func (c *Compiler) compileWithStmt(stmt *ir.WithStmt, result *[]instruction.Instruction) error {
+
+	var instrs = *result
+
+	// If the path is empty just assign the value to the target.
+	if len(stmt.Path) == 0 {
+		instrs = append(instrs, instruction.GetLocal{Index: c.local(stmt.Value)})
+		instrs = append(instrs, instruction.SetLocal{Index: c.local(stmt.Target)})
+		*result = instrs
+		return nil
+	}
+
+	lcopy := c.genLocal() // holds copy of source local
+	instrs = append(instrs, instruction.GetLocal{Index: c.local(stmt.Source)})
+	instrs = append(instrs, instruction.SetLocal{Index: lcopy})
+
+	// Initialize the target by either shallow copying the source or setting it
+	// to an empty object.
+	instrs = append(instrs, instruction.Block{
+		Instrs: []instruction.Instruction{
+			instruction.Block{Instrs: []instruction.Instruction{
+				instruction.GetLocal{Index: lcopy},
+				instruction.I32Eqz{},
+				instruction.BrIf{Index: 0},
+				instruction.GetLocal{Index: lcopy},
+				instruction.Call{Index: c.function(opaValueShallowCopy)},
+				instruction.SetLocal{Index: lcopy},
+				instruction.GetLocal{Index: lcopy},
+				instruction.SetLocal{Index: c.local(stmt.Target)},
+				instruction.Br{Index: 1},
+			}},
+			instruction.Call{Index: c.function(opaObject)},
+			instruction.SetLocal{Index: lcopy},
+			instruction.GetLocal{Index: lcopy},
+			instruction.SetLocal{Index: c.local(stmt.Target)},
+		},
+	})
+
+	// Initialize the locals that specify the path of the upsert operation.
+	lpath := make(map[int]uint32, len(stmt.Path))
+
+	for i := 0; i < len(stmt.Path); i++ {
+		lpath[i] = c.genLocal()
+		instrs = append(instrs, instruction.I32Const{Value: c.stringAddr(stmt.Path[i])})
+		instrs = append(instrs, instruction.Call{Index: c.function(opaStringTerminated)})
+		instrs = append(instrs, instruction.SetLocal{Index: lpath[i]})
+	}
+
+	// Generate a block that traverses the path of the upsert operation,
+	// shallowing copying values at each step as needed. Stop before the final
+	// segment that will only be inserted.
+	var inner []instruction.Instruction
+	ltemp := c.genLocal()
+
+	for i := 0; i < len(stmt.Path)-1; i++ {
+
+		// Lookup the next part of the path.
+		inner = append(inner, instruction.GetLocal{Index: lcopy})
+		inner = append(inner, instruction.GetLocal{Index: lpath[i]})
+		inner = append(inner, instruction.Call{Index: c.function(opaValueGet)})
+		inner = append(inner, instruction.SetLocal{Index: ltemp})
+
+		// If the next node is missing, break.
+		inner = append(inner, instruction.GetLocal{Index: ltemp})
+		inner = append(inner, instruction.I32Eqz{})
+		inner = append(inner, instruction.BrIf{Index: uint32(i)})
+
+		// Otherwise, shallow copy the next node node and insert into the copy
+		// before continuing.
+		inner = append(inner, instruction.GetLocal{Index: ltemp})
+		inner = append(inner, instruction.Call{Index: c.function(opaValueShallowCopy)})
+		inner = append(inner, instruction.SetLocal{Index: ltemp})
+		inner = append(inner, instruction.GetLocal{Index: lcopy})
+		inner = append(inner, instruction.GetLocal{Index: lpath[i]})
+		inner = append(inner, instruction.GetLocal{Index: ltemp})
+		inner = append(inner, instruction.Call{Index: c.function(opaObjectInsert)})
+		inner = append(inner, instruction.GetLocal{Index: ltemp})
+		inner = append(inner, instruction.SetLocal{Index: lcopy})
+	}
+
+	inner = append(inner, instruction.Br{Index: uint32(len(stmt.Path) - 1)})
+
+	// Generate blocks that handle missing nodes during traversal.
+	var block []instruction.Instruction
+	lval := c.genLocal()
+
+	for i := 0; i < len(stmt.Path)-1; i++ {
+		block = append(block, instruction.Block{Instrs: inner})
+		block = append(block, instruction.Call{Index: c.function(opaObject)})
+		block = append(block, instruction.SetLocal{Index: lval})
+		block = append(block, instruction.GetLocal{Index: lcopy})
+		block = append(block, instruction.GetLocal{Index: lpath[i]})
+		block = append(block, instruction.GetLocal{Index: lval})
+		block = append(block, instruction.Call{Index: c.function(opaObjectInsert)})
+		block = append(block, instruction.GetLocal{Index: lval})
+		block = append(block, instruction.SetLocal{Index: lcopy})
+		inner = block
+		block = nil
+	}
+
+	// Finish by inserting the statement's value into the shallow copied node.
+	instrs = append(instrs, instruction.Block{Instrs: inner})
+	instrs = append(instrs, instruction.GetLocal{Index: lcopy})
+	instrs = append(instrs, instruction.GetLocal{Index: lpath[len(stmt.Path)-1]})
+	instrs = append(instrs, instruction.GetLocal{Index: c.local(stmt.Value)})
+	instrs = append(instrs, instruction.Call{Index: c.function(opaObjectInsert)})
+
+	*result = instrs
+
+	return nil
+}
+
 func (c *Compiler) emitFunctionDecl(name string, tpe module.FunctionType, export bool) {
 
 	typeIndex := c.emitFunctionType(tpe)

internal/ir/ir.go

@@ -375,3 +375,14 @@ type SetAddStmt struct {
 	Value Local
 	Set   Local
 }
+
+// WithStmt creates a shallow copy of the source value with the value upserted
+// into the specified path. If the intermediate nodes in the path do not exist,
+// they will be created. If the path is empty this has the same affect as
+// assigning the value to the target.
+type WithStmt struct {
+	Source Local
+	Path   []int
+	Value  Local
+	Target Local
+}