chore(query): combine walks into walk bottomup/preorder functions

Author: barakmich

pkg/query/advisor.go

@@ -58,64 +58,43 @@ func ApplyAdvisor(co CanonicalOutline, advisor PlanAdvisor) (CanonicalOutline, e
 	}, nil
 }
 
-// applyAdvisorMutations recursively walks the outline tree bottom-up.
+// applyAdvisorMutations walks the outline tree bottom-up via WalkOutlineBottomUp.
 // For each node it calls GetMutations on the advisor and applies the returned
 // mutations in sequence. After each mutation it verifies that the resulting
 // node's ID matches the original node's ID; a mismatch is a programmer bug.
 func applyAdvisorMutations(outline Outline, co CanonicalOutline, advisor PlanAdvisor) (Outline, error) {
-	// Recurse into children first (bottom-up).
-	if len(outline.SubOutlines) > 0 {
-		newSubs := make([]Outline, len(outline.SubOutlines))
-		for i, sub := range outline.SubOutlines {
-			mutated, err := applyAdvisorMutations(sub, co, advisor)
-			if err != nil {
-				return Outline{}, err
-			}
-			newSubs[i] = mutated
-		}
-		outline = Outline{
-			Type:        outline.Type,
-			Args:        outline.Args,
-			SubOutlines: newSubs,
-			ID:          outline.ID,
+	return WalkOutlineBottomUp(outline, func(node Outline) (Outline, error) {
+		mutations, err := advisor.GetMutations(node, co)
+		if err != nil {
+			return Outline{}, err
 		}
-	}
-
-	mutations, err := advisor.GetMutations(outline, co)
-	if err != nil {
-		return Outline{}, err
-	}
-
-	result := outline
-	for _, fn := range mutations {
-		next := fn(result)
-		if next.ID != result.ID {
-			return Outline{}, spiceerrors.MustBugf(
-				"advisor mutation changed outline node ID from %d to %d; mutations must preserve the root node ID",
-				result.ID, next.ID,
-			)
+		result := node
+		for _, fn := range mutations {
+			next := fn(result)
+			if next.ID != result.ID {
+				return Outline{}, spiceerrors.MustBugf(
+					"advisor mutation changed outline node ID from %d to %d; mutations must preserve the root node ID",
+					result.ID, next.ID,
+				)
+			}
+			result = next
 		}
-		result = next
-	}
-	return result, nil
+		return result, nil
+	})
 }
 
-// collectAdvisorHints recursively walks the outline tree and calls GetHints on
-// the advisor for each node, storing any returned hints in the hints map keyed
-// by the node's ID.
+// collectAdvisorHints walks the outline tree pre-order via WalkOutlinePreOrder,
+// calling GetHints on the advisor for each node and storing any returned hints
+// in the hints map keyed by the node's ID.
 func collectAdvisorHints(outline Outline, co CanonicalOutline, advisor PlanAdvisor, hints map[OutlineNodeID][]Hint) error {
-	nodeHints, err := advisor.GetHints(outline, co)
-	if err != nil {
-		return err
-	}
-	if len(nodeHints) > 0 {
-		hints[outline.ID] = nodeHints
-	}
-
-	for _, sub := range outline.SubOutlines {
-		if err := collectAdvisorHints(sub, co, advisor, hints); err != nil {
+	return WalkOutlinePreOrder(outline, func(node Outline) error {
+		nodeHints, err := advisor.GetHints(node, co)
+		if err != nil {
 			return err
 		}
-	}
-	return nil
+		if len(nodeHints) > 0 {
+			hints[node.ID] = nodeHints
+		}
+		return nil
+	})
 }

pkg/query/mutations.go

@@ -7,25 +7,14 @@ type OutlineMutation func(Outline) Outline
 // MutateOutline performs a bottom-up traversal of the outline tree, applying
 // all the given transformation functions to each node after processing its children.
 func MutateOutline(outline Outline, fns []OutlineMutation) Outline {
-	// Recurse on children first (bottom-up)
-	if len(outline.SubOutlines) > 0 {
-		newSubs := make([]Outline, len(outline.SubOutlines))
-		for i, sub := range outline.SubOutlines {
-			newSubs[i] = MutateOutline(sub, fns)
+	// WalkOutlineBottomUp cannot fail when the callback never returns an error,
+	// so the error return is safe to ignore here.
+	result, _ := WalkOutlineBottomUp(outline, func(node Outline) (Outline, error) {
+		for _, fn := range fns {
+			node = fn(node)
 		}
-		outline = Outline{
-			Type:        outline.Type,
-			Args:        outline.Args,
-			SubOutlines: newSubs,
-			ID:          outline.ID,
-		}
-	}
-
-	// Then apply all mutation functions in sequence to current node
-	result := outline
-	for _, fn := range fns {
-		result = fn(result)
-	}
+		return node, nil
+	})
 	return result
 }
 

pkg/query/outline.go

@@ -75,6 +75,46 @@ type Outline struct {
 	ID          OutlineNodeID // Populated only inside a CanonicalOutline
 }
 
+// WalkOutlineBottomUp performs a bottom-up traversal of an Outline tree,
+// rebuilding each node with its processed children before passing it to fn.
+// fn receives the node (with already-processed children) and returns a
+// replacement node and an optional error. The traversal stops and the error
+// is propagated on the first failure.
+func WalkOutlineBottomUp(outline Outline, fn func(Outline) (Outline, error)) (Outline, error) {
+	if len(outline.SubOutlines) > 0 {
+		newSubs := make([]Outline, len(outline.SubOutlines))
+		for i, sub := range outline.SubOutlines {
+			processed, err := WalkOutlineBottomUp(sub, fn)
+			if err != nil {
+				return Outline{}, err
+			}
+			newSubs[i] = processed
+		}
+		outline = Outline{
+			Type:        outline.Type,
+			Args:        outline.Args,
+			SubOutlines: newSubs,
+			ID:          outline.ID,
+		}
+	}
+	return fn(outline)
+}
+
+// WalkOutlinePreOrder performs a pre-order traversal of an Outline tree,
+// calling fn on each node before visiting its children. fn returns an error
+// to abort the traversal early.
+func WalkOutlinePreOrder(outline Outline, fn func(Outline) error) error {
+	if err := fn(outline); err != nil {
+		return err
+	}
+	for _, sub := range outline.SubOutlines {
+		if err := WalkOutlinePreOrder(sub, fn); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // CanonicalOutline is an Outline tree that has been fully canonicalized.
 // It pairs the root Outline (with every node's ID assigned) with a map from
 // those IDs to their CanonicalKeys. Only CanonicalOutlines can be Compiled,