From 8e0ddbcff243911dba7792f92df1cc94dc833be3 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Sun, 26 Feb 2023 07:26:54 -0800 Subject: [PATCH] Support dependencies between vars defined in same statement Signed-off-by: Anand Krishnamoorthi --- src/scheduler.rs | 101 ++++++++++++++++++++++++++++++++--------- tests/scheduler/mod.rs | 39 ++++++++++++---- 2 files changed, 109 insertions(+), 31 deletions(-) diff --git a/src/scheduler.rs b/src/scheduler.rs index 3f09034..87e7946 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -1,6 +1,6 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; use anyhow::Result; @@ -57,35 +57,92 @@ pub fn schedule<'a>(infos: &mut [StmtInfo<'a>]) -> Result { let mut vars_to_process: Vec<&'a str> = defining_stmts.keys().cloned().collect(); let mut tmp = vec![]; + let mut queue = VecDeque::new(); + let mut schedule_stmt = |stmt_idx: usize| { + // Check if the statement has already been scheduled. + if scheduled[stmt_idx] { + return None; + } + + let definitions = &infos[stmt_idx].definitions; + + let can_be_scheduled = if definitions.len() == 1 { + // Handle the more common case of single definition statements optimally. + // Check if all the vars used by the definition are previously assigned. + definitions[0] + .used_vars + .iter() + .all(|uv| defined_vars.contains(uv)) + } else { + // Set of vars that can be defined in this statement. + let mut defined_in_stmt = BTreeSet::new(); + + // Add each definition to processing queue. + queue.clear(); + for defn in definitions { + queue.push_back(defn); + } + + while !queue.is_empty() { + let n = queue.len(); + for _ in 0..n { + let defn = queue.pop_front().unwrap(); + // Check if the vars used by this definition are + // 1) defined via prior assignments (or) + // 2) defined in current statement + if defn + .used_vars + .iter() + .all(|uv| defined_vars.contains(uv) || defined_in_stmt.contains(uv)) + { + defined_in_stmt.insert(defn.var); + } else { + // The definiton must be processed again. + queue.push_back(defn); + } + } + // If no definition became defined, then there is a cycle between + // the definitions in this statement. The cycle cannot be broken yet. + if n == queue.len() { + break; + } + } + + // If the vars used by all the definitions are already defined or + // can be defined by scheduling this statement, return true. + queue.is_empty() + }; + + // Schedule the var if possible. + if can_be_scheduled { + order.push(stmt_idx); + scheduled[stmt_idx] = true; + + // For each definition in the statement, mark its var as defined. + for defn in &infos[stmt_idx].definitions { + defined_vars.insert(defn.var); + } + Some(true) + } else { + Some(false) + } + }; + 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); + match schedule_stmt(stmt_idx) { + Some(true) => { stmt_scheduled = true; - } else { + } + Some(false) => { reprocess_var = true; } + None => { + // Statement has already been scheduled. + } } } diff --git a/tests/scheduler/mod.rs b/tests/scheduler/mod.rs index c24565c..186ae18 100644 --- a/tests/scheduler/mod.rs +++ b/tests/scheduler/mod.rs @@ -40,9 +40,9 @@ fn check_result(stmts: &[&str], expected: &[&str], r: SortResult) -> Result<()> fn case1() -> Result<()> { let stmts = vec![ "v = x", + "y = [1, 2, 4][_]", "x > 10", "x = y + z", - "y = [1, 2, 4][_]", "z = [4, 8][_]", "x = 5", "v = 1", @@ -60,9 +60,9 @@ fn case1() -> Result<()> { let mut infos = vec![ make_info(&[("v", &["x"]), ("x", &["v"])]), + make_info(&[("y", &[])]), make_info(&[("", &["x"])]), make_info(&[("x", &["y", "z"])]), - make_info(&[("y", &[])]), make_info(&[("z", &[])]), make_info(&[("x", &[])]), make_info(&[("v", &[])]), @@ -72,16 +72,26 @@ fn case1() -> Result<()> { } #[test] -#[ignore = "destructing needs more thought. Hoist exprs and introduce new assignments?"] +//#[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"]; + #[rustfmt::skip] + 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"]; + #[rustfmt::skip] + 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"])]), + make_info(&[("value", &["x", "p"])]), + make_info(&[("y", &[])]), ]; check_result(&stmts[..], &expected[..], schedule(&mut infos)?) @@ -89,9 +99,20 @@ fn case2() -> Result<()> { #[test] fn case2_rewritten() -> Result<()> { - let stmts = vec!["y+1 = p", "x = y", "value = x + p", "y = 5"]; + #[rustfmt::skip] + 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"]; + #[rustfmt::skip] + let expected = vec![ + "y = 5", + "y+1 = p", + "x = y", + "value = x + p" + ]; let mut infos = vec![ make_info(&[("p", &["y"])]),