Merge pull request #1702 from josephschorr/set-perf

Small micro-optimizations around set performance

internal/datastore/proxy/schemacaching/standardcache.go

@@ -125,7 +125,7 @@ func listAndCache[T schemaDefinition](
 			continue
 		}
 
-		remainingToLoad.Remove(name)
+		remainingToLoad.Delete(name)
 		loaded := loadedRaw.(*cacheEntry)
 		foundDefs = append(foundDefs, datastore.RevisionedDefinition[T]{
 			Definition:          loaded.definition.(T),

internal/datastore/proxy/schemacaching/watchingcache.go

@@ -567,7 +567,7 @@ func (swc *schemaWatchCache[T]) readDefinitionsWithNames(ctx context.Context, na
 		}
 
 		swc.definitionsReadCachedCounter.WithLabelValues(swc.kind).Inc()
-		remainingNames.Remove(name)
+		remainingNames.Delete(name)
 		if !found.wasNotFound {
 			foundDefs = append(foundDefs, found.revisionedDefinition)
 		}

internal/dispatch/graph/check_test.go

@@ -292,7 +292,7 @@ func TestCheckMetadata(t *testing.T) {
 }
 
 func addFrame(trace *v1.CheckDebugTrace, foundFrames *mapz.Set[string]) {
-	foundFrames.Add(fmt.Sprintf("%s:%s#%s", trace.Request.ResourceRelation.Namespace, strings.Join(trace.Request.ResourceIds, ","), trace.Request.ResourceRelation.Relation))
+	foundFrames.Insert(fmt.Sprintf("%s:%s#%s", trace.Request.ResourceRelation.Namespace, strings.Join(trace.Request.ResourceIds, ","), trace.Request.ResourceRelation.Relation))
 	for _, subTrace := range trace.SubProblems {
 		addFrame(subTrace, foundFrames)
 	}

internal/dispatch/graph/lookupresources_test.go

@@ -212,7 +212,7 @@ func TestSimpleLookupResourcesWithCursor(t *testing.T) {
 
 			require.Equal(1, len(stream.Results()))
 
-			found.Add(stream.Results()[0].ResolvedResource.ResourceId)
+			found.Insert(stream.Results()[0].ResolvedResource.ResourceId)
 			require.Equal(tc.expectedFirst, found.AsSlice())
 
 			cursor := stream.Results()[0].AfterResponseCursor
@@ -233,7 +233,7 @@ func TestSimpleLookupResourcesWithCursor(t *testing.T) {
 			require.NoError(err)
 
 			for _, result := range stream.Results() {
-				found.Add(result.ResolvedResource.ResourceId)
+				found.Insert(result.ResolvedResource.ResourceId)
 			}
 
 			foundResults := found.AsSlice()
@@ -585,7 +585,7 @@ func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 						}
 
 						for _, result := range stream.Results() {
-							foundResourceIDs.Add(result.ResolvedResource.ResourceId)
+							foundResourceIDs.Insert(result.ResolvedResource.ResourceId)
 							currentCursor = result.AfterResponseCursor
 						}
 

internal/dispatch/graph/reachableresources_test.go

@@ -827,7 +827,7 @@ func TestReachableResourcesCursors(t *testing.T) {
 			count := 0
 			for _, result := range stream2.Results() {
 				count++
-				foundResources.Add(result.Resource.ResourceId)
+				foundResources.Insert(result.Resource.ResourceId)
 			}
 			require.LessOrEqual(t, count, 310)
 
@@ -1262,7 +1262,7 @@ func TestReachableResourcesOverSchema(t *testing.T) {
 						}
 
 						for _, result := range stream.Results() {
-							foundResourceIDs.Add(result.Resource.ResourceId)
+							foundResourceIDs.Insert(result.Resource.ResourceId)
 							currentCursor = result.AfterResponseCursor
 						}
 
@@ -1433,12 +1433,12 @@ func TestReachableResourcesWithCachingInParallelTest(t *testing.T) {
 
 	for i := 0; i < 410; i++ {
 		if i < 250 {
-			expectedResources.Add(fmt.Sprintf("res%03d", i))
+			expectedResources.Insert(fmt.Sprintf("res%03d", i))
 			testRels = append(testRels, tuple.MustParse(fmt.Sprintf("resource:res%03d#viewer@user:tom", i)))
 		}
 
 		if i > 200 {
-			expectedResources.Add(fmt.Sprintf("res%03d", i))
+			expectedResources.Insert(fmt.Sprintf("res%03d", i))
 			testRels = append(testRels, tuple.MustParse(fmt.Sprintf("resource:res%03d#editor@user:tom", i)))
 		}
 	}
@@ -1488,7 +1488,7 @@ func TestReachableResourcesWithCachingInParallelTest(t *testing.T) {
 
 			foundResources := mapz.NewSet[string]()
 			for _, result := range stream.Results() {
-				foundResources.Add(result.Resource.ResourceId)
+				foundResources.Insert(result.Resource.ResourceId)
 			}
 
 			expectedResourcesSlice := expectedResources.AsSlice()

internal/graph/resourcesubjectsmap.go

@@ -196,9 +196,9 @@ func (rsm dispatchableResourcesSubjectMap) mapFoundResource(foundResource *v1.Re
 		}
 
 		for _, info := range infos {
-			forSubjectIDs.Add(info.subjectID)
+			forSubjectIDs.Insert(info.subjectID)
 			if !info.isCaveated {
-				nonCaveatedSubjectIDs.Add(info.subjectID)
+				nonCaveatedSubjectIDs.Insert(info.subjectID)
 			}
 		}
 	}

internal/namespace/util.go

@@ -55,7 +55,7 @@ type TypeAndRelationToCheck struct {
 func CheckNamespaceAndRelations(ctx context.Context, checks []TypeAndRelationToCheck, ds datastore.Reader) error {
 	nsNames := mapz.NewSet[string]()
 	for _, toCheck := range checks {
-		nsNames.Add(toCheck.NamespaceName)
+		nsNames.Insert(toCheck.NamespaceName)
 	}
 
 	if nsNames.IsEmpty() {
@@ -150,12 +150,12 @@ func ReadNamespaceAndTypes(
 func ListReferencedNamespaces(nsdefs []*core.NamespaceDefinition) []string {
 	referencedNamespaceNamesSet := mapz.NewSet[string]()
 	for _, nsdef := range nsdefs {
-		referencedNamespaceNamesSet.Add(nsdef.Name)
+		referencedNamespaceNamesSet.Insert(nsdef.Name)
 
 		for _, relation := range nsdef.Relation {
 			if relation.GetTypeInformation() != nil {
 				for _, allowedRel := range relation.GetTypeInformation().AllowedDirectRelations {
-					referencedNamespaceNamesSet.Add(allowedRel.GetNamespace())
+					referencedNamespaceNamesSet.Insert(allowedRel.GetNamespace())
 				}
 			}
 		}

internal/relationships/validation.go

@@ -87,10 +87,10 @@ func loadNamespacesAndCaveats(ctx context.Context, rels []*core.RelationTuple, r
 	referencedNamespaceNames := mapz.NewSet[string]()
 	referencedCaveatNamesWithContext := mapz.NewSet[string]()
 	for _, rel := range rels {
-		referencedNamespaceNames.Add(rel.ResourceAndRelation.Namespace)
-		referencedNamespaceNames.Add(rel.Subject.Namespace)
+		referencedNamespaceNames.Insert(rel.ResourceAndRelation.Namespace)
+		referencedNamespaceNames.Insert(rel.Subject.Namespace)
 		if hasNonEmptyCaveatContext(rel) {
-			referencedCaveatNamesWithContext.Add(rel.Caveat.CaveatName)
+			referencedCaveatNamesWithContext.Insert(rel.Caveat.CaveatName)
 		}
 	}
 

internal/services/integrationtesting/consistency_test.go

@@ -253,7 +253,7 @@ func validateRelationshipReads(t *testing.T, vctx validationContext) {
 
 		foundRelationshipsSet := mapz.NewSet[string]()
 		for _, rel := range foundRelationships {
-			foundRelationshipsSet.Add(tuple.MustString(rel))
+			foundRelationshipsSet.Insert(tuple.MustString(rel))
 		}
 
 		require.True(t, foundRelationshipsSet.Has(tuple.MustString(relationship)), "missing expected relationship %s in read results: %s", tuple.MustString(relationship), foundRelationshipsSet.AsSlice())

internal/services/shared/schema.go

@@ -34,7 +34,7 @@ func ValidateSchemaChanges(ctx context.Context, compiled *compiler.CompiledSchem
 			return nil, err
 		}
 
-		newCaveatDefNames.Add(caveatDef.Name)
+		newCaveatDefNames.Insert(caveatDef.Name)
 	}
 
 	// 2) Validate the namespaces defined.
@@ -58,7 +58,7 @@ func ValidateSchemaChanges(ctx context.Context, compiled *compiler.CompiledSchem
 			return nil, err
 		}
 
-		newObjectDefNames.Add(nsdef.Name)
+		newObjectDefNames.Insert(nsdef.Name)
 	}
 
 	return &ValidatedSchemaChanges{
@@ -119,7 +119,7 @@ func ApplySchemaChangesOverExisting(
 
 	for _, existingCaveat := range existingCaveats {
 		existingCaveatDefMap[existingCaveat.Name] = existingCaveat
-		existingCaveatDefNames.Add(existingCaveat.Name)
+		existingCaveatDefNames.Insert(existingCaveat.Name)
 	}
 
 	// For each caveat definition, perform a diff and ensure the changes will not result in type errors.
@@ -142,7 +142,7 @@ func ApplySchemaChangesOverExisting(
 	existingObjectDefNames := mapz.NewSet[string]()
 	for _, existingDef := range existingObjectDefs {
 		existingObjectDefMap[existingDef.Name] = existingDef
-		existingObjectDefNames.Add(existingDef.Name)
+		existingObjectDefNames.Insert(existingDef.Name)
 	}
 
 	// For each definition, perform a diff and ensure the changes will not result in any

internal/services/v1/debug_test.go

@@ -49,7 +49,7 @@ func expectDebugFrames(permissionNames ...string) rda {
 		for _, sp := range debugInfo.Check.GetSubProblems().Traces {
 			for _, permissionName := range permissionNames {
 				if sp.Permission == permissionName {
-					found.Add(permissionName)
+					found.Insert(permissionName)
 				}
 			}
 		}

internal/testutil/subjects.go

@@ -324,7 +324,7 @@ func combinatorialValues(names []string) []map[string]bool {
 // collectReferencedNames collects all referenced caveat names into the given set.
 func collectReferencedNames(expr *core.CaveatExpression, nameSet *mapz.Set[string]) {
 	if expr.GetCaveat() != nil {
-		nameSet.Add(expr.GetCaveat().CaveatName)
+		nameSet.Insert(expr.GetCaveat().CaveatName)
 		return
 	}
 

pkg/cmd/server/middleware.go

@@ -97,7 +97,7 @@ const (
 func (mc *MiddlewareChain[T]) Names() *mapz.Set[string] {
 	names := mapz.NewSet[string]()
 	for _, mw := range mc.chain {
-		names.Add(mw.Name)
+		names.Insert(mw.Name)
 	}
 	return names
 }

pkg/datastore/test/watch.go

@@ -580,17 +580,17 @@ func verifyMixedUpdates(
 
 			foundChanges := mapz.NewSet[string]()
 			for _, changedDef := range change.ChangedDefinitions {
-				foundChanges.Add("changed:" + changedDef.GetName())
+				foundChanges.Insert("changed:" + changedDef.GetName())
 			}
 			for _, deleted := range change.DeletedNamespaces {
-				foundChanges.Add("deleted-ns:" + deleted)
+				foundChanges.Insert("deleted-ns:" + deleted)
 			}
 			for _, deleted := range change.DeletedCaveats {
-				foundChanges.Add("deleted-caveat:" + deleted)
+				foundChanges.Insert("deleted-caveat:" + deleted)
 			}
 
 			for _, update := range change.RelationshipChanges {
-				foundChanges.Add("rel:" + fmt.Sprintf("OPERATION_%s(%s)", update.Operation, tuple.StringWithoutCaveat(update.Tuple)))
+				foundChanges.Insert("rel:" + fmt.Sprintf("OPERATION_%s(%s)", update.Operation, tuple.StringWithoutCaveat(update.Tuple)))
 			}
 
 			found := foundChanges.AsSlice()

pkg/diff/namespace/diff.go

@@ -2,7 +2,6 @@ package namespace
 
 import (
 	"github.com/google/go-cmp/cmp"
-	"github.com/scylladb/go-set/strset"
 	"golang.org/x/exp/slices"
 	"google.golang.org/protobuf/testing/protocmp"
 
@@ -132,16 +131,16 @@ func DiffNamespaces(existing *core.NamespaceDefinition, updated *core.NamespaceD
 
 	// Collect up relations and check.
 	existingRels := map[string]*core.Relation{}
-	existingRelNames := strset.New()
+	existingRelNames := mapz.NewSet[string]()
 
 	existingPerms := map[string]*core.Relation{}
-	existingPermNames := strset.New()
+	existingPermNames := mapz.NewSet[string]()
 
 	updatedRels := map[string]*core.Relation{}
-	updatedRelNames := strset.New()
+	updatedRelNames := mapz.NewSet[string]()
 
 	updatedPerms := map[string]*core.Relation{}
-	updatedPermNames := strset.New()
+	updatedPermNames := mapz.NewSet[string]()
 
 	for _, relation := range existing.Relation {
 		_, ok := existingRels[relation.Name]
@@ -173,35 +172,39 @@ func DiffNamespaces(existing *core.NamespaceDefinition, updated *core.NamespaceD
 		}
 	}
 
-	for _, removed := range strset.Difference(existingRelNames, updatedRelNames).List() {
+	_ = existingRelNames.Subtract(updatedRelNames).ForEach(func(removed string) error {
 		deltas = append(deltas, Delta{
 			Type:         RemovedRelation,
 			RelationName: removed,
 		})
-	}
+		return nil
+	})
 
-	for _, added := range strset.Difference(updatedRelNames, existingRelNames).List() {
+	_ = updatedRelNames.Subtract(existingRelNames).ForEach(func(added string) error {
 		deltas = append(deltas, Delta{
 			Type:         AddedRelation,
 			RelationName: added,
 		})
-	}
+		return nil
+	})
 
-	for _, removed := range strset.Difference(existingPermNames, updatedPermNames).List() {
+	_ = existingPermNames.Subtract(updatedPermNames).ForEach(func(removed string) error {
 		deltas = append(deltas, Delta{
 			Type:         RemovedPermission,
 			RelationName: removed,
 		})
-	}
+		return nil
+	})
 
-	for _, added := range strset.Difference(updatedPermNames, existingPermNames).List() {
+	_ = updatedPermNames.Subtract(existingPermNames).ForEach(func(added string) error {
 		deltas = append(deltas, Delta{
 			Type:         AddedPermission,
 			RelationName: added,
 		})
-	}
+		return nil
+	})
 
-	for _, shared := range strset.Intersection(existingPermNames, updatedPermNames).List() {
+	_ = existingPermNames.Intersect(updatedPermNames).ForEach(func(shared string) error {
 		existingPerm := existingPerms[shared]
 		updatedPerm := updatedPerms[shared]
 
@@ -222,9 +225,10 @@ func DiffNamespaces(existing *core.NamespaceDefinition, updated *core.NamespaceD
 				RelationName: shared,
 			})
 		}
-	}
+		return nil
+	})
 
-	for _, shared := range strset.Intersection(existingRelNames, updatedRelNames).List() {
+	_ = existingRelNames.Intersect(updatedRelNames).ForEach(func(shared string) error {
 		existingRel := existingRels[shared]
 		updatedRel := updatedRels[shared]
 
@@ -273,22 +277,26 @@ func DiffNamespaces(existing *core.NamespaceDefinition, updated *core.NamespaceD
 			updatedAllowedRels.Add(source)
 		}
 
-		for _, removed := range existingAllowedRels.Subtract(updatedAllowedRels).AsSlice() {
+		_ = existingAllowedRels.Subtract(updatedAllowedRels).ForEach(func(removed string) error {
 			deltas = append(deltas, Delta{
 				Type:         RelationAllowedTypeRemoved,
 				RelationName: shared,
 				AllowedType:  allowedRelsBySource[removed],
 			})
-		}
+			return nil
+		})
 
-		for _, added := range updatedAllowedRels.Subtract(existingAllowedRels).AsSlice() {
+		_ = updatedAllowedRels.Subtract(existingAllowedRels).ForEach(func(added string) error {
 			deltas = append(deltas, Delta{
 				Type:         RelationAllowedTypeAdded,
 				RelationName: shared,
 				AllowedType:  allowedRelsBySource[added],
 			})
-		}
-	}
+			return nil
+		})
+
+		return nil
+	})
 
 	return &Diff{
 		existing: existing,

pkg/genutil/mapz/countingmap.go

@@ -43,7 +43,7 @@ func (cmm *CountingMultiMap[T, Q]) Remove(key T, value Q) {
 		return
 	}
 
-	values.Remove(value)
+	values.Delete(value)
 	if values.IsEmpty() {
 		delete(cmm.valuesByKey, key)
 	}