diff --git a/omap/omap.go b/omap/omap.go index f2ae5ae..7e5c199 100644 --- a/omap/omap.go +++ b/omap/omap.go @@ -160,10 +160,9 @@ func (m Map[T, U]) Keys() []T { return nil } out := make([]T, 0, m.Len()) - m.m.Inorder(func(kv stree.KV[T, U]) bool { + for kv := range m.m.Inorder { out = append(out, kv.Key) - return true - }) + } return out } @@ -215,10 +214,10 @@ func (it *Iter[T, U]) Value() U { return it.c.Key().Value } func (it *Iter[T, U]) Seek(key T) *Iter[T, U] { it.c = nil if it.m != nil { - it.m.InorderAfter(stree.KV[T, U]{Key: key}, func(kv stree.KV[T, U]) bool { + for kv := range it.m.InorderAfter(stree.KV[T, U]{Key: key}) { it.c = it.m.Cursor(kv) - return false - }) + break + } } return it } diff --git a/oset/oset.go b/oset/oset.go index 6f1645a..09c06a8 100644 --- a/oset/oset.go +++ b/oset/oset.go @@ -152,10 +152,9 @@ func (s Set[T]) AddAll(t Set[T]) Set[T] { if t.s == nil { return s } - t.s.Inorder(func(v T) bool { + for v := range t.s.Inorder { s.s.Add(v) - return true - }) + } return s } @@ -187,10 +186,9 @@ func (s Set[T]) Slice() []T { return nil } out := make([]T, 0, s.Len()) - s.s.Inorder(func(val T) bool { + for val := range s.s.Inorder { out = append(out, val) - return true - }) + } return out } @@ -241,10 +239,10 @@ func (it *Iter[T]) Value() T { return it.c.Key() } func (it *Iter[T]) Seek(value T) *Iter[T] { it.c = nil if it.s != nil { - it.s.InorderAfter(value, func(key T) bool { + for key := range it.s.InorderAfter(value) { it.c = it.s.Cursor(key) - return false - }) + break + } } return it } diff --git a/stree/cursor.go b/stree/cursor.go index a087a71..68ee0f8 100644 --- a/stree/cursor.go +++ b/stree/cursor.go @@ -232,12 +232,9 @@ func (c *Cursor[T]) Max() *Cursor[T] { return c } -// Inorder calls f for each key of the subtree rooted at c in order. If f -// returns false, Inorder stops and returns false; otherwise it returns true -// after visiting all elements of c. -func (c *Cursor[T]) Inorder(f func(key T) bool) bool { +// Inorder is a range function over each key of the subtree at c in order. +func (c *Cursor[T]) Inorder(yield func(key T) bool) { if c.Valid() { - return c.path[len(c.path)-1].inorder(f) + c.path[len(c.path)-1].inorder(yield) } - return true } diff --git a/stree/stree.go b/stree/stree.go index 27adede..c3a4dba 100644 --- a/stree/stree.go +++ b/stree/stree.go @@ -19,6 +19,7 @@ package stree import ( "fmt" + "iter" "math" "slices" ) @@ -276,16 +277,15 @@ func (t *Tree[T]) Get(key T) (_ T, ok bool) { return } -// Inorder calls f for each key of t in order. If f returns false, Inorder -// stops and returns false; otherwise it returns true after visiting all -// elements of t. -func (t *Tree[T]) Inorder(f func(key T) bool) bool { return t.root.inorder(f) } +// Inorder is a range function that visits each key of t in order. +func (t *Tree[T]) Inorder(yield func(key T) bool) { t.root.inorder(yield) } -// InorderAfter calls f for each key greater than or equal to key, in order. -// if f returns false, InorderAfter stops and returns false. Otherwise, it -// returns true after visiting all eligible elements of t. -func (t *Tree[T]) InorderAfter(key T, f func(key T) bool) bool { - return t.root.inorderAfter(key, t.compare, f) +// InorderAfter returns a range function for each key greater than or equal to +// key, in order. +func (t *Tree[T]) InorderAfter(key T) iter.Seq[T] { + return func(yield func(T) bool) { + t.root.inorderAfter(key, t.compare, yield) + } } // Cursor constructs a cursor to the specified key, or nil if key is not diff --git a/stree/stree_test.go b/stree/stree_test.go index 23d3d50..c17d24f 100644 --- a/stree/stree_test.go +++ b/stree/stree_test.go @@ -25,11 +25,10 @@ var ( // Export all the words in tree in their stored order. func allWords(tree *stree.Tree[string]) []string { - var got []string - tree.Inorder(func(key string) bool { + got := make([]string, 0, tree.Len()) + for key := range tree.Inorder { got = append(got, key) - return true - }) + } return got } @@ -154,10 +153,9 @@ func TestInorderAfter(t *testing.T) { for _, test := range tests { want := strings.Fields(test.want) var got []string - tree.InorderAfter(test.key, func(key string) bool { + for key := range tree.InorderAfter(test.key) { got = append(got, key) - return true - }) + } if diff := gocmp.Diff(want, got, cmpopts.EquateEmpty()); diff != "" { t.Errorf("InorderAfter(%v) result differed from expected\n%s", test.key, diff) } @@ -317,10 +315,9 @@ func TestCursor(t *testing.T) { t.Run("Traverse", func(t *testing.T) { var got []string - tree.Cursor("f").Inorder(func(s string) bool { - got = append(got, s) - return true - }) + for key := range tree.Cursor("f").Inorder { + got = append(got, key) + } if diff := gocmp.Diff(got, []string{"e", "f", "g"}); diff != "" { t.Errorf("Right tree (-got, +want):\n%s", diff) } @@ -342,11 +339,10 @@ func TestKV(t *testing.T) { var gotk []string var gotv []int - st.Inorder(func(kv kv) bool { + for kv := range st.Inorder { gotk = append(gotk, kv.Key) gotv = append(gotv, kv.Value) - return true - }) + } if diff := gocmp.Diff(gotk, []string{"anybody", "hello", "here", "in", "is", "there"}); diff != "" { t.Errorf("Keys (-got, +want):\n%s", diff)