Escape stubgen names

trustfall_stubgen/src/edges_creator.rs

@@ -4,6 +4,8 @@ use maplit::btreemap;
 use quote::quote;
 use trustfall::{Schema, SchemaAdapter, TryIntoStruct};
 
+use super::util::escaped_rust_name;
+
 use super::{
     root::RustFile,
     util::{
@@ -77,7 +79,10 @@ fn make_type_edge_resolver(
     edges: Vec<(String, Vec<(String, String)>)>,
 ) -> (proc_macro2::TokenStream, proc_macro2::TokenStream) {
     let lower_type_name = to_lower_snake_case(type_name);
-    let mod_name = syn::Ident::new(&lower_type_name, proc_macro2::Span::call_site());
+    let mod_name = syn::Ident::new(
+        &escaped_rust_name(lower_type_name.clone()),
+        proc_macro2::Span::call_site(),
+    );
 
     let mut arms = proc_macro2::TokenStream::new();
     let mut edge_resolvers = proc_macro2::TokenStream::new();
@@ -135,9 +140,9 @@ fn make_edge_resolver_and_call(
         },
     );
 
-    let resolver_fn_name = to_lower_snake_case(edge_name);
+    let resolver_fn_name = escaped_rust_name(to_lower_snake_case(edge_name));
     let resolver_fn_ident = syn::Ident::new(&resolver_fn_name, proc_macro2::Span::call_site());
-    let conversion_fn_name = format!("as_{}", to_lower_snake_case(type_name));
+    let conversion_fn_name = format!("as_{}", escaped_rust_name(to_lower_snake_case(type_name)));
     let conversion_fn_ident = syn::Ident::new(&conversion_fn_name, proc_macro2::Span::call_site());
     let expect_msg = format!("conversion failed, vertex was not a {type_name}");
     let todo_msg = format!("get neighbors along edge '{edge_name}' for type '{type_name}'");

trustfall_stubgen/src/entrypoints_creator.rs

@@ -4,7 +4,10 @@ use maplit::btreemap;
 use quote::quote;
 use trustfall::{Schema, SchemaAdapter, TryIntoStruct};
 
-use crate::edges_creator::{prepare_call_parameters, FnCall};
+use crate::{
+    edges_creator::{prepare_call_parameters, FnCall},
+    util::escaped_rust_name,
+};
 
 use super::{
     root::RustFile,
@@ -70,7 +73,7 @@ fn make_entrypoint_fn(
         },
     );
 
-    let entrypoint_fn_name = to_lower_snake_case(entrypoint);
+    let entrypoint_fn_name = escaped_rust_name(to_lower_snake_case(entrypoint));
     let ident = syn::Ident::new(&entrypoint_fn_name, proc_macro2::Span::call_site());
     let todo_msg =
         format!("implement resolving starting vertices for entrypoint edge '{entrypoint}'");

trustfall_stubgen/src/root.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::{BTreeMap, BTreeSet},
+    collections::{BTreeMap, BTreeSet, HashMap},
     io::Write,
     path::Path,
     sync::{Arc, OnceLock},
@@ -10,7 +10,7 @@ use quote::quote;
 use regex::Regex;
 use trustfall::{Schema, SchemaAdapter, TryIntoStruct};
 
-use crate::util::parse_import;
+use crate::util::{escaped_rust_name, parse_import, to_lower_snake_case, upper_case_variant_name};
 
 use super::{
     adapter_creator::make_adapter_file, edges_creator::make_edges_file,
@@ -52,6 +52,9 @@ pub fn generate_rust_stub(schema: &str, target: &Path) -> anyhow::Result<()> {
 
     let mut entrypoint_match_arms = proc_macro2::TokenStream::new();
 
+    ensure_no_vertex_name_conflicts(&querying_schema, schema_adapter.clone());
+    ensure_no_field_name_conflicts_on_vertex_type(&querying_schema, schema_adapter.clone());
+
     make_vertex_file(&querying_schema, schema_adapter.clone(), &mut stub.vertex);
     make_entrypoints_file(
         &querying_schema,
@@ -124,7 +127,6 @@ impl RustFile {
         static PATTERN: OnceLock<Regex> = OnceLock::new();
         let pattern =
             PATTERN.get_or_init(|| Regex::new("([^{])\n    (pub|fn|use)").expect("invalid regex"));
-
         let pretty_item =
             prettyplease::unparse(&syn::parse_str(&item.to_string()).expect("not valid Rust"));
         let postprocessed = pattern.replace_all(&pretty_item, "$1\n\n    $2");
@@ -383,7 +385,7 @@ fn make_vertex_file(
         .collect();
     rows.sort_unstable();
     for row in rows {
-        let name = &row.name;
+        let name = &escaped_rust_name(upper_case_variant_name(&row.name));
         let ident = syn::Ident::new(name.as_str(), proc_macro2::Span::call_site());
         variants.extend(quote! {
             #ident(()),
@@ -400,3 +402,86 @@ fn make_vertex_file(
 
     vertex_file.top_level_items.push(vertex);
 }
+
+fn ensure_no_vertex_name_conflicts(querying_schema: &Schema, adapter: Arc<SchemaAdapter<'_>>) {
+    let query = r#"
+{
+    VertexType {
+        name @output
+    }
+}"#;
+    let variables: BTreeMap<String, String> = Default::default();
+
+    #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, serde::Deserialize)]
+    struct ResultRow {
+        name: String,
+    }
+
+    let mut rows: Vec<_> = trustfall::execute_query(querying_schema, adapter, query, variables)
+        .expect("invalid query")
+        .map(|x| x.try_into_struct::<ResultRow>().expect("invalid conversion"))
+        .collect();
+    rows.sort_unstable();
+
+    let mut uniq: HashMap<String, String> = HashMap::new();
+
+    for row in rows {
+        let name = row.name.clone();
+        // we normalize to lower snake case here, however in vertex name we capitalize this name instead
+        // it doesn't really matter though because the important one is just to normalize to the same capitalization scheme
+        let converted = escaped_rust_name(to_lower_snake_case(&name));
+        let v = uniq.insert(converted, name);
+        if let Some(v) = v {
+            panic!(
+                "cannot generate adapter for a schema containing both '{}' and '{}' vertices, consider renaming one of them",
+                v, &row.name
+            );
+        }
+    }
+}
+
+fn ensure_no_field_name_conflicts_on_vertex_type(
+    querying_schema: &Schema,
+    adapter: Arc<SchemaAdapter<'_>>,
+) {
+    let query = r#"
+{
+    VertexType {
+        name @output
+        edge_: edge @fold {
+            names: name @output
+        }
+        property_: property @fold {
+            names: name @output
+        }
+    }
+}"#;
+    let variables: BTreeMap<String, String> = Default::default();
+
+    #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, serde::Deserialize)]
+    struct ResultRow {
+        name: String,
+        edge_names: Vec<String>,
+        property_names: Vec<String>,
+    }
+
+    let mut rows: Vec<_> = trustfall::execute_query(querying_schema, adapter, query, variables)
+        .expect("invalid query")
+        .map(|x| x.try_into_struct::<ResultRow>().expect("invalid conversion"))
+        .collect();
+    rows.sort_unstable();
+
+    for row in &rows {
+        let mut uniq: HashMap<String, String> = HashMap::new();
+
+        for field_name in row.edge_names.iter().chain(row.property_names.iter()) {
+            let converted = escaped_rust_name(to_lower_snake_case(field_name));
+            if let Some(v) = uniq.insert(converted, field_name.clone()) {
+                panic!(
+                    "cannot generate adapter for a schema containing both '{}' and '{}' as field names on vertex '{}', consider renaming one of them",
+                    v, &field_name, &row.name
+                );
+            }
+        }
+    }
+}

trustfall_stubgen/src/tests.rs

@@ -153,3 +153,40 @@ fn test_schema(name: &str) {
 fn test_hackernews_schema() {
     test_schema("hackernews")
 }
+
+#[test]
+fn test_use_reserved_rust_names_in_schema() {
+    test_schema("use_reserved_rust_names_in_schema");
+}
+
+#[test]
+#[should_panic(
+    expected = "cannot generate adapter for a schema containing both 'Type' and 'Type_' vertices, consider renaming one of them"
+)]
+fn use_type_and_typeunderscore_edge_names() {
+    test_schema("use_type_and_typeunderscore_edge_names");
+}
+
+#[test]
+#[should_panic(
+    expected = "cannot generate adapter for a schema containing both 'Type_' and 'Type' as field names on vertex 'Type', consider renaming one of them"
+)]
+fn vertextype_with_type_and_typeunderscore_edge_and_property() {
+    test_schema("vertextype_with_type_and_typeunderscore_edge_and_property");
+}
+
+#[test]
+#[should_panic(
+    expected = "cannot generate adapter for a schema containing both 'Type' and 'Type_' as field names on vertex 'Type', consider renaming one of them"
+)]
+fn vertextype_with_type_and_typeunderscore_edges() {
+    test_schema("vertextype_with_type_and_typeunderscore_edges");
+}
+
+#[test]
+#[should_panic(
+    expected = "cannot generate adapter for a schema containing both 'Type' and 'Type_' as field names on vertex 'Type', consider renaming one of them"
+)]
+fn vertextype_with_type_and_typeunderscore_properties() {
+    test_schema("vertextype_with_type_and_typeunderscore_properties");
+}

trustfall_stubgen/src/util.rs

@@ -17,6 +17,14 @@ pub(crate) fn to_lower_snake_case(value: &str) -> String {
     result
 }
 
+pub(crate) fn upper_case_variant_name(value: &str) -> String {
+    let mut chars = value.chars();
+    let first_char = chars.next().unwrap().to_ascii_uppercase();
+    let rest_chars: String = chars.collect();
+
+    format!("{}{}", first_char, rest_chars)
+}
+
 pub(crate) fn property_resolver_fn_name(type_name: &str) -> String {
     let normalized_name = to_lower_snake_case(type_name);
     format!("resolve_{normalized_name}_property")
@@ -138,3 +146,15 @@ pub(crate) fn field_value_to_rust_type(
         #base.#suffix
     }
 }
+
+pub fn escaped_rust_name(name: String) -> String {
+    // https://doc.rust-lang.org/reference/keywords.html
+    match name.as_str() {
+        "as" | "break" | "const" | "continue" | "crate" | "else" | "enum" | "extern" | "false"
+        | "fn" | "for" | "if" | "impl" | "in" | "let" | "loop" | "match" | "mod" | "move"
+        | "mut" | "pub" | "ref" | "return" | "self" | "Self" | "static" | "struct" | "super"
+        | "trait" | "true" | "type" | "unsafe" | "use" | "where" | "while" | "async" | "await"
+        | "dyn" | "try" | "macro_rules" | "union" | "'static" => name + "_",
+        _ => name,
+    }
+}