swh_graph/
arc_iterators.rs

1// Copyright (C) 2023-2024  The Software Heritage developers
2// See the AUTHORS file at the top-level directory of this distribution
3// License: GNU General Public License version 3, or any later version
4// See top-level LICENSE file for more information
5
6//! Structures to manipulate the Software Heritage graph
7//!
8//! In order to load only what is necessary, these structures are initially created
9//! by calling [`SwhUnidirectionalGraph::new`](crate::graph::SwhUnidirectionalGraph::new) or
10//! [`SwhBidirectionalGraph::new`](crate::graph::SwhBidirectionalGraph::new), then calling methods
11//! on them to progressively load additional data (`load_properties`, `load_all_properties`,
12//! `load_labels`)
13
14#![allow(clippy::type_complexity)]
15
16use std::borrow::Borrow;
17use std::iter::Iterator;
18
19use webgraph::prelude::*;
20
21use crate::graph::{NodeId, SwhGraphWithProperties};
22use crate::labels::{EdgeLabel, UntypedEdgeLabel};
23
24pub trait IntoFlattenedLabeledArcsIterator<Label> {
25    type Flattened: IntoIterator<Item = (NodeId, Label)>;
26
27    /// Turns this `Iterator<Item=(succ, Iterator<Item=labels>)>` into an
28    /// `Iterator<ITem=(succ, label)>`.
29    fn flatten_labels(self) -> Self::Flattened;
30}
31
32pub struct LabeledSuccessorIterator<Successors: Iterator>
33where
34    <Successors as Iterator>::Item: Pair<Left = usize, Right: IntoIterator<Item: Borrow<u64>>>,
35{
36    pub(crate) successors: Successors,
37}
38
39impl<Successors: Iterator> LabeledSuccessorIterator<Successors>
40where
41    <Successors as Iterator>::Item: Pair<Left = usize, Right: IntoIterator<Item: Borrow<u64>>>,
42{
43    pub fn new(successors: Successors) -> Self {
44        LabeledSuccessorIterator { successors }
45    }
46}
47
48impl<Successors: Iterator> Iterator for LabeledSuccessorIterator<Successors>
49where
50    <Successors as Iterator>::Item: Pair<Left = usize, Right: IntoIterator<Item: Borrow<u64>>>,
51{
52    type Item = (
53        NodeId,
54        LabeledArcIterator<
55            <<<Successors as Iterator>::Item as Pair>::Right as IntoIterator>::IntoIter,
56        >,
57    );
58
59    fn next(&mut self) -> Option<Self::Item> {
60        self.successors.next().map(|pair| {
61            let (successor, arc_labels) = pair.into_pair();
62            (
63                successor,
64                LabeledArcIterator {
65                    arc_label_ids: arc_labels.into_iter(),
66                },
67            )
68        })
69    }
70}
71
72// SAFETY: if Successors is a sorted iterator of (usize, into sorted iterator),
73// then LabeledSuccessorIterator is sorted because it preserves order
74unsafe impl<Successors: SortedIterator> SortedIterator for LabeledSuccessorIterator<Successors> where
75    <Successors as Iterator>::Item:
76        Pair<Left = usize, Right: IntoIterator<Item: Borrow<u64>, IntoIter: SortedIterator>>
77{
78}
79
80impl<Successors: Iterator> IntoFlattenedLabeledArcsIterator<UntypedEdgeLabel>
81    for LabeledSuccessorIterator<Successors>
82where
83    <Successors as Iterator>::Item: Pair<Left = usize, Right: IntoIterator<Item: Borrow<u64>>>,
84{
85    type Flattened = FlattenedSuccessorsIterator<Self>;
86
87    fn flatten_labels(self) -> Self::Flattened {
88        FlattenedSuccessorsIterator::new(self)
89    }
90}
91
92pub struct FlattenedSuccessorsIterator<Successors: Iterator>
93where
94    <Successors as Iterator>::Item: Pair<Left = usize, Right: IntoIterator>,
95{
96    current_node_and_labels: Option<(
97        NodeId,
98        <<<Successors as Iterator>::Item as Pair>::Right as IntoIterator>::IntoIter,
99    )>,
100    iter: Successors,
101}
102
103impl<Successors: Iterator> FlattenedSuccessorsIterator<Successors>
104where
105    <Successors as Iterator>::Item: Pair<Left = usize, Right: IntoIterator>,
106{
107    pub fn new(successors: Successors) -> Self {
108        Self {
109            current_node_and_labels: None,
110            iter: successors,
111        }
112    }
113}
114
115impl<Successors: Iterator> Iterator for FlattenedSuccessorsIterator<Successors>
116where
117    <Successors as Iterator>::Item: Pair<Left = usize, Right: IntoIterator>,
118{
119    type Item = (
120        NodeId,
121        <<<Successors as Iterator>::Item as Pair>::Right as IntoIterator>::Item,
122    );
123
124    fn next(&mut self) -> Option<Self::Item> {
125        if self.current_node_and_labels.is_none() {
126            // first iterator, or exhausted current label list, move to the next one
127            self.current_node_and_labels = self
128                .iter
129                .next()
130                .map(Pair::into_pair)
131                .map(|(succ, labels)| (succ, labels.into_iter()))
132        }
133        let Some((current_node, ref mut current_labels)) = self.current_node_and_labels else {
134            // Reached the end of the iterator
135            return None;
136        };
137        match current_labels.next() {
138            Some(label) => Some((current_node, label)),
139            None => {
140                // exhausted this label list, move to the next one
141                self.current_node_and_labels = None;
142                self.next()
143            }
144        }
145    }
146}
147
148// SAFETY: if Successors is a sorted iterator of (usize, into sorted iterator),
149// then FlattenedSuccessorsIterator is sorted because it preserves order
150unsafe impl<Successors: SortedIterator> SortedIterator for FlattenedSuccessorsIterator<Successors> where
151    <Successors as Iterator>::Item:
152        Pair<Left = usize, Right: IntoIterator<Item: Borrow<u64>, IntoIter: SortedIterator>>
153{
154}
155
156pub struct LabeledArcIterator<T: Iterator<Item: Borrow<u64>>> {
157    arc_label_ids: T,
158}
159
160impl<T: Iterator<Item: Borrow<u64>>> Iterator for LabeledArcIterator<T> {
161    type Item = UntypedEdgeLabel;
162
163    fn next(&mut self) -> Option<Self::Item> {
164        self.arc_label_ids
165            .next()
166            .map(|label| UntypedEdgeLabel::from(*label.borrow()))
167    }
168}
169
170// SAFETY: if T is a sorted iterator of u64,
171// then LabeledArcIterator is sorted because it preserves order
172unsafe impl<T: SortedIterator<Item: Borrow<u64>>> SortedIterator for LabeledArcIterator<T> {}
173
174pub struct LabelTypingSuccessorIterator<'a, G, Successors: Iterator>
175where
176    <Successors as Iterator>::Item:
177        Pair<Left = usize, Right: IntoIterator<Item = UntypedEdgeLabel>>,
178    G: SwhGraphWithProperties,
179    <G as SwhGraphWithProperties>::Maps: crate::properties::Maps,
180{
181    pub(crate) graph: &'a G,
182    pub(crate) is_transposed: bool,
183    pub(crate) successors: Successors,
184    pub(crate) src: NodeId,
185}
186
187impl<'a, G, Successors: Iterator> Iterator for LabelTypingSuccessorIterator<'a, G, Successors>
188where
189    <Successors as Iterator>::Item:
190        Pair<Left = usize, Right: IntoIterator<Item = UntypedEdgeLabel>>,
191    G: SwhGraphWithProperties,
192    <G as SwhGraphWithProperties>::Maps: crate::properties::Maps,
193{
194    type Item = (
195        NodeId,
196        LabelTypingArcIterator<
197            'a,
198            G,
199            <<<Successors as Iterator>::Item as Pair>::Right as IntoIterator>::IntoIter,
200        >,
201    );
202
203    fn next(&mut self) -> Option<Self::Item> {
204        self.successors.next().map(|pair| {
205            let mut src = self.src;
206            let (mut dst, labels) = pair.into_pair();
207            let succ = dst;
208            if self.is_transposed {
209                (src, dst) = (dst, src)
210            }
211            (
212                succ,
213                LabelTypingArcIterator {
214                    graph: self.graph,
215                    labels: labels.into_iter(),
216                    src,
217                    dst,
218                },
219            )
220        })
221    }
222}
223
224impl<G, Successors: Iterator> IntoFlattenedLabeledArcsIterator<EdgeLabel>
225    for LabelTypingSuccessorIterator<'_, G, Successors>
226where
227    <Successors as Iterator>::Item:
228        Pair<Left = usize, Right: IntoIterator<Item = UntypedEdgeLabel>>,
229    G: SwhGraphWithProperties,
230    <G as SwhGraphWithProperties>::Maps: crate::properties::Maps,
231{
232    type Flattened = FlattenedSuccessorsIterator<Self>;
233
234    fn flatten_labels(self) -> Self::Flattened {
235        FlattenedSuccessorsIterator::new(self)
236    }
237}
238
239pub struct LabelTypingArcIterator<'a, G, Labels: Iterator<Item = UntypedEdgeLabel>>
240where
241    G: SwhGraphWithProperties,
242    <G as SwhGraphWithProperties>::Maps: crate::properties::Maps,
243{
244    graph: &'a G,
245    labels: Labels,
246    src: NodeId,
247    dst: NodeId,
248}
249
250impl<G, Labels: Iterator<Item = UntypedEdgeLabel>> Iterator
251    for LabelTypingArcIterator<'_, G, Labels>
252where
253    G: SwhGraphWithProperties,
254    <G as SwhGraphWithProperties>::Maps: crate::properties::Maps,
255{
256    type Item = EdgeLabel;
257
258    fn next(&mut self) -> Option<Self::Item> {
259        let props = self.graph.properties();
260        self.labels.next().map(move |label| {
261            label
262                .for_edge_type(
263                    props.node_type(self.src),
264                    props.node_type(self.dst),
265                    self.graph.is_transposed(),
266                )
267                .unwrap_or_else(|e| {
268                    panic!(
269                        "Unexpected edge from {} ({}) to {} ({}): {}",
270                        props.swhid(self.src),
271                        self.src,
272                        props.swhid(self.dst),
273                        self.dst,
274                        e
275                    )
276                })
277        })
278    }
279}
280
281/// Wraps an iterator of labeled successors, and yields only the successors
282pub struct DelabelingIterator<Successors: Iterator<Item: Pair<Left = usize>>> {
283    pub(crate) successors: Successors,
284}
285impl<Successors: Iterator<Item: Pair<Left = usize>>> Iterator for DelabelingIterator<Successors> {
286    type Item = NodeId;
287
288    fn next(&mut self) -> Option<Self::Item> {
289        self.successors.next().map(|pair| {
290            let (successor, _arc_labels) = pair.into_pair();
291            successor
292        })
293    }
294}
295// SAFETY: DelabelingIterator preserves order
296unsafe impl<Successors: SortedIterator<Item: Pair<Left = usize>>> SortedIterator
297    for DelabelingIterator<Successors>
298{
299}