Skip to main content

spire_core/frontier/facade/
depth.rs

1use std::future::Future;
2
3use crate::Result;
4use crate::crawler::Depth;
5use crate::frontier::seed::Seed;
6use crate::frontier::{Ack, Admission, Frontier, NextSeed, Report};
7
8/// A facade that refuses seeds deeper than a limit.
9#[derive(Debug)]
10pub struct MaxDepth<F> {
11    inner: F,
12    max: Depth,
13}
14
15impl<F> MaxDepth<F> {
16    /// Wraps `inner`, refusing seeds deeper than `max`.
17    pub const fn new(inner: F, max: u32) -> Self {
18        Self {
19            inner,
20            max: Depth::new(max),
21        }
22    }
23}
24
25impl<F: Frontier> Frontier for MaxDepth<F> {
26    async fn push(&self, seed: Seed) -> Result<Admission> {
27        if seed.depth() > self.max {
28            return Ok(Admission::TooDeep);
29        }
30        self.inner.push(seed).await
31    }
32
33    fn pop(&self) -> impl Future<Output = Result<NextSeed>> + Send {
34        self.inner.pop()
35    }
36
37    fn report(&self, report: Report) -> impl Future<Output = Result<Ack>> + Send {
38        self.inner.report(report)
39    }
40
41    fn size_hint(&self) -> Option<usize> {
42        self.inner.size_hint()
43    }
44
45    fn snapshot(&self) -> Option<Vec<Seed>> {
46        self.inner.snapshot()
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use crate::crawler::Depth;
53    use crate::frontier::test_util::seed;
54    use crate::frontier::{Admission, Frontier, FrontierExt, InMemFrontier};
55
56    #[tokio::test]
57    async fn refuses_deep_seeds() {
58        let frontier = InMemFrontier::fifo().max_depth(1);
59        let deep = |depth| seed("a").with_depth(Depth::new(depth));
60        assert_eq!(frontier.push(deep(0)).await.unwrap(), Admission::Queued);
61        assert_eq!(frontier.push(deep(1)).await.unwrap(), Admission::Queued);
62        assert_eq!(frontier.push(deep(2)).await.unwrap(), Admission::TooDeep);
63        assert_eq!(frontier.size_hint(), Some(2));
64    }
65}