From bb766608284a81a9ea2942f9bcdff4284e5df2ca Mon Sep 17 00:00:00 2001 From: Douglas Wilson Date: Thu, 26 Oct 2023 09:29:31 +0100 Subject: [PATCH 01/10] prototype walker --- src/lib.rs | 2 + src/walker.rs | 166 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 168 insertions(+) create mode 100644 src/walker.rs diff --git a/src/lib.rs b/src/lib.rs index 3b3f3671c..2a618b780 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,3 +25,5 @@ pub mod values; pub use crate::extension::Extension; pub use crate::hugr::{Direction, Hugr, HugrView, Node, Port, SimpleReplacement, Wire}; + +pub mod walker; diff --git a/src/walker.rs b/src/walker.rs new file mode 100644 index 000000000..316e65ed7 --- /dev/null +++ b/src/walker.rs @@ -0,0 +1,166 @@ +use std::rc::Rc; + +use lazy_static::__Deref; + +use crate::{ops::OpType, HugrView, Node}; + +#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone, Debug)] +pub enum WalkResult { + Advance, + Interrupt, +} + +#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone, Debug)] +pub enum WalkOrder { + Preorder, + Postorder, +} + +struct WalkerCallback<'a, T>(Box WalkResult>); + +impl<'a, T, F: 'a + Fn(Node, OpType, &mut T) -> WalkResult> From for WalkerCallback<'a, T> { + fn from(f: F) -> Self { + Self(Box::new(f)) + } +} + +pub struct Walker<'a, T> { + pre_callbacks: Vec>, + post_callbacks: Vec>, +} + +fn call_back( + n: Node, + o: OpType, + t: &mut T, + f: impl Fn(Node, O, &mut T) -> WalkResult, +) -> WalkResult +where + OpType: TryInto, +{ + match o.try_into() { + Ok(x) => f(n, x, t), + _ => WalkResult::Advance, + } +} + +enum WorkItem { + Visit(Node), + Callback(WalkOrder, Node), +} + +impl<'a, T> Walker<'a, T> { + pub fn new() -> Self { + Self { + pre_callbacks: Vec::new(), + post_callbacks: Vec::new(), + } + } + + pub fn visit WalkResult>( + &mut self, + walk_order: WalkOrder, + f: F, + ) -> &mut Self + where + OpType: TryInto, + { + let g = Rc::new(f); + let callbacks = match walk_order { + WalkOrder::Preorder => &mut self.pre_callbacks, + WalkOrder::Postorder => &mut self.post_callbacks, + }; + callbacks.push((move |n, o, t: &'_ mut _| call_back(n, o, t, g.as_ref())).into()); + self + } + + pub fn walk(&self, hugr: impl HugrView, t: &mut T) { + // We intentionally avoid recursion so that we can robustly accept very deep hugrs + let mut worklist = vec![WorkItem::Visit(hugr.root())]; + + while let Some(wi) = worklist.pop() { + match wi { + WorkItem::Visit(n) => { + worklist.push(WorkItem::Callback(WalkOrder::Postorder, n)); + // TODO we should add children in topological order + worklist.extend(hugr.children(n).map(WorkItem::Visit)); + worklist.push(WorkItem::Callback(WalkOrder::Preorder, n)); + } + WorkItem::Callback(order, n) => { + let callbacks = match order { + WalkOrder::Preorder => &self.pre_callbacks, + WalkOrder::Postorder => &self.post_callbacks, + }; + for cb in callbacks.iter() { + // this clone is unfortunate, to avoid this we would need a TryInto variant: + // try_into(&O) -> Option<&T> + if cb.0.as_ref()(n, hugr.get_optype(n).clone(), t) == WalkResult::Interrupt + { + return; + } + } + } + } + } + } +} + +#[cfg(test)] +mod test { + use std::{error::Error, iter::empty}; + + use crate::types::Signature; + use crate::{ + builder::{Container, HugrBuilder, ModuleBuilder}, + extension::{ExtensionRegistry, ExtensionSet}, + type_row, + types::FunctionType, + }; + + use super::*; + + fn test1() -> Result<(), Box> { + let mut module_builder = ModuleBuilder::new(); + let sig = Signature { + signature: FunctionType::new(type_row![], type_row![]), + input_extensions: ExtensionSet::new(), + }; + module_builder.define_function("f1", sig.clone()); + module_builder.define_function("f2", sig.clone()); + + let hugr = module_builder.finish_hugr(&ExtensionRegistry::new())?; + + let mut s = String::new(); + Walker::::new() + .visit(WalkOrder::Preorder, |_, crate::ops::Module, r| { + r.extend("pre".chars()); + r.extend(['m']); + WalkResult::Advance + }) + .visit(WalkOrder::Postorder, |_, crate::ops::Module, r| { + r.extend("post".chars()); + r.extend(['n']); + WalkResult::Advance + }) + .visit( + WalkOrder::Preorder, + |_, crate::ops::FuncDecl { ref name, .. }, r| { + r.extend("pre".chars()); + r.extend(name.chars()); + WalkResult::Advance + }, + ) + .visit( + WalkOrder::Postorder, + |_, crate::ops::FuncDecl { ref name, .. }, r| { + r.extend("post".chars()); + r.extend(name.chars()); + WalkResult::Advance + }, + ) + .walk(&hugr, &mut s); + + assert_eq!(s, "prempref1pref2postf2postf1postn"); + Ok(()) + } +} From 6b7615ccd5356012d769fa61f3757c38f2ce2718 Mon Sep 17 00:00:00 2001 From: Douglas Wilson Date: Thu, 26 Oct 2023 12:32:57 +0100 Subject: [PATCH 02/10] fix walker, rustfmt --- src/walker.rs | 90 +++++++++++++++++++++++++++------------------------ 1 file changed, 47 insertions(+), 43 deletions(-) diff --git a/src/walker.rs b/src/walker.rs index 316e65ed7..3e5f698b0 100644 --- a/src/walker.rs +++ b/src/walker.rs @@ -1,5 +1,7 @@ +#![allow(missing_docs)] use std::rc::Rc; +use itertools::Itertools; use lazy_static::__Deref; use crate::{ops::OpType, HugrView, Node}; @@ -16,31 +18,31 @@ pub enum WalkOrder { Postorder, } -struct WalkerCallback<'a, T>(Box WalkResult>); +struct WalkerCallback<'a, T, E>(Box Result>); -impl<'a, T, F: 'a + Fn(Node, OpType, &mut T) -> WalkResult> From for WalkerCallback<'a, T> { +impl<'a, T, E, F: 'a + Fn(Node, OpType, T) -> Result> From for WalkerCallback<'a, T, E> { fn from(f: F) -> Self { Self(Box::new(f)) } } -pub struct Walker<'a, T> { - pre_callbacks: Vec>, - post_callbacks: Vec>, +pub struct Walker<'a, T, E> { + pre_callbacks: Vec>, + post_callbacks: Vec>, } -fn call_back( +fn call_back( n: Node, o: OpType, - t: &mut T, - f: impl Fn(Node, O, &mut T) -> WalkResult, -) -> WalkResult + t: T, + f: &impl Fn(Node, O, T) -> Result, +) -> Result where OpType: TryInto, { match o.try_into() { Ok(x) => f(n, x, t), - _ => WalkResult::Advance, + _ => Ok(t), } } @@ -49,7 +51,7 @@ enum WorkItem { Callback(WalkOrder, Node), } -impl<'a, T> Walker<'a, T> { +impl<'a, T, E> Walker<'a, T, E> { pub fn new() -> Self { Self { pre_callbacks: Vec::new(), @@ -57,7 +59,7 @@ impl<'a, T> Walker<'a, T> { } } - pub fn visit WalkResult>( + pub fn visit Result>( &mut self, walk_order: WalkOrder, f: F, @@ -70,11 +72,11 @@ impl<'a, T> Walker<'a, T> { WalkOrder::Preorder => &mut self.pre_callbacks, WalkOrder::Postorder => &mut self.post_callbacks, }; - callbacks.push((move |n, o, t: &'_ mut _| call_back(n, o, t, g.as_ref())).into()); + callbacks.push((move |n, o, t| call_back(n, o, t, g.as_ref())).into()); self } - pub fn walk(&self, hugr: impl HugrView, t: &mut T) { + pub fn walk(&self, hugr: impl HugrView, mut t: T) -> Result { // We intentionally avoid recursion so that we can robustly accept very deep hugrs let mut worklist = vec![WorkItem::Visit(hugr.root())]; @@ -83,7 +85,8 @@ impl<'a, T> Walker<'a, T> { WorkItem::Visit(n) => { worklist.push(WorkItem::Callback(WalkOrder::Postorder, n)); // TODO we should add children in topological order - worklist.extend(hugr.children(n).map(WorkItem::Visit)); + let children = hugr.children(n).collect_vec(); + worklist.extend(children.into_iter().rev().map(WorkItem::Visit)); worklist.push(WorkItem::Callback(WalkOrder::Preorder, n)); } WorkItem::Callback(order, n) => { @@ -91,27 +94,26 @@ impl<'a, T> Walker<'a, T> { WalkOrder::Preorder => &self.pre_callbacks, WalkOrder::Postorder => &self.post_callbacks, }; + let optype = hugr.get_optype(n); for cb in callbacks.iter() { // this clone is unfortunate, to avoid this we would need a TryInto variant: // try_into(&O) -> Option<&T> - if cb.0.as_ref()(n, hugr.get_optype(n).clone(), t) == WalkResult::Interrupt - { - return; - } + t = cb.0.as_ref()(n, optype.clone(), t)?; } } } } + Ok(t) } } #[cfg(test)] mod test { - use std::{error::Error, iter::empty}; + use std::error::Error; use crate::types::Signature; use crate::{ - builder::{Container, HugrBuilder, ModuleBuilder}, + builder::{Container, HugrBuilder, ModuleBuilder, SubContainer}, extension::{ExtensionRegistry, ExtensionSet}, type_row, types::FunctionType, @@ -119,48 +121,50 @@ mod test { use super::*; + #[test] fn test1() -> Result<(), Box> { let mut module_builder = ModuleBuilder::new(); let sig = Signature { signature: FunctionType::new(type_row![], type_row![]), input_extensions: ExtensionSet::new(), }; - module_builder.define_function("f1", sig.clone()); - module_builder.define_function("f2", sig.clone()); + module_builder + .define_function("f1", sig.clone())? + .finish_sub_container()?; + module_builder + .define_function("f2", sig.clone())? + .finish_sub_container()?; let hugr = module_builder.finish_hugr(&ExtensionRegistry::new())?; - let mut s = String::new(); - Walker::::new() - .visit(WalkOrder::Preorder, |_, crate::ops::Module, r| { - r.extend("pre".chars()); - r.extend(['m']); - WalkResult::Advance + let s = Walker::<_, Box>::new() + .visit(WalkOrder::Preorder, |_, crate::ops::Module, mut r| { + r += ";prem"; + Ok(r) }) - .visit(WalkOrder::Postorder, |_, crate::ops::Module, r| { - r.extend("post".chars()); - r.extend(['n']); - WalkResult::Advance + .visit(WalkOrder::Postorder, |_, crate::ops::Module, mut r| { + r += ";postm"; + Ok(r) }) .visit( WalkOrder::Preorder, - |_, crate::ops::FuncDecl { ref name, .. }, r| { - r.extend("pre".chars()); - r.extend(name.chars()); - WalkResult::Advance + |_, crate::ops::FuncDefn { ref name, .. }, mut r| { + r += ";pre"; + r += name.as_ref(); + Ok(r) }, ) .visit( WalkOrder::Postorder, - |_, crate::ops::FuncDecl { ref name, .. }, r| { - r.extend("post".chars()); - r.extend(name.chars()); - WalkResult::Advance + |_, crate::ops::FuncDefn { ref name, .. }, mut r| { + r += ";post"; + r += name.as_ref(); + Ok(r) }, ) - .walk(&hugr, &mut s); + .walk(&hugr, String::new())?; - assert_eq!(s, "prempref1pref2postf2postf1postn"); + assert_eq!(s, ";prem;pref1;postf1;pref2;postf2;postm"); Ok(()) } } From 27503f4bbe0398ce82fca82632e0135487ff642d Mon Sep 17 00:00:00 2001 From: Douglas Wilson Date: Thu, 26 Oct 2023 13:26:38 +0100 Subject: [PATCH 03/10] tidying --- src/walker.rs | 95 ++++++++++++++++++++++++++++++--------------------- 1 file changed, 57 insertions(+), 38 deletions(-) diff --git a/src/walker.rs b/src/walker.rs index 3e5f698b0..3571ed81d 100644 --- a/src/walker.rs +++ b/src/walker.rs @@ -1,8 +1,8 @@ #![allow(missing_docs)] -use std::rc::Rc; + +use std::ops::{Deref, DerefMut}; use itertools::Itertools; -use lazy_static::__Deref; use crate::{ops::OpType, HugrView, Node}; @@ -18,9 +18,24 @@ pub enum WalkOrder { Postorder, } -struct WalkerCallback<'a, T, E>(Box Result>); +struct WalkerCallback<'a, T, E>(Box Result>); + +impl<'a, T, E> Deref for WalkerCallback<'a, T, E> { + type Target = dyn 'a + FnMut(Node, OpType, T) -> Result; + fn deref(&self) -> &Self::Target { + self.0.deref() + } +} + +impl<'a, T, E> DerefMut for WalkerCallback<'a, T, E> { + fn deref_mut(&mut self) -> &mut Self::Target { + self.0.deref_mut() + } +} -impl<'a, T, E, F: 'a + Fn(Node, OpType, T) -> Result> From for WalkerCallback<'a, T, E> { +impl<'a, T, E, F: 'a + FnMut(Node, OpType, T) -> Result> From + for WalkerCallback<'a, T, E> +{ fn from(f: F) -> Self { Self(Box::new(f)) } @@ -31,21 +46,6 @@ pub struct Walker<'a, T, E> { post_callbacks: Vec>, } -fn call_back( - n: Node, - o: OpType, - t: T, - f: &impl Fn(Node, O, T) -> Result, -) -> Result -where - OpType: TryInto, -{ - match o.try_into() { - Ok(x) => f(n, x, t), - _ => Ok(t), - } -} - enum WorkItem { Visit(Node), Callback(WalkOrder, Node), @@ -59,24 +59,44 @@ impl<'a, T, E> Walker<'a, T, E> { } } - pub fn visit Result>( + pub fn previsit Result>(&mut self, f: F) -> &mut Self + where + OpType: TryInto, + { + self.visit(WalkOrder::Preorder, f) + } + + pub fn postvisit Result>(&mut self, f: F) -> &mut Self + where + OpType: TryInto, + { + self.visit(WalkOrder::Postorder, f) + } + + fn mut_callbacks(&mut self, order: WalkOrder) -> &mut Vec> { + match order { + WalkOrder::Preorder => &mut self.pre_callbacks, + WalkOrder::Postorder => &mut self.post_callbacks, + } + } + + pub fn visit Result>( &mut self, walk_order: WalkOrder, - f: F, + mut f: F, ) -> &mut Self where OpType: TryInto, { - let g = Rc::new(f); - let callbacks = match walk_order { - WalkOrder::Preorder => &mut self.pre_callbacks, - WalkOrder::Postorder => &mut self.post_callbacks, + let cb = move |n, o: OpType, t| match o.try_into() { + Ok(x) => f(n, x, t), + _ => Ok(t), }; - callbacks.push((move |n, o, t| call_back(n, o, t, g.as_ref())).into()); + self.mut_callbacks(walk_order).push(cb.into()); self } - pub fn walk(&self, hugr: impl HugrView, mut t: T) -> Result { + pub fn walk(&mut self, hugr: impl HugrView, mut t: T) -> Result { // We intentionally avoid recursion so that we can robustly accept very deep hugrs let mut worklist = vec![WorkItem::Visit(hugr.root())]; @@ -86,19 +106,17 @@ impl<'a, T, E> Walker<'a, T, E> { worklist.push(WorkItem::Callback(WalkOrder::Postorder, n)); // TODO we should add children in topological order let children = hugr.children(n).collect_vec(); + // extend in reverse so that the first child is the top of the stack worklist.extend(children.into_iter().rev().map(WorkItem::Visit)); worklist.push(WorkItem::Callback(WalkOrder::Preorder, n)); } WorkItem::Callback(order, n) => { - let callbacks = match order { - WalkOrder::Preorder => &self.pre_callbacks, - WalkOrder::Postorder => &self.post_callbacks, - }; let optype = hugr.get_optype(n); - for cb in callbacks.iter() { - // this clone is unfortunate, to avoid this we would need a TryInto variant: + for cb in self.mut_callbacks(order).iter_mut() { + // this clone is unfortunate, to avoid this we would + // need a TryInto variant like: // try_into(&O) -> Option<&T> - t = cb.0.as_ref()(n, optype.clone(), t)?; + t = cb(n, optype.clone(), t)?; } } } @@ -115,6 +133,7 @@ mod test { use crate::{ builder::{Container, HugrBuilder, ModuleBuilder, SubContainer}, extension::{ExtensionRegistry, ExtensionSet}, + ops::{FuncDefn, Module}, type_row, types::FunctionType, }; @@ -138,17 +157,17 @@ mod test { let hugr = module_builder.finish_hugr(&ExtensionRegistry::new())?; let s = Walker::<_, Box>::new() - .visit(WalkOrder::Preorder, |_, crate::ops::Module, mut r| { + .visit(WalkOrder::Preorder, |_, Module, mut r| { r += ";prem"; Ok(r) }) - .visit(WalkOrder::Postorder, |_, crate::ops::Module, mut r| { + .visit(WalkOrder::Postorder, |_, Module, mut r| { r += ";postm"; Ok(r) }) .visit( WalkOrder::Preorder, - |_, crate::ops::FuncDefn { ref name, .. }, mut r| { + |_, FuncDefn { ref name, .. }, mut r| { r += ";pre"; r += name.as_ref(); Ok(r) @@ -156,7 +175,7 @@ mod test { ) .visit( WalkOrder::Postorder, - |_, crate::ops::FuncDefn { ref name, .. }, mut r| { + |_, FuncDefn { ref name, .. }, mut r| { r += ";post"; r += name.as_ref(); Ok(r) From 27f2595da764916694604d6b886786b4c095aa3c Mon Sep 17 00:00:00 2001 From: Douglas Wilson Date: Mon, 30 Oct 2023 13:27:51 +0000 Subject: [PATCH 04/10] derive Deref/DerefMut --- src/walker.rs | 17 ++--------------- 1 file changed, 2 insertions(+), 15 deletions(-) diff --git a/src/walker.rs b/src/walker.rs index 3571ed81d..03ccc39d0 100644 --- a/src/walker.rs +++ b/src/walker.rs @@ -1,7 +1,6 @@ #![allow(missing_docs)] -use std::ops::{Deref, DerefMut}; - +use derive_more::{Deref, DerefMut}; use itertools::Itertools; use crate::{ops::OpType, HugrView, Node}; @@ -18,21 +17,9 @@ pub enum WalkOrder { Postorder, } +#[derive(Deref, DerefMut)] struct WalkerCallback<'a, T, E>(Box Result>); -impl<'a, T, E> Deref for WalkerCallback<'a, T, E> { - type Target = dyn 'a + FnMut(Node, OpType, T) -> Result; - fn deref(&self) -> &Self::Target { - self.0.deref() - } -} - -impl<'a, T, E> DerefMut for WalkerCallback<'a, T, E> { - fn deref_mut(&mut self) -> &mut Self::Target { - self.0.deref_mut() - } -} - impl<'a, T, E, F: 'a + FnMut(Node, OpType, T) -> Result> From for WalkerCallback<'a, T, E> { From fcc78d4794978bd4eca60f9438ec66c2513961b8 Mon Sep 17 00:00:00 2001 From: Douglas Wilson Date: Mon, 30 Oct 2023 15:03:29 +0000 Subject: [PATCH 05/10] Traverse in topological order when possible --- src/walker.rs | 75 ++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 71 insertions(+), 4 deletions(-) diff --git a/src/walker.rs b/src/walker.rs index 03ccc39d0..1d19d60b8 100644 --- a/src/walker.rs +++ b/src/walker.rs @@ -1,5 +1,7 @@ #![allow(missing_docs)] +use std::collections::HashSet; + use derive_more::{Deref, DerefMut}; use itertools::Itertools; @@ -91,10 +93,31 @@ impl<'a, T, E> Walker<'a, T, E> { match wi { WorkItem::Visit(n) => { worklist.push(WorkItem::Callback(WalkOrder::Postorder, n)); - // TODO we should add children in topological order - let children = hugr.children(n).collect_vec(); + let mut pushed_children = HashSet::new(); + // We intend to only visit direct children. + // + // If the nodes children form a dataflow sibling graph we + // visit them in post dfs order. This traversal is not + // guaranteed to visit all nodes, only those reachable from + // the input node. (For example, LoadConstant nodes may not + // be reachable from the input node). So we do a second + // unordered traversal afterwards. + if let Some([input, _]) = hugr.get_io(n) { + use crate::hugr::views::PetgraphWrapper; + let wrapper = PetgraphWrapper { hugr: &hugr }; + let mut dfs = ::petgraph::visit::DfsPostOrder::new(&wrapper, input); + while let Some(x) = dfs.next(&wrapper) { + worklist.push(WorkItem::Visit(x)); + pushed_children.insert(x); + } + } + + let rest_children = hugr + .children(n) + .filter(|x| !pushed_children.contains(x)) + .collect_vec(); // extend in reverse so that the first child is the top of the stack - worklist.extend(children.into_iter().rev().map(WorkItem::Visit)); + worklist.extend(rest_children.into_iter().rev().map(WorkItem::Visit)); worklist.push(WorkItem::Callback(WalkOrder::Preorder, n)); } WorkItem::Callback(order, n) => { @@ -116,9 +139,13 @@ impl<'a, T, E> Walker<'a, T, E> { mod test { use std::error::Error; + use crate::builder::{Dataflow, DataflowHugr}; + use crate::extension::prelude::USIZE_T; + use crate::hugr::hugrmut::sealed::HugrMutInternals; + use crate::ops; use crate::types::Signature; use crate::{ - builder::{Container, HugrBuilder, ModuleBuilder, SubContainer}, + builder::{Container, FunctionBuilder, HugrBuilder, ModuleBuilder, SubContainer}, extension::{ExtensionRegistry, ExtensionSet}, ops::{FuncDefn, Module}, type_row, @@ -173,4 +200,44 @@ mod test { assert_eq!(s, ";prem;pref1;postf1;pref2;postf2;postm"); Ok(()) } + + struct Noop; + + impl TryFrom for Noop { + type Error = ops::OpType; + fn try_from(ot: ops::OpType) -> Result { + match ot { + ops::OpType::LeafOp(ops::LeafOp::Noop { .. }) => Ok(Noop), + x => Err(x), + } + } + } + #[test] + fn test2() -> Result<(), Box> { + use ops::handle::NodeHandle; + let sig = Signature { + signature: FunctionType::new(type_row![USIZE_T], type_row![USIZE_T]), + input_extensions: ExtensionSet::new(), + }; + let mut fun_builder = FunctionBuilder::new("f3", sig)?; + let [i] = fun_builder.input_wires_arr(); + let noop1 = fun_builder.add_dataflow_op(ops::LeafOp::Noop { ty: USIZE_T }, [i])?; + let noop2 = + fun_builder.add_dataflow_op(ops::LeafOp::Noop { ty: USIZE_T }, [noop1.out_wire(0)])?; + let mut h = fun_builder.finish_prelude_hugr_with_outputs([noop2.out_wire(0)])?; + h.hugr_mut() + .move_before_sibling(noop2.handle().node(), noop1.handle().node())?; + + let v = Walker::, Box>::new() + .previsit(|n, Noop, mut v| { + v.push(n); + Ok(v) + }) + .walk(&h, Vec::new())?; + assert_eq!( + &[noop1.handle().node(), noop2.handle().node()], + v.as_slice() + ); + Ok(()) + } } From 9973b72bcf3e764dad408ad5faa7589a0c5c1cbd Mon Sep 17 00:00:00 2001 From: Douglas Wilson Date: Mon, 30 Oct 2023 15:16:32 +0000 Subject: [PATCH 06/10] Callbacks need to reference the hugr, so ensure it is fixed in new --- src/walker.rs | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/src/walker.rs b/src/walker.rs index 1d19d60b8..4eb467bce 100644 --- a/src/walker.rs +++ b/src/walker.rs @@ -30,9 +30,10 @@ impl<'a, T, E, F: 'a + FnMut(Node, OpType, T) -> Result> From } } -pub struct Walker<'a, T, E> { +pub struct Walker<'a, H: HugrView, T, E> { pre_callbacks: Vec>, post_callbacks: Vec>, + hugr: &'a H, } enum WorkItem { @@ -40,11 +41,12 @@ enum WorkItem { Callback(WalkOrder, Node), } -impl<'a, T, E> Walker<'a, T, E> { - pub fn new() -> Self { +impl<'a, H: HugrView, T, E> Walker<'a, H, T, E> { + pub fn new(hugr: &'a H) -> Self { Self { pre_callbacks: Vec::new(), post_callbacks: Vec::new(), + hugr, } } @@ -85,9 +87,9 @@ impl<'a, T, E> Walker<'a, T, E> { self } - pub fn walk(&mut self, hugr: impl HugrView, mut t: T) -> Result { + pub fn walk(&mut self, mut t: T) -> Result { // We intentionally avoid recursion so that we can robustly accept very deep hugrs - let mut worklist = vec![WorkItem::Visit(hugr.root())]; + let mut worklist = vec![WorkItem::Visit(self.hugr.root())]; while let Some(wi) = worklist.pop() { match wi { @@ -102,9 +104,9 @@ impl<'a, T, E> Walker<'a, T, E> { // the input node. (For example, LoadConstant nodes may not // be reachable from the input node). So we do a second // unordered traversal afterwards. - if let Some([input, _]) = hugr.get_io(n) { + if let Some([input, _]) = self.hugr.get_io(n) { use crate::hugr::views::PetgraphWrapper; - let wrapper = PetgraphWrapper { hugr: &hugr }; + let wrapper = PetgraphWrapper { hugr: self.hugr }; let mut dfs = ::petgraph::visit::DfsPostOrder::new(&wrapper, input); while let Some(x) = dfs.next(&wrapper) { worklist.push(WorkItem::Visit(x)); @@ -112,7 +114,8 @@ impl<'a, T, E> Walker<'a, T, E> { } } - let rest_children = hugr + let rest_children = self + .hugr .children(n) .filter(|x| !pushed_children.contains(x)) .collect_vec(); @@ -121,7 +124,7 @@ impl<'a, T, E> Walker<'a, T, E> { worklist.push(WorkItem::Callback(WalkOrder::Preorder, n)); } WorkItem::Callback(order, n) => { - let optype = hugr.get_optype(n); + let optype = self.hugr.get_optype(n); for cb in self.mut_callbacks(order).iter_mut() { // this clone is unfortunate, to avoid this we would // need a TryInto variant like: @@ -170,7 +173,7 @@ mod test { let hugr = module_builder.finish_hugr(&ExtensionRegistry::new())?; - let s = Walker::<_, Box>::new() + let s = Walker::<_, _, Box>::new(&hugr) .visit(WalkOrder::Preorder, |_, Module, mut r| { r += ";prem"; Ok(r) @@ -195,7 +198,7 @@ mod test { Ok(r) }, ) - .walk(&hugr, String::new())?; + .walk(String::new())?; assert_eq!(s, ";prem;pref1;postf1;pref2;postf2;postm"); Ok(()) @@ -228,12 +231,12 @@ mod test { h.hugr_mut() .move_before_sibling(noop2.handle().node(), noop1.handle().node())?; - let v = Walker::, Box>::new() + let v = Walker::<_, Vec, Box>::new(&h) .previsit(|n, Noop, mut v| { v.push(n); Ok(v) }) - .walk(&h, Vec::new())?; + .walk(Vec::new())?; assert_eq!( &[noop1.handle().node(), noop2.handle().node()], v.as_slice() From 4134536e5e385317cfa59b77b7e76d1dc0d60c35 Mon Sep 17 00:00:00 2001 From: doug-q <141026920+doug-q@users.noreply.github.com> Date: Tue, 31 Oct 2023 16:41:51 +0000 Subject: [PATCH 07/10] Update src/walker.rs Co-authored-by: Alan Lawrence --- src/walker.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/walker.rs b/src/walker.rs index 4eb467bce..7df4c5d64 100644 --- a/src/walker.rs +++ b/src/walker.rs @@ -106,7 +106,7 @@ impl<'a, H: HugrView, T, E> Walker<'a, H, T, E> { // unordered traversal afterwards. if let Some([input, _]) = self.hugr.get_io(n) { use crate::hugr::views::PetgraphWrapper; - let wrapper = PetgraphWrapper { hugr: self.hugr }; + let wrapper = self.hugr.as_petgraph(); let mut dfs = ::petgraph::visit::DfsPostOrder::new(&wrapper, input); while let Some(x) = dfs.next(&wrapper) { worklist.push(WorkItem::Visit(x)); From e4cc9ebffe2a04dc6f7c447bbc7260b899560289 Mon Sep 17 00:00:00 2001 From: Douglas Wilson Date: Tue, 31 Oct 2023 09:27:30 +0000 Subject: [PATCH 08/10] tidy --- src/walker.rs | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/walker.rs b/src/walker.rs index 7df4c5d64..8bab54f41 100644 --- a/src/walker.rs +++ b/src/walker.rs @@ -7,12 +7,6 @@ use itertools::Itertools; use crate::{ops::OpType, HugrView, Node}; -#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone, Debug)] -pub enum WalkResult { - Advance, - Interrupt, -} - #[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone, Debug)] pub enum WalkOrder { Preorder, @@ -127,8 +121,7 @@ impl<'a, H: HugrView, T, E> Walker<'a, H, T, E> { let optype = self.hugr.get_optype(n); for cb in self.mut_callbacks(order).iter_mut() { // this clone is unfortunate, to avoid this we would - // need a TryInto variant like: - // try_into(&O) -> Option<&T> + // need a TryInto variant like: try_into(&O) -> Option<&T> t = cb(n, optype.clone(), t)?; } } From 9f2265c1209800b5cff70817bbe17659eb076fd9 Mon Sep 17 00:00:00 2001 From: Douglas Wilson Date: Thu, 2 Nov 2023 08:09:21 +0000 Subject: [PATCH 09/10] better comments, more examples --- src/walker.rs | 111 +++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 92 insertions(+), 19 deletions(-) diff --git a/src/walker.rs b/src/walker.rs index 8bab54f41..157a76ce5 100644 --- a/src/walker.rs +++ b/src/walker.rs @@ -30,11 +30,6 @@ pub struct Walker<'a, H: HugrView, T, E> { hugr: &'a H, } -enum WorkItem { - Visit(Node), - Callback(WalkOrder, Node), -} - impl<'a, H: HugrView, T, E> Walker<'a, H, T, E> { pub fn new(hugr: &'a H) -> Self { Self { @@ -82,8 +77,17 @@ impl<'a, H: HugrView, T, E> Walker<'a, H, T, E> { } pub fn walk(&mut self, mut t: T) -> Result { + enum WorkItem { + Visit(Node), + Callback(WalkOrder, Node), + } + impl From for WorkItem { + fn from(n: Node) -> Self { + WorkItem::Visit(n) + } + } // We intentionally avoid recursion so that we can robustly accept very deep hugrs - let mut worklist = vec![WorkItem::Visit(self.hugr.root())]; + let mut worklist = vec![self.hugr.root().into()]; while let Some(wi) = worklist.pop() { match wi { @@ -93,35 +97,45 @@ impl<'a, H: HugrView, T, E> Walker<'a, H, T, E> { // We intend to only visit direct children. // // If the nodes children form a dataflow sibling graph we - // visit them in post dfs order. This traversal is not - // guaranteed to visit all nodes, only those reachable from - // the input node. (For example, LoadConstant nodes may not - // be reachable from the input node). So we do a second - // unordered traversal afterwards. + // visit them in post dfs order starting from the Input + // node. Then (whether or not it's a dataflow sibling graph) + // we visit each remaining unvisited child in children() order. + // + // The second traversal is required to ensure we visit both + // nodes unreachable from Input in a dataflow sibling graph + // (e.g. LoadConstant) and the children of non dataflow + // sibling graph nodes (e.g. the children of CFG or Conditional + // nodes) if let Some([input, _]) = self.hugr.get_io(n) { - use crate::hugr::views::PetgraphWrapper; - let wrapper = self.hugr.as_petgraph(); - let mut dfs = ::petgraph::visit::DfsPostOrder::new(&wrapper, input); - while let Some(x) = dfs.next(&wrapper) { - worklist.push(WorkItem::Visit(x)); + let petgraph = self.hugr.as_petgraph(); + // Here we visit the nodes in DfsPostOrder(i.e. we have + // visited all the out_neighbors() of a node before we + // visit that node), and push a node onto the worklist + // stack when we visit it. So once we are done the stack + // will have the Input node at the top, and a nodes + // out_neighbors are always under that node on the + // worklist stack. + let mut dfs = ::petgraph::visit::DfsPostOrder::new(&petgraph, input); + while let Some(x) = dfs.next(&petgraph) { + worklist.push(x.into()); pushed_children.insert(x); } } + // Here we collect all children that were not visited by the + // DfsPostOrder traversal above, in children() order let rest_children = self .hugr .children(n) .filter(|x| !pushed_children.contains(x)) .collect_vec(); - // extend in reverse so that the first child is the top of the stack + // We extend in reverse so that the first child is the top of the stack worklist.extend(rest_children.into_iter().rev().map(WorkItem::Visit)); worklist.push(WorkItem::Callback(WalkOrder::Preorder, n)); } WorkItem::Callback(order, n) => { let optype = self.hugr.get_optype(n); for cb in self.mut_callbacks(order).iter_mut() { - // this clone is unfortunate, to avoid this we would - // need a TryInto variant like: try_into(&O) -> Option<&T> t = cb(n, optype.clone(), t)?; } } @@ -131,6 +145,31 @@ impl<'a, H: HugrView, T, E> Walker<'a, H, T, E> { } } +/// An example of use using the Walker to implement an iterator over all nodes, +/// nodes are visited in preorder where possible. More precisely, nodes are +/// visited before their children, and nodes in a dataflow sibling graph are +/// visited before their out_neighbours. +pub fn hugr_walk_iter(h: &impl HugrView) -> impl Iterator { + let mut walker = Walker::<_, Vec, std::convert::Infallible>::new(h); + walker.previsit(|n, _: OpType, mut v| { + v.push(n); + Ok(v) + }); + walker.walk(Vec::new()).unwrap().into_iter() +} + +/// An example of use using the Walker to implement a search. +/// This demonstrates terminating a walk early. +pub fn hugr_walk_find(h: &impl HugrView, mut f: impl FnMut(Node, O) -> Option) -> Option +where + OpType: TryInto, +{ + Walker::new(h) + .previsit(|n, o: O, ()| f(n, o).map_or(Ok(()), Result::Err)) + .walk(()) + .map_or_else(Option::Some, |()| None) +} + #[cfg(test)] mod test { use std::error::Error; @@ -236,4 +275,38 @@ mod test { ); Ok(()) } + + #[test] + fn leaf_op_out_degree() { + use std::collections::HashMap; + let h: crate::Hugr = todo!(); + let mut walker = Walker::new(&h); + walker.postvisit(|n, _: crate::ops::LeafOp, mut r| { + r.insert(n, h.node_outputs(n).map(|o| h.linked_ports(n, o).count())); + Ok(r) + }); + let r = walker.walk(HashMap::new()).unwrap(); + // TODO construct example and assert result of walk + } + + #[test] + fn pretty_printer() { + struct PPCtx(usize, String); + let h: crate::Hugr = todo!(); + let pp_out = Walker::<_, _, std::convert::Infallible>::new(&h) + .previsit(|n, _: OpType, PPCtx(mut indent, mut r)| { + use crate::hugr::NodeIndex; + r += format!( + "{}{}\n", + std::iter::repeat(' ').take(indent).collect::(), + n.index() + ) + .as_str(); + Ok(PPCtx(indent + 2, r)) + }) + .postvisit(|_, _: OpType, PPCtx(mut indent, r)| Ok(PPCtx(indent - 2, r))) + .walk(PPCtx(0, "".to_string())) + .unwrap(); + // TODO construct example and assert result of walk + } } From c91a02585b4332794ba658eb2392def9119c08e7 Mon Sep 17 00:00:00 2001 From: Douglas Wilson Date: Thu, 2 Nov 2023 08:57:28 +0000 Subject: [PATCH 10/10] General callbacks don't need optype --- src/walker.rs | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/walker.rs b/src/walker.rs index 157a76ce5..c06ed0bd1 100644 --- a/src/walker.rs +++ b/src/walker.rs @@ -14,11 +14,9 @@ pub enum WalkOrder { } #[derive(Deref, DerefMut)] -struct WalkerCallback<'a, T, E>(Box Result>); +struct WalkerCallback<'a, T, E>(Box Result>); -impl<'a, T, E, F: 'a + FnMut(Node, OpType, T) -> Result> From - for WalkerCallback<'a, T, E> -{ +impl<'a, T, E, F: 'a + FnMut(Node, T) -> Result> From for WalkerCallback<'a, T, E> { fn from(f: F) -> Self { Self(Box::new(f)) } @@ -68,7 +66,8 @@ impl<'a, H: HugrView, T, E> Walker<'a, H, T, E> { where OpType: TryInto, { - let cb = move |n, o: OpType, t| match o.try_into() { + let hugr = self.hugr; + let cb = move |n, t| match hugr.get_optype(n).clone().try_into() { Ok(x) => f(n, x, t), _ => Ok(t), }; @@ -134,9 +133,8 @@ impl<'a, H: HugrView, T, E> Walker<'a, H, T, E> { worklist.push(WorkItem::Callback(WalkOrder::Preorder, n)); } WorkItem::Callback(order, n) => { - let optype = self.hugr.get_optype(n); for cb in self.mut_callbacks(order).iter_mut() { - t = cb(n, optype.clone(), t)?; + t = cb(n, t)?; } } }