Create MapSet[T] and ImmutableMapSet[T] types and change Entity.Parents to be a ImmutableMapSet[EntityUID]

internal/eval/evalers.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/cedar-policy/cedar-go/internal/consts"
 	"github.com/cedar-policy/cedar-go/internal/extensions"
+	"github.com/cedar-policy/cedar-go/internal/mapset"
 	"github.com/cedar-policy/cedar-go/types"
 )
 
@@ -973,11 +974,6 @@ func newInEval(lhs, rhs Evaler) Evaler {
 	return &inEval{lhs: lhs, rhs: rhs}
 }
 
-func hasKnown(known map[types.EntityUID]struct{}, k types.EntityUID) bool {
-	_, ok := known[k]
-	return ok
-}
-
 func entityInOne(env *Env, entity types.EntityUID, parent types.EntityUID) bool {
 	key := inKey{a: entity, b: parent}
 	if cached, ok := env.inCache[key]; ok {
@@ -987,31 +983,28 @@ func entityInOne(env *Env, entity types.EntityUID, parent types.EntityUID) bool
 	env.inCache[key] = result
 	return result
 }
+
 func entityInOneWork(env *Env, entity types.EntityUID, parent types.EntityUID) bool {
 	if entity == parent {
 		return true
 	}
-	var known map[types.EntityUID]struct{}
+	var known mapset.MapSet[types.EntityUID]
 	var todo []types.EntityUID
 	var candidate = entity
 	for {
 		if fe, ok := env.Entities[candidate]; ok {
-			for _, k := range fe.Parents {
-				if k == parent {
-					return true
-				}
+			if fe.Parents.Contains(parent) {
+				return true
 			}
-			for _, k := range fe.Parents {
+			fe.Parents.Iterate(func(k types.EntityUID) bool {
 				p, ok := env.Entities[k]
-				if !ok || len(p.Parents) == 0 || k == entity || hasKnown(known, k) {
-					continue
+				if !ok || p.Parents.Len() == 0 || k == entity || known.Contains(k) {
+					return true
 				}
 				todo = append(todo, k)
-				if known == nil {
-					known = map[types.EntityUID]struct{}{}
-				}
-				known[k] = struct{}{}
-			}
+				known.Add(k)
+				return true
+			})
 		}
 		if len(todo) == 0 {
 			return false
@@ -1020,31 +1013,27 @@ func entityInOneWork(env *Env, entity types.EntityUID, parent types.EntityUID) b
 	}
 }
 
-func entityInSet(env *Env, entity types.EntityUID, parents map[types.EntityUID]struct{}) bool {
-	if _, ok := parents[entity]; ok {
+func entityInSet(env *Env, entity types.EntityUID, parents mapset.Container[types.EntityUID]) bool {
+	if parents.Contains(entity) {
 		return true
 	}
-	var known map[types.EntityUID]struct{}
+	var known mapset.MapSet[types.EntityUID]
 	var todo []types.EntityUID
 	var candidate = entity
 	for {
 		if fe, ok := env.Entities[candidate]; ok {
-			for _, k := range fe.Parents {
-				if _, ok := parents[k]; ok {
-					return true
-				}
+			if fe.Parents.Intersects(parents) {
+				return true
 			}
-			for _, k := range fe.Parents {
+			fe.Parents.Iterate(func(k types.EntityUID) bool {
 				p, ok := env.Entities[k]
-				if !ok || len(p.Parents) == 0 || k == entity || hasKnown(known, k) {
-					continue
+				if !ok || p.Parents.Len() == 0 || k == entity || known.Contains(k) {
+					return true
 				}
 				todo = append(todo, k)
-				if known == nil {
-					known = map[types.EntityUID]struct{}{}
-				}
-				known[k] = struct{}{}
-			}
+				known.Add(k)
+				return true
+			})
 		}
 		if len(todo) == 0 {
 			return false
@@ -1072,14 +1061,14 @@ func doInEval(env *Env, lhs types.EntityUID, rhs types.Value) (types.Value, erro
 	case types.EntityUID:
 		return types.Boolean(entityInOne(env, lhs, rhsv)), nil
 	case types.Set:
-		query := make(map[types.EntityUID]struct{}, rhsv.Len())
+		query := mapset.Make[types.EntityUID](rhsv.Len())
 		var err error
 		rhsv.Iterate(func(rhv types.Value) bool {
 			var e types.EntityUID
 			if e, err = ValueToEntity(rhv); err != nil {
 				return false
 			}
-			query[e] = struct{}{}
+			query.Add(e)
 			return true
 		})
 		if err != nil {

internal/eval/evalers_test.go

@@ -1833,9 +1833,9 @@ func TestEntityIn(t *testing.T) {
 		tt := tt
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
-			rhs := map[types.EntityUID]struct{}{}
+			var rhs []types.EntityUID
 			for _, v := range tt.rhs {
-				rhs[strEnt(v)] = struct{}{}
+				rhs = append(rhs, strEnt(v))
 			}
 			entityMap := types.Entities{}
 			for k, p := range tt.parents {
@@ -1846,10 +1846,10 @@ func TestEntityIn(t *testing.T) {
 				uid := strEnt(k)
 				entityMap[uid] = &types.Entity{
 					UID:     uid,
-					Parents: ps,
+					Parents: types.NewEntityUIDSet(ps...),
 				}
 			}
-			res := entityInSet(&Env{Entities: entityMap}, strEnt(tt.lhs), rhs)
+			res := entityInSet(&Env{Entities: entityMap}, strEnt(tt.lhs), types.NewEntityUIDSet(rhs...))
 			testutil.Equals(t, res, tt.result)
 		})
 	}
@@ -1860,10 +1860,10 @@ func TestEntityIn(t *testing.T) {
 
 		entityMap := types.Entities{}
 		for i := 0; i < 100; i++ {
-			p := []types.EntityUID{
+			p := types.NewEntityUIDSet(
 				types.NewEntityUID(types.EntityType(fmt.Sprint(i+1)), "1"),
 				types.NewEntityUID(types.EntityType(fmt.Sprint(i+1)), "2"),
-			}
+			)
 			uid1 := types.NewEntityUID(types.EntityType(fmt.Sprint(i)), "1")
 			entityMap[uid1] = &types.Entity{
 				UID:     uid1,
@@ -1876,7 +1876,11 @@ func TestEntityIn(t *testing.T) {
 			}
 
 		}
-		res := entityInSet(&Env{Entities: entityMap}, types.NewEntityUID("0", "1"), map[types.EntityUID]struct{}{types.NewEntityUID("0", "3"): {}})
+		res := entityInSet(
+			&Env{Entities: entityMap},
+			types.NewEntityUID("0", "1"),
+			types.NewEntityUIDSet(types.NewEntityUID("0", "3")),
+		)
 		testutil.Equals(t, res, false)
 	})
 }
@@ -2012,7 +2016,7 @@ func TestInNode(t *testing.T) {
 				uid := strEnt(k)
 				entityMap[uid] = &types.Entity{
 					UID:     uid,
-					Parents: ps,
+					Parents: types.NewEntityUIDSet(ps...),
 				}
 			}
 			ec := InitEnv(&Env{Entities: entityMap})
@@ -2152,7 +2156,7 @@ func TestIsInNode(t *testing.T) {
 				uid := strEnt(k)
 				entityMap[uid] = &types.Entity{
 					UID:     uid,
-					Parents: ps,
+					Parents: types.NewEntityUIDSet(ps...),
 				}
 			}
 			ec := InitEnv(&Env{Entities: entityMap})

internal/eval/partial.go

@@ -6,6 +6,7 @@ import (
 	"slices"
 
 	"github.com/cedar-policy/cedar-go/internal/ast"
+	"github.com/cedar-policy/cedar-go/internal/mapset"
 	"github.com/cedar-policy/cedar-go/types"
 )
 
@@ -140,10 +141,7 @@ func partialScopeEval(env *Env, ent types.Value, in ast.IsScopeNode) (evaled boo
 	case ast.ScopeTypeIn:
 		return true, entityInOne(env, e, t.Entity)
 	case ast.ScopeTypeInSet:
-		set := make(map[types.EntityUID]struct{}, len(t.Entities))
-		for _, e := range t.Entities {
-			set[e] = struct{}{}
-		}
+		set := mapset.Immutable(t.Entities...)
 		return true, entityInSet(env, e, set)
 	case ast.ScopeTypeIs:
 		return true, e.Type == t.Type

internal/mapset/immutable.go

@@ -0,0 +1,58 @@
+package mapset
+
+import (
+	"encoding/json"
+)
+
+type ImmutableMapSet[T comparable] MapSet[T]
+
+func Immutable[T comparable](args ...T) ImmutableMapSet[T] {
+	return ImmutableMapSet[T](*FromItems(args...))
+}
+
+// Contains returns whether the item exists in the set
+func (h ImmutableMapSet[T]) Contains(item T) bool {
+	return MapSet[T](h).Contains(item)
+}
+
+// Intersects returns whether any items in this set exist in o
+func (h ImmutableMapSet[T]) Intersects(o Container[T]) bool {
+	return MapSet[T](h).Intersects(o)
+}
+
+// Iterate the items in the set, calling callback for each item. If the callback returns false, iteration is halted.
+// Iteration order is undefined.
+func (h ImmutableMapSet[T]) Iterate(callback func(item T) bool) {
+	MapSet[T](h).Iterate(callback)
+}
+
+func (h ImmutableMapSet[T]) Slice() []T {
+	return MapSet[T](h).Slice()
+}
+
+// Len returns the size of the set
+func (h ImmutableMapSet[T]) Len() int {
+	return MapSet[T](h).Len()
+}
+
+// Equal returns whether the same items exist in both h and o
+func (h ImmutableMapSet[T]) Equal(o ImmutableMapSet[T]) bool {
+	om := MapSet[T](o)
+	return MapSet[T](h).Equal(&om)
+}
+
+// MarshalJSON serializes a MapSet as a JSON array. Elements are ordered lexicographically by their marshaled value.
+func (h ImmutableMapSet[T]) MarshalJSON() ([]byte, error) {
+	return MapSet[T](h).MarshalJSON()
+}
+
+// UnmarshalJSON deserializes an ImmutableMapSet from a JSON array.
+func (h *ImmutableMapSet[T]) UnmarshalJSON(b []byte) error {
+	var s MapSet[T]
+	if err := json.Unmarshal(b, &s); err != nil {
+		return err
+	}
+
+	*h = ImmutableMapSet[T](s)
+	return nil
+}