From 2bf7a6195161cbe593d891f20b4666b64805cb52 Mon Sep 17 00:00:00 2001 From: Andrew Poelstra Date: Tue, 25 Jun 2024 14:22:43 +0000 Subject: [PATCH] dag: introduce depth tracking and maximum depth to verbose preorder iter --- src/dag.rs | 47 +++++++++++++++++++++++++++++++++-------- src/node/display.rs | 2 +- src/types/final_data.rs | 2 +- src/types/mod.rs | 22 +++++++++++++++---- src/value.rs | 2 +- 5 files changed, 59 insertions(+), 16 deletions(-) diff --git a/src/dag.rs b/src/dag.rs index 4985e377..73e3df5b 100644 --- a/src/dag.rs +++ b/src/dag.rs @@ -278,13 +278,15 @@ pub trait DagLike: Sized { /// yielded, you may be better off using this iterator instead. fn verbose_pre_order_iter + Default>( self, + max_depth: Option, ) -> VerbosePreOrderIter where Self: Clone, { VerbosePreOrderIter { - stack: vec![PreOrderIterItem::initial(self, None)], + stack: vec![PreOrderIterItem::initial(self, 0, None)], index: 0, + max_depth, tracker: Default::default(), } } @@ -775,6 +777,16 @@ pub struct VerbosePreOrderIter { /// children are put onto the stack followed by their left, so that the /// appropriate one will be yielded on the next iteration. stack: Vec>, + /// Maximum depth (distance from root) that the iterator will go to. Any + /// children at a greater depth will not be yielded. However, the parent + /// of such children will still be yielded multiple times, one for each + /// (unyielded) child, with `n_children_yielded` incremented as though + /// the children had been yielded. + /// + /// To determine whether pruning has happened, you should manually check + /// the [`PreOrderIterItem::depth`] field against the maximum depth that + /// you set when constructing the iterator. + max_depth: Option, /// The index of the next item to be yielded. /// /// Note that unlike the [`PostOrderIter`], this value is not monotonic @@ -809,19 +821,31 @@ impl> Iterator for VerbosePreOrderIter< (0, 0) => {} (0, n) => { self.stack.push(top.clone().increment(n == 1)); - let child = top.node.left_child().unwrap(); - self.stack - .push(PreOrderIterItem::initial(child, Some(top.node.clone()))); + if top.depth < self.max_depth.unwrap_or(top.depth + 1) { + let child = top.node.left_child().unwrap(); + self.stack.push(PreOrderIterItem::initial( + child, + top.depth + 1, + Some(top.node.clone()), + )); + } } (1, 0) => unreachable!(), (1, 1) => {} (1, _) => { self.stack.push(top.clone().increment(true)); - let child = top.node.right_child().unwrap(); - self.stack - .push(PreOrderIterItem::initial(child, Some(top.node.clone()))); + if top.depth < self.max_depth.unwrap_or(top.depth + 1) { + let child = top.node.right_child().unwrap(); + self.stack.push(PreOrderIterItem::initial( + child, + top.depth + 1, + Some(top.node.clone()), + )); + } + } + (x, y) => { + debug_assert_eq!((x, y), (2, 2)); } - (_, _) => {} } // Then yield the element. return Some(top); @@ -840,6 +864,9 @@ pub struct PreOrderIterItem { pub parent: Option, /// The index when the element was first yielded. pub index: usize, + /// The distance of this element from the initial node. 0 for the initial + /// node itself. + pub depth: usize, /// How many of this item's children have been yielded. /// /// This can also be interpreted as a count of how many times this @@ -853,12 +880,13 @@ impl PreOrderIterItem { /// Creates a `PreOrderIterItem` which yields a given element for the first time. /// /// Marks the index as 0. The index must be manually set before yielding. - fn initial(d: D, parent: Option) -> Self { + fn initial(d: D, depth: usize, parent: Option) -> Self { PreOrderIterItem { is_complete: matches!(d.as_dag_node(), Dag::Nullary), node: d, parent, index: 0, + depth, n_children_yielded: 0, } } @@ -868,6 +896,7 @@ impl PreOrderIterItem { PreOrderIterItem { node: self.node, index: self.index, + depth: self.depth, parent: self.parent, n_children_yielded: self.n_children_yielded + 1, is_complete, diff --git a/src/node/display.rs b/src/node/display.rs index dd237760..5d01e4a3 100644 --- a/src/node/display.rs +++ b/src/node/display.rs @@ -46,7 +46,7 @@ where &'a Node: DagLike, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - for data in self.0.verbose_pre_order_iter::() { + for data in self.0.verbose_pre_order_iter::(None) { match data.n_children_yielded { 1 => match data.node.inner() { Inner::Comp(..) => f.write_str("; ")?, diff --git a/src/types/final_data.rs b/src/types/final_data.rs index 8088916c..1bfc06e9 100644 --- a/src/types/final_data.rs +++ b/src/types/final_data.rs @@ -77,7 +77,7 @@ impl fmt::Debug for Final { impl fmt::Display for Final { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { let mut skipping: Option = None; - for data in self.verbose_pre_order_iter::() { + for data in self.verbose_pre_order_iter::(None) { if let Some(skip) = skipping { if data.is_complete && data.node.tmr == skip { skipping = None; diff --git a/src/types/mod.rs b/src/types/mod.rs index 625c2616..a82317d9 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -294,10 +294,18 @@ impl Bound { } } +const MAX_DISPLAY_DEPTH: usize = 64; + impl fmt::Debug for Bound { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { let arc = Arc::new(self.shallow_clone()); - for data in arc.verbose_pre_order_iter::() { + for data in arc.verbose_pre_order_iter::(Some(MAX_DISPLAY_DEPTH)) { + if data.depth == MAX_DISPLAY_DEPTH { + if data.n_children_yielded == 0 { + f.write_str("...")?; + } + continue; + } match (&*data.node, data.n_children_yielded) { (Bound::Free(ref s), _) => f.write_str(s)?, (Bound::Complete(ref comp), _) => fmt::Debug::fmt(comp, f)?, @@ -314,13 +322,19 @@ impl fmt::Debug for Bound { impl fmt::Display for Bound { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { let arc = Arc::new(self.shallow_clone()); - for data in arc.verbose_pre_order_iter::() { + for data in arc.verbose_pre_order_iter::(Some(MAX_DISPLAY_DEPTH)) { + if data.depth == MAX_DISPLAY_DEPTH { + if data.n_children_yielded == 0 { + f.write_str("...")?; + } + continue; + } match (&*data.node, data.n_children_yielded) { (Bound::Free(ref s), _) => f.write_str(s)?, (Bound::Complete(ref comp), _) => fmt::Display::fmt(comp, f)?, (Bound::Sum(..), 0) | (Bound::Product(..), 0) => { if data.index > 0 { - f.write_str("(")? + f.write_str("(")?; } } (Bound::Sum(..), 2) | (Bound::Product(..), 2) => { @@ -457,7 +471,7 @@ impl Type { // First, do occurs-check to ensure that we have no infinitely sized types. let mut occurs_check = HashSet::new(); - for data in bound.verbose_pre_order_iter::() { + for data in bound.verbose_pre_order_iter::(None) { if data.is_complete { occurs_check.remove(&(data.node.as_ref() as *const _)); } else if data.n_children_yielded == 0 diff --git a/src/value.rs b/src/value.rs index 17f86a43..4e2ca079 100644 --- a/src/value.rs +++ b/src/value.rs @@ -339,7 +339,7 @@ impl fmt::Debug for Value { impl fmt::Display for Value { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - for data in self.verbose_pre_order_iter::() { + for data in self.verbose_pre_order_iter::(None) { match data.node { Value::Unit => { if data.n_children_yielded == 0