Propagate some source locations from cedar syntax schema parser to JSON syntax structures

cedar-policy-validator/Cargo.toml

@@ -26,6 +26,7 @@ arbitrary = { version = "1", features = ["derive"], optional = true }
 lalrpop-util = { version = "0.22.0", features = ["lexer", "unicode"] }
 lazy_static = "1.4.0"
 nonempty = "0.10.0"
+educe = "0.6.0"
 
 # wasm dependencies
 serde-wasm-bindgen = { version = "0.6", optional = true }

cedar-policy-validator/src/cedar_schema/test.rs

@@ -343,6 +343,7 @@ mod demo_tests {
             applies_to: None,
             member_of: None,
             annotations: Annotations::new(),
+            loc: None,
         };
         let namespace =
             json_schema::NamespaceDefinition::new(empty(), once(("foo".to_smolstr(), action)));
@@ -445,6 +446,7 @@ namespace Baz {action "Foo" appliesTo {
                     shape: json_schema::AttributesOrContext::default(),
                     tags: None,
                     annotations: Annotations::new(),
+                    loc: None,
                 },
             )]),
             actions: BTreeMap::from([(
@@ -458,6 +460,7 @@ namespace Baz {action "Foo" appliesTo {
                     }),
                     member_of: None,
                     annotations: Annotations::new(),
+                    loc: None,
                 },
             )]),
             annotations: Annotations::new(),

cedar-policy-validator/src/cedar_schema/to_json_schema.rs

@@ -193,16 +193,17 @@ impl TryFrom<Annotated<Namespace>> for json_schema::NamespaceDefinition<RawName>
         let common_types = common_types
             .into_iter()
             .map(|decl| {
-                let name_loc = decl.data.name.loc.clone();
-                let id = UnreservedId::try_from(decl.data.name.node)
+                let name_loc = decl.data.node.name.loc.clone();
+                let id = UnreservedId::try_from(decl.data.node.name.node)
                     .map_err(|e| ToJsonSchemaError::reserved_name(e.name(), name_loc.clone()))?;
                 let ctid = json_schema::CommonTypeId::new(id)
                     .map_err(|e| ToJsonSchemaError::reserved_keyword(&e.id, name_loc))?;
                 Ok((
                     ctid,
                     CommonType {
-                        ty: cedar_type_to_json_type(decl.data.def),
+                        ty: cedar_type_to_json_type(decl.data.node.def),
                         annotations: decl.annotations.into(),
+                        loc: Some(decl.data.loc),
                     },
                 ))
             })
@@ -219,13 +220,13 @@ impl TryFrom<Annotated<Namespace>> for json_schema::NamespaceDefinition<RawName>
 
 /// Converts action type decls
 fn convert_action_decl(
-    a: Annotated<ActionDecl>,
+    a: Annotated<Node<ActionDecl>>,
 ) -> Result<impl Iterator<Item = (SmolStr, json_schema::ActionType<RawName>)>, ToJsonSchemaErrors> {
     let ActionDecl {
         names,
         parents,
         app_decls,
-    } = a.data;
+    } = a.data.node;
     // Create the internal type from the 'applies_to' clause and 'member_of'
     let applies_to = app_decls
         .map(|decls| convert_app_decls(&names.first().node, &names.first().loc, decls))
@@ -241,6 +242,7 @@ fn convert_action_decl(
         applies_to: Some(applies_to),
         member_of,
         annotations: a.annotations.into(),
+        loc: Some(a.data.loc),
     };
     // Then map that type across all of the bound names
     Ok(names.into_iter().map(move |name| (name.node, ty.clone())))
@@ -356,7 +358,7 @@ fn convert_id(node: Node<Id>) -> Result<UnreservedId, ToJsonSchemaError> {
 
 /// Convert Entity declarations
 fn convert_entity_decl(
-    e: Annotated<EntityDecl>,
+    e: Annotated<Node<EntityDecl>>,
 ) -> Result<
     impl Iterator<Item = (UnreservedId, json_schema::EntityType<RawName>)>,
     ToJsonSchemaErrors,
@@ -365,24 +367,21 @@ fn convert_entity_decl(
     let etype = json_schema::EntityType {
         member_of_types: e
             .data
+            .node
             .member_of_types
             .into_iter()
             .map(RawName::from)
             .collect(),
-        shape: convert_attr_decls(e.data.attrs),
-        tags: e.data.tags.map(cedar_type_to_json_type),
+        shape: convert_attr_decls(e.data.node.attrs),
+        tags: e.data.node.tags.map(cedar_type_to_json_type),
         annotations: e.annotations.into(),
+        loc: Some(e.data.loc),
     };
 
     // Then map over all of the bound names
-    collect_all_errors(
-        e.data
-            .names
-            .into_iter()
-            .map(move |name| -> Result<_, ToJsonSchemaErrors> {
-                Ok((convert_id(name)?, etype.clone()))
-            }),
-    )
+    collect_all_errors(e.data.node.names.into_iter().map(
+        move |name| -> Result<_, ToJsonSchemaErrors> { Ok((convert_id(name)?, etype.clone())) },
+    ))
 }
 
 /// Create a [`json_schema::AttributesOrContext`] from a series of `AttrDecl`s
@@ -642,32 +641,123 @@ fn partition_decls(
 }
 
 fn into_partition_decls(
-    decls: Vec<Annotated<Node<Declaration>>>,
+    decls: impl IntoIterator<Item = Annotated<Node<Declaration>>>,
 ) -> (
-    Vec<Annotated<EntityDecl>>,
-    Vec<Annotated<ActionDecl>>,
-    Vec<Annotated<TypeDecl>>,
+    Vec<Annotated<Node<EntityDecl>>>,
+    Vec<Annotated<Node<ActionDecl>>>,
+    Vec<Annotated<Node<TypeDecl>>>,
 ) {
     let mut entities = vec![];
     let mut actions = vec![];
     let mut types = vec![];
 
     for decl in decls.into_iter() {
+        let loc = decl.data.loc;
         match decl.data.node {
             Declaration::Entity(e) => entities.push(Annotated {
-                data: e,
+                data: Node { node: e, loc },
                 annotations: decl.annotations,
             }),
             Declaration::Action(a) => actions.push(Annotated {
-                data: a,
+                data: Node { node: a, loc },
                 annotations: decl.annotations,
             }),
             Declaration::Type(t) => types.push(Annotated {
-                data: t,
+                data: Node { node: t, loc },
                 annotations: decl.annotations,
             }),
         }
     }
 
     (entities, actions, types)
 }
+
+#[cfg(test)]
+mod preserves_source_locations {
+    use super::*;
+    use cool_asserts::assert_matches;
+
+    #[test]
+    fn entity_action_and_common_type_decls() {
+        let (schema, _) = json_schema::Fragment::from_cedarschema_str(
+            r#"
+        namespace NS {
+            type S = String;
+            entity A;
+            entity B in A;
+            entity C in A {
+                bool: Bool,
+                s: S,
+                a: Set<A>,
+                b: { inner: B },
+            };
+            type AA = A;
+            action Read, Write;
+            action List in Read appliesTo {
+                principal: [A],
+                resource: [B, C],
+                context: {
+                    s: Set<S>,
+                    ab: { a: AA, b: B },
+                }
+            };
+        }
+        "#,
+            &Extensions::all_available(),
+        )
+        .unwrap();
+        let ns = schema
+            .0
+            .get(&Some(Name::parse_unqualified_name("NS").unwrap()))
+            .expect("couldn't find namespace NS");
+
+        let entityA = ns
+            .entity_types
+            .get(&"A".parse().unwrap())
+            .expect("couldn't find entity A");
+        let entityB = ns
+            .entity_types
+            .get(&"B".parse().unwrap())
+            .expect("couldn't find entity B");
+        let entityC = ns
+            .entity_types
+            .get(&"C".parse().unwrap())
+            .expect("couldn't find entity C");
+        let ctypeS = ns
+            .common_types
+            .get(&json_schema::CommonTypeId::new("S".parse().unwrap()).unwrap())
+            .expect("couldn't find common type S");
+        let ctypeAA = ns
+            .common_types
+            .get(&json_schema::CommonTypeId::new("AA".parse().unwrap()).unwrap())
+            .expect("couldn't find common type AA");
+        let actionRead = ns.actions.get("Read").expect("couldn't find action Read");
+        let actionWrite = ns.actions.get("Write").expect("couldn't find action Write");
+        let actionList = ns.actions.get("List").expect("couldn't find action List");
+
+        assert_matches!(&entityA.loc, Some(loc) => assert_matches!(loc.snippet(),
+            Some("entity A;")
+        ));
+        assert_matches!(&entityB.loc, Some(loc) => assert_matches!(loc.snippet(),
+            Some("entity B in A;")
+        ));
+        assert_matches!(&entityC.loc, Some(loc) => assert_matches!(loc.snippet(),
+            Some("entity C in A {\n                bool: Bool,\n                s: S,\n                a: Set<A>,\n                b: { inner: B },\n            };")
+        ));
+        assert_matches!(&ctypeS.loc, Some(loc) => assert_matches!(loc.snippet(),
+            Some("type S = String;")
+        ));
+        assert_matches!(&ctypeAA.loc, Some(loc) => assert_matches!(loc.snippet(),
+            Some("type AA = A;")
+        ));
+        assert_matches!(&actionRead.loc, Some(loc) => assert_matches!(loc.snippet(),
+            Some("action Read, Write;")
+        ));
+        assert_matches!(&actionWrite.loc, Some(loc) => assert_matches!(loc.snippet(),
+            Some("action Read, Write;")
+        ));
+        assert_matches!(&actionList.loc, Some(loc) => assert_matches!(loc.snippet(),
+            Some("action List in Read appliesTo {\n                principal: [A],\n                resource: [B, C],\n                context: {\n                    s: Set<S>,\n                    ab: { a: AA, b: B },\n                }\n            };")
+        ));
+    }
+}