ROX-36764: Add rules for default work arguments

src/main.rs

@@ -62,19 +62,21 @@ struct Args {
 fn run_script(script_path: String) -> Vec<(i32, u64)> {
     info!("Loading script: {:?}", script_path);
 
-    let ast: Vec<Node> =
+    let nodes: Vec<Node> =
         parse_instructions(&std::fs::read_to_string(script_path).unwrap())
             .unwrap();
 
+    let prepared_nodes = apply_rules(nodes);
+
     let (machine, works): (Vec<_>, Vec<_>) =
-        ast.iter().partition_map(|node| match node {
+        prepared_nodes.iter().partition_map(|node| match node {
             Node::Work { .. } => Either::Right(node),
             Node::Machine { .. } => Either::Left(node),
         });
 
     let _ = apply(machine);
 
-    apply_rules(works)
+    works
         .into_iter()
         .flat_map(|node| {
             debug!("AST NODE: {:?}", node);
@@ -83,19 +85,11 @@ fn run_script(script_path: String) -> Vec<(i32, u64)> {
                 unreachable!()
             };
 
-            let workers: u32 = args
-                .get("workers")
-                .cloned()
-                .unwrap_or(String::from("0"))
-                .parse()
-                .unwrap();
+            let workers: u32 =
+                args.get("workers").cloned().unwrap().parse().unwrap();
 
-            let duration: u64 = args
-                .get("duration")
-                .cloned()
-                .unwrap_or(String::from("0"))
-                .parse()
-                .unwrap();
+            let duration: u64 =
+                args.get("duration").cloned().unwrap().parse().unwrap();
 
             (0..workers)
                 .filter_map(|_| {
@@ -339,4 +333,46 @@ mod tests {
         new_script_worker(ast[1].clone()).run_payload().unwrap();
         new_script_worker(ast[2].clone()).run_payload().unwrap();
     }
+
+    #[test]
+    fn test_default_work_args() {
+        let input = r#"
+            main () {
+              task(stub);
+            }
+        "#;
+
+        let nodes: Vec<Node> = parse_instructions(input).unwrap();
+        assert_eq!(nodes.len(), 1);
+
+        let prepared_nodes = apply_rules(nodes);
+
+        let Node::Work { ref args, .. } = prepared_nodes[0] else {
+            unreachable!()
+        };
+
+        assert_eq!(args.get("workers").cloned().unwrap(), "1".to_string());
+        assert_eq!(args.get("duration").cloned().unwrap(), "0".to_string());
+    }
+
+    #[test]
+    fn test_custom_work_args() {
+        let input = r#"
+            main (workers = 2, duration = 10) {
+              task(stub);
+            }
+        "#;
+
+        let nodes: Vec<Node> = parse_instructions(input).unwrap();
+        assert_eq!(nodes.len(), 1);
+
+        let prepared_nodes = apply_rules(nodes);
+
+        let Node::Work { ref args, .. } = prepared_nodes[0] else {
+            unreachable!()
+        };
+
+        assert_eq!(args.get("workers").cloned().unwrap(), "2".to_string());
+        assert_eq!(args.get("duration").cloned().unwrap(), "10".to_string());
+    }
 }

src/script/rules.rs

@@ -1,12 +1,73 @@
 use log::debug;
 
-use crate::script::ast::Node;
+use crate::script::ast::{Instruction, Node};
+use std::collections::HashMap;
 
-/// Contains a list of transformation to apply after parsing
+fn apply_instruction_rules(
+    instructions: &[Instruction],
+    _node: &Node,
+) -> Vec<Instruction> {
+    instructions.to_vec()
+}
+
+fn apply_arg_rules(
+    args: &HashMap<String, String>,
+    _node: &Node,
+) -> HashMap<String, String> {
+    let mut new_args = args.clone();
+
+    new_args.entry("workers".to_string()).or_insert_with(|| {
+        debug!("Applying default number of workers");
+        "1".to_string()
+    });
+
+    new_args.entry("duration".to_string()).or_insert_with(|| {
+        debug!("Applying default duration");
+        "0".to_string()
+    });
+
+    new_args
+}
+
+fn apply_work_rules(work: Node) -> Node {
+    let Node::Work {
+        ref name,
+        ref args,
+        ref instructions,
+        ref dist,
+    } = work
+    else {
+        unreachable!()
+    };
+
+    Node::Work {
+        name: name.clone(),
+        args: apply_arg_rules(args, &work),
+        instructions: apply_instruction_rules(instructions, &work),
+        dist: dist.clone(),
+    }
+}
+
+fn apply_machine_rules(machine: Node) -> Node {
+    machine
+}
+
+fn apply_node_rules(node: Node) -> Node {
+    match node {
+        Node::Work { .. } => apply_work_rules(node),
+        Node::Machine { .. } => apply_machine_rules(node),
+    }
+}
+
+/// Contains a list of transformation to apply after parsing.
+/// Note that transformation does not update AST in place, but
+/// rather provides an isolated copy of it. This may introduce
+/// some parsing overhead of course, and has to be re-evaluated
+/// every now and then.
+///
 /// TODO: Add following rules:
 /// - add path if directory is expected
-/// - add default worker arguments
-pub fn apply_rules(works: Vec<&Node>) -> Vec<&Node> {
+pub fn apply_rules(nodes: Vec<Node>) -> Vec<Node> {
     debug!("Applying rules");
-    works
+    nodes.into_iter().map(apply_node_rules).collect::<Vec<_>>()
 }