github.com/charypar/monobuild@v0.0.0-20211122220434-fd884ed50212/rs/src/graph/iter.rs (about)

     1  use std::collections::{btree_set, BTreeSet};
     2  use std::{iter, slice};
     3  
     4  use super::{Edge, Graph, Subgraph, Vertex};
     5  
     6  // Iterator over the graph adjacency lists
     7  pub struct GraphIter<'g, 's, V, E>
     8  where
     9      V: Vertex,
    10      E: Edge,
    11      's: 'g,
    12  {
    13      pub(super) graph: &'g Graph<V, E>,
    14      pub(super) iter: iter::Enumerate<slice::Iter<'g, V>>,
    15      pub(super) masks: Option<(&'s Vec<bool>, &'s BTreeSet<(usize, usize)>)>,
    16  }
    17  
    18  impl<'g, 's, V, E> Iterator for GraphIter<'g, 's, V, E>
    19  where
    20      V: Vertex,
    21      E: Edge,
    22      's: 'g,
    23  {
    24      type Item = (&'g V, Edges<'g, V, E>);
    25  
    26      fn next<'i>(&'i mut self) -> Option<Self::Item> {
    27          let edge_mask = self.masks.map(|m| m.1);
    28  
    29          match self.masks {
    30              Some((vec_mask, _)) => self
    31                  .iter
    32                  .find(|(i, _)| *vec_mask.get(*i).expect("Vector mask size is wrong")),
    33              None => self.iter.next(),
    34          }
    35          .map(|(i, v)| (v, Edges::new(self.graph, i, edge_mask)))
    36      }
    37  }
    38  
    39  // Iterator over edges originating in a vertex
    40  pub struct Edges<'g, V, E>
    41  where
    42      V: Vertex,
    43      E: Edge,
    44  {
    45      vertex: usize,
    46      graph: &'g Graph<V, E>,
    47      mask: Option<&'g BTreeSet<(usize, usize)>>,
    48      iter: btree_set::Iter<'g, (usize, E)>,
    49  }
    50  
    51  impl<'g, V, E> Edges<'g, V, E>
    52  where
    53      V: Vertex,
    54      E: Edge,
    55  {
    56      fn new(
    57          graph: &'g Graph<V, E>,
    58          index: usize,
    59          mask: Option<&'g BTreeSet<(usize, usize)>>,
    60      ) -> Self {
    61          Self {
    62              vertex: index,
    63              graph,
    64              iter: graph.edges[index].iter(),
    65              mask,
    66          }
    67      }
    68  }
    69  
    70  impl<'g, V, E> Iterator for Edges<'g, V, E>
    71  where
    72      V: Vertex,
    73      E: Edge,
    74  {
    75      type Item = (&'g V, E);
    76  
    77      fn next(&mut self) -> Option<Self::Item> {
    78          match self.mask {
    79              Some(mask) => {
    80                  let vertex = self.vertex;
    81                  self.iter.find(|(dest, _)| mask.contains(&(vertex, *dest)))
    82              }
    83              None => self.iter.next(),
    84          }
    85          .map(|(vid, e)| {
    86              (
    87                  self.graph
    88                      .vertices
    89                      .get(*vid)
    90                      .expect("vertex not found in graph<"),
    91                  *e,
    92              )
    93          })
    94      }
    95  }
    96  
    97  // Iterator over the vertices of the graph
    98  pub struct Vertices<'g, 's, V>
    99  where
   100      V: 'g,
   101  {
   102      pub(super) iter: std::iter::Enumerate<std::slice::Iter<'g, V>>,
   103      pub(super) mask: Option<&'s Vec<bool>>,
   104  }
   105  
   106  impl<'g, 's, V> Iterator for Vertices<'g, 's, V> {
   107      type Item = &'g V;
   108  
   109      fn next(&mut self) -> Option<Self::Item> {
   110          match self.mask {
   111              Some(mask) => self.iter.find(|(i, _)| mask[*i]).map(|(_, v)| v),
   112              None => self.iter.next().map(|(_, v)| v),
   113          }
   114      }
   115  }
   116  
   117  // Converting into iterator
   118  
   119  impl<'g, V, E> IntoIterator for &'g Graph<V, E>
   120  where
   121      V: Vertex,
   122      E: Edge,
   123  {
   124      type Item = (&'g V, Edges<'g, V, E>);
   125      type IntoIter = GraphIter<'g, 'g, V, E>;
   126  
   127      fn into_iter(self) -> Self::IntoIter {
   128          self.iter()
   129      }
   130  }
   131  
   132  impl<'g, V, E> IntoIterator for &'g Subgraph<'g, V, E>
   133  where
   134      V: Vertex,
   135      E: Edge,
   136  {
   137      type Item = (&'g V, Edges<'g, V, E>);
   138      type IntoIter = GraphIter<'g, 'g, V, E>;
   139  
   140      fn into_iter(self) -> Self::IntoIter {
   141          self.iter()
   142      }
   143  }