Refactor TPE (#1822)

Signed-off-by: Shaobo He <[email protected]>

Author: shaobo-he-aws

cedar-policy-core/src/tpe/batched_evaluator.rs renamed to cedar-policy-core/src/batched_evaluator.rs

@@ -16,23 +16,22 @@
 
 //! This module contains the batched evaluator implementation and the (internal) definition of [`EntityLoader`]
 
+pub mod err;
+
 use std::collections::{HashMap, HashSet};
 use std::iter;
 use std::sync::Arc;
 
-use crate::ast::{Entity, EntityUID, EntityUIDEntry, Expr, PolicyID, Request};
+use crate::ast::{Entity, EntityUID, EntityUIDEntry, Request};
 use crate::authorizer::Decision;
+use crate::batched_evaluator::err::{BatchedEvalError, InsufficientIterationsError};
 use crate::entities::TCComputation;
 use crate::tpe::entities::PartialEntity;
-use crate::tpe::err::{
-    BatchedEvalError, InsufficientIterationsError, NonstaticPolicyError, PartialRequestError,
-    TPEError,
-};
+use crate::tpe::err::{PartialRequestError, TPEError};
+use crate::tpe::policy_expr_map;
 use crate::tpe::request::{PartialEntityUID, PartialRequest};
 use crate::tpe::residual::Residual;
 use crate::tpe::response::{ResidualPolicy, Response};
-use crate::validator::typecheck::{PolicyCheck, Typechecker};
-use crate::validator::types::Type;
 use crate::validator::ValidatorSchema;
 use crate::{ast::PolicySet, extensions::Extensions};
 
@@ -47,38 +46,6 @@ pub trait EntityLoader {
     fn load_entities(&mut self, uids: &HashSet<EntityUID>) -> HashMap<EntityUID, Option<Entity>>;
 }
 
-pub(crate) fn policy_expr_map<'a>(
-    request: &'a PartialRequest,
-    ps: &'a PolicySet,
-    schema: &ValidatorSchema,
-) -> std::result::Result<HashMap<&'a PolicyID, Expr<Option<Type>>>, TPEError> {
-    let mut exprs = HashMap::new();
-    let tc = Typechecker::new(schema, crate::validator::ValidationMode::Strict);
-    let env = request.find_request_env(schema)?;
-    for p in ps.policies() {
-        if !p.is_static() {
-            return Err(NonstaticPolicyError.into());
-        }
-        let t = p.template();
-        match tc.typecheck_by_single_request_env(t, &env) {
-            PolicyCheck::Success(expr) => {
-                exprs.insert(p.id(), expr);
-            }
-            PolicyCheck::Fail(errs) => {
-                return Err(TPEError::Validation(errs));
-            }
-            PolicyCheck::Irrelevant(errs, expr) => {
-                if errs.is_empty() {
-                    exprs.insert(p.id(), expr);
-                } else {
-                    return Err(TPEError::Validation(errs));
-                }
-            }
-        }
-    }
-    Ok(exprs)
-}
-
 fn concrete_request_to_partial(
     request: &Request,
     schema: &ValidatorSchema,

cedar-policy-core/src/batched_evaluator/err.rs

@@ -0,0 +1,56 @@
+/*
+ * Copyright Cedar Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//! This module contains possible errors thrown by batched evaluation
+
+use thiserror::Error;
+
+use crate::ast::PartialValueToValueError;
+use crate::tpe::err::{EntitiesError, MissingEntitiesError, PartialRequestError, TPEError};
+use crate::validator::RequestValidationError;
+
+/// Errors for Batched Evaluation
+#[derive(Debug, Error)]
+#[non_exhaustive]
+pub enum BatchedEvalError {
+    /// Error thrown by TPE
+    #[error(transparent)]
+    TPE(#[from] TPEError),
+    /// Error when the request is not valid
+    #[error(transparent)]
+    RequestValidation(#[from] RequestValidationError),
+    /// Error when the request is partial
+    #[error(transparent)]
+    PartialRequest(#[from] PartialRequestError),
+    /// Error when the loaded entities are not valid
+    #[error(transparent)]
+    Entities(#[from] EntitiesError),
+    /// Error thrown when a entity loader provided entity was partial instead of fully concrete
+    #[error(transparent)]
+    PartialValueToValue(#[from] PartialValueToValueError),
+    /// Error the entity loader failed to load all requested entities
+    #[error(transparent)]
+    MissingEntities(#[from] MissingEntitiesError),
+    /// Error when batched evaluation did not converge due to the iteration limit
+    #[error(transparent)]
+    InsufficientIterations(#[from] InsufficientIterationsError),
+}
+
+/// Batched evaluation may not return an answer when the maximum
+/// iterations is too low.
+#[derive(Debug, Error)]
+#[error("Batched evaluation failed: insufficient iteration limit.")]
+pub struct InsufficientIterationsError {}

cedar-policy-core/src/lib.rs

@@ -28,6 +28,8 @@ pub use from_normalized_str::*;
 pub mod entities;
 #[macro_use]
 mod error_macros;
+#[cfg(feature = "tpe")]
+pub mod batched_evaluator;
 pub mod est;
 pub mod evaluator;
 pub mod expr_builder;

cedar-policy-core/src/tpe.rs

@@ -16,24 +16,56 @@
 
 //! This module contains the type-aware partial evaluator.
 
-pub mod batched_evaluator;
 pub mod entities;
 pub mod err;
 pub mod evaluator;
 pub mod request;
 pub mod residual;
 pub mod response;
 
-use std::sync::Arc;
+use std::{collections::HashMap, sync::Arc};
 
-use crate::tpe::batched_evaluator::policy_expr_map;
-use crate::tpe::err::TPEError;
+use crate::ast::{Expr, PolicyID};
+use crate::tpe::err::{NonstaticPolicyError, TPEError};
 use crate::tpe::response::{ResidualPolicy, Response};
-use crate::validator::ValidatorSchema;
+use crate::validator::types::Type;
+use crate::validator::{typecheck::PolicyCheck, typecheck::Typechecker, ValidatorSchema};
 use crate::{ast::PolicySet, extensions::Extensions};
 
 use crate::tpe::{entities::PartialEntities, evaluator::Evaluator, request::PartialRequest};
 
+pub(crate) fn policy_expr_map<'a>(
+    request: &'a PartialRequest,
+    ps: &'a PolicySet,
+    schema: &ValidatorSchema,
+) -> std::result::Result<HashMap<&'a PolicyID, Expr<Option<Type>>>, TPEError> {
+    let mut exprs = HashMap::new();
+    let tc = Typechecker::new(schema, crate::validator::ValidationMode::Strict);
+    let env = request.find_request_env(schema)?;
+    for p in ps.policies() {
+        if !p.is_static() {
+            return Err(NonstaticPolicyError.into());
+        }
+        let t = p.template();
+        match tc.typecheck_by_single_request_env(t, &env) {
+            PolicyCheck::Success(expr) => {
+                exprs.insert(p.id(), expr);
+            }
+            PolicyCheck::Fail(errs) => {
+                return Err(TPEError::Validation(errs));
+            }
+            PolicyCheck::Irrelevant(errs, expr) => {
+                if errs.is_empty() {
+                    exprs.insert(p.id(), expr);
+                } else {
+                    return Err(TPEError::Validation(errs));
+                }
+            }
+        }
+    }
+    Ok(exprs)
+}
+
 /// Type-aware partial-evaluation on a `PolicySet`.
 /// Both `request` and `entities` should be valid and hence be constructed
 /// using their safe constructors.

cedar-policy-core/src/tpe/err.rs

@@ -23,7 +23,7 @@ use smol_str::SmolStr;
 use thiserror::Error;
 
 use crate::{
-    ast::{Eid, EntityType, EntityUID, PartialValueToValueError},
+    ast::{Eid, EntityType, EntityUID},
     entities::{
         conformance::err::{EntitySchemaConformanceError, InvalidEnumEntityError},
         err::Duplicate,
@@ -109,39 +109,6 @@ pub enum TPEError {
     ExprToResidualError(#[from] ExprToResidualError),
 }
 
-/// Errors for Batched Evaluation
-#[derive(Debug, Error)]
-#[non_exhaustive]
-pub enum BatchedEvalError {
-    /// Error thrown by TPE
-    #[error(transparent)]
-    TPE(#[from] TPEError),
-    /// Error when the request is not valid
-    #[error(transparent)]
-    RequestValidation(#[from] RequestValidationError),
-    /// Error when the request is partial
-    #[error(transparent)]
-    PartialRequest(#[from] PartialRequestError),
-    /// Error when the loaded entities are not valid
-    #[error(transparent)]
-    Entities(#[from] EntitiesError),
-    /// Error thrown when a entity loader provided entity was partial instead of fully concrete
-    #[error(transparent)]
-    PartialValueToValue(#[from] PartialValueToValueError),
-    /// Error the entity loader failed to load all requested entities
-    #[error(transparent)]
-    MissingEntities(#[from] MissingEntitiesError),
-    /// Error when batched evaluation did not converge due to the iteration limit
-    #[error(transparent)]
-    InsufficientIterations(#[from] InsufficientIterationsError),
-}
-
-/// Batched evaluation may not return an answer when the maximum
-/// iterations is too low.
-#[derive(Debug, Error)]
-#[error("Batched evaluation failed: insufficient iteration limit.")]
-pub struct InsufficientIterationsError {}
-
 /// Residuals require fully typed expressions without
 /// unknowns or parse errors.
 #[derive(Debug, Error)]

cedar-policy/src/api.rs

@@ -5117,11 +5117,11 @@ mod tpe {
 
     use cedar_policy_core::ast;
     use cedar_policy_core::authorizer::Decision;
-    use cedar_policy_core::tpe;
-    use cedar_policy_core::tpe::batched_evaluator::is_authorized_batched;
-    use cedar_policy_core::tpe::{
-        batched_evaluator::EntityLoader as EntityLoaderInternal, err::BatchedEvalError,
+    use cedar_policy_core::batched_evaluator::is_authorized_batched;
+    use cedar_policy_core::batched_evaluator::{
+        err::BatchedEvalError, EntityLoader as EntityLoaderInternal,
     };
+    use cedar_policy_core::tpe;
     use cedar_policy_core::{
         entities::conformance::EntitySchemaConformanceChecker, extensions::Extensions,
         validator::CoreSchema,
@@ -5154,8 +5154,8 @@ mod tpe {
     }
 
     /// A partial [`Request`]
-    /// Its principal and resource could have unknown [`EntityId`]; its action
-    /// must be know; and its context could be unknown
+    /// Its principal/resource types and action must be known and its context
+    /// must either be fully known or unknown
     #[repr(transparent)]
     #[derive(Debug, Clone, RefCast)]
     pub struct PartialRequest(pub(crate) tpe::request::PartialRequest);
@@ -5296,6 +5296,8 @@ mod tpe {
 
     impl PartialEntities {
         /// Construct [`PartialEntities`] from a JSON value
+        /// The `parent`, `attrs`, `tags` field must be either fully known or
+        /// unknown. And parent entities cannot have unknown parents.
         pub fn from_json_value(
             value: serde_json::Value,
             schema: &Schema,

cedar-policy/src/test/test.rs

@@ -8783,7 +8783,8 @@ unless
             str::FromStr,
         };
 
-        use cedar_policy_core::tpe::err::{BatchedEvalError, TPEError};
+        use cedar_policy_core::batched_evaluator::err::BatchedEvalError;
+        use cedar_policy_core::tpe::err::TPEError;
         use cool_asserts::assert_matches;
         use itertools::Itertools;