Changes needed by RFC 53 (#524)

Signed-off-by: Shaobo He <shaobohe@amazon.com>

Author: shaobo-he-aws

cedar-drt/fuzz/src/schemas.rs

@@ -16,7 +16,9 @@
 
 use cedar_policy_core::ast::{Id, InternalName, Name};
 use cedar_policy_validator::json_schema;
+use cedar_policy_validator::json_schema::EntityTypeKind;
 use cedar_policy_validator::RawName;
+use cedar_policy_validator::ValidatorEntityTypeKind;
 use itertools::Itertools;
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 
@@ -228,36 +230,67 @@ impl<N: Clone + PartialEq + Debug + Display + TypeName + Ord> Equiv for json_sch
     fn equiv(lhs: &Self, rhs: &Self) -> Result<(), String> {
         Equiv::equiv(&lhs.annotations, &rhs.annotations)
             .map_err(|e| format!("mismatch in entity annotations: {e}"))?;
-        Equiv::equiv(
-            &lhs.member_of_types.iter().collect::<BTreeSet<_>>(),
-            &rhs.member_of_types.iter().collect::<BTreeSet<_>>(),
-        )
-        .map_err(|e| format!("memberOfTypes are not equal: {e}"))?;
-        Equiv::equiv(&lhs.shape, &rhs.shape).map_err(|e| format!("mismatched types: {e}"))?;
-        match (&lhs.tags, &rhs.tags) {
-            (Some(ts1), Some(ts2)) => {
-                Equiv::equiv(ts1, ts2).map_err(|msg| format!("mismatched entity tags: {msg}"))
+        match (&lhs.kind, &rhs.kind) {
+            (EntityTypeKind::Enum { choices: c1 }, EntityTypeKind::Enum { choices: c2 }) => {
+                if c1 != c2 {
+                    Err(format!(
+                        "enumerated entity types have different eid choices: {c1:?} and {c2:?}"
+                    ))
+                } else {
+                    Ok(())
+                }
+            }
+            (EntityTypeKind::Standard(lhs), EntityTypeKind::Standard(rhs)) => {
+                Equiv::equiv(
+                    &lhs.member_of_types.iter().collect::<BTreeSet<_>>(),
+                    &rhs.member_of_types.iter().collect::<BTreeSet<_>>(),
+                )
+                .map_err(|e| format!("memberOfTypes are not equal: {e}"))?;
+                Equiv::equiv(&lhs.shape, &rhs.shape)
+                    .map_err(|e| format!("mismatched types: {e}"))?;
+                match (&lhs.tags, &rhs.tags) {
+                    (Some(ts1), Some(ts2)) => Equiv::equiv(ts1, ts2)
+                        .map_err(|msg| format!("mismatched entity tags: {msg}")),
+                    (None, None) => Ok(()),
+                    (Some(ts), None) | (None, Some(ts)) => {
+                        Err(format!("only one side has tags: {ts}"))
+                    }
+                }
             }
-            (None, None) => Ok(()),
-            (Some(ts), None) | (None, Some(ts)) => Err(format!("only one side has tags: {ts}")),
+            (k1, k2) => Err(format!("different entity type kind: {:?} and {:?}", k1, k2)),
         }
     }
 }
 
 impl Equiv for cedar_policy_validator::ValidatorEntityType {
     fn equiv(lhs: &Self, rhs: &Self) -> Result<(), String> {
-        Equiv::equiv(&lhs.descendants, &rhs.descendants)?;
-        Equiv::equiv(
-            &lhs.attributes().collect::<HashMap<_, _>>(),
-            &rhs.attributes().collect::<HashMap<_, _>>(),
-        )?;
-        if lhs.tag_type() != rhs.tag_type() {
-            return Err(format!(
-                "encountered different tags types: {:?} and {:?}",
-                lhs.tag_type(),
-                rhs.tag_type()
-            ));
-        }
+        match (&lhs.kind, &rhs.kind) {
+            (ValidatorEntityTypeKind::Enum(c1), ValidatorEntityTypeKind::Enum(c2)) => {
+                if c1 != c2 {
+                    return Err(format!(
+                        "enumerated entity types have different eid choices: {c1:?} and {c2:?}"
+                    ));
+                }
+            }
+            (ValidatorEntityTypeKind::Standard(_), ValidatorEntityTypeKind::Standard(_)) => {
+                Equiv::equiv(&lhs.descendants, &rhs.descendants)?;
+                Equiv::equiv(
+                    &lhs.attributes().iter().collect::<HashMap<_, _>>(),
+                    &rhs.attributes().iter().collect::<HashMap<_, _>>(),
+                )?;
+                if lhs.tag_type() != rhs.tag_type() {
+                    return Err(format!(
+                        "encountered different tags types: {:?} and {:?}",
+                        lhs.tag_type(),
+                        rhs.tag_type()
+                    ));
+                }
+            }
+            (k1, k2) => {
+                return Err(format!("different entity type kind: {:?} and {:?}", k1, k2));
+            }
+        };
+
         Ok(())
     }
 }

cedar-lean/CedarProto/ValidatorEntityType.lean

@@ -32,6 +32,7 @@ structure ValidatorEntityType where
   descendants : Array Spec.EntityTypeProto
   attrs : RecordType
   tags : Option CedarType
+  enums: Array String
 deriving Inhabited
 
 -- the Tag message in Validator.proto
@@ -91,12 +92,19 @@ def mergeTags (result : ValidatorEntityType) (x : TagMessage) : ValidatorEntityT
     tags := mergeOption result.tags (some x.optional_type)
   }
 
+@[inline]
+def mergeEnums (result : ValidatorEntityType) (x : Array String) : ValidatorEntityType :=
+  {result with
+    enums := result.enums ++ x
+  }
+
 @[inline]
 def merge (x y : ValidatorEntityType) : ValidatorEntityType :=
   {
     descendants := y.descendants ++ x.descendants
     attrs := Field.merge x.attrs y.attrs
     tags := mergeOption x.tags y.tags
+    enums := x.enums ++ y.enums
   }
 
 @[inline]
@@ -111,6 +119,9 @@ def parseField (t : Tag) : BParsec (MergeFn ValidatorEntityType) := do
     | 5 =>
       let x : TagMessage ← Field.guardedParse t
       pure (pure $ mergeTags · x)
+    | 6 =>
+      let x : Repeated String ← Field.guardedParse t
+      pure (pure $ mergeEnums · x)
     | _ =>
       t.wireType.skip
       pure ignore

cedar-lean/CedarProto/ValidatorSchema.lean

@@ -71,11 +71,19 @@ def toEntitySchema (ets : EntityTypeWithTypesMap) : EntitySchema :=
   let ancestorMap := descendantsToAncestors descendantMap
   Data.Map.make (ets.map
     (λ (k,v) => (k,
+      if v.enums.isEmpty then
       {
         ancestors := ancestorMap.find! k,
         attrs := Data.Map.make v.attrs.kvs,
         tags := v.tags,
-      })))
+      }
+      else
+      {
+        ancestors := ancestorMap.find! k,
+        attrs := Data.Map.empty,
+        tags := none,
+      }
+      )))
 end Cedar.Validation.Proto.EntityTypeWithTypesMap
 
 namespace Cedar.Validation.Proto.EntityUidWithActionsIdMap

cedar-lean/DiffTest/Parser.lean

@@ -508,14 +508,22 @@ partial def jsonToCedarType (json : Lean.Json) : ParseResult CedarType := do
       .ok (.ext name)
     | tag => .error s!"jsonToCedarType: unknown tag {tag}"
 
+partial def jsonToEntityTypeKind (json: Lean.Json) : ParseResult (RecordType × Option CedarType) := do
+  let (tag, value) ← unpackJsonSum json
+  match tag with
+    | "Enum" => .ok (Map.empty, none)
+    | "Standard" =>
+      let attrs ← getJsonField value "attributes" >>= (getJsonField · "attrs") >>= jsonToRecordType
+      let tags ← -- the "tags" field may be absent
+      match getJsonField value "tags" with
+        | .ok jty  => (jsonToCedarType jty).map λ ty => (attrs, some ty)
+        | .error _ => .ok (attrs, none)
+    | tag => .error s!"jsonToEntityTypeKind: unknown tag {tag}"
+
 partial def jsonToEntityTypeEntry (json : Lean.Json) : ParseResult JsonEntitySchemaEntry := do
   let descendants_json ← getJsonField json "descendants" >>= jsonToArray
   let descendants ← List.mapM jsonToName descendants_json.toList
-  let attrs ← getJsonField json "attributes" >>= (getJsonField · "attrs") >>= jsonToRecordType
-  let tags ← -- the "tags" field may be absent
-    match getJsonField json "tags" with
-    | .ok jty  => (jsonToCedarType jty).map .some
-    | .error _ => .ok .none
+  let (attrs, tags) ← getJsonField json "kind" >>= jsonToEntityTypeKind
   .ok {
     descendants := Set.make descendants,
     attrs := attrs,

cedar-policy-generators/Cargo.toml

@@ -17,6 +17,7 @@ rand = "0.8.5"
 anyhow = "1.0.72"
 serde_with = "3.4.0"
 thiserror = "2.0"
+nonempty = { version = "0.10", features = ["arbitrary"] }
 
 [dev.dependencies]
 rand = "0.8.5"

cedar-policy-generators/src/expr.rs

@@ -28,7 +28,7 @@ use crate::{accum, gen, gen_inner, uniform};
 use arbitrary::{Arbitrary, MaxRecursionReached, Unstructured};
 use cedar_policy_core::ast::{self, UnreservedId};
 use cedar_policy_core::est::Annotations;
-use cedar_policy_validator::json_schema;
+use cedar_policy_validator::json_schema::{self, EntityTypeKind, StandardEntityType};
 use smol_str::SmolStr;
 use std::collections::BTreeMap;
 
@@ -590,11 +590,14 @@ impl<'a> ExprGenerator<'a> {
                                         .expect("Failed to select entity index."),
                                 )
                                 .expect("Failed to select entity from map.");
-                            let attr_names: Vec<&SmolStr> =
-                                attrs_from_attrs_or_context(&self.schema.schema, &entity_type.shape)
-                                    .attrs
-                                    .keys()
-                                    .collect::<Vec<_>>();
+                            let attr_names: Vec<&SmolStr> = match &entity_type.kind {
+                                // Generate an empty vec here so that we fail fast
+                                EntityTypeKind::Enum { .. } => vec![],
+                                EntityTypeKind::Standard(StandardEntityType { shape, ..}) => attrs_from_attrs_or_context(&self.schema.schema, shape)
+                                .attrs
+                                .keys()
+                                .collect::<Vec<_>>()
+                            };
                             let attr_name = SmolStr::clone(u.choose(&attr_names)?);
                             Ok(ast::Expr::has_attr(
                                 self.generate_expr_for_schematype(
@@ -2559,7 +2562,7 @@ impl<'a> ExprGenerator<'a> {
         u: &mut Unstructured<'_>,
     ) -> Result<ast::EntityUID> {
         match self.hierarchy {
-            None => generate_uid_with_type(ty.clone(), u),
+            None => generate_uid_with_type(ty.clone(), &self.schema.get_uid_choices(ty), u),
             Some(hierarchy) => hierarchy.arbitrary_uid_with_type(ty, u),
         }
     }