containers: Implement an interval tree

Hits a bug in the Go 1.19.1 compiler, you must use a dev build of Go.

2 files changed, 223 insertions, 0 deletions

diff --git a/lib/containers/intervaltree.go b/lib/containers/intervaltree.go
new file mode 100644
index 0000000..424b297
--- /dev/null
+++ b/lib/containers/intervaltree.go
@@ -0,0 +1,142 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package containers
+
+type intervalKey[K Ordered[K]] struct {
+	Min, Max K
+}
+
+func (ival intervalKey[K]) ContainsFn(fn func(K) int) bool {
+	return fn(ival.Min) >= 0 && fn(ival.Max) <= 0
+}
+
+func (a intervalKey[K]) Cmp(b intervalKey[K]) int {
+	if d := a.Min.Cmp(b.Min); d != 0 {
+		return d
+	}
+	return a.Max.Cmp(b.Max)
+}
+
+type intervalValue[K Ordered[K], V any] struct {
+	Val            V
+	SpanOfChildren intervalKey[K]
+}
+
+type IntervalTree[K Ordered[K], V any] struct {
+	MinFn func(V) K
+	MaxFn func(V) K
+	inner RBTree[intervalKey[K], intervalValue[K, V]]
+}
+
+func (t *IntervalTree[K, V]) keyFn(v intervalValue[K, V]) intervalKey[K] {
+	return intervalKey[K]{
+		Min: t.MinFn(v.Val),
+		Max: t.MaxFn(v.Val),
+	}
+}
+
+func (t *IntervalTree[K, V]) attrFn(node *RBNode[intervalValue[K, V]]) {
+	max := t.MaxFn(node.Value.Val)
+	if node.Left != nil && node.Left.Value.SpanOfChildren.Max.Cmp(max) > 0 {
+		max = node.Left.Value.SpanOfChildren.Max
+	}
+	if node.Right != nil && node.Right.Value.SpanOfChildren.Max.Cmp(max) > 0 {
+		max = node.Right.Value.SpanOfChildren.Max
+	}
+	node.Value.SpanOfChildren.Max = max
+
+	min := t.MinFn(node.Value.Val)
+	if node.Left != nil && node.Left.Value.SpanOfChildren.Min.Cmp(min) < 0 {
+		min = node.Left.Value.SpanOfChildren.Min
+	}
+	if node.Right != nil && node.Right.Value.SpanOfChildren.Min.Cmp(min) < 0 {
+		min = node.Right.Value.SpanOfChildren.Min
+	}
+	node.Value.SpanOfChildren.Min = min
+}
+
+func (t *IntervalTree[K, V]) init() {
+	if t.inner.KeyFn == nil {
+		t.inner.KeyFn = t.keyFn
+		t.inner.AttrFn = t.attrFn
+	}
+}
+
+func (t *IntervalTree[K, V]) Delete(min, max K) {
+	t.init()
+	t.inner.Delete(intervalKey[K]{
+		Min: min,
+		Max: max,
+	})
+}
+
+func (t *IntervalTree[K, V]) Equal(u *IntervalTree[K, V]) bool {
+	return t.inner.Equal(&u.inner)
+}
+
+func (t *IntervalTree[K, V]) Insert(val V) {
+	t.init()
+	t.inner.Insert(intervalValue[K, V]{Val: val})
+}
+
+func (t *IntervalTree[K, V]) Min() (K, bool) {
+	if t.inner.root == nil {
+		var zero K
+		return zero, false
+	}
+	return t.inner.root.Value.SpanOfChildren.Min, true
+}
+
+func (t *IntervalTree[K, V]) Max() (K, bool) {
+	if t.inner.root == nil {
+		var zero K
+		return zero, false
+	}
+	return t.inner.root.Value.SpanOfChildren.Max, true
+}
+
+func (t *IntervalTree[K, V]) Lookup(k K) (V, bool) {
+	return t.Search(k.Cmp)
+}
+
+func (t *IntervalTree[K, V]) Search(fn func(K) int) (V, bool) {
+	node := t.inner.root
+	for node != nil {
+		switch {
+		case t.keyFn(node.Value).ContainsFn(fn):
+			return node.Value.Val, true
+		case node.Left != nil && node.Left.Value.SpanOfChildren.ContainsFn(fn):
+			node = node.Left
+		case node.Right != nil && node.Right.Value.SpanOfChildren.ContainsFn(fn):
+			node = node.Right
+		default:
+			node = nil
+		}
+	}
+	var zero V
+	return zero, false
+}
+
+func (t *IntervalTree[K, V]) searchAll(fn func(K) int, node *RBNode[intervalValue[K, V]], ret *[]V) {
+	if node == nil {
+		return
+	}
+	if !node.Value.SpanOfChildren.ContainsFn(fn) {
+		return
+	}
+	t.searchAll(fn, node.Left, ret)
+	if t.keyFn(node.Value).ContainsFn(fn) {
+		*ret = append(*ret, node.Value.Val)
+	}
+	t.searchAll(fn, node.Right, ret)
+}
+
+func (t *IntervalTree[K, V]) SearchAll(fn func(K) int) []V {
+	var ret []V
+	t.searchAll(fn, t.inner.root, &ret)
+	return ret
+}
+
+//func (t *IntervalTree[K, V]) Walk(fn func(*RBNode[V]) error) error
diff --git a/lib/containers/intervaltree_test.go b/lib/containers/intervaltree_test.go
new file mode 100644
index 0000000..7a4689b
--- /dev/null
+++ b/lib/containers/intervaltree_test.go
@@ -0,0 +1,81 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package containers
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func (t *IntervalTree[K, V]) ASCIIArt() string {
+	return t.inner.ASCIIArt()
+}
+
+func (v intervalValue[K, V]) String() string {
+	return fmt.Sprintf("%v) ([%v,%v]",
+		v.Val,
+		v.SpanOfChildren.Min,
+		v.SpanOfChildren.Max)
+}
+
+func (v NativeOrdered[T]) String() string {
+	return fmt.Sprintf("%v", v.Val)
+}
+
+type SimpleInterval struct {
+	Min, Max int
+}
+
+func (ival SimpleInterval) String() string {
+	return fmt.Sprintf("[%v,%v]", ival.Min, ival.Max)
+}
+
+func TestIntervalTree(t *testing.T) {
+	tree := IntervalTree[NativeOrdered[int], SimpleInterval]{
+		MinFn: func(ival SimpleInterval) NativeOrdered[int] { return NativeOrdered[int]{ival.Min} },
+		MaxFn: func(ival SimpleInterval) NativeOrdered[int] { return NativeOrdered[int]{ival.Max} },
+	}
+
+	// CLRS Figure 14.4
+	// level 0
+	tree.Insert(SimpleInterval{16, 21})
+	// level 1
+	tree.Insert(SimpleInterval{8, 9})
+	tree.Insert(SimpleInterval{25, 30})
+	// level 2
+	tree.Insert(SimpleInterval{5, 8})
+	tree.Insert(SimpleInterval{15, 23})
+	tree.Insert(SimpleInterval{17, 19})
+	tree.Insert(SimpleInterval{26, 26})
+	// level 3
+	tree.Insert(SimpleInterval{0, 3})
+	tree.Insert(SimpleInterval{6, 10})
+	tree.Insert(SimpleInterval{19, 20})
+
+	t.Log(tree.ASCIIArt())
+
+	// find intervals that touch [9,20]
+	intervals := tree.SearchAll(func(k NativeOrdered[int]) int {
+		if k.Val < 9 {
+			return 1
+		}
+		if k.Val > 20 {
+			return -1
+		}
+		return 0
+	})
+	assert.Equal(t,
+		[]SimpleInterval{
+			{6, 10},
+			{8, 9},
+			{15, 23},
+			{16, 21},
+			{17, 19},
+			{19, 20},
+		},
+		intervals)
+}