Add recursive iter() on ParseResult and NodeRef

Cargo.toml

@@ -13,18 +13,21 @@ repository = "https://github.com/pganalyze/pg_query.rs"
 
 [dependencies]
 itertools = "0.10.3"
-prost = "0.10.4"
+prost = "0.13.0"
 serde = { version = "1.0.139", features = ["derive"] }
 serde_json = "1.0.82"
 thiserror = "1.0.31"
 
 [build-dependencies]
 bindgen = "0.66.1"
 clippy = { version = "0.0.302", optional = true }
+prost = "0.13.0"
 prost-build = "0.10.4"
+prost-types = "0.13.0"
 fs_extra = "1.2.0"
 cc = "1.0.83"
 glob = "0.3.1"
+heck = "0.4"
 
 [dev-dependencies]
 easy-parallel = "3.2.0"

build.rs

@@ -3,12 +3,128 @@
 
 use fs_extra::dir::CopyOptions;
 use glob::glob;
+use heck::ToUpperCamelCase;
+use prost::Message;
+use prost_types::field_descriptor_proto::Type;
+use std::collections::{BTreeSet, HashMap, HashSet};
 use std::env;
 use std::path::{Path, PathBuf};
 
 static SOURCE_DIRECTORY: &str = "libpg_query";
 static LIBRARY_NAME: &str = "pg_query";
 
+type Cardinality = prost_types::field_descriptor_proto::Label;
+
+struct Edge {
+    field: String,
+    message: usize,
+    cardinality: Cardinality,
+}
+
+/// Represents a directed labeled multigraph of Message types. Each vertex represents a message
+/// type. An edge A->B is a tuple (field_name: String, type: FieldType), that states that
+/// Message A has a field (with name equal to `field_name`) of Message type B.
+struct MessageGraph {
+    messages: HashMap<String, usize>,
+
+    /// For each vertex A, the list of edges from A to other vertices, and a set of vertices B such that there is at least one edge B->A
+    edges: Vec<(String, Vec<Edge>, BTreeSet<usize>)>,
+}
+
+impl MessageGraph {
+    fn new() -> Self {
+        Self { messages: HashMap::new(), edges: Vec::new() }
+    }
+
+    /// Get the ID for a given `type_name` if it exists, or generate a new one if it doesn't
+    fn id_for(&mut self, type_name: &str) -> usize {
+        if let Some(id) = self.messages.get(type_name) {
+            *id
+        } else {
+            let id = self.edges.len();
+            self.edges.push((type_name.to_string(), Vec::new(), BTreeSet::new()));
+            self.messages.insert(type_name.to_string(), id);
+            id
+        }
+    }
+
+    /// Parse protobuf files and populate the graph with its Messages and corresponding edges
+    fn make(&mut self, fds: prost_types::FileDescriptorSet) {
+        for fd in fds.file {
+            let package = fd.package().to_string();
+            for msg in fd.message_type {
+                let full_name = format!(".{}.{}", package, msg.name());
+                let id = self.id_for(&full_name);
+
+                // We use this to check for duplicate fields
+                let mut fields: HashSet<String> = HashSet::new();
+
+                if msg.name() != "Node" && msg.name() != "A_Const" {
+                    for field in &msg.field {
+                        if field.r#type() != Type::Message {
+                            continue;
+                        }
+
+                        if field.oneof_index.is_some() {
+                            panic!("No support for enums: field {} of message {}", field.name(), msg.name());
+                        }
+
+                        if !fields.insert(field.name().to_string()) {
+                            panic!("Duplicate field: {}", field.name());
+                        }
+
+                        let message_id = self.id_for(field.type_name());
+                        self.edges[id].1.push(Edge { field: field.name().to_string(), message: message_id, cardinality: field.label() });
+                        self.edges[message_id].2.insert(id);
+                    }
+                }
+            }
+        }
+    }
+
+    /// Set `filter[x] = true` for all vertices `x` with a path to vertex `id`
+    fn filter_incoming(&self, id: usize, filter: &mut Vec<bool>) {
+        if !filter[id] {
+            filter[id] = true;
+            for nb in self.edges[id].2.iter() {
+                self.filter_incoming(*nb, filter);
+            }
+        }
+    }
+
+    /// Generate code for `unpack` impls for all Message types
+    fn write(&self, buf: &mut String) {
+        let mut filter = vec![false; self.messages.len()];
+        self.filter_incoming(*self.messages.get(".pg_query.Node").unwrap(), &mut filter);
+        for (id, (name, edges, _incoming)) in self.edges.iter().enumerate() {
+            let filtered = filter[id];
+            let short_name = &name[name.rfind(".").unwrap() + 1..].to_upper_camel_case();
+            if short_name == "Node" || short_name == "ParseResult" || short_name == "ScanResult" || short_name == "ScanToken" {
+                continue;
+            }
+
+            buf.push_str(&format!("impl<'a> Unpack<'a> for protobuf::{} {{\n", short_name));
+            if filtered && edges.iter().any(|e| filter[e.message]) {
+                buf.push_str("    fn unpack(&'a self, vec: &mut VecDeque<NodeRef<'a>>) {\n");
+                for edge in edges.iter() {
+                    if filter[edge.message] {
+                        match edge.cardinality {
+                            Cardinality::Repeated => buf.push_str(&format!("        self.{}.iter().for_each(|n| n.unpack(vec));\n", edge.field)),
+                            Cardinality::Required => buf.push_str(&format!("        vec.push_back(self.{});\n", edge.field)),
+                            Cardinality::Optional => {
+                                buf.push_str(&format!("        if let Some(ref e) = self.{} {{ e.unpack(vec); }}\n", edge.field))
+                            }
+                        }
+                    }
+                }
+                buf.push_str("    }\n}\n\n");
+            } else {
+                buf.push_str("    fn unpack(&'a self, _vec: &mut VecDeque<NodeRef<'a>>) { }\n}\n");
+            }
+        }
+    }
+}
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let out_dir = PathBuf::from(env::var("OUT_DIR")?);
     let build_path = Path::new(".").join(SOURCE_DIRECTORY);
@@ -65,6 +181,18 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .write_to_file(out_dir.join("bindings.rs"))?;
 
     // Generate the protobuf definition
+    let mut config = prost_build::Config::new();
+    let fds_path = out_dir.join("./file_descriptor_set.bin");
+    config.file_descriptor_set_path(fds_path.clone());
+    config.compile_protos(&[&out_protobuf_path.join(LIBRARY_NAME).with_extension("proto")], &[&out_protobuf_path])?;
+
+    let mut buf = String::new();
+    let fds = prost_types::FileDescriptorSet::decode(std::fs::read(fds_path)?.as_slice())?;
+    let mut graph = MessageGraph::new();
+    graph.make(fds);
+    graph.write(&mut buf);
+    std::fs::write(out_dir.join("./unpack.rs"), buf)?;
+
     prost_build::compile_protos(&[&out_protobuf_path.join(LIBRARY_NAME).with_extension("proto")], &[&out_protobuf_path])?;
 
     Ok(())

src/bindings.rs

@@ -3,4 +3,5 @@
 #![allow(non_snake_case)]
 #![allow(unused)]
 #![allow(clippy::all)]
+
 include!(concat!(env!("OUT_DIR"), "/bindings.rs"));

src/node_enum.rs

@@ -25,6 +25,7 @@ impl NodeEnum {
         while !iter.is_empty() {
             let (node, depth, context, has_filter_columns) = iter.remove(0);
             let depth = depth + 1;
+
             match node {
                 //
                 // The following statement types do not modify tables