Speed Improvement: Make incoming entities Arc<Entities> like schema entities (#1296)

Signed-off-by: Thomas Hill <[email protected]>

Author: tomlikestorock

cedar-policy-core/src/ast/entity.rs

@@ -30,6 +30,7 @@ use serde_with::{serde_as, TryFromInto};
 use smol_str::SmolStr;
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::str::FromStr;
+use std::sync::Arc;
 use thiserror::Error;
 
 /// The entity type that Actions must have
@@ -605,6 +606,25 @@ impl TCNode<EntityUID> for Entity {
     }
 }
 
+impl TCNode<EntityUID> for Arc<Entity> {
+    fn get_key(&self) -> EntityUID {
+        self.uid().clone()
+    }
+
+    fn add_edge_to(&mut self, k: EntityUID) {
+        // Use Arc::make_mut to get a mutable reference to the inner value
+        Arc::make_mut(self).add_ancestor(k)
+    }
+
+    fn out_edges(&self) -> Box<dyn Iterator<Item = &EntityUID> + '_> {
+        Box::new(self.ancestors())
+    }
+
+    fn has_edge_to(&self, e: &EntityUID) -> bool {
+        self.is_descendant_of(e)
+    }
+}
+
 impl std::fmt::Display for Entity {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(

cedar-policy-core/src/entities.rs

@@ -58,7 +58,7 @@ pub struct Entities {
     /// Important internal invariant: for any `Entities` object that exists, the
     /// the `ancestor` relation is transitively closed.
     #[serde_as(as = "Vec<(_, _)>")]
-    entities: HashMap<EntityUID, Entity>,
+    entities: HashMap<EntityUID, Arc<Entity>>,
 
     /// The mode flag determines whether this store functions as a partial store or
     /// as a fully concrete store.
@@ -109,7 +109,7 @@ impl Entities {
 
     /// Iterate over the `Entity`s in the `Entities`
     pub fn iter(&self) -> impl Iterator<Item = &Entity> {
-        self.entities.values()
+        self.entities.values().map(|e| e.as_ref())
     }
 
     /// Adds the [`crate::ast::Entity`]s in the iterator to this [`Entities`].
@@ -125,7 +125,7 @@ impl Entities {
     /// responsible for ensuring that TC and DAG hold before calling this method.
     pub fn add_entities(
         mut self,
-        collection: impl IntoIterator<Item = Entity>,
+        collection: impl IntoIterator<Item = Arc<Entity>>,
         schema: Option<&impl Schema>,
         tc_computation: TCComputation,
         extensions: &Extensions<'_>,
@@ -174,7 +174,7 @@ impl Entities {
         tc_computation: TCComputation,
         extensions: &Extensions<'_>,
     ) -> Result<Self> {
-        let mut entity_map = create_entity_map(entities.into_iter())?;
+        let mut entity_map = create_entity_map(entities.into_iter().map(Arc::new))?;
         if let Some(schema) = schema {
             // Validate non-action entities against schema.
             // We do this before adding the actions, because we trust the
@@ -213,7 +213,7 @@ impl Entities {
                 schema
                     .action_entities()
                     .into_iter()
-                    .map(|e| (e.uid().clone(), Arc::unwrap_or_clone(e))),
+                    .map(|e: Arc<Entity>| (e.uid().clone(), e)),
             );
         }
         Ok(Self {
@@ -252,6 +252,7 @@ impl Entities {
     fn to_ejsons(&self) -> Result<Vec<EntityJson>> {
         self.entities
             .values()
+            .map(Arc::as_ref)
             .map(EntityJson::from_entity)
             .collect::<std::result::Result<_, JsonSerializationError>>()
             .map_err(Into::into)
@@ -322,7 +323,9 @@ impl Entities {
 }
 
 /// Create a map from EntityUids to Entities, erroring if there are any duplicates
-fn create_entity_map(es: impl Iterator<Item = Entity>) -> Result<HashMap<EntityUID, Entity>> {
+fn create_entity_map(
+    es: impl Iterator<Item = Arc<Entity>>,
+) -> Result<HashMap<EntityUID, Arc<Entity>>> {
     let mut map = HashMap::new();
     for e in es {
         match map.entry(e.uid().clone()) {
@@ -338,10 +341,13 @@ fn create_entity_map(es: impl Iterator<Item = Entity>) -> Result<HashMap<EntityU
 impl IntoIterator for Entities {
     type Item = Entity;
 
-    type IntoIter = hash_map::IntoValues<EntityUID, Entity>;
+    type IntoIter = std::iter::Map<
+        std::collections::hash_map::IntoValues<EntityUID, Arc<Entity>>,
+        fn(Arc<Entity>) -> Entity,
+    >;
 
     fn into_iter(self) -> Self::IntoIter {
-        self.entities.into_values()
+        self.entities.into_values().map(Arc::unwrap_or_clone)
     }
 }
 
@@ -497,7 +503,8 @@ mod json_parsing_tests {
 
         let addl_entities = parser
             .iter_from_json_value(new)
-            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)));
+            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)))
+            .map(Arc::new);
         let err = simple_entities(&parser).add_entities(
             addl_entities,
             None::<&NoEntitiesSchema>,
@@ -537,7 +544,8 @@ mod json_parsing_tests {
 
         let addl_entities = parser
             .iter_from_json_value(new)
-            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)));
+            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)))
+            .map(Arc::new);
         let err = simple_entities(&parser).add_entities(
             addl_entities,
             None::<&NoEntitiesSchema>,
@@ -576,7 +584,8 @@ mod json_parsing_tests {
 
         let addl_entities = parser
             .iter_from_json_value(new)
-            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)));
+            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)))
+            .map(Arc::new);
         let err = simple_entities(&parser).add_entities(
             addl_entities,
             None::<&NoEntitiesSchema>,
@@ -619,7 +628,8 @@ mod json_parsing_tests {
 
         let addl_entities = parser
             .iter_from_json_value(new)
-            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)));
+            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)))
+            .map(Arc::new);
         let es = simple_entities(&parser)
             .add_entities(
                 addl_entities,
@@ -658,7 +668,8 @@ mod json_parsing_tests {
 
         let addl_entities = parser
             .iter_from_json_value(new)
-            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)));
+            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)))
+            .map(Arc::new);
         let es = simple_entities(&parser)
             .add_entities(
                 addl_entities,
@@ -699,7 +710,8 @@ mod json_parsing_tests {
 
         let addl_entities = parser
             .iter_from_json_value(new)
-            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)));
+            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)))
+            .map(Arc::new);
         let es = simple_entities(&parser)
             .add_entities(
                 addl_entities,
@@ -739,7 +751,8 @@ mod json_parsing_tests {
 
         let addl_entities = parser
             .iter_from_json_value(new)
-            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)));
+            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)))
+            .map(Arc::new);
         let es = simple_entities(&parser)
             .add_entities(
                 addl_entities,
@@ -766,7 +779,8 @@ mod json_parsing_tests {
 
         let addl_entities = parser
             .iter_from_json_value(new)
-            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)));
+            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)))
+            .map(Arc::new);
         let err = simple_entities(&parser)
             .add_entities(
                 addl_entities,
@@ -787,7 +801,8 @@ mod json_parsing_tests {
         let new = serde_json::json!([{"uid":{ "type": "Test", "id": "alice" }, "attrs" : {}, "parents" : []}]);
         let addl_entities = parser
             .iter_from_json_value(new)
-            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)));
+            .unwrap_or_else(|e| panic!("{:?}", &miette::Report::new(e)))
+            .map(Arc::new);
         let err = simple_entities(&parser).add_entities(
             addl_entities,
             None::<&NoEntitiesSchema>,

cedar-policy/src/api.rs

@@ -58,6 +58,7 @@ use smol_str::SmolStr;
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use std::io::Read;
 use std::str::FromStr;
+use std::sync::Arc;
 
 // PANIC SAFETY: `CARGO_PKG_VERSION` should return a valid SemVer version string
 #[allow(clippy::unwrap_used)]
@@ -407,7 +408,7 @@ impl Entities {
     ) -> Result<Self, EntitiesError> {
         Ok(Self(
             self.0.add_entities(
-                entities.into_iter().map(|e| e.0),
+                entities.into_iter().map(|e| Arc::new(e.0)),
                 schema
                     .map(|s| cedar_policy_validator::CoreSchema::new(&s.0))
                     .as_ref(),
@@ -446,7 +447,7 @@ impl Entities {
             Extensions::all_available(),
             cedar_policy_core::entities::TCComputation::ComputeNow,
         );
-        let new_entities = eparser.iter_from_json_str(json)?;
+        let new_entities = eparser.iter_from_json_str(json)?.map(Arc::new);
         Ok(Self(self.0.add_entities(
             new_entities,
             schema.as_ref(),
@@ -484,7 +485,7 @@ impl Entities {
             Extensions::all_available(),
             cedar_policy_core::entities::TCComputation::ComputeNow,
         );
-        let new_entities = eparser.iter_from_json_value(json)?;
+        let new_entities = eparser.iter_from_json_value(json)?.map(Arc::new);
         Ok(Self(self.0.add_entities(
             new_entities,
             schema.as_ref(),
@@ -523,7 +524,7 @@ impl Entities {
             Extensions::all_available(),
             cedar_policy_core::entities::TCComputation::ComputeNow,
         );
-        let new_entities = eparser.iter_from_json_file(json)?;
+        let new_entities = eparser.iter_from_json_file(json)?.map(Arc::new);
         Ok(Self(self.0.add_entities(
             new_entities,
             schema.as_ref(),