Add feature testing

juno_scheduler/src/compile.rs

@@ -22,6 +22,7 @@ pub enum ScheduleCompilerError {
         actual: usize,
         loc: Location,
     },
+    SemanticError(String, Location),
 }
 
 impl fmt::Display for ScheduleCompilerError {
@@ -46,6 +47,11 @@ impl fmt::Display for ScheduleCompilerError {
                 "({}, {}) -- ({}, {}): Expected {} arguments, found {}",
                 loc.0 .0, loc.0 .1, loc.1 .0, loc.1 .1, expected, actual
             ),
+            ScheduleCompilerError::SemanticError(msg, loc) => write!(
+                f,
+                "({}, {}) -- ({}, {}): {}",
+                loc.0 .0, loc.0 .1, loc.1 .0, loc.1 .1, msg,
+            ),
         }
     }
 }
@@ -76,6 +82,8 @@ enum Appliable {
     // DeleteUncalled requires special handling because it changes FunctionIDs, so it is not
     // treated like a pass
     DeleteUncalled,
+    // Test whether a feature is enabled
+    Feature,
     Schedule(Schedule),
     Device(Device),
 }
@@ -85,6 +93,8 @@ impl Appliable {
     fn is_valid_num_args(&self, num: usize) -> bool {
         match self {
             Appliable::Pass(pass) => pass.is_valid_num_args(num),
+            // Testing whether a feature is enabled takes the feature instead of a selection, so it
+            // has 0 arguments
             // Delete uncalled, Schedules, and devices do not take arguments
             _ => num == 0,
         }
@@ -158,6 +168,8 @@ impl FromStr for Appliable {
             "serialize" => Ok(Appliable::Pass(ir::Pass::Serialize)),
             "write-predication" => Ok(Appliable::Pass(ir::Pass::WritePredication)),
 
+            "feature" => Ok(Appliable::Feature),
+
             "print" => Ok(Appliable::Pass(ir::Pass::Print)),
 
             "cpu" | "llvm" => Ok(Appliable::Device(Device::LLVM)),
@@ -409,6 +421,17 @@ fn compile_expr(
                         on: selection,
                     }))
                 }
+                Appliable::Feature => match selection {
+                    ir::Selector::Selection(mut args) if args.len() == 1 => {
+                        Ok(ExprResult::Expr(ir::ScheduleExp::Feature {
+                            feature: Box::new(args.pop().unwrap()),
+                        }))
+                    }
+                    _ => Err(ScheduleCompilerError::SemanticError(
+                        "feature requires exactly one argument as its selection".to_string(),
+                        lexer.line_col(span),
+                    )),
+                },
                 Appliable::Schedule(sched) => Ok(ExprResult::Stmt(ir::ScheduleStmt::AddSchedule {
                     sched,
                     on: selection,

juno_scheduler/src/ir.rs

@@ -121,6 +121,9 @@ pub enum ScheduleExp {
     DeleteUncalled {
         on: Selector,
     },
+    Feature {
+        feature: Box<ScheduleExp>,
+    },
     Record {
         fields: Vec<(String, ScheduleExp)>,
     },

juno_scheduler/src/pm.rs

@@ -16,6 +16,7 @@ use juno_utils::stringtab::StringTable;
 
 use std::cell::RefCell;
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
+use std::env;
 use std::fmt;
 use std::fs::File;
 use std::io::Write;
@@ -1460,6 +1461,23 @@ fn interp_expr(
                 changed,
             ))
         }
+        ScheduleExp::Feature { feature } => {
+            let (feature, modified) = interp_expr(pm, &*feature, stringtab, env, functions)?;
+            let Value::String { val } = feature else {
+                return Err(SchedulerError::SemanticError(
+                    "Feature expects a single string argument (instead of a selection)".to_string(),
+                ));
+            };
+            // To test for features, the scheduler needs to be invoked from a build script so that
+            // Cargo provides the enabled features via environment variables
+            let key = val.to_uppercase().replace("-", "_");
+            Ok((
+                Value::Boolean {
+                    val: env::var(key).is_ok(),
+                },
+                modified,
+            ))
+        }
         ScheduleExp::Record { fields } => {
             let mut result = HashMap::new();
             let mut changed = false;