Fix SubjectSetByType to properly *merge* mapped relations

internal/datasets/subjectset.go

@@ -36,6 +36,10 @@ func (ss SubjectSet) MustUnionWithSet(other SubjectSet) {
 	ss.BaseSubjectSet.MustUnionWithSet(other.BaseSubjectSet)
 }
 
+func (ss SubjectSet) Clone() SubjectSet {
+	return SubjectSet{ss.BaseSubjectSet.Clone()}
+}
+
 func (ss SubjectSet) UnionWithSet(other SubjectSet) error {
 	return ss.BaseSubjectSet.UnionWithSet(other.BaseSubjectSet)
 }

internal/datasets/subjectsetbytype.go

@@ -69,7 +69,17 @@ func (s *SubjectByTypeSet) Map(mapper func(rr *core.RelationReference) (*core.Re
 		if updatedType == nil {
 			continue
 		}
-		mapped.byType[tuple.JoinRelRef(updatedType.Namespace, updatedType.Relation)] = subjectset
+
+		key := tuple.JoinRelRef(updatedType.Namespace, updatedType.Relation)
+		if existing, ok := mapped.byType[key]; ok {
+			cloned := subjectset.Clone()
+			if err := cloned.UnionWithSet(existing); err != nil {
+				return nil, err
+			}
+			mapped.byType[key] = cloned
+		} else {
+			mapped.byType[key] = subjectset
+		}
 	}
 	return mapped, nil
 }

internal/datasets/subjectsetbytype_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 	"github.com/authzed/spicedb/pkg/tuple"
 )
@@ -104,3 +105,29 @@ func TestSubjectSetByTypeWithCaveats(t *testing.T) {
 		tom.GetCaveatExpression(),
 	)
 }
+
+func TestSubjectSetMapOverSameSubjectDifferentRelation(t *testing.T) {
+	set := NewSubjectByTypeSet()
+	require.True(t, set.IsEmpty())
+
+	err := set.AddSubjectOf(tuple.MustParse("document:foo#folder@folder:folder1"))
+	require.NoError(t, err)
+
+	err = set.AddSubjectOf(tuple.MustParse("document:foo#folder@folder:folder2#parent"))
+	require.NoError(t, err)
+
+	mapped, err := set.Map(func(rr *core.RelationReference) (*core.RelationReference, error) {
+		return &core.RelationReference{
+			Namespace: rr.Namespace,
+			Relation:  "shared",
+		}, nil
+	})
+	require.NoError(t, err)
+
+	foundSubjectIDs := mapz.NewSet[string]()
+	for _, sub := range mapped.byType["folder#shared"].AsSlice() {
+		foundSubjectIDs.Add(sub.SubjectId)
+	}
+
+	require.ElementsMatch(t, []string{"folder1", "folder2"}, foundSubjectIDs.AsSlice())
+}

internal/dispatch/graph/lookupsubjects_test.go

@@ -640,6 +640,73 @@ func TestCaveatedLookupSubjects(t *testing.T) {
 				},
 			},
 		},
+		{
+			"arrow over different relations of the same subject",
+			`definition user {}
+	
+			 definition folder {
+				relation parent: folder
+				relation viewer: user
+				permission view = viewer
+			 }
+
+		 	 definition document {
+				relation folder: folder | folder#parent
+				permission view = folder->view
+  		 }`,
+			[]*corev1.RelationTuple{
+				tuple.MustParse("folder:folder1#viewer@user:tom"),
+				tuple.MustParse("folder:folder2#viewer@user:fred"),
+				tuple.MustParse("document:somedoc#folder@folder:folder1"),
+				tuple.MustParse("document:somedoc#folder@folder:folder2#parent"),
+			},
+			ONR("document", "somedoc", "view"),
+			RR("user", "..."),
+			[]*v1.FoundSubject{
+				{
+					SubjectId: "tom",
+				},
+				{
+					SubjectId: "fred",
+				},
+			},
+		},
+		{
+			"caveated arrow over different relations of the same subject",
+			`definition user {}
+	
+			 caveat somecaveat(somecondition int) {
+				somecondition == 42
+			 }
+
+			 definition folder {
+				relation parent: folder
+				relation viewer: user
+				permission view = viewer
+			 }
+
+		 	 definition document {
+				relation folder: folder | folder#parent with somecaveat
+				permission view = folder->view
+  		 }`,
+			[]*corev1.RelationTuple{
+				tuple.MustParse("folder:folder1#viewer@user:tom"),
+				tuple.MustParse("folder:folder2#viewer@user:fred"),
+				tuple.MustParse("document:somedoc#folder@folder:folder1"),
+				tuple.MustWithCaveat(tuple.MustParse("document:somedoc#folder@folder:folder2#parent"), "somecaveat"),
+			},
+			ONR("document", "somedoc", "view"),
+			RR("user", "..."),
+			[]*v1.FoundSubject{
+				{
+					SubjectId: "tom",
+				},
+				{
+					SubjectId:        "fred",
+					CaveatExpression: caveatexpr("somecaveat"),
+				},
+			},
+		},
 	}
 
 	for _, tc := range testCases {

internal/services/integrationtesting/testconfigs/arrowoversametype.yaml

@@ -0,0 +1,28 @@
+---
+schema: |+
+  definition user {}
+
+  definition folder {
+    relation parent: folder
+
+    relation viewer: user
+    permission view = viewer
+  }
+
+  definition document {
+    relation folder: folder#parent | folder
+    permission view = folder->view
+  }
+
+relationships: >-
+  document:firstdoc#folder@folder:folder1
+
+  document:firstdoc#folder@folder:folder2#parent
+
+  folder:folder1#viewer@user:tom
+
+  folder:folder2#viewer@user:fred
+assertions:
+  assertTrue:
+    - "document:firstdoc#view@user:tom#..."
+    - "document:firstdoc#view@user:fred#..."