Add get_entity_literals to api.rs (#1149)

Signed-off-by: Andrew Wells <[email protected]>

Author: andrewmwells-amazon

cedar-policy/CHANGELOG.md

@@ -14,6 +14,7 @@ Cedar Language Version: TBD
 
 ### Added
 
+- Added `get_entity_literals` API (#1149).
 - Implemented [RFC 82](https://github.com/cedar-policy/rfcs/pull/82), adding
   entity tags to the Cedar language under experimental flag `entity-tags` (#1204, #1207, more coming)
 - Implemented [RFC 74](https://github.com/cedar-policy/rfcs/pull/74): A new experimental API (`compute_entity_manifest`)

cedar-policy/src/api.rs

@@ -2919,6 +2919,23 @@ impl Policy {
         get_valid_request_envs(self.ast.template(), s)
     }
 
+    /// Get all entity literals occuring in a `Policy`
+    pub fn entity_literals(&self) -> Vec<EntityUid> {
+        self.ast
+            .condition()
+            .subexpressions()
+            .filter_map(|e| match e.expr_kind() {
+                cedar_policy_core::ast::ExprKind::Lit(l) => match l {
+                    cedar_policy_core::ast::Literal::EntityUID(euid) => {
+                        Some(EntityUid((*euid).as_ref().clone()))
+                    }
+                    _ => None,
+                },
+                _ => None,
+            })
+            .collect()
+    }
+
     fn from_est(id: Option<PolicyId>, est: est::Policy) -> Result<Self, PolicyFromJsonError> {
         Ok(Self {
             ast: est.clone().try_into_ast_policy(id.map(PolicyId::into))?,

cedar-policy/src/tests.rs

@@ -5715,3 +5715,29 @@ mod context_tests {
         );
     }
 }
+
+mod policy_manipulation_functions_tests {
+    use super::*;
+
+    #[test]
+    fn empty_policy() {
+        let policy_str = r###"permit(principal, action, resource);
+        "###;
+        let policy = Policy::from_str(policy_str).expect("should succeed");
+        assert_eq!(policy.entity_literals(), vec![]);
+    }
+
+    #[test]
+    fn non_empty_policy() {
+        let policy_str = r###"permit(principal == User::"Bob", action == Action::"view", resource) when {
+            !resource.private && resource.owner != User::"Alice"
+        };
+        "###;
+        let policy = Policy::from_str(policy_str).expect("should succeed");
+        let res = policy.entity_literals();
+        assert_eq!(res.len(), 3);
+        assert!(res.contains(&EntityUid::from_str("User::\"Bob\"").expect("should parse")));
+        assert!(res.contains(&EntityUid::from_str("Action::\"view\"").expect("should parse")));
+        assert!(res.contains(&EntityUid::from_str("User::\"Alice\"").expect("should parse")));
+    }
+}