intersect.go 1.1 KB
Newer Older
Petar Maymounkov's avatar
Petar Maymounkov committed
1 2 3 4 5
package trie

import (
	"github.com/libp2p/go-libp2p-xor/key"
)
Petar Maymounkov's avatar
Petar Maymounkov committed
6 7 8

// Intersect computes the intersection of the keys in p and q.
// p and q must be non-nil. The returned trie is never nil.
9 10
func Intersect(p, q *Trie) *Trie {
	return IntersectAtDepth(0, p, q)
Petar Maymounkov's avatar
Petar Maymounkov committed
11 12
}

13
func IntersectAtDepth(depth int, p, q *Trie) *Trie {
Petar Maymounkov's avatar
Petar Maymounkov committed
14
	switch {
15 16 17
	case p.IsLeaf() && q.IsLeaf():
		if p.IsEmpty() || q.IsEmpty() {
			return &Trie{} // empty set
Petar Maymounkov's avatar
Petar Maymounkov committed
18
		} else {
Petar Maymounkov's avatar
Petar Maymounkov committed
19
			if key.Equal(p.Key, q.Key) {
20
				return &Trie{Key: p.Key} // singleton
Petar Maymounkov's avatar
Petar Maymounkov committed
21
			} else {
22
				return &Trie{} // empty set
Petar Maymounkov's avatar
Petar Maymounkov committed
23 24
			}
		}
25 26 27
	case p.IsLeaf() && !q.IsLeaf():
		if p.IsEmpty() {
			return &Trie{} // empty set
Petar Maymounkov's avatar
Petar Maymounkov committed
28
		} else {
29 30
			if _, found := q.FindAtDepth(depth, p.Key); found {
				return &Trie{Key: p.Key}
Petar Maymounkov's avatar
Petar Maymounkov committed
31
			} else {
32
				return &Trie{} // empty set
Petar Maymounkov's avatar
Petar Maymounkov committed
33 34
			}
		}
35
	case !p.IsLeaf() && q.IsLeaf():
Petar Maymounkov's avatar
Petar Maymounkov committed
36
		return Intersect(q, p)
37 38 39 40 41
	case !p.IsLeaf() && !q.IsLeaf():
		disjointUnion := &Trie{
			Branch: [2]*Trie{
				IntersectAtDepth(depth+1, p.Branch[0], q.Branch[0]),
				IntersectAtDepth(depth+1, p.Branch[1], q.Branch[1]),
Petar Maymounkov's avatar
Petar Maymounkov committed
42 43 44 45 46 47 48
			},
		}
		disjointUnion.shrink()
		return disjointUnion
	}
	panic("unreachable")
}