From 6cbb011597cdf61b2e37655817625f88b5fb2c37 Mon Sep 17 00:00:00 2001
From: Michal 'vorner' Vaner <vorner@vorner.cz>
Date: Sat, 7 Dec 2019 19:32:51 +0100
Subject: [PATCH] Iterators for the IpNetwork

Version-agnostic iterator in addition to the version specific ones.
---
 src/lib.rs | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 24abf7a..09377d4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -16,6 +16,8 @@ mod ipv6;
 pub use crate::common::IpNetworkError;
 pub use crate::ipv4::{ipv4_mask_to_prefix, Ipv4Network};
 pub use crate::ipv6::{ipv6_mask_to_prefix, Ipv6Network};
+use crate::ipv4::Ipv4NetworkIterator;
+use crate::ipv6::Ipv6NetworkIterator;
 
 /// Represents a generic network range. This type can have two variants:
 /// the v4 and the v6 case.
@@ -247,6 +249,17 @@ impl IpNetwork {
             IpNetwork::V6(ref ip) => NetworkSize::V6(ip.size()),
         }
     }
+
+    /// Returns an iterator over the addresses contained in the network.
+    ///
+    /// This lists all the addresses in the network range, in ascending order.
+    pub fn iter(&self) -> IpNetworkIterator {
+        let inner = match self {
+            IpNetwork::V4(ip) => IpNetworkIteratorInner::V4(ip.iter()),
+            IpNetwork::V6(ip) => IpNetworkIteratorInner::V6(ip.iter()),
+        };
+        IpNetworkIterator { inner }
+    }
 }
 
 /// Tries to parse the given string into a `IpNetwork`. Will first try to parse
@@ -306,6 +319,34 @@ impl fmt::Display for IpNetwork {
     }
 }
 
+#[derive(Clone, Debug)]
+enum IpNetworkIteratorInner {
+    V4(Ipv4NetworkIterator),
+    V6(Ipv6NetworkIterator),
+}
+
+pub struct IpNetworkIterator {
+    inner: IpNetworkIteratorInner,
+}
+
+impl Iterator for IpNetworkIterator {
+    type Item = IpAddr;
+    fn next(&mut self) -> Option<IpAddr> {
+        match &mut self.inner {
+            IpNetworkIteratorInner::V4(iter) => iter.next().map(IpAddr::V4),
+            IpNetworkIteratorInner::V6(iter) => iter.next().map(IpAddr::V6),
+        }
+    }
+}
+
+impl IntoIterator for &'_ IpNetwork {
+    type IntoIter = IpNetworkIterator;
+    type Item = IpAddr;
+    fn into_iter(self) -> IpNetworkIterator {
+        self.iter()
+    }
+}
+
 /// Converts a `IpAddr` network mask into a prefix.
 /// If the mask is invalid this will return an `IpNetworkError::InvalidPrefix`.
 pub fn ip_mask_to_prefix(mask: IpAddr) -> Result<u8, IpNetworkError> {