Implementation of RFC 34 (#430)

Co-authored-by: Kesha Hietala <[email protected]>

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

@@ -15,13 +15,17 @@
  */
 
 use crate::ast::*;
+use crate::evaluator::{EvaluationError, RestrictedEvaluator};
+use crate::extensions::Extensions;
 use crate::parser::err::ParseErrors;
 use crate::transitive_closure::TCNode;
 use crate::FromNormalizedStr;
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
+use serde_with::{serde_as, TryFromInto};
 use smol_str::SmolStr;
 use std::collections::{HashMap, HashSet};
+use thiserror::Error;
 
 /// We support two types of entities. The first is a nominal type (e.g., User, Action)
 /// and the second is an unspecified type, which is used (internally) to represent cases
@@ -216,8 +220,8 @@ pub struct Entity {
     /// Internal HashMap of attributes.
     ///
     /// In the serialized form of `Entity`, attribute values appear as
-    /// `RestrictedExpr`s.
-    attrs: HashMap<SmolStr, RestrictedExpr>,
+    /// `RestrictedExpr`s, for mostly historical reasons.
+    attrs: HashMap<SmolStr, PartialValueSerializedAsExpr>,
 
     /// Set of ancestors of this `Entity` (i.e., all direct and transitive
     /// parents), as UIDs
@@ -230,6 +234,53 @@ impl Entity {
         uid: EntityUID,
         attrs: HashMap<SmolStr, RestrictedExpr>,
         ancestors: HashSet<EntityUID>,
+        extensions: &Extensions<'_>,
+    ) -> Result<Self, EntityAttrEvaluationError> {
+        let evaluator = RestrictedEvaluator::new(extensions);
+        let evaluated_attrs = attrs
+            .into_iter()
+            .map(|(k, v)| {
+                let attr_val = evaluator
+                    .partial_interpret(v.as_borrowed())
+                    .map_err(|err| EntityAttrEvaluationError {
+                        uid: uid.clone(),
+                        attr: k.clone(),
+                        err,
+                    })?;
+                Ok((k, attr_val.into()))
+            })
+            .collect::<Result<_, EntityAttrEvaluationError>>()?;
+        Ok(Entity {
+            uid,
+            attrs: evaluated_attrs,
+            ancestors,
+        })
+    }
+
+    /// Create a new `Entity` with this UID, attributes, and ancestors.
+    ///
+    /// Unlike in `Entity::new()`, in this constructor, attributes are expressed
+    /// as `PartialValue`.
+    pub fn new_with_attr_partial_value(
+        uid: EntityUID,
+        attrs: HashMap<SmolStr, PartialValue>,
+        ancestors: HashSet<EntityUID>,
+    ) -> Self {
+        Entity {
+            uid,
+            attrs: attrs.into_iter().map(|(k, v)| (k, v.into())).collect(), // TODO: can we do this without disassembling and reassembling the HashMap
+            ancestors,
+        }
+    }
+
+    /// Create a new `Entity` with this UID, attributes, and ancestors.
+    ///
+    /// Unlike in `Entity::new()`, in this constructor, attributes are expressed
+    /// as `PartialValueSerializedAsExpr`.
+    pub fn new_with_attr_partial_value_serialized_as_expr(
+        uid: EntityUID,
+        attrs: HashMap<SmolStr, PartialValueSerializedAsExpr>,
+        ancestors: HashSet<EntityUID>,
     ) -> Self {
         Entity {
             uid,
@@ -244,8 +295,8 @@ impl Entity {
     }
 
     /// Get the value for the given attribute, or `None` if not present
-    pub fn get(&self, attr: &str) -> Option<&RestrictedExpr> {
-        self.attrs.get(attr)
+    pub fn get(&self, attr: &str) -> Option<&PartialValue> {
+        self.attrs.get(attr).map(|v| v.as_ref())
     }
 
     /// Is this `Entity` a descendant of `e` in the entity hierarchy?
@@ -259,10 +310,8 @@ impl Entity {
     }
 
     /// Iterate over this entity's attributes
-    pub fn attrs(&self) -> impl Iterator<Item = (&str, BorrowedRestrictedExpr<'_>)> {
-        self.attrs
-            .iter()
-            .map(|(k, v)| (k.as_str(), v.as_borrowed()))
+    pub fn attrs(&self) -> impl Iterator<Item = (&SmolStr, &PartialValue)> {
+        self.attrs.iter().map(|(k, v)| (k, v.as_ref()))
     }
 
     /// Create an `Entity` with the given UID, no attributes, and no parents.
@@ -274,12 +323,6 @@ impl Entity {
         }
     }
 
-    /// Read-only access the internal `attrs` map of String to RestrictedExpr.
-    /// This function is available only inside Core.
-    pub(crate) fn attrs_map(&self) -> &HashMap<SmolStr, RestrictedExpr> {
-        &self.attrs
-    }
-
     /// Test if two `Entity` objects are deep/structurally equal.
     /// That is, not only do they have the same UID, but also the same
     /// attributes, attribute values, and ancestors.
@@ -290,8 +333,15 @@ impl Entity {
     /// Set the given attribute to the given value.
     // Only used for convenience in some tests and when fuzzing
     #[cfg(any(test, fuzzing))]
-    pub fn set_attr(&mut self, attr: SmolStr, val: RestrictedExpr) {
-        self.attrs.insert(attr, val);
+    pub fn set_attr(
+        &mut self,
+        attr: SmolStr,
+        val: RestrictedExpr,
+        extensions: &Extensions<'_>,
+    ) -> Result<(), EvaluationError> {
+        let val = RestrictedEvaluator::new(extensions).partial_interpret(val.as_borrowed())?;
+        self.attrs.insert(attr, val.into());
+        Ok(())
     }
 
     /// Mark the given `UID` as an ancestor of this `Entity`.
@@ -354,6 +404,61 @@ impl std::fmt::Display for Entity {
     }
 }
 
+/// `PartialValue`, but serialized as a `RestrictedExpr`.
+///
+/// (Extension values can't be directly serialized, but can be serialized as
+/// `RestrictedExpr`)
+#[serde_as]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+pub struct PartialValueSerializedAsExpr(
+    #[serde_as(as = "TryFromInto<RestrictedExpr>")] PartialValue,
+);
+
+impl AsRef<PartialValue> for PartialValueSerializedAsExpr {
+    fn as_ref(&self) -> &PartialValue {
+        &self.0
+    }
+}
+
+impl std::ops::Deref for PartialValueSerializedAsExpr {
+    type Target = PartialValue;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl From<PartialValue> for PartialValueSerializedAsExpr {
+    fn from(value: PartialValue) -> PartialValueSerializedAsExpr {
+        PartialValueSerializedAsExpr(value)
+    }
+}
+
+impl From<PartialValueSerializedAsExpr> for PartialValue {
+    fn from(value: PartialValueSerializedAsExpr) -> PartialValue {
+        value.0
+    }
+}
+
+impl std::fmt::Display for PartialValueSerializedAsExpr {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+/// Error type for evaluation errors when evaluating an entity attribute.
+/// Contains some extra contextual information and the underlying
+/// `EvaluationError`.
+#[derive(Debug, Error)]
+#[error("failed to evaluate attribute `{attr}` of `{uid}`: {err}")]
+pub struct EntityAttrEvaluationError {
+    /// UID of the entity where the error was encountered
+    pub uid: EntityUID,
+    /// Attribute of the entity where the error was encountered
+    pub attr: SmolStr,
+    /// Underlying evaluation error
+    pub err: EvaluationError,
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

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

@@ -175,6 +175,15 @@ impl From<Value> for Expr {
     }
 }
 
+impl From<PartialValue> for Expr {
+    fn from(pv: PartialValue) -> Self {
+        match pv {
+            PartialValue::Value(v) => Expr::from(v),
+            PartialValue::Residual(expr) => expr,
+        }
+    }
+}
+
 impl<T> Expr<T> {
     fn new(expr_kind: ExprKind<T>, source_info: Option<SourceInfo>, data: T) -> Self {
         Self {
@@ -643,8 +652,8 @@ impl Unknown {
 
 impl std::fmt::Display for Unknown {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        // Like the Display impl for Expr, we delegate to the EST
-        // pretty-printer, to avoid code duplication
+        // Like the Display impl for Expr, we delegate to the EST pretty-printer,
+        // to avoid code duplication
         write!(f, "{}", crate::est::Expr::from(Expr::unknown(self.clone())))
     }
 }
@@ -936,10 +945,14 @@ impl<T> ExprBuilder<T> {
 
     /// Create an `Expr` which calls the extension function with the given
     /// `Name` on `args`
-    pub fn call_extension_fn(self, fn_name: Name, args: Vec<Expr<T>>) -> Expr<T> {
+    pub fn call_extension_fn(
+        self,
+        fn_name: Name,
+        args: impl IntoIterator<Item = Expr<T>>,
+    ) -> Expr<T> {
         self.with_expr_kind(ExprKind::ExtensionFunctionApp {
             fn_name,
-            args: Arc::new(args),
+            args: Arc::new(args.into_iter().collect()),
         })
     }
 

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

@@ -44,7 +44,7 @@ impl<'a, T> Iterator for ExprIterator<'a, T> {
         let next_expr = self.expression_stack.pop()?;
         match next_expr.expr_kind() {
             ExprKind::Lit(_) => (),
-            ExprKind::Unknown { .. } => (),
+            ExprKind::Unknown(_) => (),
             ExprKind::Slot(_) => (),
             ExprKind::Var(_) => (),
             ExprKind::If {

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

@@ -344,15 +344,32 @@ impl<V: ExtensionValue> StaticallyTyped for V {
 }
 
 #[derive(Debug, Clone)]
-/// Object container for extension values, also stores the fully reduced AST
-/// for the arguments
+/// Object container for extension values, also stores the constructor-and-args
+/// that can reproduce the value (important for converting the value back to
+/// `RestrictedExpr` for instance)
 pub struct ExtensionValueWithArgs {
     value: Arc<dyn InternalExtensionValue>,
-    args: Vec<Expr>,
-    constructor: Name,
+    pub(crate) constructor: Name,
+    /// Args are stored in `RestrictedExpr` form, just because that's most
+    /// convenient for reconstructing a `RestrictedExpr` that reproduces this
+    /// extension value
+    pub(crate) args: Vec<RestrictedExpr>,
 }
 
 impl ExtensionValueWithArgs {
+    /// Create a new `ExtensionValueWithArgs`
+    pub fn new(
+        value: Arc<dyn InternalExtensionValue + Send + Sync>,
+        constructor: Name,
+        args: Vec<RestrictedExpr>,
+    ) -> Self {
+        Self {
+            value,
+            constructor,
+            args,
+        }
+    }
+
     /// Get the internal value
     pub fn value(&self) -> &(dyn InternalExtensionValue) {
         self.value.as_ref()
@@ -363,23 +380,15 @@ impl ExtensionValueWithArgs {
         self.value.typename()
     }
 
-    /// Constructor
-    pub fn new(
-        value: Arc<dyn InternalExtensionValue + Send + Sync>,
-        args: Vec<Expr>,
-        constructor: Name,
-    ) -> Self {
-        Self {
-            value,
-            args,
-            constructor,
-        }
+    /// Get the constructor and args that can reproduce this value
+    pub fn constructor_and_args(&self) -> (&Name, &[RestrictedExpr]) {
+        (&self.constructor, &self.args)
     }
 }
 
 impl From<ExtensionValueWithArgs> for Expr {
     fn from(val: ExtensionValueWithArgs) -> Self {
-        ExprBuilder::new().call_extension_fn(val.constructor, val.args)
+        ExprBuilder::new().call_extension_fn(val.constructor, val.args.into_iter().map(Into::into))
     }
 }