[fud2] Serialize Plans into Assignment Lists

fud2/fud-core/Cargo.toml

@@ -14,9 +14,10 @@ description = "Library for building declarative build tools"
 argh.workspace = true
 cranelift-entity = "0.103.0"
 serde.workspace = true
+serde_json.workspace = true
 figment = { version = "0.10.12", features = ["toml"] }
 pathdiff = { version = "0.2.1", features = ["camino"] }
-camino = "1.1.6"
+camino = { version = "1.1.6", features = ["serde1"] }
 anyhow.workspace = true
 log.workspace = true
 env_logger.workspace = true
@@ -31,6 +32,7 @@ libc = "0.2.174"
 
 [dev-dependencies]
 rand_chacha = "0.3.1"
+insta = "1.43.1"
 
 [features]
 egg_planner = ["dep:egg"]

fud2/fud-core/src/ast_converter.rs

@@ -0,0 +1,135 @@
+use std::{collections::HashMap, ops};
+
+use camino::Utf8PathBuf;
+use cranelift_entity::PrimaryMap;
+
+use crate::{
+    exec::{IO, OpRef, Operation},
+    flang::ast::{
+        Assignment, AssignmentList, Op, Visitable, Visitor, VisitorResult,
+    },
+};
+
+pub fn steps_to_ast(
+    plan: &Vec<(OpRef, Vec<IO>, Vec<IO>)>,
+    ops: &PrimaryMap<OpRef, Operation>,
+) -> AssignmentList {
+    let mut ast = AssignmentList { assigns: vec![] };
+    for step in plan {
+        let vars = step
+            .1
+            .iter()
+            .map(|v| match v {
+                IO::StdIO(utf8_path_buf) => utf8_path_buf,
+                IO::File(utf8_path_buf) => utf8_path_buf,
+            })
+            .cloned()
+            .collect();
+        let args = step
+            .2
+            .iter()
+            .map(|v| match v {
+                IO::StdIO(utf8_path_buf) => utf8_path_buf,
+                IO::File(utf8_path_buf) => utf8_path_buf,
+            })
+            .cloned()
+            .collect();
+
+        let fun = Op {
+            name: ops[step.0].name.clone(),
+            args,
+        };
+
+        let assignment = Assignment { vars, value: fun };
+        ast.assigns.push(assignment);
+    }
+
+    ast
+}
+
+/// A struct to convert a flang AST into the steps of a `Plan`.
+struct ASTToStepList {
+    step_list: Vec<(OpRef, Vec<IO>, Vec<IO>)>,
+    name_to_op_ref: HashMap<String, OpRef>,
+}
+
+impl ASTToStepList {
+    fn from_ops(ops: &PrimaryMap<OpRef, Operation>) -> Self {
+        let name_to_op_ref =
+            ops.iter().map(|(k, v)| (v.name.clone(), k)).collect();
+        ASTToStepList {
+            step_list: vec![],
+            name_to_op_ref,
+        }
+    }
+
+    fn step_list_from_ast(
+        &mut self,
+        ast: &AssignmentList,
+    ) -> Vec<(OpRef, Vec<IO>, Vec<IO>)> {
+        self.step_list = vec![];
+        let _ = ast.visit(self);
+        self.step_list.clone()
+    }
+}
+
+impl Visitor for ASTToStepList {
+    type Result = ops::ControlFlow<()>;
+
+    fn visit_assignment(&mut self, a: &Assignment) -> Self::Result {
+        let vars = a.vars.iter().map(|s| IO::File(s.clone())).collect();
+        let args = a.value.args.iter().map(|s| IO::File(s.clone())).collect();
+        let op_ref = self.name_to_op_ref[&a.value.name];
+
+        self.step_list.push((op_ref, vars, args));
+        Self::Result::output()
+    }
+}
+
+/// Given a flang AST and a set of ops, returns the steps of a `Plan` which the flang AST
+/// represents.
+pub fn ast_to_steps(
+    ast: &AssignmentList,
+    ops: &PrimaryMap<OpRef, Operation>,
+) -> Vec<(OpRef, Vec<IO>, Vec<IO>)> {
+    ASTToStepList::from_ops(ops).step_list_from_ast(ast)
+}
+
+#[derive(Default)]
+struct ASTToString {
+    assigns: Vec<String>,
+}
+
+impl ASTToString {
+    fn new() -> Self {
+        ASTToString { assigns: vec![] }
+    }
+
+    fn string_from_ast(&mut self, ast: &AssignmentList) -> String {
+        self.assigns = vec![];
+        let _ = ast.visit(self);
+        self.assigns.join("\n")
+    }
+}
+
+impl Visitor for ASTToString {
+    type Result = ops::ControlFlow<()>;
+
+    fn visit_assignment(&mut self, a: &Assignment) -> Self::Result {
+        let var_vec: Vec<String> =
+            a.vars.iter().map(Utf8PathBuf::to_string).collect();
+        let vars = var_vec.join(", ");
+        let arg_vec: Vec<String> =
+            a.value.args.iter().map(Utf8PathBuf::to_string).collect();
+        let args = arg_vec.join(", ");
+        let assign_string = format!("{} = {}({});", vars, a.value.name, args);
+        self.assigns.push(assign_string);
+        Self::Result::output()
+    }
+}
+
+/// Returns a pretty printed string from a flang AST. The returned string will be valid flang
+/// syntax.
+pub fn ast_to_string(ast: &AssignmentList) -> String {
+    ASTToString::new().string_from_ast(ast)
+}

fud2/fud-core/src/cli.rs

@@ -18,6 +18,8 @@ enum Mode {
     Generate,
     Run,
     Cmds,
+    FlangPlan,
+    JsonPlan,
 }
 
 impl FromStr for Mode {
@@ -31,6 +33,8 @@ impl FromStr for Mode {
             "run" => Ok(Mode::Run),
             "dot" => Ok(Mode::ShowDot),
             "cmds" => Ok(Mode::Cmds),
+            "flang-plan" => Ok(Mode::FlangPlan),
+            "json-plan" => Ok(Mode::JsonPlan),
             _ => Err("unknown mode".to_string()),
         }
     }
@@ -45,6 +49,8 @@ impl Display for Mode {
             Mode::Run => write!(f, "run"),
             Mode::ShowDot => write!(f, "dot"),
             Mode::Cmds => write!(f, "cmds"),
+            Mode::FlangPlan => write!(f, "flang-plan"),
+            Mode::JsonPlan => write!(f, "json-plan"),
         }
     }
 }
@@ -57,6 +63,8 @@ enum Planner {
     #[cfg(feature = "egg_planner")]
     Egg,
     Enumerate,
+    FromFlang,
+    FromJson,
 }
 
 impl FromStr for Planner {
@@ -68,6 +76,8 @@ impl FromStr for Planner {
             #[cfg(feature = "egg_planner")]
             "egg" => Ok(Planner::Egg),
             "enumerate" => Ok(Planner::Enumerate),
+            "flang" => Ok(Planner::FromFlang),
+            "json" => Ok(Planner::FromJson),
             _ => Err("unknown planner".to_string()),
         }
     }
@@ -80,6 +90,8 @@ impl Display for Planner {
             #[cfg(feature = "egg_planner")]
             Planner::Egg => write!(f, "egg"),
             Planner::Enumerate => write!(f, "enumerate"),
+            Planner::FromFlang => write!(f, "flang"),
+            Planner::FromJson => write!(f, "json"),
         }
     }
 }
@@ -166,7 +178,7 @@ pub struct FudArgs<T: CliExt> {
     #[argh(option)]
     to: Vec<String>,
 
-    /// execution mode (run, plan, emit, gen, dot)
+    /// execution mode (run, plan, emit, gen, dot, flang-plan, json-plan)
     #[argh(option, short = 'm', default = "Mode::Run")]
     mode: Mode,
 
@@ -305,6 +317,8 @@ fn get_request<T: CliExt>(
             #[cfg(feature = "egg_planner")]
             Planner::Egg => Box::new(plan::EggPlanner {}),
             Planner::Enumerate => Box::new(plan::EnumeratePlanner {}),
+            Planner::FromFlang => Box::new(plan::FlangPlanner {}),
+            Planner::FromJson => Box::new(plan::JsonPlanner {}),
         },
         timing_csv: args.timing_csv.clone(),
     })
@@ -523,6 +537,8 @@ fn cli_ext<T: CliExt>(
             args.force_rebuild,
             csv_path,
         )?,
+        Mode::FlangPlan => run.show_ops_flang(),
+        Mode::JsonPlan => run.show_ops_json(),
     }
 
     Ok(())

fud2/fud-core/src/exec/driver.rs

@@ -1,9 +1,22 @@
-use super::{OpRef, Operation, Request, Setup, SetupRef, State, StateRef};
-use crate::{run, script, utils};
+use super::{
+    OpRef, Operation, Request, Setup, SetupRef, State, StateRef,
+    plan::PlannerType,
+};
+use crate::{
+    ast_converter::ast_to_steps,
+    flang::{error::Wrap, session::ParseSession},
+    run, script, utils,
+};
 use camino::{Utf8Path, Utf8PathBuf};
 use cranelift_entity::PrimaryMap;
 use rand::distributions::{Alphanumeric, DistString};
-use std::{collections::HashMap, error::Error, ffi::OsStr, fmt::Display};
+use std::{
+    collections::HashMap,
+    error::Error,
+    ffi::OsStr,
+    fmt::Display,
+    io::{self, Read},
+};
 
 type FileData = HashMap<&'static str, &'static [u8]>;
 
@@ -168,6 +181,53 @@ impl Driver {
     /// This works by searching for a path through the available operations from the input state
     /// to the output state. If no such path exists in the operation graph, we return None.
     pub fn plan(&self, req: &Request) -> Option<Plan> {
+        // Special case if the planner is the one which reads from stdin.
+        let planner_ty = req.planner.ty();
+        if let PlannerType::FromFlang | PlannerType::FromJson = planner_ty {
+            let mut stdin = io::stdin().lock();
+            let mut input = String::new();
+            let res = stdin.read_to_string(&mut input);
+            if let Err(e) = res {
+                eprintln!("error: {e}");
+                return None;
+            }
+
+            let p = ParseSession::with_str_buf(&input);
+            let ast = if matches!(planner_ty, PlannerType::FromFlang) {
+                p.parse()
+            } else {
+                serde_json::from_str(&input)
+                    .map_err(|e| Wrap::new(e) as Wrap<dyn Error>)
+            };
+            match ast {
+                Err(e) => {
+                    eprintln!("error: {e}");
+                    return None;
+                }
+                Ok(ast) => {
+                    let steps = ast_to_steps(&ast, &self.ops);
+                    let results = self.gen_names(
+                        &req.end_states,
+                        req,
+                        false,
+                        &req.end_states,
+                    );
+                    let inputs = self.gen_names(
+                        &req.start_states,
+                        req,
+                        true,
+                        &req.start_states,
+                    );
+                    return Some(Plan {
+                        steps,
+                        inputs,
+                        results,
+                        workdir: req.workdir.clone(),
+                    });
+                }
+            }
+        }
+
         // Find a plan through the states.
         let path = req.planner.find_plan(
             &req.start_states,

fud2/fud-core/src/exec/plan/egg_planner.rs

@@ -55,7 +55,7 @@ use crate::exec::State;
 
 use super::{
     super::{OpRef, Operation, StateRef},
-    FindPlan, Step,
+    FindPlan, PlannerType, Step,
 };
 use cranelift_entity::PrimaryMap;
 use egg::{
@@ -301,4 +301,8 @@ impl FindPlan for EggPlanner {
             })
             .filter(|steps| !steps.is_empty())
     }
+
+    fn ty(&self) -> PlannerType {
+        PlannerType::Egg
+    }
 }

fud2/fud-core/src/exec/plan/enumerative_planner.rs

@@ -2,7 +2,7 @@ use crate::exec::State;
 
 use super::{
     super::{OpRef, Operation, StateRef},
-    FindPlan, Step,
+    FindPlan, PlannerType, Step,
 };
 use cranelift_entity::PrimaryMap;
 
@@ -177,4 +177,8 @@ impl FindPlan for EnumeratePlanner {
     ) -> Option<Vec<Step>> {
         Self::find_plan(start, end, through, ops)
     }
+
+    fn ty(&self) -> PlannerType {
+        PlannerType::Enumerative
+    }
 }