diff --git a/src/lib.rs b/src/lib.rs index b889378..12f5b71 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,7 @@ pub mod builtins; pub mod interpreter; pub mod lexer; pub mod parser; +pub mod scheduler; pub mod value; pub use ast::*; diff --git a/src/scheduler.rs b/src/scheduler.rs new file mode 100644 index 0000000..3f09034 --- /dev/null +++ b/src/scheduler.rs @@ -0,0 +1,125 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +use std::collections::{BTreeMap, BTreeSet}; + +use anyhow::Result; + +#[derive(Debug)] +pub struct Definition<'a> { + // The variable being defined. + // This can be an empty string to indicate that + // no variable is being defined. + pub var: &'a str, + + // Other variables in the same scope used to compute + // the value of this variable. + pub used_vars: Vec<&'a str>, +} + +#[derive(Debug)] +pub struct StmtInfo<'a> { + // A statement can define multiple variables. + // A variable can also be defined by multiple statement. + pub definitions: Vec>, +} + +#[derive(Debug)] +pub enum SortResult { + // The order in which statements must be executed. + Order(Vec), + // List of statements comprising a cycle for a given var. + Cycle(String, Vec), +} + +pub fn schedule<'a>(infos: &mut [StmtInfo<'a>]) -> Result { + // Mapping from each var to the list of statements that define it. + let mut defining_stmts: BTreeMap<&'a str, Vec> = BTreeMap::new(); + + // For each statement, interate through its definitions and add the + // statement (index) to the var's defining-statements list. + for (idx, info) in infos.iter().enumerate() { + for defn in &info.definitions { + defining_stmts.entry(defn.var).or_default().push(idx); + } + } + + // Order of execution for statements. + let mut order = vec![]; + order.reserve(infos.len()); + + // Keep track of whether a var has been defined or not. + let mut defined_vars = BTreeSet::new(); + + // Keep track of whether a statement has been scheduled or not. + let mut scheduled = vec![false; infos.len()]; + + // List of vars to be processed. + let mut vars_to_process: Vec<&'a str> = defining_stmts.keys().cloned().collect(); + let mut tmp = vec![]; + + let mut process_var = |var| { + let mut stmt_scheduled = false; + let mut reprocess_var = false; + // Loop through each statement that defines the var. + for stmt_idx in defining_stmts.entry(var).or_default().iter().cloned() { + // If the statement has already been scheduled, skip it. + if scheduled[stmt_idx] { + continue; + } + + // In the statement, find the defn for the var. + for defn in &infos[stmt_idx].definitions { + if defn.var != var { + continue; + } + + // If all the vars used by the definition are defined, + // then the statement can be scheduled. + if defn.used_vars.iter().all(|v| defined_vars.contains(v)) { + // Schedule the statement. + order.push(stmt_idx); + scheduled[stmt_idx] = true; + + // Mark the var as defined. + defined_vars.insert(var); + stmt_scheduled = true; + } else { + reprocess_var = true; + } + } + } + + (stmt_scheduled, reprocess_var) + }; + + let mut done = false; + while !done { + done = true; + + // Swap with temporary vec. + std::mem::swap(&mut vars_to_process, &mut tmp); + + // Loop through each unscheduled var. + for var in tmp.iter().cloned() { + let (stmt_scheduled, reprocess_var) = process_var(var); + + if stmt_scheduled { + done = false; + + // If a statement has been scheduled, it means that the + // var has been defined. Process "" (statements that don't define any var) + // to see if any statements that depend on var can be scheduled. + // Doing so allows statements like `x > 10` to be scheduled immediately after x has been defined. + // TODO: Also schedule statements like `y = x > 10` immediately. + process_var(""); + } + + if reprocess_var { + vars_to_process.push(var); + } + } + } + + // TODO: determine cycles. + Ok(SortResult::Order(order)) +} diff --git a/tests/scheduler/mod.rs b/tests/scheduler/mod.rs new file mode 100644 index 0000000..c24565c --- /dev/null +++ b/tests/scheduler/mod.rs @@ -0,0 +1,104 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use anyhow::{bail, Result}; + +use regorus::scheduler::*; + +fn make_info<'a>(definitions: &[(&'a str, &[&'a str])]) -> StmtInfo<'a> { + StmtInfo { + definitions: definitions + .iter() + .map(|d| Definition { + var: d.0, + used_vars: d.1.to_vec(), + }) + .collect(), + } +} + +fn print_stmts(stmts: &[&str], order: &[usize]) { + for idx in order.iter().cloned() { + println!("{}", stmts[idx]); + } +} + +fn check_result(stmts: &[&str], expected: &[&str], r: SortResult) -> Result<()> { + match r { + SortResult::Order(order) => { + print_stmts(stmts, &order); + for (i, o) in order.iter().cloned().enumerate() { + assert_eq!(stmts[o], expected[i]); + } + Ok(()) + } + _ => bail!("scheduling failed"), + } +} + +#[test] +fn case1() -> Result<()> { + let stmts = vec![ + "v = x", + "x > 10", + "x = y + z", + "y = [1, 2, 4][_]", + "z = [4, 8][_]", + "x = 5", + "v = 1", + ]; + + let expected = vec![ + "v = 1", + "v = x", + "x = 5", + "x > 10", + "y = [1, 2, 4][_]", + "z = [4, 8][_]", + "x = y + z", + ]; + + let mut infos = vec![ + make_info(&[("v", &["x"]), ("x", &["v"])]), + make_info(&[("", &["x"])]), + make_info(&[("x", &["y", "z"])]), + make_info(&[("y", &[])]), + make_info(&[("z", &[])]), + make_info(&[("x", &[])]), + make_info(&[("v", &[])]), + ]; + + check_result(&stmts[..], &expected[..], schedule(&mut infos)?) +} + +#[test] +#[ignore = "destructing needs more thought. Hoist exprs and introduce new assignments?"] +fn case2() -> Result<()> { + let stmts = vec!["[x, y+1] = [y, p]", "value = x + p", "y = 5"]; + + let expected = vec!["y = 5", "[x, y+1] = [y, p]", "value = x + p"]; + + let mut infos = vec![ + make_info(&[("y", &[])]), + make_info(&[("value", &["x", "p"])]), + make_info(&[("x", &["y"]), ("y", &["x"]), ("p", &["y"])]), + ]; + + check_result(&stmts[..], &expected[..], schedule(&mut infos)?) +} + +#[test] +fn case2_rewritten() -> Result<()> { + let stmts = vec!["y+1 = p", "x = y", "value = x + p", "y = 5"]; + + let expected = vec!["y = 5", "y+1 = p", "x = y", "value = x + p"]; + + let mut infos = vec![ + make_info(&[("p", &["y"])]), + make_info(&[("x", &["y"])]), + make_info(&[("value", &["x", "p"])]), + make_info(&[("y", &[])]), + ]; + + check_result(&stmts[..], &expected[..], schedule(&mut infos)?) +} diff --git a/tests/tests.rs b/tests/tests.rs index 3ec840a..c735c9a 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -4,4 +4,5 @@ mod interpreter; mod lexer; mod parser; +mod scheduler; mod value;