diff --git a/storage/bindings.go b/storage/bindings.go index d506f3e1f9..be6cefe6b1 100644 --- a/storage/bindings.go +++ b/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) diff --git a/topdown/aggregates.go b/topdown/aggregates.go index 81fa61ead7..5d6cfd8352 100644 --- a/topdown/aggregates.go +++ b/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) diff --git a/topdown/arithmetic.go b/topdown/arithmetic.go index cd22afa471..2ec360252d 100644 --- a/topdown/arithmetic.go +++ b/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) diff --git a/topdown/casts.go b/topdown/casts.go index 34eab860ff..776762cb2b 100644 --- a/topdown/casts.go +++ b/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) diff --git a/topdown/eq.go b/topdown/eq.go index 277ff1869d..9a16f5c085 100644 --- a/topdown/eq.go +++ b/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 } diff --git a/topdown/topdown.go b/topdown/topdown.go index 2b145fd2eb..e8dafb606d 100644 --- a/topdown/topdown.go +++ b/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[]". - 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 }) } diff --git a/topdown/topdown_test.go b/topdown/topdown_test.go index d3155a1e05..d4bbf8c71e 100644 --- a/topdown/topdown_test.go +++ b/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 { diff --git a/util/hashmap.go b/util/hashmap.go index a04af8e18b..3ba40208fb 100644 --- a/util/hashmap.go +++ b/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 diff --git a/util/hashmap_test.go b/util/hashmap_test.go index 725372dcb5..9836353d63 100644 --- a/util/hashmap_test.go +++ b/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"