Optimization: Remove copying of bindings

Instead of copying the bindings each time a binding is added, return an object
that can be used to undo the binding at the end of the proof.

storage/bindings.go

@@ -65,6 +65,11 @@ func (b *Bindings) Put(k, v ast.Value) {
 	b.hashMap.Put(k, v)
 }
 
+// Delete removes a key/value pair.
+func (b *Bindings) Delete(k ast.Value) {
+	b.hashMap.Delete(k)
+}
+
 // Update returns new bindings that are the union of these bindings and the other bindings.
 func (b *Bindings) Update(other *Bindings) *Bindings {
 	new := b.hashMap.Update(other.hashMap)

topdown/aggregates.go

@@ -40,8 +40,7 @@ func evalReduce(f reduceFunc) BuiltinFunc {
 
 		switch dst := dst.(type) {
 		case ast.Var:
-			ctx = ctx.BindValue(dst, y)
-			return iter(ctx)
+			return Continue(ctx, dst, y, iter)
 		default:
 			if dst.Equal(y) {
 				return iter(ctx)

topdown/arithmetic.go

@@ -59,8 +59,7 @@ func evalArithArity1(f arithArity1) BuiltinFunc {
 
 		switch b := b.(type) {
 		case ast.Var:
-			ctx = ctx.BindValue(b, r)
-			return iter(ctx)
+			return Continue(ctx, b, r, iter)
 		default:
 			if b.Equal(r) {
 				return iter(ctx)
@@ -93,8 +92,7 @@ func evalArithArity2(f arithArity2) BuiltinFunc {
 
 		switch cv := cv.(type) {
 		case ast.Var:
-			ctx = ctx.BindValue(cv, c)
-			return iter(ctx)
+			return Continue(ctx, cv, c, iter)
 		default:
 			if cv.Equal(c) {
 				return iter(ctx)

topdown/casts.go

@@ -44,8 +44,7 @@ func evalToNumber(ctx *Context, expr *ast.Expr, iter Iterator) error {
 
 	switch b := b.(type) {
 	case ast.Var:
-		ctx = ctx.BindValue(b, n)
-		return iter(ctx)
+		return Continue(ctx, b, n, iter)
 	default:
 		if n.Equal(b) {
 			return iter(ctx)

topdown/eq.go

@@ -17,7 +17,9 @@ func evalEq(ctx *Context, expr *ast.Expr, iter Iterator) error {
 	a := operands[1].Value
 	b := operands[2].Value
 
-	return evalEqUnify(ctx, a, b, iter)
+	undo, err := evalEqUnify(ctx, a, b, nil, iter)
+	ctx.Unbind(undo)
+	return err
 }
 
 func evalEqGround(ctx *Context, a ast.Value, b ast.Value, iter Iterator) error {
@@ -51,7 +53,7 @@ func evalEqGround(ctx *Context, a ast.Value, b ast.Value, iter Iterator) error {
 //
 // In cases involving references, OPA assumes that the references are ground at this stage.
 // As a result, references are just special cases of the normal scalar/composite unification.
-func evalEqUnify(ctx *Context, a ast.Value, b ast.Value, iter Iterator) error {
+func evalEqUnify(ctx *Context, a ast.Value, b ast.Value, prev *Undo, iter Iterator) (*Undo, error) {
 
 	// Plug bindings into both terms because this will be called recursively and there may be
 	// new bindings that have been made as part of unification.
@@ -60,53 +62,53 @@ func evalEqUnify(ctx *Context, a ast.Value, b ast.Value, iter Iterator) error {
 
 	switch a := a.(type) {
 	case ast.Var:
-		return evalEqUnifyVar(ctx, a, b, iter)
+		return evalEqUnifyVar(ctx, a, b, prev, iter)
 	case ast.Object:
-		return evalEqUnifyObject(ctx, a, b, iter)
+		return evalEqUnifyObject(ctx, a, b, prev, iter)
 	case ast.Array:
-		return evalEqUnifyArray(ctx, a, b, iter)
+		return evalEqUnifyArray(ctx, a, b, prev, iter)
 	default:
 		switch b := b.(type) {
 		case ast.Var:
-			return evalEqUnifyVar(ctx, b, a, iter)
+			return evalEqUnifyVar(ctx, b, a, prev, iter)
 		case ast.Array:
-			return evalEqUnifyArray(ctx, b, a, iter)
+			return evalEqUnifyArray(ctx, b, a, prev, iter)
 		case ast.Object:
-			return evalEqUnifyObject(ctx, b, a, iter)
+			return evalEqUnifyObject(ctx, b, a, prev, iter)
 		default:
-			return evalEqGround(ctx, a, b, iter)
+			return prev, evalEqGround(ctx, a, b, iter)
 		}
 	}
 
 }
 
-func evalEqUnifyArray(ctx *Context, a ast.Array, b ast.Value, iter Iterator) error {
+func evalEqUnifyArray(ctx *Context, a ast.Array, b ast.Value, prev *Undo, iter Iterator) (*Undo, error) {
 	switch b := b.(type) {
 	case ast.Var:
-		return evalEqUnifyVar(ctx, b, a, iter)
+		return evalEqUnifyVar(ctx, b, a, prev, iter)
 	case ast.Ref:
-		return evalEqUnifyArrayRef(ctx, a, b, iter)
+		return evalEqUnifyArrayRef(ctx, a, b, prev, iter)
 	case ast.Array:
-		return evalEqUnifyArrays(ctx, a, b, iter)
+		return evalEqUnifyArrays(ctx, a, b, prev, iter)
 	default:
-		return nil
+		return prev, nil
 	}
 }
 
-func evalEqUnifyArrayRef(ctx *Context, a ast.Array, b ast.Ref, iter Iterator) error {
+func evalEqUnifyArrayRef(ctx *Context, a ast.Array, b ast.Ref, prev *Undo, iter Iterator) (*Undo, error) {
 
 	r, err := ctx.DataStore.GetRef(b)
 	if err != nil {
-		return err
+		return prev, err
 	}
 
 	slice, ok := r.([]interface{})
 	if !ok {
-		return nil
+		return prev, nil
 	}
 
 	if len(a) != len(slice) {
-		return nil
+		return prev, nil
 	}
 
 	for i := range a {
@@ -114,128 +116,131 @@ func evalEqUnifyArrayRef(ctx *Context, a ast.Array, b ast.Ref, iter Iterator) er
 		child := make(ast.Ref, len(b), len(b)+1)
 		copy(child, b)
 		child = append(child, ast.NumberTerm(float64(i)))
-		err := evalEqUnify(ctx, a[i].Value, child, func(ctx *Context) error {
+		p, err := evalEqUnify(ctx, a[i].Value, child, prev, func(ctx *Context) error {
 			tmp = ctx
 			return nil
 		})
+		prev = p
 		if err != nil {
-			return err
+			return nil, err
 		}
 		if tmp == nil {
-			return nil
+			return nil, nil
 		}
 		ctx = tmp
 	}
-	return iter(ctx)
+	return prev, iter(ctx)
 }
 
-func evalEqUnifyArrays(ctx *Context, a ast.Array, b ast.Array, iter Iterator) error {
+func evalEqUnifyArrays(ctx *Context, a ast.Array, b ast.Array, prev *Undo, iter Iterator) (*Undo, error) {
 	aLen := len(a)
 	bLen := len(b)
 	if aLen != bLen {
-		return nil
+		return nil, nil
 	}
 	for i := 0; i < aLen; i++ {
 		ai := a[i].Value
 		bi := b[i].Value
 		var tmp *Context
-		err := evalEqUnify(ctx, ai, bi, func(ctx *Context) error {
+		p, err := evalEqUnify(ctx, ai, bi, prev, func(ctx *Context) error {
 			tmp = ctx
 			return nil
 		})
+		prev = p
 		if err != nil {
-			return err
+			return nil, err
 		}
 		if tmp == nil {
-			return nil
+			return nil, nil
 		}
 		ctx = tmp
 	}
-	return iter(ctx)
+	return prev, iter(ctx)
 }
 
 // evalEqUnifyObject attempts to unify the object "a" with some other value "b".
 // TODO(tsandal): unification of object keys (or unordered sets in general) is not
 // supported because it would be too expensive. We may revisit this in the future.
-func evalEqUnifyObject(ctx *Context, a ast.Object, b ast.Value, iter Iterator) error {
+func evalEqUnifyObject(ctx *Context, a ast.Object, b ast.Value, prev *Undo, iter Iterator) (*Undo, error) {
 	switch b := b.(type) {
 	case ast.Var:
-		return evalEqUnifyVar(ctx, b, a, iter)
+		return evalEqUnifyVar(ctx, b, a, prev, iter)
 	case ast.Ref:
-		return evalEqUnifyObjectRef(ctx, a, b, iter)
+		return evalEqUnifyObjectRef(ctx, a, b, prev, iter)
 	case ast.Object:
-		return evalEqUnifyObjects(ctx, a, b, iter)
+		return evalEqUnifyObjects(ctx, a, b, prev, iter)
 	default:
-		return nil
+		return nil, nil
 	}
 }
 
-func evalEqUnifyObjectRef(ctx *Context, a ast.Object, b ast.Ref, iter Iterator) error {
+func evalEqUnifyObjectRef(ctx *Context, a ast.Object, b ast.Ref, prev *Undo, iter Iterator) (*Undo, error) {
 
 	r, err := ctx.DataStore.GetRef(b)
 
 	if err != nil {
-		return err
+		return prev, err
 	}
 
 	for i := range a {
 		if !a[i][0].IsGround() {
-			return fmt.Errorf("illegal variable object key: %v", a[i][0])
+			return prev, fmt.Errorf("illegal variable object key: %v", a[i][0])
 		}
 	}
 
 	obj, ok := r.(map[string]interface{})
 	if !ok {
-		return nil
+		return prev, nil
 	}
 
 	if len(obj) != len(a) {
-		return nil
+		return prev, nil
 	}
 
 	for i := range a {
 		// TODO(tsandall): support non-string keys in storage.
 		k, ok := a[i][0].Value.(ast.String)
 		if !ok {
-			return fmt.Errorf("illegal object key type %T: %v", a[i][0], a[i][0])
+			return prev, fmt.Errorf("illegal object key type %T: %v", a[i][0], a[i][0])
 		}
 
 		_, ok = obj[string(k)]
 		if !ok {
-			return nil
+			return nil, nil
 		}
 
 		child := make(ast.Ref, len(b), len(b)+1)
 		copy(child, b)
 		child = append(child, a[i][0])
 		var tmp *Context
-		err := evalEqUnify(ctx, a[i][1].Value, child, func(ctx *Context) error {
+		p, err := evalEqUnify(ctx, a[i][1].Value, child, prev, func(ctx *Context) error {
 			tmp = ctx
 			return nil
 		})
+		prev = p
 		if err != nil {
-			return err
+			return nil, err
 		}
 		if tmp == nil {
-			return nil
+			return nil, nil
 		}
 		ctx = tmp
 	}
-	return iter(ctx)
+	return prev, iter(ctx)
 }
 
-func evalEqUnifyObjects(ctx *Context, a ast.Object, b ast.Object, iter Iterator) error {
+func evalEqUnifyObjects(ctx *Context, a ast.Object, b ast.Object, prev *Undo, iter Iterator) (*Undo, error) {
 
 	if len(a) != len(b) {
-		return nil
+		return nil, nil
 	}
 
 	for i := range a {
 		if !a[i][0].IsGround() {
-			return fmt.Errorf("illegal variable object key: %v", a[i][0])
+			return prev, fmt.Errorf("illegal variable object key: %v", a[i][0])
 		}
 		if !b[i][0].IsGround() {
-			return fmt.Errorf("illegal variable object key: %v", b[i][0])
+			return prev, fmt.Errorf("illegal variable object key: %v", b[i][0])
 		}
 	}
 
@@ -243,12 +248,13 @@ func evalEqUnifyObjects(ctx *Context, a ast.Object, b ast.Object, iter Iterator)
 		var tmp *Context
 		for j := range b {
 			if b[j][0].Equal(a[i][0]) {
-				err := evalEqUnify(ctx, a[i][1].Value, b[j][1].Value, func(ctx *Context) error {
+				p, err := evalEqUnify(ctx, a[i][1].Value, b[j][1].Value, prev, func(ctx *Context) error {
 					tmp = ctx
 					return nil
 				})
+				prev = p
 				if err != nil {
-					return err
+					return nil, err
 				}
 				if tmp == nil {
 					break
@@ -256,18 +262,16 @@ func evalEqUnifyObjects(ctx *Context, a ast.Object, b ast.Object, iter Iterator)
 			}
 		}
 		if tmp == nil {
-			return nil
+			return nil, nil
 		}
 		ctx = tmp
 	}
 
-	return iter(ctx)
+	return prev, iter(ctx)
 }
 
-func evalEqUnifyVar(ctx *Context, a ast.Var, b ast.Value, iter Iterator) error {
-	ctx = ctx.BindValue(a, b)
-	if ctx == nil {
-		return nil
-	}
-	return iter(ctx)
+func evalEqUnifyVar(ctx *Context, a ast.Var, b ast.Value, prev *Undo, iter Iterator) (*Undo, error) {
+	undo := ctx.Bind(a, b, prev)
+	err := iter(ctx)
+	return undo, err
 }

topdown/topdown.go

@@ -14,11 +14,6 @@ import (
 )
 
 // Context contains the state of the evaluation process.
-//
-// TODO(tsandall): profile perf/memory usage with current approach;
-// the current approach copies the Context structure for each
-// step and binding. This avoids the need to undo steps and bindings
-// each time the proof fails but this may be too expensive.
 type Context struct {
 	Query     ast.Body
 	Globals   *storage.Bindings
@@ -50,12 +45,30 @@ func (ctx *Context) Binding(k ast.Value) ast.Value {
 	return nil
 }
 
-// BindValue returns a new Context with bindings that map the key to the value.
-func (ctx *Context) BindValue(key ast.Value, value ast.Value) *Context {
-	cpy := *ctx
-	cpy.Locals = ctx.Locals.Copy()
-	cpy.Locals.Put(key, value)
-	return &cpy
+// Undo represents a binding that can be undone.
+type Undo struct {
+	Key   ast.Value
+	Value ast.Value
+	Prev  *Undo
+}
+
+// Bind updates the context to include a binding from the key to the value. The return
+// value is used to return the context to the state before the binding was added.
+func (ctx *Context) Bind(key ast.Value, value ast.Value, prev *Undo) *Undo {
+	o := ctx.Locals.Get(key)
+	ctx.Locals.Put(key, value)
+	return &Undo{key, o, prev}
+}
+
+// Unbind updates the context by removing the binding represented by the undo.
+func (ctx *Context) Unbind(undo *Undo) {
+	for u := undo; u != nil; u = u.Prev {
+		if u.Value != nil {
+			ctx.Locals.Put(u.Key, u.Value)
+		} else {
+			ctx.Locals.Delete(u.Key)
+		}
+	}
 }
 
 // Child returns a new context to evaluate a query that was referenced by this context.
@@ -90,11 +103,11 @@ func (ctx *Context) trace(f string, a ...interface{}) {
 }
 
 func (ctx *Context) traceEval() {
-	ctx.trace("Eval %v", ctx.Current(), ctx.Locals)
+	ctx.trace("Eval %v", ctx.Current())
 }
 
 func (ctx *Context) traceTry(expr *ast.Expr) {
-	ctx.trace(" Try %v", expr, ctx.Locals)
+	ctx.trace(" Try %v", expr)
 }
 
 func (ctx *Context) traceSuccess(expr *ast.Expr) {
@@ -157,6 +170,29 @@ func conflictErr(query interface{}, kind string, rule *ast.Rule) error {
 // Iterator is the interface for processing contexts.
 type Iterator func(*Context) error
 
+// Continue binds the key to the value in the current context and invokes the iterator.
+// This is a helper function for simple cases where a single value (e.g., a variable) needs
+// to be bound to a value in order for the evaluation the proceed.
+func Continue(ctx *Context, key, value ast.Value, iter Iterator) error {
+	undo := ctx.Bind(key, value, nil)
+	err := iter(ctx)
+	ctx.Unbind(undo)
+	return err
+}
+
+// ContinueN binds N keys to N values. The key/value pairs are passed in as alternating pairs, e.g.,
+// key-1, value-1, key-2, value-2, ..., key-N, value-N.
+func ContinueN(ctx *Context, iter Iterator, x ...ast.Value) error {
+	var prev *Undo
+	for i := 0; i < len(x)/2; i++ {
+		offset := i * 2
+		prev = ctx.Bind(x[offset], x[offset+1], prev)
+	}
+	err := iter(ctx)
+	ctx.Unbind(prev)
+	return err
+}
+
 // Eval runs the evaluation algorithm on the contxet and calls the iterator
 // foreach context that contains bindings that satisfy all of the expressions
 // inside the body.
@@ -549,16 +585,21 @@ func evalRef(ctx *Context, ref, path ast.Ref, iter Iterator) error {
 	}
 
 	return evalRef(ctx, n, ast.Ref{}, func(ctx *Context) error {
+		var undo *Undo
 		if b := ctx.Binding(n); b == nil {
 			p := PlugValue(n, ctx).(ast.Ref)
 			v, err := lookupValue(ctx.DataStore, p)
 			if err != nil {
 				return err
 			}
-			ctx = ctx.BindValue(n, v)
+			undo = ctx.Bind(n, v, nil)
 		}
 		tmp := append(path, head)
-		return evalRef(ctx, tail, tmp, iter)
+		err := evalRef(ctx, tail, tmp, iter)
+		if undo != nil {
+			ctx.Unbind(undo)
+		}
+		return err
 	})
 }
 
@@ -627,9 +668,10 @@ func evalRefRecWalkColl(ctx *Context, path, tail ast.Ref, iter Iterator) error {
 	switch node := node.(type) {
 	case map[string]interface{}:
 		for key := range node {
-			cpy := ctx.BindValue(head, ast.String(key))
+			undo := ctx.Bind(head, ast.String(key), nil)
 			path = append(path, ast.StringTerm(key))
-			err := evalRefRec(cpy, path, tail, iter)
+			err := evalRefRec(ctx, path, tail, iter)
+			ctx.Unbind(undo)
 			if err != nil {
 				return err
 			}
@@ -638,9 +680,10 @@ func evalRefRecWalkColl(ctx *Context, path, tail ast.Ref, iter Iterator) error {
 		return nil
 	case []interface{}:
 		for i := range node {
-			cpy := ctx.BindValue(head, ast.Number(i))
+			undo := ctx.Bind(head, ast.Number(i), nil)
 			path = append(path, ast.NumberTerm(float64(i)))
-			err := evalRefRec(cpy, path, tail, iter)
+			err := evalRefRec(ctx, path, tail, iter)
+			ctx.Unbind(undo)
 			if err != nil {
 				return err
 			}
@@ -753,8 +796,10 @@ func evalRefRulePartialObjectDoc(ctx *Context, ref ast.Ref, path ast.Ref, rule *
 			}
 			value = PlugValue(value, child)
 
-			ctx = ctx.BindValue(suffix[0].Value.(ast.Var), key)
-			return evalRefRuleResult(ctx, ref, ref[len(path)+1:], value, iter)
+			undo := ctx.Bind(suffix[0].Value.(ast.Var), key, nil)
+			err := evalRefRuleResult(ctx, ref, ref[len(path)+1:], value, iter)
+			ctx.Unbind(undo)
+			return err
 		})
 	}
 
@@ -802,8 +847,7 @@ func evalRefRulePartialObjectDocFull(ctx *Context, ref ast.Ref, rules []*ast.Rul
 		}
 	}
 
-	ctx = ctx.BindValue(ref, result)
-	return iter(ctx)
+	return Continue(ctx, ref, result, iter)
 }
 
 func evalRefRulePartialSetDoc(ctx *Context, ref ast.Ref, path ast.Ref, rule *ast.Rule, iter Iterator) error {
@@ -833,9 +877,7 @@ func evalRefRulePartialSetDoc(ctx *Context, ref ast.Ref, path ast.Ref, rule *ast
 			// so that expression will be defined. E.g., given a simple rule:
 			// "p = true :- q[x]", we say that "p" should be defined if "q"
 			// is defined for some value "x".
-			ctx = ctx.BindValue(key.(ast.Var), value)
-			ctx = ctx.BindValue(ref[:len(path)+1], ast.Boolean(true))
-			return iter(ctx)
+			return ContinueN(ctx, iter, key, value, ref[:len(path)+1], ast.Boolean(true))
 		})
 	}
 
@@ -845,8 +887,7 @@ func evalRefRulePartialSetDoc(ctx *Context, ref ast.Ref, path ast.Ref, rule *ast
 
 	return Eval(child, func(child *Context) error {
 		// See comment above for explanation of why the reference is bound to true.
-		ctx = ctx.BindValue(ref[:len(path)+1], ast.Boolean(true))
-		return iter(ctx)
+		return Continue(ctx, ref[:len(path)+1], ast.Boolean(true), iter)
 	})
 
 }
@@ -864,12 +905,17 @@ func evalRefRuleResult(ctx *Context, ref ast.Ref, suffix ast.Ref, result ast.Val
 		// needing to be processed. This is done by substituting the prefix of the original reference
 		// with the binding. In this case, the prefix is "q[k]" and the binding value would be
 		// "a[<some key>]".
-		var binding ast.Ref
-		binding = append(binding, result...)
-		binding = append(binding, suffix...)
+		offset := len(result)
+		binding := make(ast.Ref, offset+len(suffix))
+		for i := range result {
+			binding[i] = result[i]
+		}
+		for i := range suffix {
+			binding[i+offset] = suffix[i]
+		}
 		return evalRefRec(ctx, result, suffix, func(ctx *Context) error {
-			ctx = ctx.BindValue(ref, PlugValue(binding, ctx))
-			return iter(ctx)
+			v := PlugValue(binding, ctx)
+			return Continue(ctx, ref, v, iter)
 		})
 
 	case ast.Array:
@@ -879,15 +925,16 @@ func evalRefRuleResult(ctx *Context, ref ast.Ref, suffix ast.Ref, result ast.Val
 				pluggedSuffix = append(pluggedSuffix, PlugTerm(t, ctx))
 			}
 			return result.Query(pluggedSuffix, func(keys map[ast.Var]ast.Value, value ast.Value) error {
-				ctx = ctx.BindValue(ref, value)
+				prev := ctx.Bind(ref, value, nil)
 				for k, v := range keys {
-					ctx = ctx.BindValue(k, v)
+					prev = ctx.Bind(k, v, prev)
 				}
-				return iter(ctx)
+				err := iter(ctx)
+				ctx.Unbind(prev)
+				return err
 			})
 		}
-		ctx = ctx.BindValue(ref, result)
-		return iter(ctx)
+		return Continue(ctx, ref, result, iter)
 
 	case ast.Object:
 		if len(suffix) > 0 {
@@ -896,23 +943,23 @@ func evalRefRuleResult(ctx *Context, ref ast.Ref, suffix ast.Ref, result ast.Val
 				pluggedSuffix = append(pluggedSuffix, PlugTerm(t, ctx))
 			}
 			return result.Query(pluggedSuffix, func(keys map[ast.Var]ast.Value, value ast.Value) error {
-				ctx = ctx.BindValue(ref, value)
+				prev := ctx.Bind(ref, value, nil)
 				for k, v := range keys {
-					ctx = ctx.BindValue(k, v)
+					prev = ctx.Bind(k, v, prev)
 				}
-				return iter(ctx)
+				err := iter(ctx)
+				ctx.Unbind(prev)
+				return err
 			})
 		}
-		ctx = ctx.BindValue(ref, result)
-		return iter(ctx)
+		return Continue(ctx, ref, result, iter)
 
 	default:
 		if len(suffix) > 0 {
 			// This is not defined because it attempts to dereference a scalar.
 			return nil
 		}
-		ctx = ctx.BindValue(ref, result)
-		return iter(ctx)
+		return Continue(ctx, ref, result, iter)
 	}
 }
 
@@ -979,8 +1026,7 @@ func evalTermsComprehension(ctx *Context, comp ast.Value, iter Iterator) error {
 		if err != nil {
 			return err
 		}
-		ctx = ctx.BindValue(comp, r)
-		return iter(ctx)
+		return Continue(ctx, comp, r, iter)
 	default:
 		panic(fmt.Sprintf("illegal argument: %v %v", ctx, comp))
 	}
@@ -1006,8 +1052,14 @@ func evalTermsIndexed(ctx *Context, iter Iterator, indexed ast.Ref, nonIndexed *
 		// Iterate the bindings for the indexed term that when applied to the reference
 		// would locate the non-indexed value obtained above.
 		return index.Iter(nonIndexedValue, func(bindings *storage.Bindings) error {
-			ctx.Locals = ctx.Locals.Update(bindings)
-			return iter(ctx)
+			var prev *Undo
+			bindings.Iter(func(k, v ast.Value) bool {
+				prev = ctx.Bind(k, v, prev)
+				return false
+			})
+			err := iter(ctx)
+			ctx.Unbind(prev)
+			return err
 		})
 
 	}

topdown/topdown_test.go

@@ -192,16 +192,16 @@ func TestPlugValue(t *testing.T) {
 	world := ast.String("world")
 
 	ctx1 := &Context{Locals: storage.NewBindings(), Globals: storage.NewBindings()}
-	ctx1 = ctx1.BindValue(a, b)
-	ctx1 = ctx1.BindValue(b, cs)
-	ctx1 = ctx1.BindValue(c, ks)
-	ctx1 = ctx1.BindValue(k, hello)
+	ctx1.Bind(a, b, nil)
+	ctx1.Bind(b, cs, nil)
+	ctx1.Bind(c, ks, nil)
+	ctx1.Bind(k, hello, nil)
 
 	ctx2 := &Context{Locals: storage.NewBindings(), Globals: storage.NewBindings()}
-	ctx2 = ctx2.BindValue(a, b)
-	ctx2 = ctx2.BindValue(b, cs)
-	ctx2 = ctx2.BindValue(c, vs)
-	ctx2 = ctx2.BindValue(v, world)
+	ctx2.Bind(a, b, nil)
+	ctx2.Bind(b, cs, nil)
+	ctx2.Bind(c, vs, nil)
+	ctx2.Bind(v, world, nil)
 
 	expected := ast.MustParseTerm(`[{"hello": "world"}]`).Value
 
@@ -221,8 +221,8 @@ func TestPlugValue(t *testing.T) {
 	n := ast.MustParseTerm("a.b[x.y[i]]").Value
 
 	ctx3 := &Context{Locals: storage.NewBindings(), Globals: storage.NewBindings()}
-	ctx3 = ctx3.BindValue(ast.Var("i"), ast.Number(1))
-	ctx3 = ctx3.BindValue(ast.MustParseTerm("x.y[i]").Value, ast.Number(1))
+	ctx3.Bind(ast.Var("i"), ast.Number(1), nil)
+	ctx3.Bind(ast.MustParseTerm("x.y[i]").Value, ast.Number(1), nil)
 
 	expected = ast.MustParseTerm("a.b[1]").Value
 
@@ -909,10 +909,10 @@ func TestExample(t *testing.T) {
     `)
 
 	assertTopDown(t, store, 0, "violations", []string{"opa", "example", "violations"}, "{}", `
-        [
-            {"id": "s4", "name": "dev", "protocols": ["http"], "ports": ["p1", "p2"]}
-        ]
-    `)
+	    [
+	        {"id": "s4", "name": "dev", "protocols": ["http"], "ports": ["p1", "p2"]}
+	    ]
+	`)
 }
 
 func compileModules(input []string) map[string]*ast.Module {

util/hashmap.go

@@ -70,6 +70,24 @@ func (h *HashMap) Get(k T) (T, bool) {
 	return nil, false
 }
 
+// Delete removes the the key k.
+func (h *HashMap) Delete(k T) {
+	hash := h.hash(k)
+	var prev *hashEntry
+	for entry := h.table[hash]; entry != nil; entry = entry.next {
+		if h.eq(entry.k, k) {
+			if prev != nil {
+				prev.next = entry.next
+			} else {
+				h.table[hash] = entry.next
+			}
+			h.size--
+			return
+		}
+		prev = entry
+	}
+}
+
 // Hash returns the hash code for this hash map.
 func (h *HashMap) Hash() int {
 	var hash int

util/hashmap_test.go

@@ -11,6 +11,26 @@ import (
 	"testing"
 )
 
+func TestHashMapPutDelete(t *testing.T) {
+	m := stringHashMap()
+	m.Put("a", "b")
+	m.Put("b", "c")
+	m.Delete("b")
+	r, _ := m.Get("a")
+	if r != "b" {
+		t.Fatal("Expected a to be intact")
+	}
+	r, ok := m.Get("b")
+	if ok {
+		t.Fatalf("Expected b to be removed: %v", r)
+	}
+	m.Delete("b")
+	r, _ = m.Get("a")
+	if r != "b" {
+		t.Fatal("Expected a to be intact")
+	}
+}
+
 func TestHashMapOverwrite(t *testing.T) {
 	m := stringHashMap()
 	key := "hello"