From 27a91ac9ecce41c76102110f26078ac25e4a5c37 Mon Sep 17 00:00:00 2001
From: Ayaka Yorihiro <ayakayorihiro@gmail.com>
Date: Mon, 21 Oct 2024 14:30:45 -0400
Subject: [PATCH] Add new tool to obtain names of groups

---
 Cargo.lock                         | 13 ++++
 Cargo.toml                         |  1 +
 tools/component_groups/Cargo.toml  | 28 +++++++++
 tools/component_groups/src/main.rs | 99 ++++++++++++++++++++++++++++++
 4 files changed, 141 insertions(+)
 create mode 100644 tools/component_groups/Cargo.toml
 create mode 100644 tools/component_groups/src/main.rs

diff --git a/Cargo.lock b/Cargo.lock
index 79563a1ba..56c81de6e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -735,6 +735,19 @@ dependencies = [
  "serde_json",
 ]
 
+[[package]]
+name = "component_groups"
+version = "0.7.1"
+dependencies = [
+ "argh",
+ "calyx-frontend",
+ "calyx-ir",
+ "calyx-opt",
+ "calyx-utils",
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "console"
 version = "0.15.8"
diff --git a/Cargo.toml b/Cargo.toml
index d509d9a1f..aa8b637a4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,6 +19,7 @@ members = [
     "tools/calyx-pass-explorer",
     "tools/cider-data-converter",
     "tools/component_cells",
+    "tools/component_groups",
     "tools/yxi",
     "tools/calyx-writer",
 ]
diff --git a/tools/component_groups/Cargo.toml b/tools/component_groups/Cargo.toml
new file mode 100644
index 000000000..cbc2c2abb
--- /dev/null
+++ b/tools/component_groups/Cargo.toml
@@ -0,0 +1,28 @@
+[package]
+name = "component_groups"
+authors.workspace = true
+license-file.workspace = true
+keywords.workspace = true
+repository.workspace = true
+readme.workspace = true
+description.workspace = true
+categories.workspace = true
+homepage.workspace = true
+edition.workspace = true
+version.workspace = true
+rust-version.workspace = true
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+serde.workspace = true
+argh.workspace = true
+serde_json = "1.0.79"
+
+calyx-utils = { path = "../../calyx-utils" }
+calyx-frontend = { path = "../../calyx-frontend" }
+calyx-opt = { path = "../../calyx-opt" }
+
+[dependencies.calyx-ir]
+path = "../../calyx-ir"
+features = ["serialize"]
diff --git a/tools/component_groups/src/main.rs b/tools/component_groups/src/main.rs
new file mode 100644
index 000000000..d3e501c14
--- /dev/null
+++ b/tools/component_groups/src/main.rs
@@ -0,0 +1,99 @@
+use argh::FromArgs;
+use calyx_frontend as frontend;
+use calyx_ir::{self as ir, Id};
+use calyx_utils::{CalyxResult, OutputFile};
+use serde::Serialize;
+use std::path::{Path, PathBuf};
+use std::{collections::HashSet, io};
+
+#[derive(FromArgs)]
+/// Path for library and path for file to read from
+struct Args {
+    /// file path to read data from
+    #[argh(positional, from_str_fn(read_path))]
+    file_path: Option<PathBuf>,
+
+    /// library path
+    #[argh(option, short = 'l', default = "Path::new(\".\").into()")]
+    pub lib_path: PathBuf,
+
+    /// output file
+    #[argh(option, short = 'o', default = "OutputFile::Stdout")]
+    pub output: OutputFile,
+}
+
+fn read_path(path: &str) -> Result<PathBuf, String> {
+    Ok(Path::new(path).into())
+}
+
+fn main() -> CalyxResult<()> {
+    let p: Args = argh::from_env();
+
+    let ws = frontend::Workspace::construct(&p.file_path, &p.lib_path)?;
+
+    let ctx: ir::Context = ir::from_ast::ast_to_ir(ws)?;
+
+    let main_comp = ctx.entrypoint();
+
+    let mut component_info: HashSet<ComponentInfo> = HashSet::new();
+    gen_component_info(&ctx, main_comp, true, &mut component_info);
+    write_json(component_info.clone(), p.output)?;
+    Ok(())
+}
+
+fn id_serialize_passthrough<S>(id: &Id, ser: S) -> Result<S::Ok, S::Error>
+where
+    S: serde::Serializer,
+{
+    id.to_string().serialize(ser)
+}
+
+#[derive(PartialEq, Eq, Hash, Clone, Serialize)]
+struct ComponentInfo {
+    #[serde(serialize_with = "id_serialize_passthrough")]
+    pub component: Id,
+    pub is_main_component: bool,
+    pub groups: Vec<Id>, // list of all groups in the component
+}
+
+/// Accumulates a set of components to the groups that they contain
+/// in the program with entrypoint `main_comp`.
+fn gen_component_info(
+    ctx: &ir::Context,
+    comp: &ir::Component,
+    is_main_comp: bool,
+    component_info: &mut HashSet<ComponentInfo>,
+) {
+    let curr_comp_info = ComponentInfo {
+        component: comp.name,
+        is_main_component: is_main_comp,
+        groups: comp
+            .groups
+            .into_iter()
+            .map(|g| g.borrow().name())
+            .collect::<Vec<_>>(),
+    };
+    // recurse into any other cells
+    for cell in comp.cells.iter() {
+        let cell_ref = cell.borrow();
+        if let ir::CellType::Component { name } = cell_ref.prototype {
+            let component = ctx
+                .components
+                .iter()
+                .find(|comp| comp.name == name)
+                .unwrap();
+            gen_component_info(ctx, component, false, component_info);
+        }
+    }
+    component_info.insert(curr_comp_info);
+}
+
+/// Write the collected set of component information to a JSON file.
+fn write_json(
+    component_info: HashSet<ComponentInfo>,
+    file: OutputFile,
+) -> Result<(), io::Error> {
+    let created_vec: Vec<ComponentInfo> = component_info.into_iter().collect();
+    serde_json::to_writer_pretty(file.get_write(), &created_vec)?;
+    Ok(())
+}