feat(qp): prune branches that cannot lead to the subject type of the … (#2968)

Author: miparnisari

CHANGELOG.md

@@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 - Updated CI so that Postgres tests run against v18 which is GA and not against v13 which is EOL (https://github.com/authzed/spicedb/pull/2926)
 - Added tracing to request validation (https://github.com/authzed/spicedb/pull/2950)
+- Query Planner optimization: in Check requests, prune branches that cannot lead to the subject type specified (https://github.com/authzed/spicedb/pull/2968)
 
 ### Fixed
 - Regression introduced in 1.49.2: missing spans in ReadSchema calls (https://github.com/authzed/spicedb/pull/2947)

internal/services/integrationtesting/query_plan_consistency_test.go

@@ -135,7 +135,7 @@ func runQueryPlanAssertions(t *testing.T, handle *queryPlanConsistencyHandle) {
 
 									// Apply static optimizations if requested
 									if optimizationMode.optimize {
-										co, err = queryopt.ApplyOptimizations(co, queryopt.StandardOptimzations)
+										co, err = queryopt.ApplyOptimizations(co, queryopt.StandardOptimzations, queryopt.RequestParams{SubjectType: rel.Subject.ObjectType, SubjectRelation: rel.Subject.Relation})
 										require.NoError(err)
 									}
 

internal/services/v1/permissions_queryplan.go

@@ -58,7 +58,10 @@ func (ps *permissionServer) checkPermissionWithQueryPlan(ctx context.Context, re
 		return nil, ps.rewriteError(ctx, err)
 	}
 
-	optimized, err := queryopt.ApplyOptimizations(co, queryopt.StandardOptimzations)
+	optimized, err := queryopt.ApplyOptimizations(co, queryopt.StandardOptimzations, queryopt.RequestParams{
+		SubjectType:     req.Subject.Object.ObjectType,
+		SubjectRelation: req.Subject.OptionalRelation,
+	})
 	if err != nil {
 		return nil, ps.rewriteError(ctx, err)
 	}

pkg/query/mutations.go

@@ -18,6 +18,53 @@ func MutateOutline(outline Outline, fns []OutlineMutation) Outline {
 	return result
 }
 
+// NullPropagation is an OutlineMutation that propagates NullIteratorType nodes
+// upward through the tree according to each node type's semantics:
+//   - Union: null if ALL children are null
+//   - Intersection: null if ANY child is null
+//   - Arrow/IntersectionArrow: null if the right child is null
+//   - Exclusion: null if the left child is null
+//   - Caveat/Alias/Recursive: null if the only child is null
+//
+// This is intended to run after a mutation that nullifies leaf nodes (e.g.
+// reachability pruning), so that the nulls cascade correctly through the tree
+// during the bottom-up walk.
+func NullPropagation(outline Outline) Outline {
+	switch outline.Type {
+	case UnionIteratorType:
+		for _, sub := range outline.SubOutlines {
+			if sub.Type != NullIteratorType {
+				return outline
+			}
+		}
+		return Outline{Type: NullIteratorType, ID: outline.ID}
+
+	case IntersectionIteratorType:
+		for _, sub := range outline.SubOutlines {
+			if sub.Type == NullIteratorType {
+				return Outline{Type: NullIteratorType, ID: outline.ID}
+			}
+		}
+
+	case ArrowIteratorType, IntersectionArrowIteratorType:
+		if len(outline.SubOutlines) == 2 && outline.SubOutlines[1].Type == NullIteratorType {
+			return Outline{Type: NullIteratorType, ID: outline.ID}
+		}
+
+	case ExclusionIteratorType:
+		if len(outline.SubOutlines) == 2 && outline.SubOutlines[0].Type == NullIteratorType {
+			return Outline{Type: NullIteratorType, ID: outline.ID}
+		}
+
+	case CaveatIteratorType, AliasIteratorType, RecursiveIteratorType:
+		if len(outline.SubOutlines) == 1 && outline.SubOutlines[0].Type == NullIteratorType {
+			return Outline{Type: NullIteratorType, ID: outline.ID}
+		}
+	}
+
+	return outline
+}
+
 // ReorderMutation returns an OutlineMutation that reorders SubOutlines according
 // to order, where order[i] is the index of the child to place at position i.
 // If order has a different length than the node's SubOutlines, it is a no-op.

pkg/query/queryopt/caveat_pushdown.go

@@ -9,7 +9,11 @@ func init() {
 		Pushes caveat evalution to the lowest point in the tree.
 		Cannot push through intersection arrows
 		`,
-		Mutation: caveatPushdown,
+		NewTransform: func(_ RequestParams) OutlineTransform {
+			return func(outline query.Outline) query.Outline {
+				return query.MutateOutline(outline, []query.OutlineMutation{caveatPushdown})
+			}
+		},
 	})
 }
 

pkg/query/queryopt/reachability_pruning.go

@@ -0,0 +1,91 @@
+package queryopt
+
+import (
+	"github.com/authzed/spicedb/pkg/query"
+)
+
+func init() {
+	MustRegisterOptimization(Optimizer{
+		Name: "reachability-pruning",
+		Description: `
+		Replaces subtrees with NullIteratorType nodes when they can never
+		produce the target subject type of the request.
+		`,
+		NewTransform: func(params RequestParams) OutlineTransform {
+			return reachabilityPruning(params)
+		},
+	})
+}
+
+func reachabilityPruning(params RequestParams) func(outline query.Outline) query.Outline {
+	return func(outline query.Outline) query.Outline {
+		if params.SubjectType == "" || (params.SubjectRelation != "" && params.SubjectRelation != "...") {
+			// do not mutate if subjectType is empty or if subjectRelation is non-empty
+			return outline
+		}
+		// Pre-pass: find all node IDs inside arrow left subtrees.
+		// These are intermediate hops whose subject type does not need
+		// to match the target, so they must be excluded from pruning.
+		immune := collectArrowLeftSubtreeIDs(outline)
+		return query.MutateOutline(outline, []query.OutlineMutation{
+			leafSubjectTypePruner(params.SubjectType, immune),
+			query.NullPropagation,
+		})
+	}
+}
+
+// leafSubjectTypePruner returns an OutlineMutation that replaces leaf outline
+// nodes (DatastoreIteratorType, SelfIteratorType) with NullIteratorType when
+// their subject type does not match the target. Nodes whose IDs appear in the
+// immune set are skipped — these are nodes inside arrow left subtrees whose
+// subject types are intermediate hops, not final outputs.
+//
+// Null propagation through compound nodes (unions, intersections, arrows, etc.)
+// is handled by query.NullPropagation, which should run after this mutation in
+// the same MutateOutline call.
+func leafSubjectTypePruner(targetSubjectType string, immune map[query.OutlineNodeID]bool) query.OutlineMutation {
+	return func(outline query.Outline) query.Outline {
+		if immune[outline.ID] {
+			return outline
+		}
+
+		switch outline.Type {
+		case query.DatastoreIteratorType:
+			if outline.Args != nil && outline.Args.Relation != nil && outline.Args.Relation.Type() == targetSubjectType {
+				return outline
+			}
+			return query.Outline{Type: query.NullIteratorType, ID: outline.ID}
+
+		case query.SelfIteratorType:
+			if outline.Args != nil && outline.Args.DefinitionName == targetSubjectType {
+				return outline
+			}
+			return query.Outline{Type: query.NullIteratorType, ID: outline.ID}
+
+		default:
+			return outline
+		}
+	}
+}
+
+// collectArrowLeftSubtreeIDs walks the outline tree and collects the IDs of all
+// nodes that appear inside the left subtree of an arrow (or intersection arrow).
+// These nodes have intermediate subject types that should not be pruned.
+func collectArrowLeftSubtreeIDs(outline query.Outline) map[query.OutlineNodeID]bool {
+	ids := make(map[query.OutlineNodeID]bool)
+	_ = query.WalkOutlinePreOrder(outline, func(node query.Outline) error {
+		if (node.Type == query.ArrowIteratorType || node.Type == query.IntersectionArrowIteratorType) && len(node.SubOutlines) == 2 {
+			markAllIDs(node.SubOutlines[0], ids)
+		}
+		return nil
+	})
+	return ids
+}
+
+// markAllIDs recursively adds the ID of every node in the subtree to the set.
+func markAllIDs(outline query.Outline, ids map[query.OutlineNodeID]bool) {
+	ids[outline.ID] = true
+	for _, sub := range outline.SubOutlines {
+		markAllIDs(sub, ids)
+	}
+}