4 files changed, 68 insertions, 101 deletions
diff --git a/plugin/file/tree/all.go b/plugin/file/tree/all.go
index fd806365f..e1fc5b392 100644
--- a/plugin/file/tree/all.go
+++ b/plugin/file/tree/all.go
@@ -1,6 +1,6 @@
 package tree
 
-// All traverses tree and returns all elements
+// All traverses tree and returns all elements.
 func (t *Tree) All() []*Elem {
 	if t.Root == nil {
 		return nil
@@ -19,30 +19,3 @@ func (n *Node) all(found []*Elem) []*Elem {
 	}
 	return found
 }
-
-// Do performs fn on all values stored in the tree. A boolean is returned indicating whether the
-// Do traversal was interrupted by an Operation returning true. If fn alters stored values' sort
-// relationships, future tree operation behaviors are undefined.
-func (t *Tree) Do(fn func(e *Elem) bool) bool {
-	if t.Root == nil {
-		return false
-	}
-	return t.Root.do(fn)
-}
-
-func (n *Node) do(fn func(e *Elem) bool) (done bool) {
-	if n.Left != nil {
-		done = n.Left.do(fn)
-		if done {
-			return
-		}
-	}
-	done = fn(n.Elem)
-	if done {
-		return
-	}
-	if n.Right != nil {
-		done = n.Right.do(fn)
-	}
-	return
-}
diff --git a/plugin/file/tree/elem.go b/plugin/file/tree/elem.go
index 6317cc912..c1909649d 100644
--- a/plugin/file/tree/elem.go
+++ b/plugin/file/tree/elem.go
@@ -15,20 +15,34 @@ func newElem(rr dns.RR) *Elem {
 	return &e
 }
 
-// Types returns the RRs with type qtype from e. If qname is given (only the
-// first one is used), the RR are copied and the owner is replaced with qname[0].
-func (e *Elem) Types(qtype uint16, qname ...string) []dns.RR {
+// Types returns the types of the records in e. The returned list is not sorted.
+func (e *Elem) Types() []uint16 {
+	t := make([]uint16, len(e.m))
+	i := 0
+	for ty := range e.m {
+		t[i] = ty
+		i++
+	}
+	return t
+}
+
+// Type returns the RRs with type qtype from e.
+func (e *Elem) Type(qtype uint16) []dns.RR { return e.m[qtype] }
+
+// TypeForWildcard returns the RRs with type qtype from e. The ownername returned is set to qname.
+func (e *Elem) TypeForWildcard(qtype uint16, qname string) []dns.RR {
 	rrs := e.m[qtype]
 
-	if rrs != nil && len(qname) > 0 {
-		copied := make([]dns.RR, len(rrs))
-		for i := range rrs {
-			copied[i] = dns.Copy(rrs[i])
-			copied[i].Header().Name = qname[0]
-		}
-		return copied
+	if rrs == nil {
+		return nil
+	}
+
+	copied := make([]dns.RR, len(rrs))
+	for i := range rrs {
+		copied[i] = dns.Copy(rrs[i])
+		copied[i].Header().Name = qname
 	}
-	return rrs
+	return copied
 }
 
 // All returns all RRs from e, regardless of type.
@@ -52,13 +66,10 @@ func (e *Elem) Name() string {
 	return ""
 }
 
-// Empty returns true is e does not contain any RRs, i.e. is an
-// empty-non-terminal.
-func (e *Elem) Empty() bool {
-	return len(e.m) == 0
-}
+// Empty returns true is e does not contain any RRs, i.e. is an empty-non-terminal.
+func (e *Elem) Empty() bool { return len(e.m) == 0 }
 
-// Insert inserts rr into e. If rr is equal to existing rrs this is a noop.
+// Insert inserts rr into e. If rr is equal to existing RRs, the RR will be added anyway.
 func (e *Elem) Insert(rr dns.RR) {
 	t := rr.Header().Rrtype
 	if e.m == nil {
@@ -71,66 +82,20 @@ func (e *Elem) Insert(rr dns.RR) {
 		e.m[t] = []dns.RR{rr}
 		return
 	}
-	for _, er := range rrs {
-		if equalRdata(er, rr) {
-			return
-		}
-	}
 
 	rrs = append(rrs, rr)
 	e.m[t] = rrs
 }
 
-// Delete removes rr from e. When e is empty after the removal the returned bool is true.
-func (e *Elem) Delete(rr dns.RR) (empty bool) {
+// Delete removes all RRs of type rr.Header().Rrtype from e.
+func (e *Elem) Delete(rr dns.RR) {
 	if e.m == nil {
-		return true
-	}
-
-	t := rr.Header().Rrtype
-	rrs, ok := e.m[t]
-	if !ok {
 		return
 	}
 
-	for i, er := range rrs {
-		if equalRdata(er, rr) {
-			rrs = removeFromSlice(rrs, i)
-			e.m[t] = rrs
-			empty = len(rrs) == 0
-			if empty {
-				delete(e.m, t)
-			}
-			return
-		}
-	}
-	return
+	t := rr.Header().Rrtype
+	delete(e.m, t)
 }
 
 // Less is a tree helper function that calls less.
 func Less(a *Elem, name string) int { return less(name, a.Name()) }
-
-// Assuming the same type and name this will check if the rdata is equal as well.
-func equalRdata(a, b dns.RR) bool {
-	switch x := a.(type) {
-	// TODO(miek): more types, i.e. all types. + tests for this.
-	case *dns.A:
-		return x.A.Equal(b.(*dns.A).A)
-	case *dns.AAAA:
-		return x.AAAA.Equal(b.(*dns.AAAA).AAAA)
-	case *dns.MX:
-		if x.Mx == b.(*dns.MX).Mx && x.Preference == b.(*dns.MX).Preference {
-			return true
-		}
-	}
-	return false
-}
-
-// removeFromSlice removes index i from the slice.
-func removeFromSlice(rrs []dns.RR, i int) []dns.RR {
-	if i >= len(rrs) {
-		return rrs
-	}
-	rrs = append(rrs[:i], rrs[i+1:]...)
-	return rrs
-}
diff --git a/plugin/file/tree/tree.go b/plugin/file/tree/tree.go
index ed33c09a4..3aeeba4d5 100644
--- a/plugin/file/tree/tree.go
+++ b/plugin/file/tree/tree.go
@@ -275,7 +275,8 @@ func (n *Node) deleteMax() (root *Node, d int) {
 	return
 }
 
-// Delete removes rr from the tree, is the node turns empty, that node is deleted with DeleteNode.
+// Delete removes all RRs of type rr.Header().Rrtype from e. If after the deletion of rr the node is empty the
+// entire node is deleted.
 func (t *Tree) Delete(rr dns.RR) {
 	if t.Root == nil {
 		return
@@ -283,15 +284,13 @@ func (t *Tree) Delete(rr dns.RR) {
 
 	el, _ := t.Search(rr.Header().Name)
 	if el == nil {
-		t.deleteNode(rr)
 		return
 	}
-	// Delete from this element.
-	empty := el.Delete(rr)
-	if empty {
+	el.Delete(rr)
+	if el.Empty() {
 		t.deleteNode(rr)
-		return
 	}
+	return
 }
 
 // DeleteNode deletes the node that matches rr according to Less().
diff --git a/plugin/file/tree/walk.go b/plugin/file/tree/walk.go
new file mode 100644
index 000000000..00accbafa
--- /dev/null
+++ b/plugin/file/tree/walk.go
@@ -0,0 +1,30 @@
+package tree
+
+import "github.com/miekg/dns"
+
+// Walk performs fn on all values stored in the tree. If a non-nil error is returned the
+// Walk was interrupted by an fn returning that error. If fn alters stored values' sort
+// relationships, future tree operation behaviors are undefined.
+func (t *Tree) Walk(fn func(e *Elem, rrs map[uint16][]dns.RR) error) error {
+	if t.Root == nil {
+		return nil
+	}
+	return t.Root.walk(fn)
+}
+
+func (n *Node) walk(fn func(e *Elem, rrs map[uint16][]dns.RR) error) error {
+	if n.Left != nil {
+		if err := n.Left.walk(fn); err != nil {
+			return err
+		}
+	}
+	if err := fn(n.Elem, n.Elem.m); err != nil {
+		return err
+	}
+	if n.Right != nil {
+		if err := n.Right.walk(fn); err != nil {
+			return err
+		}
+	}
+	return nil
+}