Add --output-directory/-d, child index info

Author: JuanCortezEclyp

treeedb/src/cli.rs

@@ -1,6 +1,7 @@
 use std::fs;
 use std::io;
 use std::io::Read;
+use std::path::{Path, PathBuf};
 use std::process;
 
 use anyhow::{Context, Result};
@@ -44,14 +45,15 @@ fn handle_parse_errors(path: &str, tree: &Tree, on_parse_error: &OnParseError) {
 /// Generate Datalog facts from source code
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
-// TODO(lb): Output directory, default is current directory
-// #[arg(short, long, default_value = None)]
-// pub output: Option<String>,
 pub struct Args {
     /// Behavior on parse errors
     #[arg(long, default_value_t = OnParseError::Warn, value_name = "CHOICE")]
     on_parse_error: OnParseError,
 
+    /// Output directory
+    #[arg(short, long, default_value = ".", value_name = "OUT_DIR")]
+    pub output_directory: PathBuf,
+
     /// Source code to consume; if empty, parse from stdin
     #[arg(value_name = "SRC_FILE")]
     pub source_files: Vec<String>,
@@ -75,13 +77,19 @@ fn stdin_string() -> Result<String> {
     Ok(stdin_str)
 }
 
+fn create_consumer(output_directory: &Path) -> Result<super::wide::WideCsvConsumer> {
+    // TODO(lb): Create consumer based on config
+    // For now, just use the wide CSV consumer as the default
+    Ok(super::wide::WideCsvConsumer::new(
+        output_directory.join("node.csv"),
+        output_directory.join("field.csv"),
+        output_directory.join("child.csv"),
+    )?)
+}
+
 pub fn main(language: tree_sitter::Language) -> Result<()> {
     let args = Args::parse();
-    let mut fc = super::wide::WideCsvConsumer::new(
-        "node.csv".into(),
-        "field.csv".into(),
-        "child.csv".into(),
-    )?;
+    let mut fc = create_consumer(&args.output_directory)?;
     if args.source_files.is_empty() {
         let content = stdin_string()?;
         let tree = parse(language, &content)?;

treeedb/src/consumer.rs

@@ -22,7 +22,7 @@ pub trait FactConsumer {
         child: &Node<'_>,
     ) -> Result<(), Self::Err>;
 
-    fn child(&mut self, parent: &Node<'_>, child: &Node<'_>) -> Result<(), Self::Err>;
+    fn child(&mut self, parent: &Node<'_>, index: u32, child: &Node<'_>) -> Result<(), Self::Err>;
 
     fn node(&mut self, node: &Node<'_>, source: &[u8]) -> Result<(), Self::Err>;
 }

treeedb/src/lib.rs

@@ -21,8 +21,8 @@ pub fn facts<E>(
                 fc.field(&node, name, &child)?;
             }
         }
-        for child in node.named_children(&mut cursor) {
-            fc.child(&node, &child)?;
+        for (i, child) in node.named_children(&mut cursor).enumerate() {
+            fc.child(&node, i as u32, &child)?;
             nodes.push(child);
         }
     }

treeedb/src/narrow.rs

@@ -35,7 +35,12 @@ impl FactConsumer for NarrowCsvConsumer {
         Ok(())
     }
 
-    fn child(&mut self, _parent: &Node<'_>, _child: &Node<'_>) -> Result<(), Self::Err> {
+    fn child(
+        &mut self,
+        _parent: &Node<'_>,
+        _index: u32,
+        _child: &Node<'_>,
+    ) -> Result<(), Self::Err> {
         Ok(())
     }
 

treeedb/src/wide.rs

@@ -41,9 +41,12 @@ impl FactConsumer for WideCsvConsumer {
         Ok(())
     }
 
-    fn child(&mut self, parent: &Node<'_>, child: &Node<'_>) -> Result<(), Self::Err> {
-        self.child
-            .write_record([&parent.id().to_string(), &child.id().to_string()])?;
+    fn child(&mut self, parent: &Node<'_>, index: u32, child: &Node<'_>) -> Result<(), Self::Err> {
+        self.child.write_record([
+            &parent.id().to_string(),
+            &index.to_string(),
+            &child.id().to_string(),
+        ])?;
         Ok(())
     }
 

treeedbgen-souffle/src/gen.rs

@@ -107,7 +107,7 @@ fn node_with_fields(
             )?;
             writeln!(
                 w,
-                "{}(x, as(y, {})) :- {}{}(x), {}child(x, y).",
+                "{}(x, as(y, {})) :- {}{}(x), {}child(x, _, y).",
                 child_relation_name,
                 child_type_name,
                 config.relation_prefix,
@@ -219,6 +219,7 @@ fn declare_node(config: &PrivGenConfig, w: &mut impl Write) -> Result<(), GenErr
     writeln!(w, ".type {}EndRow <: number", config.type_prefix)?;
     writeln!(w, ".type {}EndCol <: number", config.type_prefix)?;
     writeln!(w, ".type {}NodeText <: symbol", config.type_prefix)?;
+    writeln!(w, ".type {}NodeIndex <: number", config.type_prefix)?;
     writeln!(
         w,
         ".decl {}node({})",
@@ -294,6 +295,7 @@ fn declare_child(config: &PrivGenConfig, w: &mut impl Write) -> Result<(), GenEr
         config.relation_prefix,
         [
             format!("parent: {}Node", config.type_prefix),
+            format!("index: {}NodeIndex", config.type_prefix),
             format!("child: {}Node", config.type_prefix),
         ]
         .join(", ")