Update to generate is operator

cedar-policy-generators/src/expr.rs

@@ -146,6 +146,12 @@ impl<'a> ExprGenerator<'a> {
                             Err(Error::LikeDisabled)
                         }
                     },
+                    1 => {
+                            Ok(ast::Expr::is_entity_type(
+                                self.generate_expr(max_depth - 1, u)?,
+                                u.choose(&self.schema.entity_types)?.clone(),
+                            ))
+                    },
                     1 => {
                         let mut l = Vec::new();
                         u.arbitrary_loop(Some(0), Some(self.settings.max_width as u32), |u| {
@@ -476,6 +482,17 @@ impl<'a> ExprGenerator<'a> {
                                 Err(Error::LikeDisabled)
                             }
                         },
+                        // is
+                        2 => {
+                                Ok(ast::Expr::is_entity_type(
+                                    self.generate_expr_for_type(
+                                        &Type::entity(),
+                                        max_depth - 1,
+                                        u,
+                                    )?,
+                                    u.choose(&self.schema.entity_types)?.clone(),
+                                ))
+                        },
                         // extension function that returns bool
                         2 => self.generate_ext_func_call_for_type(
                             &Type::bool(),

cedar-policy-generators/src/hierarchy.rs

@@ -65,6 +65,9 @@ pub struct Hierarchy {
     /// We keep this in sync with the `entities` HashMap too.
     /// This is to make arbitrary_uid_with_type() fast.
     uids_by_type: HashMap<ast::Name, EntityUIDs>,
+    /// Vec of all entity types used by entities in the hierarchy, again kept in
+    /// sync with the `entities` HashMap. Makes `arbitrary_entity_type()` fast.
+    entity_types: Vec<ast::Name>,
 }
 
 impl Hierarchy {
@@ -76,16 +79,19 @@ impl Hierarchy {
             .values()
             .flat_map(|uids_inner| uids_inner.iter().cloned())
             .collect();
+        let uids_by_type: HashMap<_, _> = uids_by_type
+            .into_iter()
+            .map(|(n, uids)| (n, EntityUIDs::from_iter(uids)))
+            .collect();
+        let entity_types: Vec<_> = uids_by_type.keys().cloned().collect();
         Self {
             entities: uids
                 .iter()
                 .map(|uid| (uid.clone(), Entity::with_uid(uid.clone())))
                 .collect(),
             uids,
-            uids_by_type: uids_by_type
-                .into_iter()
-                .map(|(n, uids)| (n, EntityUIDs::from_iter(uids)))
-                .collect(),
+            uids_by_type,
+            entity_types,
         }
     }
 
@@ -147,6 +153,19 @@ impl Hierarchy {
         )
     }
 
+    /// Generate an entity type, usually picking on that's used by some entity in
+    /// the hierarchy.
+    pub fn arbitrary_entity_type(&self, u: &mut Unstructured<'_>) -> Result<ast::Name> {
+        // entity type that is used by some entity or isn't. 90% of the time
+        // pick one that is used.
+        if u.ratio::<u8>(9, 10)? {
+            let ety = u.choose(&self.entity_types)?;
+            Ok(ety.clone())
+        } else {
+            Ok(u.arbitrary()?)
+        }
+    }
+
     /// Get an Entity object by UID
     pub fn entity(&self, uid: &EntityUID) -> Option<&Entity> {
         self.entities.get(uid)
@@ -198,6 +217,7 @@ impl Hierarchy {
             entities: new_entities,
             uids: self.uids,
             uids_by_type: self.uids_by_type,
+            entity_types: self.entity_types,
         }
     }
 
@@ -248,13 +268,15 @@ impl From<Entities> for Hierarchy {
             uids_by_type.entry(etype).or_default().insert(e.uid());
             uids.push(e.uid());
         }
+        let entity_types: Vec<_> = uids_by_type.keys().cloned().collect();
         Hierarchy {
             uids,
             uids_by_type: uids_by_type
                 .into_iter()
                 .map(|(k, v)| (k, EntityUIDs::from_iter(v.into_iter())))
                 .collect(),
             entities: entities.into_iter().map(|e| (e.uid(), e)).collect(),
+            entity_types,
         }
     }
 }

cedar-policy-generators/src/policy.rs

@@ -4,7 +4,7 @@ use crate::hierarchy::Hierarchy;
 use crate::size_hint_utils::size_hint_for_ratio;
 use arbitrary::{Arbitrary, Unstructured};
 use cedar_policy_core::ast::{
-    Effect, EntityUID, Expr, Id, Policy, PolicyID, PolicySet, StaticPolicy, Template,
+    Effect, EntityUID, Expr, Id, Name, Policy, PolicyID, PolicySet, StaticPolicy, Template,
 };
 use cedar_policy_core::{ast, est};
 use serde::Serialize;
@@ -189,6 +189,12 @@ pub enum PrincipalOrResourceConstraint {
     In(EntityUID),
     /// Eg `principal in ?principal`
     InSlot,
+    /// Eg `principal is User`
+    IsType(Name),
+    /// Eg `principal is User in Group::"123"`
+    IsTypeIn(Name, EntityUID),
+    /// Eg `principal is User in ?principal`
+    IsTypeInSlot(Name),
 }
 
 impl PrincipalOrResourceConstraint {
@@ -199,6 +205,9 @@ impl PrincipalOrResourceConstraint {
             PrincipalOrResourceConstraint::EqSlot => true,
             PrincipalOrResourceConstraint::In(_) => false,
             PrincipalOrResourceConstraint::InSlot => true,
+            PrincipalOrResourceConstraint::IsType(_) => false,
+            PrincipalOrResourceConstraint::IsTypeIn(_, _) => false,
+            PrincipalOrResourceConstraint::IsTypeInSlot(_) => true,
         }
     }
 }
@@ -211,6 +220,15 @@ impl From<PrincipalOrResourceConstraint> for ast::PrincipalConstraint {
             PrincipalOrResourceConstraint::EqSlot => ast::PrincipalConstraint::is_eq_slot(),
             PrincipalOrResourceConstraint::In(euid) => ast::PrincipalConstraint::is_in(euid),
             PrincipalOrResourceConstraint::InSlot => ast::PrincipalConstraint::is_in_slot(),
+            PrincipalOrResourceConstraint::IsType(ty) => {
+                ast::PrincipalConstraint::is_entity_type(ty)
+            }
+            PrincipalOrResourceConstraint::IsTypeIn(ty, euid) => {
+                ast::PrincipalConstraint::is_entity_type_in(ty, euid)
+            }
+            PrincipalOrResourceConstraint::IsTypeInSlot(ty) => {
+                ast::PrincipalConstraint::is_entity_type_in_slot(ty)
+            }
         }
     }
 }
@@ -223,6 +241,15 @@ impl From<PrincipalOrResourceConstraint> for ast::ResourceConstraint {
             PrincipalOrResourceConstraint::EqSlot => ast::ResourceConstraint::is_eq_slot(),
             PrincipalOrResourceConstraint::In(euid) => ast::ResourceConstraint::is_in(euid),
             PrincipalOrResourceConstraint::InSlot => ast::ResourceConstraint::is_in_slot(),
+            PrincipalOrResourceConstraint::IsType(ty) => {
+                ast::ResourceConstraint::is_entity_type(ty)
+            }
+            PrincipalOrResourceConstraint::IsTypeIn(ty, euid) => {
+                ast::ResourceConstraint::is_entity_type_in(ty, euid)
+            }
+            PrincipalOrResourceConstraint::IsTypeInSlot(ty) => {
+                ast::ResourceConstraint::is_entity_type_in_slot(ty)
+            }
         }
     }
 }
@@ -238,6 +265,9 @@ impl std::fmt::Display for PrincipalOrResourceConstraint {
             Self::EqSlot => write!(f, " == ?"),
             Self::In(uid) => write!(f, " in {uid}"),
             Self::InSlot => write!(f, " in ?"),
+            Self::IsType(ty) => write!(f, " is {ty}"),
+            Self::IsTypeIn(ty, euid) => write!(f, " is {ty} in {euid}"),
+            Self::IsTypeInSlot(ty) => write!(f, "is {ty} in ?"),
         }
     }
 }
@@ -248,26 +278,28 @@ impl PrincipalOrResourceConstraint {
         allow_slots: bool,
         u: &mut Unstructured<'_>,
     ) -> Result<Self> {
-        // 20% of the time, NoConstraint; 40%, Eq; 40%, In
+        // 20% of the time, NoConstraint
         if u.ratio(1, 5)? {
             Ok(Self::NoConstraint)
         } else {
-            // choose Eq or In
-            let use_eq = u.ratio(1, 2)?;
             // If slots are allowed, then generate a slot 50% of the time.
             if allow_slots && u.ratio(1, 2)? {
-                if use_eq {
-                    Ok(Self::EqSlot)
-                } else {
-                    Ok(Self::InSlot)
-                }
+                // 40% Eq, 40% In or IsIn.
+                // Don't generate `Is` on its own because it can't have a slot.
+                gen!(u,
+                    2 => Ok(Self::EqSlot),
+                    1 => Ok(Self::InSlot),
+                    1 => Ok(Self::IsTypeInSlot(hierarchy.arbitrary_entity_type(u)?))
+                )
             } else {
+                // 32% Eq, 32% In (16% with and without Is), 32% Is (16% with and without In)
                 let uid = hierarchy.arbitrary_uid(u)?;
-                if use_eq {
-                    Ok(Self::Eq(uid))
-                } else {
-                    Ok(Self::In(uid))
-                }
+                gen!(u,
+                    2 => Ok(Self::Eq(uid)),
+                    1 => Ok(Self::In(uid)),
+                    1 => Ok(Self::IsType(hierarchy.arbitrary_entity_type(u)?)),
+                    1 => Ok(Self::IsTypeIn(hierarchy.arbitrary_entity_type(u)?, uid))
+                )
             }
         }
     }