Support dependencies between vars defined in same statement

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2023-02-26 07:26:54 -08:00
committed by Anand Krishnamoorthi
parent 34eee1f17f
commit 8e0ddbcff2
2 changed files with 109 additions and 31 deletions

View File

@@ -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<SortResult> {
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.
}
}
}

View File

@@ -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"])]),