1 files changed, 248 insertions, 0 deletions
diff --git a/vendor/github.com/authzed/spicedb/internal/graph/resourcesubjectsmap2.go b/vendor/github.com/authzed/spicedb/internal/graph/resourcesubjectsmap2.go
new file mode 100644
index 0000000..4e41955
--- /dev/null
+++ b/vendor/github.com/authzed/spicedb/internal/graph/resourcesubjectsmap2.go
@@ -0,0 +1,248 @@
+package graph
+
+import (
+	"sort"
+	"sync"
+
+	"github.com/authzed/spicedb/pkg/genutil/mapz"
+	core "github.com/authzed/spicedb/pkg/proto/core/v1"
+	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
+	"github.com/authzed/spicedb/pkg/spiceerrors"
+	"github.com/authzed/spicedb/pkg/tuple"
+)
+
+type syncONRSet struct {
+	sync.Mutex
+	items map[string]struct{} // GUARDED_BY(Mutex)
+}
+
+func (s *syncONRSet) Add(onr *core.ObjectAndRelation) bool {
+	key := tuple.StringONR(tuple.FromCoreObjectAndRelation(onr))
+	s.Lock()
+	_, existed := s.items[key]
+	if !existed {
+		s.items[key] = struct{}{}
+	}
+	s.Unlock()
+	return !existed
+}
+
+func NewSyncONRSet() *syncONRSet {
+	return &syncONRSet{items: make(map[string]struct{})}
+}
+
+// resourcesSubjectMap2 is a multimap which tracks mappings from found resource IDs
+// to the subject IDs (may be more than one) for each, as well as whether the mapping
+// is conditional due to the use of a caveat on the relationship which formed the mapping.
+type resourcesSubjectMap2 struct {
+	resourceType         *core.RelationReference
+	resourcesAndSubjects *mapz.MultiMap[string, subjectInfo2]
+}
+
+// subjectInfo2 is the information about a subject contained in a resourcesSubjectMap2.
+type subjectInfo2 struct {
+	subjectID                string
+	missingContextParameters []string
+}
+
+func newResourcesSubjectMap2(resourceType *core.RelationReference) resourcesSubjectMap2 {
+	return resourcesSubjectMap2{
+		resourceType:         resourceType,
+		resourcesAndSubjects: mapz.NewMultiMap[string, subjectInfo2](),
+	}
+}
+
+func newResourcesSubjectMap2WithCapacity(resourceType *core.RelationReference, capacity uint32) resourcesSubjectMap2 {
+	return resourcesSubjectMap2{
+		resourceType:         resourceType,
+		resourcesAndSubjects: mapz.NewMultiMapWithCap[string, subjectInfo2](capacity),
+	}
+}
+
+func subjectIDsToResourcesMap2(resourceType *core.RelationReference, subjectIDs []string) resourcesSubjectMap2 {
+	rsm := newResourcesSubjectMap2(resourceType)
+	for _, subjectID := range subjectIDs {
+		rsm.addSubjectIDAsFoundResourceID(subjectID)
+	}
+	return rsm
+}
+
+// addRelationship adds the relationship to the resource subject map, recording a mapping from
+// the resource of the relationship to the subject, as well as whether the relationship was caveated.
+func (rsm resourcesSubjectMap2) addRelationship(rel tuple.Relationship, missingContextParameters []string) error {
+	spiceerrors.DebugAssert(func() bool {
+		return rel.Resource.ObjectType == rsm.resourceType.Namespace && rel.Resource.Relation == rsm.resourceType.Relation
+	}, "invalid relationship for addRelationship. expected: %v, found: %v", rsm.resourceType, rel.Resource)
+
+	spiceerrors.DebugAssert(func() bool {
+		return len(missingContextParameters) == 0 || rel.OptionalCaveat != nil
+	}, "missing context parameters must be empty if there is no caveat")
+
+	rsm.resourcesAndSubjects.Add(rel.Resource.ObjectID, subjectInfo2{rel.Subject.ObjectID, missingContextParameters})
+	return nil
+}
+
+// withAdditionalMissingContextForDispatchedResourceID adds additional missing context parameters
+// to the existing missing context parameters for the dispatched resource ID.
+func (rsm resourcesSubjectMap2) withAdditionalMissingContextForDispatchedResourceID(
+	resourceID string,
+	additionalMissingContext []string,
+) {
+	if len(additionalMissingContext) == 0 {
+		return
+	}
+
+	subjectInfo2s, _ := rsm.resourcesAndSubjects.Get(resourceID)
+	updatedInfos := make([]subjectInfo2, 0, len(subjectInfo2s))
+	for _, info := range subjectInfo2s {
+		info.missingContextParameters = append(info.missingContextParameters, additionalMissingContext...)
+		updatedInfos = append(updatedInfos, info)
+	}
+	rsm.resourcesAndSubjects.Set(resourceID, updatedInfos)
+}
+
+// addSubjectIDAsFoundResourceID adds a subject ID directly as a found subject for itself as the resource,
+// with no associated caveat.
+func (rsm resourcesSubjectMap2) addSubjectIDAsFoundResourceID(subjectID string) {
+	rsm.resourcesAndSubjects.Add(subjectID, subjectInfo2{subjectID, nil})
+}
+
+// asReadOnly returns a read-only dispatchableResourcesSubjectMap2 for dispatching for the
+// resources in this map (if any).
+func (rsm resourcesSubjectMap2) asReadOnly() dispatchableResourcesSubjectMap2 {
+	return dispatchableResourcesSubjectMap2{rsm}
+}
+
+func (rsm resourcesSubjectMap2) len() int {
+	return rsm.resourcesAndSubjects.Len()
+}
+
+// dispatchableResourcesSubjectMap2 is a read-only, frozen version of the resourcesSubjectMap2 that
+// can be used for mapping conditionals once calls have been dispatched. This is read-only due to
+// its use by concurrent callers.
+type dispatchableResourcesSubjectMap2 struct {
+	resourcesSubjectMap2
+}
+
+func (rsm dispatchableResourcesSubjectMap2) len() int {
+	return rsm.resourcesAndSubjects.Len()
+}
+
+func (rsm dispatchableResourcesSubjectMap2) isEmpty() bool {
+	return rsm.resourcesAndSubjects.IsEmpty()
+}
+
+func (rsm dispatchableResourcesSubjectMap2) resourceIDs() []string {
+	return rsm.resourcesAndSubjects.Keys()
+}
+
+// filterSubjectIDsToDispatch returns the set of subject IDs that have not yet been
+// dispatched, by adding them to the dispatched set.
+func (rsm dispatchableResourcesSubjectMap2) filterSubjectIDsToDispatch(dispatched *syncONRSet, dispatchSubjectType *core.RelationReference) []string {
+	resourceIDs := rsm.resourceIDs()
+	filtered := make([]string, 0, len(resourceIDs))
+	for _, resourceID := range resourceIDs {
+		if dispatched.Add(&core.ObjectAndRelation{
+			Namespace: dispatchSubjectType.Namespace,
+			ObjectId:  resourceID,
+			Relation:  dispatchSubjectType.Relation,
+		}) {
+			filtered = append(filtered, resourceID)
+		}
+	}
+
+	return filtered
+}
+
+// cloneAsMutable returns a mutable clone of this dispatchableResourcesSubjectMap2.
+func (rsm dispatchableResourcesSubjectMap2) cloneAsMutable() resourcesSubjectMap2 {
+	return resourcesSubjectMap2{
+		resourceType:         rsm.resourceType,
+		resourcesAndSubjects: rsm.resourcesAndSubjects.Clone(),
+	}
+}
+
+func (rsm dispatchableResourcesSubjectMap2) asPossibleResources() []*v1.PossibleResource {
+	resources := make([]*v1.PossibleResource, 0, rsm.resourcesAndSubjects.Len())
+
+	// Sort for stability.
+	sortedResourceIds := rsm.resourcesAndSubjects.Keys()
+	sort.Strings(sortedResourceIds)
+
+	for _, resourceID := range sortedResourceIds {
+		subjectInfo2s, _ := rsm.resourcesAndSubjects.Get(resourceID)
+		subjectIDs := make([]string, 0, len(subjectInfo2s))
+		allCaveated := true
+		nonCaveatedSubjectIDs := make([]string, 0, len(subjectInfo2s))
+		missingContextParameters := mapz.NewSet[string]()
+
+		for _, info := range subjectInfo2s {
+			subjectIDs = append(subjectIDs, info.subjectID)
+			if len(info.missingContextParameters) == 0 {
+				allCaveated = false
+				nonCaveatedSubjectIDs = append(nonCaveatedSubjectIDs, info.subjectID)
+			} else {
+				missingContextParameters.Extend(info.missingContextParameters)
+			}
+		}
+
+		// Sort for stability.
+		sort.Strings(subjectIDs)
+
+		// If all the incoming edges are caveated, then the entire status has to be marked as a check
+		// is required. Otherwise, if there is at least *one* non-caveated incoming edge, then we can
+		// return the existing status as a short-circuit for those non-caveated found subjects.
+		if allCaveated {
+			resources = append(resources, &v1.PossibleResource{
+				ResourceId:           resourceID,
+				ForSubjectIds:        subjectIDs,
+				MissingContextParams: missingContextParameters.AsSlice(),
+			})
+		} else {
+			resources = append(resources, &v1.PossibleResource{
+				ResourceId:    resourceID,
+				ForSubjectIds: nonCaveatedSubjectIDs,
+			})
+		}
+	}
+	return resources
+}
+
+func (rsm dispatchableResourcesSubjectMap2) mapPossibleResource(foundResource *v1.PossibleResource) (*v1.PossibleResource, error) {
+	forSubjectIDs := mapz.NewSet[string]()
+	nonCaveatedSubjectIDs := mapz.NewSet[string]()
+	missingContextParameters := mapz.NewSet[string]()
+
+	for _, forSubjectID := range foundResource.ForSubjectIds {
+		// Map from the incoming subject ID to the subject ID(s) that caused the dispatch.
+		infos, ok := rsm.resourcesAndSubjects.Get(forSubjectID)
+		if !ok {
+			return nil, spiceerrors.MustBugf("missing for subject ID")
+		}
+
+		for _, info := range infos {
+			forSubjectIDs.Insert(info.subjectID)
+			if len(info.missingContextParameters) == 0 {
+				nonCaveatedSubjectIDs.Insert(info.subjectID)
+			} else {
+				missingContextParameters.Extend(info.missingContextParameters)
+			}
+		}
+	}
+
+	// If there are some non-caveated IDs, return those and mark as the parent status.
+	if nonCaveatedSubjectIDs.Len() > 0 {
+		return &v1.PossibleResource{
+			ResourceId:    foundResource.ResourceId,
+			ForSubjectIds: nonCaveatedSubjectIDs.AsSlice(),
+		}, nil
+	}
+
+	// Otherwise, everything is caveated, so return the full set of subject IDs and mark
+	// as a check is required.
+	return &v1.PossibleResource{
+		ResourceId:           foundResource.ResourceId,
+		ForSubjectIds:        forSubjectIDs.AsSlice(),
+		MissingContextParams: missingContextParameters.AsSlice(),
+	}, nil
+}