From 6b0151257a76b5f2f6aa4bbcdc027fce57cc6170 Mon Sep 17 00:00:00 2001 From: JP Appel Date: Fri, 13 Jun 2025 14:26:14 -0400 Subject: Add bfs and dfs iterators for clause trees --- pkg/query/parser.go | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/pkg/query/parser.go b/pkg/query/parser.go index e8fccfe..5de803c 100644 --- a/pkg/query/parser.go +++ b/pkg/query/parser.go @@ -2,6 +2,7 @@ package query import ( "fmt" + "iter" "os" "slices" "strings" @@ -283,6 +284,60 @@ func (c Clause) buildString(b *strings.Builder, level int) { b.WriteString(")\n") } +// Depth of tree via recursion +func (root Clause) Depth() int { + maxHeight := 1 + for _, child := range root.Clauses { + maxHeight = max(maxHeight, child.Depth()+1) + } + return maxHeight +} + +func (root *Clause) DFS() iter.Seq[*Clause] { + return func(yield func(*Clause) bool) { + stack := make([]*Clause, 0, len(root.Clauses)) + stack = append(stack, root) + + for len(stack) != 0 { + node := stack[len(stack)-1] + + if !yield(node) { + return + } + + stack := stack[:len(stack)-1] + stack = append(stack, node.Clauses...) + } + } +} + +func (root *Clause) BFS() iter.Seq[*Clause] { + return func(yield func(*Clause) bool) { + queue := make([]*Clause, 64*len(root.Clauses)) + queue = append(queue, root) + read := 0 + write := 1 + size := 1 + + // FIXME: can potentially discard values if queue is too small + for size != 0 { + node := queue[read] + + if !yield(node) { + return + } + read = (read + 1) % len(queue) + size -= 1 + + for _, child := range node.Clauses { + queue[write] = child + write = (write + 1) % len(queue) + size += 1 + } + } + } +} + func Parse(tokens []Token) (*Clause, error) { stack := make([]*Clause, 0, 10) -- cgit v1.2.3