Statement scheduler (WIP)

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

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::*;

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<Definition<'a>>,
+}
+
+#[derive(Debug)]
+pub enum SortResult {
+    // The order in which statements must be executed.
+    Order(Vec<usize>),
+    // List of statements comprising a cycle for a given var.
+    Cycle(String, Vec<usize>),
+}
+
+pub fn schedule<'a>(infos: &mut [StmtInfo<'a>]) -> Result<SortResult> {
+    // Mapping from each var to the list of statements that define it.
+    let mut defining_stmts: BTreeMap<&'a str, Vec<usize>> = 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))
+}

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)?)
+}

tests/tests.rs

@@ -4,4 +4,5 @@
 mod interpreter;
 mod lexer;
 mod parser;
+mod scheduler;
 mod value;