Skip to content

Commit 86f1717

Browse files
sean-rdleal
authored andcommitted
interval: add InOrderTraverse() visitor
1 parent 178e917 commit 86f1717

File tree

1 file changed

+76
-0
lines changed

1 file changed

+76
-0
lines changed

interval/search.go

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -478,3 +478,79 @@ func maxEnd[V, T any](n *node[V, T], searchEnd T, cmp CmpFunc[T], visit func(*no
478478
maxEnd(n.Right, searchEnd, cmp, visit)
479479
}
480480
}
481+
482+
// StopTraversal is used as a return value from [VisitFunc] to indicate that the iteration is to be stopped.
483+
// It is not returned as an error by any function.
484+
var StopTraversal = errors.New("stop tree traversal")
485+
486+
// VisitFunc is called on all values. Returning non-nil error will stop iteration.
487+
// If the returned error is [StopTraversal], the iteration is interrupted, but no error is returned to the caller.
488+
type VisitFunc[V, T any] func(V, T) error
489+
490+
// InOrderTraverse traverses the tree in order and applies VisitFunc to each node. It's safe for concurrent use. To prevent deadlock, avoid calling other tree methods within visitFunc.
491+
func (st *SearchTree[V, T]) InOrderTraverse(visitFunc VisitFunc[V, T]) error {
492+
tree.mu.RLock()
493+
defer tree.mu.RUnlock()
494+
495+
var inOrder func(n *node[V, T]) error
496+
inOrder = func(n *node[V, T]) error {
497+
if n == nil {
498+
return nil
499+
}
500+
501+
// Visit left child
502+
if err := inOrder(n.Left); err != nil {
503+
return err
504+
}
505+
506+
// Visit current node
507+
err := visitFunc(n.Interval.Val, n.Interval.Start)
508+
if err != nil {
509+
return err
510+
}
511+
512+
// Visit right child
513+
return inOrder(n.Right)
514+
}
515+
516+
err := inOrder(tree.root)
517+
// Do not percolate StopTraversal error to the caller.
518+
if errors.Is(err, StopTraversal) {
519+
return nil
520+
}
521+
return err
522+
}
523+
524+
// InOrderTraverse traverses the tree in order and applies VisitFunc to each node. It's safe for concurrent use. To prevent deadlock, avoid calling other tree methods within visitFunc.
525+
func (st *MultiValueSearchTree[V, T]) InOrderTraverse(visitFunc VisitFunc[V, T]) error {
526+
tree.mu.RLock()
527+
defer tree.mu.RUnlock()
528+
529+
var inOrder func(n *node[V, T]) error
530+
inOrder = func(n *node[V, T]) error {
531+
if n == nil {
532+
return nil
533+
}
534+
535+
// Visit left child
536+
if err := inOrder(n.Left); err != nil {
537+
return err
538+
}
539+
540+
// Visit current node
541+
err := visitFunc(n.Interval.Vals, n.Interval.Start)
542+
if err != nil {
543+
return err
544+
}
545+
546+
// Visit right child
547+
return inOrder(n.Right)
548+
}
549+
550+
err := inOrder(tree.root)
551+
// Do not percolate StopTraversal error to the caller.
552+
if errors.Is(err, StopTraversal) {
553+
return nil
554+
}
555+
return err
556+
}

0 commit comments

Comments
 (0)