Add iter() on ParseResult and NodeRef

This allows walking over all nodes in the AST, instead of just a limited
subset as in the current `nodes()` function. See pganalyze#31

The implementation uses static code generation in `build.rs`. The
protobuf definitions are parsed, and a graph of all Message types is
constructed. All NodeRef types are given an `unpack()` function, that
recursively calls `unpack()` on all relevant fields (i.e., the fields
that have a Node type, or that have a type that eventually has a Node
type as a field).

The result is guaranteed to visit all nodes. The code generation
mechanism is maybe also useful to replace parts of the codebase that
currently need to be manually hardcoded.

Adds prost, prost-types and heck to the build dependencies, and updates
the prost dependency version.

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