github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/pkg/tuple/onrbytypeset.go (about)

     1  package tuple
     2  
     3  import (
     4  	"github.com/samber/lo"
     5  
     6  	core "github.com/authzed/spicedb/pkg/proto/core/v1"
     7  )
     8  
     9  // ONRByTypeSet is a set of ObjectAndRelation's, grouped by namespace+relation.
    10  type ONRByTypeSet struct {
    11  	byType map[string][]string
    12  }
    13  
    14  // NewONRByTypeSet creates and returns a new ONRByTypeSet.
    15  func NewONRByTypeSet() *ONRByTypeSet {
    16  	return &ONRByTypeSet{
    17  		byType: map[string][]string{},
    18  	}
    19  }
    20  
    21  // Add adds the specified ObjectAndRelation to the set.
    22  func (s *ONRByTypeSet) Add(onr *core.ObjectAndRelation) {
    23  	key := JoinRelRef(onr.Namespace, onr.Relation)
    24  	if _, ok := s.byType[key]; !ok {
    25  		s.byType[key] = []string{}
    26  	}
    27  
    28  	s.byType[key] = append(s.byType[key], onr.ObjectId)
    29  }
    30  
    31  // ForEachType invokes the handler for each type of ObjectAndRelation found in the set, along
    32  // with all IDs of objects of that type.
    33  func (s *ONRByTypeSet) ForEachType(handler func(rr *core.RelationReference, objectIds []string)) {
    34  	for key, objectIds := range s.byType {
    35  		ns, rel := MustSplitRelRef(key)
    36  		handler(&core.RelationReference{
    37  			Namespace: ns,
    38  			Relation:  rel,
    39  		}, lo.Uniq(objectIds))
    40  	}
    41  }
    42  
    43  // Map runs the mapper function over each type of object in the set, returning a new ONRByTypeSet with
    44  // the object type replaced by that returned by the mapper function.
    45  func (s *ONRByTypeSet) Map(mapper func(rr *core.RelationReference) (*core.RelationReference, error)) (*ONRByTypeSet, error) {
    46  	mapped := NewONRByTypeSet()
    47  	for key, objectIds := range s.byType {
    48  		ns, rel := MustSplitRelRef(key)
    49  		updatedType, err := mapper(&core.RelationReference{
    50  			Namespace: ns,
    51  			Relation:  rel,
    52  		})
    53  		if err != nil {
    54  			return nil, err
    55  		}
    56  		if updatedType == nil {
    57  			continue
    58  		}
    59  		mapped.byType[JoinRelRef(updatedType.Namespace, updatedType.Relation)] = lo.Uniq(objectIds)
    60  	}
    61  	return mapped, nil
    62  }
    63  
    64  // IsEmpty returns true if the set is empty.
    65  func (s *ONRByTypeSet) IsEmpty() bool {
    66  	return len(s.byType) == 0
    67  }
    68  
    69  // KeyLen returns the number of keys in the set.
    70  func (s *ONRByTypeSet) KeyLen() int {
    71  	return len(s.byType)
    72  }
    73  
    74  // ValueLen returns the number of values in the set.
    75  func (s *ONRByTypeSet) ValueLen() int {
    76  	var total int
    77  	for _, vals := range s.byType {
    78  		total += len(vals)
    79  	}
    80  
    81  	return total
    82  }