Internal data structures for Generalized templates

cedar-language-server/src/schema/fold.rs

@@ -80,7 +80,7 @@ pub(crate) fn fold_schema(schema_info: &SchemaInfo) -> Option<Vec<FoldingRange>>
         .filter_map(|et| et.loc.as_loc_ref());
     let action_locs = validator.action_ids().filter_map(|a| a.loc());
     let common_types = validator
-        .common_types()
+        .common_types_extended()
         .filter_map(|ct| ct.type_loc.as_loc_ref());
 
     // Combine all locations and create folding ranges

cedar-language-server/src/schema/symbols.rs

@@ -125,7 +125,7 @@ pub(crate) fn schema_symbols(schema_info: &SchemaInfo) -> Option<Vec<DocumentSym
 
     // Create common type symbols
     let common_type_symbols: Vec<DocumentSymbol> = validator
-        .common_types()
+        .common_types_extended()
         .filter_map(|ct| {
             ct.name_loc
                 .as_ref()

cedar-policy-core/src/ast.rs

@@ -54,5 +54,7 @@ mod expr_iterator;
 pub use expr_iterator::*;
 mod annotation;
 pub use annotation::*;
+mod generalized_slots_annotation;
+pub use generalized_slots_annotation::*;
 mod expr_visitor;
 pub use expr_visitor::*;

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

@@ -293,7 +293,7 @@ impl<T> Expr<T> {
         self.subexpressions()
             .filter_map(|exp| match &exp.expr_kind {
                 ExprKind::Slot(slotid) => Some(Slot {
-                    id: *slotid,
+                    id: slotid.clone(),
                     loc: exp.source_loc().into_maybe_loc(),
                 }),
                 _ => None,
@@ -1842,7 +1842,7 @@ mod test {
         let e = Expr::slot(SlotId::principal());
         let p = SlotId::principal();
         let r = SlotId::resource();
-        let set: HashSet<SlotId> = HashSet::from_iter([p]);
+        let set: HashSet<SlotId> = HashSet::from_iter([p.clone()]);
         assert_eq!(set, e.slots().map(|slot| slot.id).collect::<HashSet<_>>());
         let e = Expr::or(
             Expr::slot(SlotId::principal()),

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

@@ -55,7 +55,7 @@ pub trait ExprVisitor {
         match expr.expr_kind() {
             ExprKind::Lit(lit) => self.visit_literal(lit, loc),
             ExprKind::Var(var) => self.visit_var(*var, loc),
-            ExprKind::Slot(slot) => self.visit_slot(*slot, loc),
+            ExprKind::Slot(slot) => self.visit_slot(slot.clone(), loc),
             ExprKind::Unknown(unknown) => self.visit_unknown(unknown, loc),
             ExprKind::If {
                 test_expr,

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

@@ -0,0 +1,122 @@
+/*
+ * 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.
+ */
+
+use std::collections::BTreeMap;
+
+use crate::ast::SlotId;
+use crate::extensions::Extensions;
+use crate::validator::{
+    json_schema::Type as JSONSchemaType, types::Type as ValidatorType, RawName, SchemaError,
+    ValidatorSchema,
+};
+use serde::{Deserialize, Serialize};
+
+/// Struct storing the pairs of SlotId's and their corresponding type
+#[derive(Clone, Eq, PartialEq, PartialOrd, Ord, Debug, Hash, Serialize, Deserialize)]
+pub struct GeneralizedSlotsAnnotation(BTreeMap<SlotId, JSONSchemaType<RawName>>);
+
+impl GeneralizedSlotsAnnotation {
+    /// Create a new empty `GeneralizedSlotsAnnotation` (with no slots)
+    pub fn new() -> Self {
+        Self(BTreeMap::new())
+    }
+
+    /// Get the type of the slot by key
+    pub fn get(&self, key: &SlotId) -> Option<&JSONSchemaType<RawName>> {
+        self.0.get(key)
+    }
+
+    /// Iterate over all pairs of slots and their types
+    pub fn iter(&self) -> impl Iterator<Item = (&SlotId, &JSONSchemaType<RawName>)> {
+        self.0.iter()
+    }
+
+    /// Tell if it's empty
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    /// Converts the types of generalized slots annotation to
+    /// use validator types so that they can be used by the typechecker
+    pub fn into_validator_generalized_slots_annotation(
+        self,
+        schema: &ValidatorSchema,
+    ) -> Result<ValidatorGeneralizedSlotsAnnotation, SchemaError> {
+        let validator_generalized_slots_annotation: Result<BTreeMap<_, _>, SchemaError> = self
+            .0
+            .into_iter()
+            .map(|(k, ty)| -> Result<_, SchemaError> {
+                Ok((
+                    k,
+                    schema.json_schema_type_to_validator_type(ty, Extensions::all_available())?,
+                ))
+            })
+            .collect();
+        Ok(validator_generalized_slots_annotation?.into())
+    }
+}
+
+impl Default for GeneralizedSlotsAnnotation {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl FromIterator<(SlotId, JSONSchemaType<RawName>)> for GeneralizedSlotsAnnotation {
+    fn from_iter<T: IntoIterator<Item = (SlotId, JSONSchemaType<RawName>)>>(iter: T) -> Self {
+        Self(BTreeMap::from_iter(iter))
+    }
+}
+
+impl From<BTreeMap<SlotId, JSONSchemaType<RawName>>> for GeneralizedSlotsAnnotation {
+    fn from(value: BTreeMap<SlotId, JSONSchemaType<RawName>>) -> Self {
+        Self(value)
+    }
+}
+
+/// Struct storing the pairs of SlotId's and their corresponding validator types
+#[derive(Clone, Eq, PartialEq, PartialOrd, Ord, Debug, Hash)]
+pub struct ValidatorGeneralizedSlotsAnnotation(BTreeMap<SlotId, ValidatorType>);
+
+impl FromIterator<(SlotId, ValidatorType)> for ValidatorGeneralizedSlotsAnnotation {
+    fn from_iter<T: IntoIterator<Item = (SlotId, ValidatorType)>>(iter: T) -> Self {
+        Self(BTreeMap::from_iter(iter))
+    }
+}
+
+impl From<BTreeMap<SlotId, ValidatorType>> for ValidatorGeneralizedSlotsAnnotation {
+    fn from(value: BTreeMap<SlotId, ValidatorType>) -> Self {
+        Self(value)
+    }
+}
+
+impl Default for ValidatorGeneralizedSlotsAnnotation {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl ValidatorGeneralizedSlotsAnnotation {
+    /// Create a new empty `ValidatorGeneralizedSlotsAnnotation` (with no slots)
+    pub fn new() -> Self {
+        Self(BTreeMap::new())
+    }
+
+    /// Get the validator type of the slot by key
+    pub fn get(&self, slot: &SlotId) -> Option<&ValidatorType> {
+        self.0.get(slot)
+    }
+}

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

@@ -283,7 +283,7 @@ impl<'de> Deserialize<'de> for InternalName {
 /// Clone is O(1).
 // This simply wraps a separate enum -- currently [`ValidSlotId`] -- in case we
 // want to generalize later
-#[derive(Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 #[serde(transparent)]
 pub struct SlotId(pub(crate) ValidSlotId);
 
@@ -298,6 +298,11 @@ impl SlotId {
         Self(ValidSlotId::Resource)
     }
 
+    /// Create a `generalized slot`
+    pub fn generalized_slot(id: Id) -> Self {
+        Self(ValidSlotId::GeneralizedSlot(id))
+    }
+
     /// Check if a slot represents a principal
     pub fn is_principal(&self) -> bool {
         matches!(self, Self(ValidSlotId::Principal))
@@ -307,6 +312,11 @@ impl SlotId {
     pub fn is_resource(&self) -> bool {
         matches!(self, Self(ValidSlotId::Resource))
     }
+
+    /// Check if a slot represents a generalized slot
+    pub fn is_generalized_slot(&self) -> bool {
+        matches!(self, Self(ValidSlotId::GeneralizedSlot(_)))
+    }
 }
 
 impl From<PrincipalOrResource> for SlotId {
@@ -324,20 +334,22 @@ impl std::fmt::Display for SlotId {
     }
 }
 
-/// Two possible variants for Slots
-#[derive(Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+/// Three possible variants for Slots
+#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub(crate) enum ValidSlotId {
     #[serde(rename = "?principal")]
     Principal,
     #[serde(rename = "?resource")]
     Resource,
+    GeneralizedSlot(Id), // Slots for generalized templates, for more info see [RFC 98](https://github.com/cedar-policy/rfcs/pull/98).
 }
 
 impl std::fmt::Display for ValidSlotId {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let s = match self {
             ValidSlotId::Principal => "principal",
             ValidSlotId::Resource => "resource",
+            ValidSlotId::GeneralizedSlot(id) => id.as_ref(),
         };
         write!(f, "?{s}")
     }