fix: deduplicate through alias as well as union, and refactor those functions

Author: barakmich

pkg/query/alias.go

@@ -39,7 +39,8 @@ func (a *Alias) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRel
 				return nil, err
 			}
 
-			return func(yield func(Path, error) bool) {
+			// Create combined sequence with self-edge and rewritten paths
+			combined := func(yield func(Path, error) bool) {
 				// Yield the self-edge first
 				if !yield(selfPath, nil) {
 					return
@@ -57,7 +58,10 @@ func (a *Alias) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRel
 						return
 					}
 				}
-			}, nil
+			}
+
+			// Wrap with deduplication to handle duplicate paths after rewriting
+			return DeduplicatePathSeq(combined), nil
 		}
 	}
 
@@ -67,7 +71,7 @@ func (a *Alias) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRel
 		return nil, err
 	}
 
-	return func(yield func(Path, error) bool) {
+	rewritten := func(yield func(Path, error) bool) {
 		for path, err := range subSeq {
 			if err != nil {
 				yield(Path{}, err)
@@ -79,7 +83,10 @@ func (a *Alias) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRel
 				return
 			}
 		}
-	}, nil
+	}
+
+	// Wrap with deduplication to handle duplicate paths after rewriting
+	return DeduplicatePathSeq(rewritten), nil
 }
 
 func (a *Alias) IterSubjectsImpl(ctx *Context, resource Object) (PathSeq, error) {

pkg/query/path.go

@@ -46,6 +46,12 @@ func (p Path) Key() string {
 	return fmt.Sprintf("%s#%s@%s", p.Resource.Key(), p.Relation, ObjectAndRelationKey(p.Subject))
 }
 
+// EndpointsKey returns a unique string key for this Path based on its resource and subject only,
+// excluding the relation. This matches the semantics of EqualsEndpoints.
+func (p Path) EndpointsKey() string {
+	return fmt.Sprintf("%s@%s", p.Resource.Key(), ObjectAndRelationKey(p.Subject))
+}
+
 // MergeOr combines the paths, ORing the caveats and expiration and metadata together.
 // Returns a new Path with the merged values.
 func (p Path) MergeOr(other Path) (Path, error) {
@@ -300,3 +306,39 @@ func CollectAll(seq PathSeq) ([]Path, error) {
 	}
 	return out, nil
 }
+
+// DeduplicatePathSeq returns a new PathSeq that deduplicates paths based on their
+// endpoints (resource and subject, excluding relation). Paths with the same endpoints
+// are merged using OR semantics (caveats are OR'd, no caveat wins over caveat).
+// This collects all paths first, deduplicates with merging, then yields results.
+func DeduplicatePathSeq(seq PathSeq) PathSeq {
+	return func(yield func(Path, error) bool) {
+		seen := make(map[string]Path)
+		for path, err := range seq {
+			if err != nil {
+				yield(Path{}, err)
+				return
+			}
+
+			key := path.EndpointsKey()
+			if existing, exists := seen[key]; !exists {
+				seen[key] = path
+			} else {
+				// Merge with existing path using OR semantics
+				merged, err := existing.MergeOr(path)
+				if err != nil {
+					yield(Path{}, err)
+					return
+				}
+				seen[key] = merged
+			}
+		}
+
+		// Yield all deduplicated paths
+		for _, path := range seen {
+			if !yield(path, nil) {
+				return
+			}
+		}
+	}
+}

pkg/query/union.go

@@ -21,62 +21,37 @@ func (u *Union) addSubIterator(subIt Iterator) {
 }
 
 func (u *Union) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRelation) (PathSeq, error) {
-	var out []Path
-	// Collect paths from all sub-iterators
 	ctx.TraceStep(u, "processing %d sub-iterators with %d resources", len(u.subIts), len(resources))
 
-	for iterIdx, it := range u.subIts {
-		ctx.TraceStep(u, "processing sub-iterator %d", iterIdx)
+	// Create a concatenated sequence from all sub-iterators
+	combinedSeq := func(yield func(Path, error) bool) {
+		for iterIdx, it := range u.subIts {
+			ctx.TraceStep(u, "processing sub-iterator %d", iterIdx)
 
-		pathSeq, err := ctx.Check(it, resources, subject)
-		if err != nil {
-			return nil, err
-		}
-		paths, err := CollectAll(pathSeq)
-		if err != nil {
-			return nil, err
-		}
-
-		ctx.TraceStep(u, "sub-iterator %d returned %d paths", iterIdx, len(paths))
-		out = append(out, paths...)
-	}
-
-	ctx.TraceStep(u, "collected %d total paths before deduplication", len(out))
-
-	// Deduplicate paths based on resource for CheckImpl
-	// Since the subject is fixed in CheckImpl, we only need to deduplicate by resource
-	seen := make(map[string]Path)
-	for _, path := range out {
-		// Use resource object (type + id) as key for deduplication, not the full resource with relation
-		key := path.Resource.Key()
-		if existing, exists := seen[key]; !exists {
-			seen[key] = path
-		} else {
-			// If we already have a path for this resource,
-			// merge it with the new one using OR semantics
-			merged, err := existing.MergeOr(path)
+			pathSeq, err := ctx.Check(it, resources, subject)
 			if err != nil {
-				return nil, err
+				yield(Path{}, err)
+				return
 			}
-			seen[key] = merged
-		}
-	}
 
-	// Convert map to slice
-	deduplicatedSlice := make([]Path, 0, len(seen))
-	for _, path := range seen {
-		deduplicatedSlice = append(deduplicatedSlice, path)
-	}
-
-	ctx.TraceStep(u, "deduplicated to %d paths", len(deduplicatedSlice))
-
-	return func(yield func(Path, error) bool) {
-		for _, path := range deduplicatedSlice {
-			if !yield(path, nil) {
-				return
+			pathCount := 0
+			for path, err := range pathSeq {
+				if err != nil {
+					yield(Path{}, err)
+					return
+				}
+				pathCount++
+				if !yield(path, nil) {
+					return
+				}
 			}
+
+			ctx.TraceStep(u, "sub-iterator %d returned %d paths", iterIdx, pathCount)
 		}
-	}, nil
+	}
+
+	// Wrap with deduplication
+	return DeduplicatePathSeq(combinedSeq), nil
 }
 
 func (u *Union) IterSubjectsImpl(ctx *Context, resource Object) (PathSeq, error) {