juno_scheduler/src/pm.rs

@@ -281,6 +281,13 @@ impl ModuleSelection {
     }
 }
 
+#[derive(Debug, Clone)]
+pub enum SetOp {
+    Difference,
+    Union,
+    Intersection,
+}
+
 #[derive(Debug, Clone)]
 pub enum Value {
     Label {
@@ -303,7 +310,7 @@ pub enum Value {
         selection: Vec<Value>,
     },
     SetOp {
-        op: parser::SetOp,
+        op: SetOp,
         lhs: Box<Value>,
         rhs: Box<Value>,
     },
@@ -364,9 +371,9 @@ impl Value {
                 let lhs = lhs.as_selection(pm, funcs)?;
                 let rhs = rhs.as_selection(pm, funcs)?;
                 let res = match op {
-                    parser::SetOp::Union => lhs.union(rhs),
-                    parser::SetOp::Intersection => lhs.intersection(rhs),
-                    parser::SetOp::Difference => lhs.difference(rhs),
+                    SetOp::Union => lhs.union(rhs),
+                    SetOp::Intersection => lhs.intersection(rhs),
+                    SetOp::Difference => lhs.difference(rhs),
                 };
                 selection.add(res);
             }
@@ -1301,17 +1308,62 @@ fn interp_expr(
         ScheduleExp::Integer { val } => Ok((Value::Integer { val: *val }, false)),
         ScheduleExp::Boolean { val } => Ok((Value::Boolean { val: *val }, false)),
         ScheduleExp::String { val } => Ok((Value::String { val: val.clone() }, false)),
-        ScheduleExp::SetOp { op, lhs, rhs } => {
+        ScheduleExp::UnaryOp { op, exp } => {
+            let (exp, modified) = interp_expr(pm, exp, stringtab, env, functions)?;
+            match op {
+                parser::UnaryOp::Not => {
+                    let Value::Boolean { val } = exp else {
+                        return Err(SchedulerError::SemanticError(
+                            "Not can only be applied to boolean values".to_string(),
+                        ));
+                    };
+                    Ok((Value::Boolean { val: !val }, modified))
+                }
+            }
+        }
+        ScheduleExp::BinaryOp { op, lhs, rhs } => {
             let (lhs, lhs_mod) = interp_expr(pm, lhs, stringtab, env, functions)?;
             let (rhs, rhs_mod) = interp_expr(pm, rhs, stringtab, env, functions)?;
-            Ok((
-                Value::SetOp {
-                    op: *op,
-                    lhs: Box::new(lhs),
-                    rhs: Box::new(rhs),
-                },
-                lhs_mod || rhs_mod,
-            ))
+            match op {
+                parser::BinaryOp::Difference
+                | parser::BinaryOp::Union
+                | parser::BinaryOp::Intersection => Ok((
+                    Value::SetOp {
+                        op: match op {
+                            parser::BinaryOp::Difference => SetOp::Difference,
+                            parser::BinaryOp::Union => SetOp::Union,
+                            parser::BinaryOp::Intersection => SetOp::Intersection,
+                            _ => panic!(),
+                        },
+                        lhs: Box::new(lhs),
+                        rhs: Box::new(rhs),
+                    },
+                    lhs_mod || rhs_mod,
+                )),
+                parser::BinaryOp::Or | parser::BinaryOp::And => {
+                    let (Value::Boolean { val: lhs }, Value::Boolean { val: rhs }) = (lhs, rhs)
+                    else {
+                        return Err(SchedulerError::SemanticError(format!(
+                            "{} can only be applied to boolean values",
+                            match op {
+                                parser::BinaryOp::Or => "or",
+                                parser::BinaryOp::And => "and",
+                                _ => panic!(),
+                            },
+                        )));
+                    };
+                    Ok((
+                        Value::Boolean {
+                            val: match op {
+                                parser::BinaryOp::Or => lhs || rhs,
+                                parser::BinaryOp::And => lhs && rhs,
+                                _ => panic!(),
+                            },
+                        },
+                        lhs_mod || rhs_mod,
+                    ))
+                }
+            }
         }
         ScheduleExp::Field { collect, field } => {
             let (lhs, changed) = interp_expr(pm, collect, stringtab, env, functions)?;
@@ -2780,7 +2832,7 @@ fn run_pass(
             }
             pm.delete_gravestones();
             pm.clear_analyses();
-        },
+        }
         Pass::ForkFission => {
             assert!(args.len() == 1);
 
@@ -2794,7 +2846,6 @@ fn run_pass(
                 });
             };
 
-
             pm.make_fork_join_maps();
             pm.make_loops();
             pm.make_reduce_cycles();
@@ -2817,18 +2868,18 @@ fn run_pass(
 
             let fork_label = fork_labels[0].label;
 
-            for ((((func, fork_join_map), loop_tree), reduce_cycles), nodes_in_fork_joins,
-            ) in build_selection(pm, selection, false)
-                .into_iter()
-                .zip(fork_join_maps.iter())
-                .zip(loops.iter())
-                .zip(reduce_cycles.iter())
-                .zip(nodes_in_fork_joins.iter())
+            for ((((func, fork_join_map), loop_tree), reduce_cycles), nodes_in_fork_joins) in
+                build_selection(pm, selection, false)
+                    .into_iter()
+                    .zip(fork_join_maps.iter())
+                    .zip(loops.iter())
+                    .zip(reduce_cycles.iter())
+                    .zip(nodes_in_fork_joins.iter())
             {
                 let Some(mut func) = func else {
                     continue;
                 };
-                
+
                 let new_forks = fork_fission(
                     &mut func,
                     nodes_in_fork_joins,
@@ -2837,12 +2888,12 @@ fn run_pass(
                     fork_join_map,
                     fork_label,
                 );
-                
+
                 let created_fork_joins = &mut created_fork_joins[func.func_id().idx()];
 
                 for f in new_forks {
                     created_fork_joins.push(f);
-                };
+                }
 
                 changed |= func.modified();
             }
@@ -2899,8 +2950,7 @@ fn run_pass(
             result = Value::Record {
                 fields: new_fork_joins,
             };
-
-        },
+        }
         Pass::ForkFissionBufferize => {
             assert!(args.len() == 1 || args.len() == 2);
             let Some(Value::Label {