From 696a7d192e5eefa53230168a4b200ec0560c8a10 Mon Sep 17 00:00:00 2001 From: Luke Shumaker Date: Sun, 5 Feb 2023 00:31:29 -0700 Subject: containers: Rethink the RBTree interface to be simpler --- lib/containers/intervaltree.go | 132 +++++++++++++++++++++-------------------- 1 file changed, 68 insertions(+), 64 deletions(-) (limited to 'lib/containers/intervaltree.go') diff --git a/lib/containers/intervaltree.go b/lib/containers/intervaltree.go index b7ff866..7b96526 100644 --- a/lib/containers/intervaltree.go +++ b/lib/containers/intervaltree.go @@ -4,81 +4,80 @@ package containers -type intervalKey[K Ordered[K]] struct { +type interval[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]) Compare(b intervalKey[K]) int { +// Compare implements Ordered. +func (a interval[K]) Compare(b interval[K]) int { if d := a.Min.Compare(b.Min); d != 0 { return d } return a.Max.Compare(b.Max) } +// ContainsFn returns whether this interval contains the range matched +// by the given function. +func (ival interval[K]) ContainsFn(fn func(K) int) bool { + return fn(ival.Min) >= 0 && fn(ival.Max) <= 0 +} + type intervalValue[K Ordered[K], V any] struct { - Val V - SpanOfChildren intervalKey[K] + Val V + ValSpan interval[K] + ChildSpan interval[K] +} + +// Compare implements Ordered. +func (a intervalValue[K, V]) Compare(b intervalValue[K, V]) int { + return a.ValSpan.Compare(b.ValSpan) } 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), - } + inner RBTree[intervalValue[K, V]] } 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.Compare(max) > 0 { - max = node.Left.Value.SpanOfChildren.Max + max := node.Value.ValSpan.Max + if node.Left != nil && node.Left.Value.ChildSpan.Max.Compare(max) > 0 { + max = node.Left.Value.ChildSpan.Max } - if node.Right != nil && node.Right.Value.SpanOfChildren.Max.Compare(max) > 0 { - max = node.Right.Value.SpanOfChildren.Max + if node.Right != nil && node.Right.Value.ChildSpan.Max.Compare(max) > 0 { + max = node.Right.Value.ChildSpan.Max } - node.Value.SpanOfChildren.Max = max + node.Value.ChildSpan.Max = max - min := t.MinFn(node.Value.Val) - if node.Left != nil && node.Left.Value.SpanOfChildren.Min.Compare(min) < 0 { - min = node.Left.Value.SpanOfChildren.Min + min := node.Value.ValSpan.Min + if node.Left != nil && node.Left.Value.ChildSpan.Min.Compare(min) < 0 { + min = node.Left.Value.ChildSpan.Min } - if node.Right != nil && node.Right.Value.SpanOfChildren.Min.Compare(min) < 0 { - min = node.Right.Value.SpanOfChildren.Min + if node.Right != nil && node.Right.Value.ChildSpan.Min.Compare(min) < 0 { + min = node.Right.Value.ChildSpan.Min } - node.Value.SpanOfChildren.Min = min + node.Value.ChildSpan.Min = min } func (t *IntervalTree[K, V]) init() { - if t.inner.KeyFn == nil { - t.inner.KeyFn = t.keyFn + if t.inner.AttrFn == nil { 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}) + t.inner.Insert(intervalValue[K, V]{ + Val: val, + ValSpan: interval[K]{ + Min: t.MinFn(val), + Max: t.MaxFn(val), + }, + }) } func (t *IntervalTree[K, V]) Min() (K, bool) { @@ -86,7 +85,7 @@ func (t *IntervalTree[K, V]) Min() (K, bool) { var zero K return zero, false } - return t.inner.root.Value.SpanOfChildren.Min, true + return t.inner.root.Value.ChildSpan.Min, true } func (t *IntervalTree[K, V]) Max() (K, bool) { @@ -94,22 +93,18 @@ func (t *IntervalTree[K, V]) Max() (K, bool) { 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.Compare) + return t.inner.root.Value.ChildSpan.Max, true } 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): + case node.Value.ValSpan.ContainsFn(fn): return node.Value.Val, true - case node.Left != nil && node.Left.Value.SpanOfChildren.ContainsFn(fn): + case node.Left != nil && node.Left.Value.ChildSpan.ContainsFn(fn): node = node.Left - case node.Right != nil && node.Right.Value.SpanOfChildren.ContainsFn(fn): + case node.Right != nil && node.Right.Value.ChildSpan.ContainsFn(fn): node = node.Right default: node = nil @@ -119,24 +114,33 @@ func (t *IntervalTree[K, V]) Search(fn func(K) int) (V, bool) { return zero, false } -func (t *IntervalTree[K, V]) searchAll(fn func(K) int, node *RBNode[intervalValue[K, V]], ret *[]V) { +func (t *IntervalTree[K, V]) Range(fn func(V) bool) { + t.inner.Range(func(node *RBNode[intervalValue[K, V]]) bool { + return fn(node.Value.Val) + }) +} + +func (t *IntervalTree[K, V]) Subrange(rangeFn func(K) int, handleFn func(V) bool) { + t.subrange(t.inner.root, rangeFn, handleFn) +} + +func (t *IntervalTree[K, V]) subrange(node *RBNode[intervalValue[K, V]], rangeFn func(K) int, handleFn func(V) bool) bool { if node == nil { - return + return true } - if !node.Value.SpanOfChildren.ContainsFn(fn) { - return + if !node.Value.ChildSpan.ContainsFn(rangeFn) { + return true } - t.searchAll(fn, node.Left, ret) - if t.keyFn(node.Value).ContainsFn(fn) { - *ret = append(*ret, node.Value.Val) + if !t.subrange(node.Left, rangeFn, handleFn) { + return false } - 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 + if node.Value.ValSpan.ContainsFn(rangeFn) { + if !handleFn(node.Value.Val) { + return false + } + } + if !t.subrange(node.Right, rangeFn, handleFn) { + return false + } + return true } - -// TODO: func (t *IntervalTree[K, V]) Walk(fn func(*RBNode[V]) error) error -- cgit v1.2.3-2-g168b