-
Notifications
You must be signed in to change notification settings - Fork 490
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Refactor topo-order alg out of sierra-gas.
In preparation for ap-change solve. commit-id:eb28e7c1
- Loading branch information
Showing
5 changed files
with
137 additions
and
104 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
pub mod topological_order; |
110 changes: 110 additions & 0 deletions
110
crates/cairo-lang-sierra/src/algorithm/topological_order.rs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,110 @@ | ||
use crate::program::StatementIdx; | ||
|
||
/// The status of a node during the topological ordering finding algorithm. | ||
#[derive(Clone, Debug)] | ||
enum TopologicalOrderStatus { | ||
/// The computation for that statement did not start. | ||
NotStarted, | ||
/// The computation is in progress. | ||
InProgress, | ||
/// The computation was completed, and all the children were visited. | ||
Done, | ||
} | ||
|
||
/// Returns the topological ordering. | ||
/// `detect_cycles` - if true, the function will return an error if a cycle is detected. | ||
/// `roots` - the roots of the graph. | ||
/// `node_count` - the number of nodes in the graph. | ||
/// `get_children` - a function that returns the children of a node. | ||
/// `out_of_bounds_err` - a function that returns an error for a node is out of bounds. | ||
/// `cycle_err` - a function that returns an error for a node that is part of a cycle. | ||
/// Note: Will only work properly if the nodes are in the range [0, node_count). | ||
pub fn get_topological_ordering<E>( | ||
detect_cycles: bool, | ||
roots: impl Iterator<Item = StatementIdx>, | ||
node_count: usize, | ||
get_children: impl Fn(StatementIdx) -> Result<Vec<StatementIdx>, E>, | ||
out_of_bounds_err: impl Fn(StatementIdx) -> E, | ||
cycle_err: impl Fn(StatementIdx) -> E, | ||
) -> Result<Vec<StatementIdx>, E> { | ||
let mut ordering = vec![]; | ||
let mut status = vec![TopologicalOrderStatus::NotStarted; node_count]; | ||
for root in roots { | ||
calculate_topological_ordering( | ||
detect_cycles, | ||
&mut ordering, | ||
&mut status, | ||
root, | ||
&get_children, | ||
&out_of_bounds_err, | ||
&cycle_err, | ||
)?; | ||
} | ||
Ok(ordering) | ||
} | ||
|
||
/// Calculates the topological ordering starting from `root`. For more info see | ||
/// `get_topological_ordering`. | ||
fn calculate_topological_ordering<E>( | ||
detect_cycles: bool, | ||
ordering: &mut Vec<StatementIdx>, | ||
status: &mut [TopologicalOrderStatus], | ||
root: StatementIdx, | ||
get_children: &impl Fn(StatementIdx) -> Result<Vec<StatementIdx>, E>, | ||
out_of_bounds_err: &impl Fn(StatementIdx) -> E, | ||
cycle_err: &impl Fn(StatementIdx) -> E, | ||
) -> Result<(), E> { | ||
// A stack of statements to visit. | ||
// When the pair is popped out of the stack, `is_done=true` means that we've already visited | ||
// all of its children, and we just need to add it to the ordering. | ||
let mut stack = vec![root]; | ||
|
||
while let Some(idx) = stack.pop() { | ||
match status.get(idx.0) { | ||
Some(TopologicalOrderStatus::NotStarted) => { | ||
// Mark the statement as `InProgress`. | ||
status[idx.0] = TopologicalOrderStatus::InProgress; | ||
|
||
// Push the statement back to the stack, so that after visiting all | ||
// of its children, we would add it to the ordering. | ||
// Add the missing children on top of it. | ||
stack.push(idx); | ||
for child in get_children(idx)? { | ||
match status.get(child.0) { | ||
Some(TopologicalOrderStatus::InProgress) => { | ||
if detect_cycles { | ||
return Err(cycle_err(child)); | ||
} | ||
continue; | ||
} | ||
Some(TopologicalOrderStatus::Done) => { | ||
continue; | ||
} | ||
Some(TopologicalOrderStatus::NotStarted) => { | ||
stack.push(child); | ||
} | ||
None => { | ||
return Err(out_of_bounds_err(child)); | ||
} | ||
} | ||
} | ||
} | ||
Some(TopologicalOrderStatus::InProgress) => { | ||
// Mark the statement as `Done`. | ||
status[idx.0] = TopologicalOrderStatus::Done; | ||
|
||
// Add the element to the ordering after visiting all its children. | ||
// This gives us reverse topological ordering. | ||
ordering.push(idx); | ||
} | ||
Some(TopologicalOrderStatus::Done) => { | ||
// Do nothing. | ||
} | ||
None => { | ||
return Err(out_of_bounds_err(idx)); | ||
} | ||
} | ||
} | ||
|
||
Ok(()) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters