From 35977adc8103cfc232bc95fbc32f4e34f2b6a6d7 Mon Sep 17 00:00:00 2001
From: Kieran <kieran@harkin.me>
Date: Fri, 28 Feb 2025 16:00:56 +0000
Subject: [PATCH] Iterator size hints (#217)

* feat: iterator size hints

* fix: add size_hint to IpNetworkIterator
---
 src/ipv4.rs | 28 ++++++++++++++++++++++++++++
 src/ipv6.rs | 21 +++++++++++++++++++++
 src/lib.rs  |  6 ++++++
 3 files changed, 55 insertions(+)

diff --git a/src/ipv4.rs b/src/ipv4.rs
index c977f69..086cfc4 100644
--- a/src/ipv4.rs
+++ b/src/ipv4.rs
@@ -368,6 +368,15 @@ impl Iterator for Ipv4NetworkIterator {
         };
         Some(next.into())
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if let Some(n) = self.next {
+            let elms = (self.end - n + 1) as usize;
+            (elms, Some(elms))
+        } else {
+            (0, None)
+        }
+    }
 }
 
 impl IntoIterator for &'_ Ipv4Network {
@@ -582,6 +591,25 @@ mod test {
         assert_eq!(None, iter.next());
     }
 
+    #[test]
+    fn iterator_v4_size_hint() {
+        let cidr: Ipv4Network = "192.168.0.0/24".parse().unwrap();
+        let mut iter = cidr.iter();
+        assert_eq!((256, Some(256)), iter.size_hint());
+        iter.next();
+        assert_eq!((255, Some(255)), iter.size_hint());
+
+        let cidr: Ipv4Network = "192.168.0.0/32".parse().unwrap();
+        let mut iter = cidr.iter();
+        assert_eq!((1, Some(1)), iter.size_hint());
+        iter.next();
+        assert_eq!((0, None), iter.size_hint());
+
+        let cidr: Ipv4Network = "192.168.0.0/0".parse().unwrap();
+        let iter = cidr.iter();
+        assert_eq!((4294967295, Some(4294967295)), iter.size_hint());
+    }
+
     #[test]
     fn v4_mask_to_prefix() {
         let mask = Ipv4Addr::new(255, 255, 255, 128);
diff --git a/src/ipv6.rs b/src/ipv6.rs
index de07288..baf89d2 100644
--- a/src/ipv6.rs
+++ b/src/ipv6.rs
@@ -372,6 +372,15 @@ impl Iterator for Ipv6NetworkIterator {
         };
         Some(next.into())
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if let Some(n) = self.next {
+            let elms = (self.end - n + 1) as usize;
+            (elms, Some(elms))
+        } else {
+            (0, None)
+        }
+    }
 }
 
 impl IntoIterator for &'_ Ipv6Network {
@@ -612,6 +621,18 @@ mod test {
         assert_eq!(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 2), iter.next().unwrap());
     }
 
+    #[test]
+    fn iterator_v6_size_hint() {
+        let cidr: Ipv6Network = "2001:db8::/128".parse().unwrap();
+        let mut iter = cidr.iter();
+        assert_eq!((1, Some(1)), iter.size_hint());
+        assert_eq!(
+            Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0),
+            iter.next().unwrap()
+        );
+        assert_eq!((0, None), iter.size_hint());
+    }
+
     #[test]
     fn network_v6() {
         let cidr: Ipv6Network = "2001:db8::0/96".parse().unwrap();
diff --git a/src/lib.rs b/src/lib.rs
index 1b8702a..24ec499 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -414,6 +414,12 @@ impl Iterator for IpNetworkIterator {
             IpNetworkIteratorInner::V6(iter) => iter.next().map(IpAddr::V6),
         }
     }
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        match &self.inner {
+            IpNetworkIteratorInner::V4(iter) => iter.size_hint(),
+            IpNetworkIteratorInner::V6(iter) => iter.size_hint(),
+        }
+    }
 }
 
 impl IntoIterator for &'_ IpNetwork {