Repository: cedar-policy/cedar
Pull request #1480 (2025-02-18 21:19:09)

Manifest cherry pick
## Description of changes

Cherry-pick #1429 and #1462 fixes to entity manifest feature

## Issue #, if available

## Checklist for requesting a review

The change in this PR is (choose one, and delete the other options):

- [ ] A breaking change requiring a major version bump to `cedar-policy` (e.g., changes to the signature of an existing API).
- [ ] A backwards-compatible change requiring a minor version bump to `cedar-policy` (e.g., addition of a new API).
- [ ] A bug fix or other functionality change requiring a patch to `cedar-policy`.
- [ ] A change "invisible" to users (e.g., documentation, changes to "internal" crates like `cedar-policy-core`, `cedar-validator`, etc.)
- [ ] A change (breaking or otherwise) that only impacts unreleased or experimental code.

I confirm that this PR (choose one, and delete the other options):

- [ ] Updates the "Unreleased" section of the CHANGELOG with a description of my change (required for major/minor version bumps).
- [ ] Does not update the CHANGELOG because my change does not significantly impact released code.

I confirm that [`cedar-spec`](https://github.com/cedar-policy/cedar-spec) (choose one, and delete the other options):

- [ ] Does not require updates because my change does not impact the Cedar formal model or DRT infrastructure.
- [ ] Requires updates, and I have made / will make these updates myself. (Please include in your description a timeline or link to the relevant PR in `cedar-spec`, and how you have tested that your updates are correct.)
- [ ] Requires updates, but I do not plan to make them in the near future. (Make sure that your changes are hidden behind a feature flag to mark them as experimental.)
- [ ] I'm not sure how my change impacts `cedar-spec`. (Post your PR anyways, and we'll discuss in the comments.)

I confirm that [`docs.cedarpolicy.com`](https://docs.cedarpolicy.com/) (choose one, and delete the other options):

- [ ] Does not require updates because my change does not impact the Cedar language specification.
- [ ] Requires updates, and I have made / will make these updates myself. (Please include in your description a timeline or link to the relevant PR in [`cedar-docs`](https://github.com/cedar-policy/cedar-docs). PRs should be targeted at a `staging-X.Y` branch, not `main`.)
- [ ] I'm not sure how my change impacts the documentation. (Post your PR anyways, and we'll discuss in the comments.)


john-h-kastner-aws: > Didn't review in detail but I assume this is a clean cherry-pick

Yes. The only merge conflict was in the changelog 


Patch:
diff --git a/cedar-policy-validator/src/entity_manifest/loader.rs b/cedar-policy-validator/src/entity_manifest/loader.rs
index bafc2e1c..e575d6b5 100644
--- a/cedar-policy-validator/src/entity_manifest/loader.rs
+++ b/cedar-policy-validator/src/entity_manifest/loader.rs
@@ -18,7 +18,7 @@
 //! Loads entities based on the entity manifest.
 
 use std::{
-    collections::{BTreeMap, HashMap, HashSet},
+    collections::{btree_map, hash_map, BTreeMap, HashMap, HashSet},
     sync::Arc,
 };
 
@@ -258,7 +258,20 @@ pub(crate) fn load_entities(
                     entity_request.access_trie,
                     &mut Default::default(),
                 )?);
-                entities.insert(entity_request.entity_id, loaded);
+                match entities.entry(entity_request.entity_id) {
+                    hash_map::Entry::Occupied(o) => {
+                        // If the entity is already present in the slice, then
+                        // we need to be careful not to clobber its existing
+                        // attributes.  This can happen when an entity is
+                        // referenced by both an entity literal and a variable.
+                        let (k, v) = o.remove_entry();
+                        let merged = merge_entities(v, loaded);
+                        entities.insert(k, merged);
+                    }
+                    hash_map::Entry::Vacant(v) => {
+                        v.insert(loaded);
+                    }
+                }
             }
         }
 
@@ -299,6 +312,110 @@ pub(crate) fn load_entities(
     }
 }
 
+/// Merge the contents of two entities in the slice. Combines the attributes
+/// records for both entities, recursively merging any attribute that exist in
+/// both. If one entity is referenced by multiple entity roots in the slice,
+/// then we need to be sure that we don't clobber the attribute for the first
+/// when inserting the second into the slice.
+// INVARIANT: `e1` and `e2` must be the result of slicing the same original
+// entity using the same entity manifest and request. I.e., they may differ only in
+// what attributes they contain. When an attribute exists in both, the
+// attributes may differ only if they are records, and then only in what nested
+// attributes they contain.
+fn merge_entities(e1: Entity, e2: Entity) -> Entity {
+    let (uid1, mut attrs1, ancestors1, tags1) = e1.into_inner();
+    let (uid2, attrs2, ancestors2, tags2) = e2.into_inner();
+
+    assert_eq!(
+        uid1, uid2,
+        "attempting to merge entities with different uids!"
+    );
+    assert_eq!(
+        ancestors1, ancestors2,
+        "attempting to merge entities with different ancestors!"
+    );
+    assert!(
+        tags1.is_empty() && tags2.is_empty(),
+        "attempting to merge entities with tags!"
+    );
+
+    for (k, v2) in attrs2 {
+        match attrs1.entry(k) {
+            hash_map::Entry::Occupied(occupied) => {
+                let (k, v1) = occupied.remove_entry();
+                match (v1, v2) {
+                    (PartialValue::Value(v1), PartialValue::Value(v2)) => {
+                        let merged_v = merge_values(v1, v2);
+                        attrs1.insert(k, PartialValue::Value(merged_v));
+                    }
+                    (PartialValue::Residual(e1), PartialValue::Residual(e2)) => {
+                        assert_eq!(e1, e2, "attempting to merge different residuals!");
+                        attrs1.insert(k, PartialValue::Residual(e1));
+                    }
+                    // PANIC SAFETY: We're merging sliced copies of the same entity, so the attribute must be the same
+                    #[allow(clippy::panic)]
+                    (PartialValue::Value(_), PartialValue::Residual(_))
+                    | (PartialValue::Residual(_), PartialValue::Value(_)) => {
+                        panic!("attempting to merge a value with a residual")
+                    }
+                };
+            }
+            hash_map::Entry::Vacant(vacant) => {
+                vacant.insert(v2);
+            }
+        }
+    }
+
+    Entity::new_with_attr_partial_value(uid1, attrs1, ancestors1)
+}
+
+/// Merge two value for corresponding attributes in the slice.
+// INVARIANT: `v1` and `v2` must be the result of slicing the same original
+// value using the same entity manifest and request. I.e., they must be
+// identical, except for the attributes they contain when the values are a
+// records. When an attribute exists in both records, the attributes must be
+// recursively identical, with the same exception.
+fn merge_values(v1: Value, v2: Value) -> Value {
+    match (v1.value, v2.value) {
+        (ValueKind::Record(r1), ValueKind::Record(r2)) => {
+            let mut r1 = Arc::unwrap_or_clone(r1);
+            for (k, v2) in Arc::unwrap_or_clone(r2) {
+                match r1.entry(k) {
+                    btree_map::Entry::Occupied(occupied) => {
+                        let (k, v1) = occupied.remove_entry();
+                        let merged_v = merge_values(v1, v2);
+                        r1.insert(k, merged_v);
+                    }
+                    btree_map::Entry::Vacant(vacant) => {
+                        vacant.insert(v2);
+                    }
+                }
+            }
+            Value::new(ValueKind::Record(Arc::new(r1)), v1.loc)
+        }
+        (ValueKind::Lit(l1), ValueKind::Lit(l2)) => {
+            assert_eq!(l1, l2, "attempting to merge different literals!");
+            Value::new(l1, v1.loc)
+        }
+        (vk1 @ ValueKind::ExtensionValue(_), vk2 @ ValueKind::ExtensionValue(_))
+        | (vk1 @ ValueKind::Set(_), vk2 @ ValueKind::Set(_)) => {
+            // It might seem that we should recur into the sets and extensions
+            // values, but `AccessTrie::slice_val` doesn't, so the merge
+            // function can stop here too.
+            assert_eq!(
+                vk1, vk2,
+                "attempting to merge different sets or extensions!"
+            );
+            Value::new(vk1, v1.loc)
+        }
+        // PANIC SAFETY: We're merging sliced copies of the same entity, so the attribute must be the same
+        #[allow(clippy::panic)]
+        _ => {
+            panic!("attempting to merge values of different kinds!")
+        }
+    }
+}
+
 /// Given a context value and an access trie, find all of the remaining
 /// entities in the context.
 /// Also keep track of required ancestors when encountering the `is_ancestor` flag.
@@ -458,3 +575,60 @@ fn compute_ancestors_request(
         entity_id,
     })
 }
+
+#[cfg(test)]
+mod test {
+    use cedar_policy_core::ast::Value;
+    use smol_str::ToSmolStr;
+
+    use super::merge_values;
+
+    #[test]
+    fn test_merge_values() {
+        assert_eq!(
+            merge_values(Value::new(1, None), Value::new(1, None)),
+            Value::new(1, None),
+        );
+        assert_eq!(
+            merge_values(
+                Value::set([Value::new(1, None), Value::new(2, None)], None),
+                Value::set([Value::new(1, None), Value::new(2, None)], None),
+            ),
+            Value::set([Value::new(1, None), Value::new(2, None)], None),
+        );
+        assert_eq!(
+            merge_values(
+                Value::record([("a".to_smolstr(), Value::new(1, None))], None),
+                Value::record([("a".to_smolstr(), Value::new(1, None))], None),
+            ),
+            Value::record([("a".to_smolstr(), Value::new(1, None))], None),
+        );
+        assert_eq!(
+            merge_values(
+                Value::empty_record(None),
+                Value::record([("a".to_smolstr(), Value::new(1, None))], None),
+            ),
+            Value::record([("a".to_smolstr(), Value::new(1, None))], None),
+        );
+        assert_eq!(
+            merge_values(
+                Value::record([("a".to_smolstr(), Value::new(1, None))], None),
+                Value::empty_record(None),
+            ),
+            Value::record([("a".to_smolstr(), Value::new(1, None))], None),
+        );
+        assert_eq!(
+            merge_values(
+                Value::record([("a".to_smolstr(), Value::new(1, None))], None),
+                Value::record([("b".to_smolstr(), Value::new(2, None))], None),
+            ),
+            Value::record(
+                [
+                    ("a".to_smolstr(), Value::new(1, None)),
+                    ("b".to_smolstr(), Value::new(2, None))
+                ],
+                None
+            ),
+        );
+    }
+}
diff --git a/cedar-policy-validator/src/entity_manifest/slicing.rs b/cedar-policy-validator/src/entity_manifest/slicing.rs
index 6d4b5577..adcc5a9b 100644
--- a/cedar-policy-validator/src/entity_manifest/slicing.rs
+++ b/cedar-policy-validator/src/entity_manifest/slicing.rs
@@ -244,13 +244,14 @@ impl AccessTrie {
 
 #[cfg(test)]
 mod entity_slice_tests {
+    use similar_asserts::assert_eq;
     use std::collections::BTreeSet;
 
     use cedar_policy_core::{
         ast::{Context, PolicyID, PolicySet},
         entities::{EntityJsonParser, TCComputation},
         extensions::Extensions,
-        parser::parse_policy,
+        parser::{self, parse_policy},
     };
 
     use crate::{entity_manifest::compute_entity_manifest, CoreSchema, ValidatorSchema};
@@ -336,6 +337,7 @@ action Read appliesTo {
         .0
     }
 
+    #[track_caller]
     fn expect_entity_slice_to(
         original: serde_json::Value,
         expected: serde_json::Value,
@@ -925,4 +927,262 @@ action BeSad appliesTo {
         let entity_manifest = compute_entity_manifest(&schema, &pset).expect("Should succeed");
         assert_eq!(entity_manifest, entity_manifest);
     }
+
+    #[test]
+    fn test_slice_with_entity_alias() {
+        let schema = schema();
+        let entities_json = serde_json::json!([{
+            "uid" : { "type" : "User", "id" : "oliver"},
+            "parents": [],
+            "attrs": { "name": "oliver" }
+        }]);
+
+        // Only lit is accessed
+        let pset = parser::parse_policyset(
+            r#"permit(principal in User::"oliver", action, resource) when { User::"oliver".name == "oliver" };"#,
+        ).unwrap();
+        let manifest = compute_entity_manifest(&schema, &pset).unwrap();
+        expect_entity_slice_to(
+            entities_json.clone(),
+            entities_json.clone(),
+            &schema,
+            &manifest,
+        );
+
+        // Only var is accessed
+        let pset = parser::parse_policyset(
+            r#"permit(principal in User::"oliver", action, resource) when { principal.name == "oliver" };"#,
+        ).unwrap();
+        let manifest = compute_entity_manifest(&schema, &pset).unwrap();
+        expect_entity_slice_to(
+            entities_json.clone(),
+            entities_json.clone(),
+            &schema,
+            &manifest,
+        );
+
+        // Both are accessed
+        let pset = parser::parse_policyset(
+            r#"permit(principal in User::"oliver", action, resource) when { principal.name == User::"oliver".name };"#,
+        ).unwrap();
+        let manifest = compute_entity_manifest(&schema, &pset).unwrap();
+        expect_entity_slice_to(
+            entities_json.clone(),
+            entities_json.clone(),
+            &schema,
+            &manifest,
+        );
+    }
+
+    #[test]
+    fn test_slice_in_const_true_guard() {
+        let schema = ValidatorSchema::from_cedarschema_str(
+            "
+entity User = {
+  foo: String,
+  bar: String,
+  baz: String,
+};
+
+entity Document;
+
+action Read appliesTo {
+  principal: [User],
+  resource: [Document]
+};",
+            Extensions::all_available(),
+        )
+        .unwrap()
+        .0;
+
+        let pset = parser::parse_policyset(
+            r#"permit(principal, action, resource) when {
+                if (principal.foo == "foo") || true then
+                    principal.bar == "bar"
+                else
+                    principal.baz == "baz"
+            };"#,
+        )
+        .unwrap();
+        let manifest = compute_entity_manifest(&schema, &pset).unwrap();
+        let entities_json = serde_json::json!([{
+            "uid" : { "type" : "User", "id" : "oliver"},
+            "parents": [],
+            "attrs": {
+                "foo": "foo",
+                "bar": "bar",
+                "baz": "baz",
+            }
+        }]);
+        // We need `foo` to evaluate the guard even though we know it's always
+        // `true`. We need `bar` to evaluate the `then` branch. We don't need
+        // `baz` because we'll never evaluate the `else` branch.
+        let expected_json = serde_json::json!([{
+            "uid" : { "type" : "User", "id" : "oliver"},
+            "parents": [],
+            "attrs": {
+                "foo": "foo",
+                "bar": "bar",
+            }
+        }]);
+        expect_entity_slice_to(
+            entities_json.clone(),
+            expected_json.clone(),
+            &schema,
+            &manifest,
+        );
+    }
+
+    #[test]
+    fn test_slice_with_entity_alias_with_two_attrs() {
+        let schema = ValidatorSchema::from_cedarschema_str(
+            "
+entity User = {
+  foo: String,
+  bar: String,
+};
+
+entity Document;
+
+action Read appliesTo {
+  principal: [User],
+  resource: [Document]
+};",
+            Extensions::all_available(),
+        )
+        .unwrap()
+        .0;
+
+        let entities_json = serde_json::json!([{
+            "uid" : { "type" : "User", "id" : "oliver"},
+            "parents": [],
+            "attrs": {
+                "foo": "1",
+                "bar": "1",
+            }
+        }]);
+
+        // principal and literal alias, but access different attrs. Slice needs to take both
+        let pset = parser::parse_policyset(
+            r#"permit(principal in User::"oliver", action, resource) when { principal.foo == User::"oliver".bar };"#,
+        ).unwrap();
+        let manifest = compute_entity_manifest(&schema, &pset).unwrap();
+        expect_entity_slice_to(
+            entities_json.clone(),
+            entities_json.clone(),
+            &schema,
+            &manifest,
+        );
+    }
+
+    #[test]
+    fn test_slice_with_entity_alias_with_nested_record() {
+        let schema = ValidatorSchema::from_cedarschema_str(
+            "
+entity User = {
+  foo: {
+    bar: String,
+    baz: String,
+  },
+};
+
+entity Document;
+
+action Read appliesTo {
+  principal: [User],
+  resource: [Document]
+};",
+            Extensions::all_available(),
+        )
+        .unwrap()
+        .0;
+
+        let entities_json = serde_json::json!([{
+            "uid" : { "type" : "User", "id" : "oliver"},
+            "parents": [],
+            "attrs": { "foo": {
+                "bar": "1",
+                "baz": "1",
+            }}
+        }]);
+
+        // principal and literal alias, but access different attrs of a nested record. Slice needs to take both
+        let pset = parser::parse_policyset(
+            r#"permit(principal in User::"oliver", action, resource) when { principal.foo.bar == User::"oliver".foo.baz };"#,
+        ).unwrap();
+        let manifest = compute_entity_manifest(&schema, &pset).unwrap();
+        expect_entity_slice_to(
+            entities_json.clone(),
+            entities_json.clone(),
+            &schema,
+            &manifest,
+        );
+    }
+
+    #[test]
+    fn test_slice_with_entity_alias_ancestors() {
+        let schema = ValidatorSchema::from_cedarschema_str(
+            "
+entity User in Group = {
+  name: String
+};
+
+entity Group;
+
+entity Document;
+
+action Read appliesTo {
+  principal: [User],
+  resource: [Document]
+};",
+            Extensions::all_available(),
+        )
+        .unwrap()
+        .0;
+        let entities_json = serde_json::json!([{
+            "uid" : { "type" : "User", "id" : "oliver"},
+            "parents": [ { "type" : "Group", "id" : "oliver"}, ],
+            "attrs": { "name": "oliver" }
+        }]);
+
+        // The `principal` alias needs to load ancestors, but the lit alias does not. Slicing still needs to load ancestors
+        let pset = parser::parse_policyset(
+            r#"permit(principal in Group::"oliver", action, resource) when {User::"oliver".name == "oliver"};"#,
+        ).unwrap();
+        let manifest = compute_entity_manifest(&schema, &pset).unwrap();
+        expect_entity_slice_to(
+            entities_json.clone(),
+            entities_json.clone(),
+            &schema,
+            &manifest,
+        );
+
+        // Lit wants ancestors, `principal` does not
+        let pset = parser::parse_policyset(
+            r#"permit(principal, action, resource) when { User::"oliver" in Group::"oliver" && principal.name == "oliver"};"#,
+        ).unwrap();
+        let manifest = compute_entity_manifest(&schema, &pset).unwrap();
+        expect_entity_slice_to(
+            entities_json.clone(),
+            entities_json.clone(),
+            &schema,
+            &manifest,
+        );
+
+        // Both need ancestors
+        let pset = parser::parse_policyset(
+            r#"permit(principal, action, resource) when { User::"oliver" in Group::"oliver" && principal in Group::"oliver" };"#,
+        ).unwrap();
+        let manifest = compute_entity_manifest(&schema, &pset).unwrap();
+        expect_entity_slice_to(
+            entities_json,
+            serde_json::json!([{
+                "uid" : { "type" : "User", "id" : "oliver"},
+                "parents": [ { "type" : "Group", "id" : "oliver"}, ],
+                "attrs": {}
+            }]),
+            &schema,
+            &manifest,
+        );
+    }
 }
diff --git a/cedar-policy-validator/src/typecheck.rs b/cedar-policy-validator/src/typecheck.rs
index 9918f485..33b51493 100644
--- a/cedar-policy-validator/src/typecheck.rs
+++ b/cedar-policy-validator/src/typecheck.rs
@@ -497,7 +497,16 @@ impl<'a> SingleEnvTypechecker<'a> {
 
                         ans_then.then_typecheck(|typ_then, then_capability| {
                             TypecheckAnswer::success_with_capability(
-                                typ_then,
+                                ExprBuilder::with_data(typ_then.data().clone())
+                                    .with_same_source_loc(e)
+                                    .ite(
+                                        typ_test,
+                                        typ_then.clone(),
+                                        // The type of the test expression is `True`, so we know the `else` branch
+                                        // will never be evaluated. We still need to put something here, so we use
+                                        // a copy of the `then` branch.
+                                        typ_then,
+                                    ),
                                 // The output capability of the whole `if` expression also
                                 // needs to contain the capability of the condition.
                                 then_capability.union(&test_capability),
@@ -511,7 +520,19 @@ impl<'a> SingleEnvTypechecker<'a> {
                         let ans_else = self.typecheck(prior_capability, else_expr, type_errors);
 
                         ans_else.then_typecheck(|typ_else, else_capability| {
-                            TypecheckAnswer::success_with_capability(typ_else, else_capability)
+                            TypecheckAnswer::success_with_capability(
+                                ExprBuilder::with_data(typ_else.data().clone())
+                                    .with_same_source_loc(e)
+                                    .ite(
+                                        typ_test,
+                                        // The type of the test expression is `False`, so we know the `then` branch
+                                        // will never be evaluated. We still need to put something here, so we use
+                                        // a copy of the `else` branch.
+                                        typ_else.clone(),
+                                        typ_else,
+                                    ),
+                                else_capability,
+                            )
                         })
                     } else {
                         // When we don't short circuit, the `then` and `else`
@@ -575,15 +596,15 @@ impl<'a> SingleEnvTypechecker<'a> {
                 );
                 ans_left.then_typecheck(|typ_left, capability_left| {
                     match typ_left.data() {
-                        // First argument is false, so short circuit the `&&` to
-                        // false _without_ typechecking the second argument.
-                        // Since the type of the `&&` is `false`, it is known to
-                        // always evaluate to `false` at run time. The `&&`
-                        // expression typechecks with an empty capability rather
-                        // than the capability of the lhs.
-                        // The right operand is not typechecked, so it is not
-                        // included in the type annotated AST.
-                        Some(Type::False) => TypecheckAnswer::success(typ_left),
+                        // LHS argument is false, so short circuit the `&&` to `False` _without_
+                        // typechecking the RHS.  We still need to build an annotated `&&` with
+                        // some RHS, so we use a literal `true`. The `&&` expression typechecks
+                        // with an empty capability rather than the capability of the lhs.
+                        Some(Type::False) => TypecheckAnswer::success(
+                            ExprBuilder::with_data(Some(Type::False))
+                                .with_same_source_loc(e)
+                                .and(typ_left, ExprBuilder::new().val(true)),
+                        ),
                         _ => {
                             // Similar to the `then` branch of an `if`
                             // expression, the rhs of an `&&` is typechecked
@@ -670,14 +691,18 @@ impl<'a> SingleEnvTypechecker<'a> {
                     |_| None,
                 );
                 ans_left.then_typecheck(|ty_expr_left, capability_left| match ty_expr_left.data() {
-                    // Contrary to `&&` where short circuiting did not permit
-                    // any capability, a capability can be maintained when short
-                    // circuiting `||`. We know the left operand is `true`, so
-                    // its capability is maintained. The right operand is not
-                    // evaluated, so its capability does not need to be considered.
-                    // The right operand is not typechecked, so it is not
-                    // included in the type annotated AST.
-                    Some(Type::True) => TypecheckAnswer::success(ty_expr_left),
+                    // LHS argument is true, so short circuit the `|| to `True` _without_
+                    // typechecking the RHS. We still need to build an annotated `||` with
+                    // some RHS, so we use a literal `true`.  Contrary to `&&`, we keep a
+                    // capability  when short circuiting `||`. The left operand is `true`,
+                    // so its capability is maintained. The right operand is not evaluated,
+                    // so its capability is not considered.
+                    Some(Type::True) => TypecheckAnswer::success_with_capability(
+                        ExprBuilder::with_data(Some(Type::True))
+                            .with_same_source_loc(e)
+                            .or(ty_expr_left, ExprBuilder::new().val(true)),
+                        capability_left,
+                    ),
                     _ => {
                         // The right operand of an `||` cannot be typechecked
                         // using the capability learned from the left because the
@@ -1718,12 +1743,7 @@ impl<'a> SingleEnvTypechecker<'a> {
         context: LubContext,
     ) -> TypecheckAnswer<'b> {
         match annotated_expr.data() {
-            Some(Type::False) => {
-                TypecheckAnswer::success(ExprBuilder::with_data(Some(Type::False)).val(false))
-            }
-            Some(Type::True) => {
-                TypecheckAnswer::success(ExprBuilder::with_data(Some(Type::True)).val(true))
-            }
+            Some(Type::True | Type::False) => TypecheckAnswer::success(annotated_expr),
             _ => match (lhs_ty, rhs_ty) {
                 (Some(lhs_ty), Some(rhs_ty)) => {
                     if let Err(lub_hint) =
@@ -1952,12 +1972,8 @@ impl<'a> SingleEnvTypechecker<'a> {
                 .then_typecheck(|type_of_in, _| {
                     if !self.mode.is_strict() {
                         TypecheckAnswer::success(type_of_in)
-                    } else if matches!(type_of_in.data(), Some(Type::False)) {
-                        TypecheckAnswer::success(
-                            ExprBuilder::with_data(Some(Type::False)).val(false),
-                        )
-                    } else if matches!(type_of_in.data(), Some(Type::True)) {
-                        TypecheckAnswer::success(ExprBuilder::with_data(Some(Type::True)).val(true))
+                    } else if matches!(type_of_in.data(), Some(Type::True | Type::False)) {
+                        TypecheckAnswer::success(type_of_in)
                     } else {
                         match (lhs_ty, rhs_ty) {
                             (Some(lhs_ty), Some(rhs_ty)) => {


Test patch:
diff --git a/cedar-policy-validator/src/typecheck/test/type_annotation.rs b/cedar-policy-validator/src/typecheck/test/type_annotation.rs
index 752205cd..dd7f578d 100644
--- a/cedar-policy-validator/src/typecheck/test/type_annotation.rs
+++ b/cedar-policy-validator/src/typecheck/test/type_annotation.rs
@@ -16,6 +16,7 @@
 
 use cool_asserts::assert_matches;
 use serde_json::json;
+use similar_asserts::assert_eq;
 use std::collections::HashSet;
 
 use cedar_policy_core::{
@@ -68,6 +69,14 @@ fn expr_typechecks_with_correct_annotation() {
             ExprBuilder::with_data(Some(Type::singleton_boolean(false))).val(false),
         ),
     );
+    assert_expr_has_annotated_ast(
+        &Expr::and(Expr::val(false), Expr::val(1)),
+        &ExprBuilder::with_data(Some(Type::singleton_boolean(false))).and(
+            ExprBuilder::with_data(Some(Type::singleton_boolean(false))).val(false),
+            // Doesn't contain `1` because we didn't typecheck this side
+            ExprBuilder::new().val(true),
+        ),
+    );
     assert_expr_has_annotated_ast(
         &Expr::or(Expr::greater(Expr::val(1), Expr::val(1)), Expr::val(true)),
         &ExprBuilder::with_data(Some(Type::singleton_boolean(true))).or(
@@ -78,6 +87,14 @@ fn expr_typechecks_with_correct_annotation() {
             ExprBuilder::with_data(Some(Type::singleton_boolean(true))).val(true),
         ),
     );
+    assert_expr_has_annotated_ast(
+        &Expr::or(Expr::val(true), Expr::val(1)),
+        &ExprBuilder::with_data(Some(Type::singleton_boolean(true))).or(
+            ExprBuilder::with_data(Some(Type::singleton_boolean(true))).val(true),
+            // Doesn't contain `1` because we didn't typecheck this side
+            ExprBuilder::new().val(true),
+        ),
+    );
     assert_expr_has_annotated_ast(
         &Expr::ite(
             Expr::less(Expr::val(1), Expr::val(0)),
@@ -93,6 +110,24 @@ fn expr_typechecks_with_correct_annotation() {
             ExprBuilder::with_data(Some(Type::primitive_string())).val("foo"),
         ),
     );
+    assert_expr_has_annotated_ast(
+        &Expr::ite(Expr::val(true), Expr::val("bar"), Expr::val("foo")),
+        &ExprBuilder::with_data(Some(Type::primitive_string())).ite(
+            ExprBuilder::with_data(Some(Type::True)).val(true),
+            ExprBuilder::with_data(Some(Type::primitive_string())).val("bar"),
+            // Contains `bar` instead of `foo` because we didn't typecheck the `else` branch
+            ExprBuilder::with_data(Some(Type::primitive_string())).val("bar"),
+        ),
+    );
+    assert_expr_has_annotated_ast(
+        &Expr::ite(Expr::val(false), Expr::val("bar"), Expr::val("foo")),
+        &ExprBuilder::with_data(Some(Type::primitive_string())).ite(
+            ExprBuilder::with_data(Some(Type::False)).val(false),
+            // Contains `foo` instead of `bar` because we didn't typecheck the `else` branch
+            ExprBuilder::with_data(Some(Type::primitive_string())).val("foo"),
+            ExprBuilder::with_data(Some(Type::primitive_string())).val("foo"),
+        ),
+    );
     assert_expr_has_annotated_ast(
         &Expr::not(Expr::val(false)),
         &ExprBuilder::with_data(Some(Type::singleton_boolean(true)))