refactor: migrate get_location to codegen (not working yet)

crates/codegen/src/get_children.rs

@@ -0,0 +1,133 @@
+// use pg_query_proto_parser::{FieldType, Node, ProtoParser};
+// use proc_macro2::{Ident, TokenStream};
+// use quote::{format_ident, quote};
+//
+// // todo: get_children should only return a Vec<NestedNode> with location being an Option<i32>
+// // we then pass the results into a resolve_locations function that takes a Vec<NestedNode> and returns a Vec<NestedNode>, but where location is an i32
+//
+// pub fn get_children_mod(_item: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
+// let parser = ProtoParser::new("./libpg_query/protobuf/pg_query.proto");
+// let proto_file = parser.parse();
+//
+// let manual_node_names = manual_node_names();
+//
+// let node_identifiers = node_identifiers(&proto_file.nodes, &manual_node_names);
+// let node_handlers = node_handlers(&proto_file.nodes, &manual_node_names);
+//
+// quote! {
+// use pg_query::NodeEnum;
+// use std::collections::VecDeque;
+//
+// #[derive(Debug, Clone)]
+// pub struct ChildrenNode {
+// pub node: NodeEnum,
+// pub depth: i32,
+// pub location: Option<i32>,
+// pub path: String,
+// }
+//
+// /// Returns all children of the node, recursively
+// pub fn get_children(node: &NodeEnum, text: String, current_depth: i32) -> Vec<ChildrenNode> {
+// let mut nodes: Vec<ChildrenNode> = vec![];
+// // Node, depth, path
+// let mut stack: VecDeque<(NodeEnum, i32, String)> =
+// VecDeque::from(vec![(node.to_owned(), current_depth, "0".to_string())]);
+// while !stack.is_empty() {
+// let (node, depth, path) = stack.pop_front().unwrap();
+// let current_depth = depth + 1;
+// let mut child_ctr: i32 = 0;
+// let mut handle_child = |c: NodeEnum| {
+// let location = get_location(&c);
+// let path = path.clone() + "." + child_ctr.to_string().as_str();
+// child_ctr = child_ctr + 1;
+// stack.push_back((c.to_owned(), current_depth, path.clone()));
+// nodes.push(ChildrenNode {
+// node: c,
+// depth: current_depth,
+// location,
+// path: path.clone(),
+// });
+// };
+// match &node {
+// // `AConst` is the only node with a `one of` property, so we handle it manually
+// // if you need to handle other nodes manually, add them to the `manual_node_names` function below
+// NodeEnum::AConst(n) => {
+// if n.val.is_some() {
+// handle_child(match n.val.to_owned().unwrap() {
+// pg_query::protobuf::a_const::Val::Ival(v) => NodeEnum::Integer(v),
+// pg_query::protobuf::a_const::Val::Fval(v) => NodeEnum::Float(v),
+// pg_query::protobuf::a_const::Val::Boolval(v) => NodeEnum::Boolean(v),
+// pg_query::protobuf::a_const::Val::Sval(v) => NodeEnum::String(v),
+// pg_query::protobuf::a_const::Val::Bsval(v) => NodeEnum::BitString(v),
+// });
+// }
+// }
+// #(NodeEnum::#node_identifiers(n) => {
+// #node_handlers
+// }),*,
+// };
+// }
+// nodes
+// }
+// }
+// }
+//
+// fn manual_node_names() -> Vec<&'static str> {
+// vec!["AConst"]
+// }
+//
+// fn node_identifiers(nodes: &[Node], exclude_nodes: &[&str]) -> Vec<Ident> {
+// nodes
+// .iter()
+// .filter(|node| !exclude_nodes.contains(&node.name.as_str()))
+// .map(|node| format_ident!("{}", &node.name))
+// .collect()
+// }
+//
+// fn node_handlers(nodes: &[Node], exclude_nodes: &[&str]) -> Vec<TokenStream> {
+// nodes
+// .iter()
+// .filter(|node| !exclude_nodes.contains(&node.name.as_str()))
+// .map(|node| {
+// let property_handlers = property_handlers(&node);
+// quote! {
+// #(#property_handlers)*
+// }
+// })
+// .collect()
+// }
+//
+// fn property_handlers(node: &Node) -> Vec<TokenStream> {
+// node.fields
+// .iter()
+// .map(|field| {
+// if field.field_type == FieldType::Node && field.repeated {
+// let field_name = field.name.as_str();
+// quote! {
+// n.#field_name
+// .iter()
+// .for_each(|x| handle_child(x.node.as_ref().unwrap().to_owned()));
+// }
+// } else if field.field_type == FieldType::Node && field.is_one_of == false {
+// if field.node_name == Some("Node".to_owned()) {
+// let field_name = field.name.as_str();
+// quote! {
+// if n.#field_name.is_some() {
+// handle_child(n.#field_name.to_owned().unwrap().node.unwrap());
+// }
+// }
+// } else {
+// let enum_variant_name = field.enum_variant_name.as_ref().unwrap();
+// let field_name = field.name.as_str();
+// quote! {
+// if n.#field_name.is_some() {
+// handle_child(NodeEnum::#enum_variant_name(n.#field_name.to_owned().unwrap()));
+// }
+// }
+// }
+// } else {
+// panic!("Unhandled field type: {:?}", field);
+// }
+// })
+// .collect()
+// }

crates/codegen/src/get_location.rs

@@ -0,0 +1,72 @@
+use pg_query_proto_parser::{FieldType, Node, ProtoParser};
+use proc_macro2::{Ident, TokenStream};
+use quote::{format_ident, quote};
+
+pub fn get_location_mod(_item: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
+ let parser = ProtoParser::new("./libpg_query/protobuf/pg_query.proto");
+ let proto_file = parser.parse();
+
+ let manual_node_names = manual_node_names();
+
+ let node_identifiers = node_identifiers(&proto_file.nodes, &manual_node_names);
+ let location_idents = location_idents(&proto_file.nodes, &manual_node_names);
+
+ quote! {
+ use pg_query::NodeEnum;
+
+ //! Returns the location of a node
+ pub fn get_location(node: &NodeEnum) -> Option<i32> {
+ let location = match node {
+ // for some nodes, the location of the node itself is after their childrens location.
+ // we implement the logic for those nodes manually.
+ // if you add one, make sure to add its name to `manual_node_names()`.
+ NodeEnum::BoolExpr(n) => {
+ let a = n.args.iter().min_by(|a, b| {
+ let loc_a = get_location(&a.node.as_ref().unwrap());
+ let loc_b = get_location(&b.node.as_ref().unwrap());
+ loc_a.cmp(&loc_b)
+ });
+ get_location(&a.unwrap().node.as_ref().unwrap())
+ },
+ NodeEnum::AExpr(n) => get_location(&n.lexpr.as_ref().unwrap().node.as_ref().unwrap()),
+ #(NodeEnum::#node_identifiers(n) => #location_idents),*
+ };
+ if location.is_some() && location.unwrap() < 0 {
+ None
+ } else {
+ location
+ }
+ }
+ }
+}
+
+fn manual_node_names() -> Vec<&'static str> {
+ vec!["BoolExpr", "AExpr"]
+}
+
+fn location_idents(nodes: &[Node], exclude_nodes: &[&str]) -> Vec<TokenStream> {
+ nodes
+ .iter()
+ .filter(|n| !exclude_nodes.contains(&n.name.as_str()))
+ .map(|node| {
+ if node
+ .fields
+ .iter()
+ .find(|n| n.name == "location" && n.field_type == FieldType::Int32)
+ .is_some()
+ {
+ quote! { Some(n.location) }
+ } else {
+ quote! { None }
+ }
+ })
+ .collect()
+}
+
+fn node_identifiers(nodes: &[Node], exclude_nodes: &[&str]) -> Vec<Ident> {
+ nodes
+ .iter()
+ .filter(|n| !exclude_nodes.contains(&n.name.as_str()))
+ .map(|node| format_ident!("{}", &node.name))
+ .collect()
+}

crates/codegen/src/lib.rs

@@ -1,8 +1,22 @@
+// mod get_children;
+mod get_location;
 mod syntax_kind;
 
+// use get_children::get_children_mod;
+use get_location::get_location_mod;
 use syntax_kind::syntax_kind_mod;
 
+// #[proc_macro]
+// pub fn get_children(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
+// get_children_mod(item.into()).into()
+// }
+
 #[proc_macro]
 pub fn syntax_kind(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
  syntax_kind_mod(item.into()).into()
 }
+
+#[proc_macro]
+pub fn get_location(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
+ get_location_mod(item.into()).into()
+}

crates/parser/src/get_location_codegen.rs

@@ -0,0 +1,3 @@
+use codegen::get_location;
+
+get_location!();

crates/parser/src/lib.rs

@@ -16,6 +16,7 @@
 //! To see how these drawbacks are mitigated, see the `statement.rs` and the `source_file.rs` module.
 
 mod ast_node;
+mod get_location_codegen;
 mod parser;
 mod pg_query_utils_generated;
 mod pg_query_utils_generated_test;