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

import (
tavit ohanian's avatar
tavit ohanian committed
4
	"gitlab.dms3.io/p2p/go-p2p-xor/key"
Petar Maymounkov's avatar
Petar Maymounkov committed
5
)
Petar Maymounkov's avatar
Petar Maymounkov committed
6

7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
func IntersectKeySlices(p, q []key.Key) []key.Key {
	hat := []key.Key{}
	for _, p := range p {
		if keyIsIn(p, q) && !keyIsIn(p, hat) {
			hat = append(hat, p)
		}
	}
	return hat
}

func keyIsIn(q key.Key, s []key.Key) bool {
	for _, s := range s {
		if key.Equal(q, s) {
			return true
		}
	}
	return false
}

Petar Maymounkov's avatar
Petar Maymounkov committed
26 27
// Intersect computes the intersection of the keys in p and q.
// p and q must be non-nil. The returned trie is never nil.
28 29
func Intersect(p, q *Trie) *Trie {
	return IntersectAtDepth(0, p, q)
Petar Maymounkov's avatar
Petar Maymounkov committed
30 31
}

32
func IntersectAtDepth(depth int, p, q *Trie) *Trie {
Petar Maymounkov's avatar
Petar Maymounkov committed
33
	switch {
34 35 36
	case p.IsLeaf() && q.IsLeaf():
		if p.IsEmpty() || q.IsEmpty() {
			return &Trie{} // empty set
Petar Maymounkov's avatar
Petar Maymounkov committed
37
		} else {
Petar Maymounkov's avatar
Petar Maymounkov committed
38
			if key.Equal(p.Key, q.Key) {
39
				return &Trie{Key: p.Key} // singleton
Petar Maymounkov's avatar
Petar Maymounkov committed
40
			} else {
41
				return &Trie{} // empty set
Petar Maymounkov's avatar
Petar Maymounkov committed
42 43
			}
		}
44 45 46
	case p.IsLeaf() && !q.IsLeaf():
		if p.IsEmpty() {
			return &Trie{} // empty set
Petar Maymounkov's avatar
Petar Maymounkov committed
47
		} else {
48 49
			if _, found := q.FindAtDepth(depth, p.Key); found {
				return &Trie{Key: p.Key}
Petar Maymounkov's avatar
Petar Maymounkov committed
50
			} else {
51
				return &Trie{} // empty set
Petar Maymounkov's avatar
Petar Maymounkov committed
52 53
			}
		}
54
	case !p.IsLeaf() && q.IsLeaf():
Petar Maymounkov's avatar
Petar Maymounkov committed
55
		return IntersectAtDepth(depth, q, p)
56 57 58 59 60
	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
61 62 63 64 65 66 67
			},
		}
		disjointUnion.shrink()
		return disjointUnion
	}
	panic("unreachable")
}