Skip to content

Commit

Permalink
Add iter() on ParseResult and NodeRef
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
absporl committed Jul 23, 2024
1 parent 5562e4a commit 2d05243
Show file tree
Hide file tree
Showing 7 changed files with 518 additions and 20 deletions.
60 changes: 47 additions & 13 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 4 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
127 changes: 127 additions & 0 deletions build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,127 @@

use fs_extra::dir::CopyOptions;
use glob::glob;
use prost::Message;
use prost_types::field_descriptor_proto::Type;
use std::collections::{BTreeSet, HashMap, HashSet};
use std::env;
use std::path::{Path, PathBuf};
//use prost_build::ident::to_snake;
use heck::ToUpperCamelCase;

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() }
}

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
}
}

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);

// buf.push_str(&format!("Message: {}\n", msg.name()));
// let my_edges: Vec<(String, FieldType)> = Vec::new();
//
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);
}
}
}
}
}

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);
}
}
}

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> 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);
Expand Down Expand Up @@ -65,6 +180,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(())
Expand Down
3 changes: 3 additions & 0 deletions src/bindings.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,4 +3,7 @@
#![allow(non_snake_case)]
#![allow(unused)]
#![allow(clippy::all)]

use prost::Message;

include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
1 change: 1 addition & 0 deletions src/node_enum.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading

0 comments on commit 2d05243

Please sign in to comment.