From 99203160b9e3d51e8262d6e4ec28795da3cf9026 Mon Sep 17 00:00:00 2001
From: Igor Bykov <con29rus@icloud.com>
Date: Sat, 8 Jul 2017 22:28:00 +0300
Subject: [PATCH] Add an iterator for ipv6 network. (#49)

This is feature gated since i128 is not stable yet
---
 .travis.yml | 14 ++++++---
 Cargo.toml  |  1 +
 src/ipv6.rs | 82 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs  |  1 +
 4 files changed, 94 insertions(+), 4 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index edc522f..e9e94e3 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -6,12 +6,18 @@ rust:
 env:
   global:
     - secure: gokQ7xIWwmAuEUW3IyS5B/pbZxdFSSDBto5beJ+4ACGcRMDqJ/eCPf1ekSVXME4TWM46uUCjxcdUjYhIhQ6sG4zfWck4u45qRJ5JbIoTvR+ykxhN1j3Zi5x9ptP3ALDbHn2i3v6t9xohORfQpz3dVND5c7thbYDyKP2ZR1sez5c=
+    - FEATURES: default
+
+matrix:
+  include:
+    - rust: nightly
+      env: FEATURES=ipv6-iterator
 
 script:
-  - cargo build --verbose
-  - cargo test --verbose
-  - cargo build --release --verbose
-  - cargo doc --verbose
+  - cargo build --features $FEATURES --verbose
+  - cargo test --features $FEATURES --verbose
+  - cargo build --features $FEATURES --release --verbose
+  - cargo doc --features $FEATURES --verbose
 
 after_script:
   - curl http://www.rust-ci.org/artifacts/put?t=$RUSTCI_TOKEN | sh
diff --git a/Cargo.toml b/Cargo.toml
index d401aa1..6860ab7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,3 +16,4 @@ clippy = {version = "0.0.104", optional = true}
 [features]
 default = []
 dev = ["clippy"]
+ipv6-iterator = []
diff --git a/src/ipv6.rs b/src/ipv6.rs
index f46f84d..5b08519 100644
--- a/src/ipv6.rs
+++ b/src/ipv6.rs
@@ -29,6 +29,29 @@ impl Ipv6Network {
         }
     }
 
+    /// Returns an iterator over `Ipv6Network`. Each call to `next` will return the next
+    /// `Ipv6Addr` in the given network. `None` will be returned when there are no more
+    /// addresses.
+    #[cfg(feature = "ipv6-iterator")]
+    pub fn iter(&self) -> Ipv6NetworkIterator {
+
+        let dec = u128::from(self.addr);
+        let max = u128::max_value();
+        let prefix = self.prefix;
+
+        let mask = max.checked_shl((IPV6_BITS - prefix) as u32).unwrap_or(0);
+        let start: u128 = dec & mask;
+
+        let mask = max.checked_shr(prefix as u32).unwrap_or(0);
+        let end: u128 = dec | mask;
+
+        Ipv6NetworkIterator{
+            next: start,
+            end: end,
+        }
+
+    }
+
     pub fn ip(&self) -> Ipv6Addr {
         self.addr
     }
@@ -90,6 +113,33 @@ impl FromStr for Ipv6Network {
     }
 }
 
+#[cfg(feature = "ipv6-iterator")]
+pub struct Ipv6NetworkIterator {
+    next: u128,
+    end: u128,
+}
+
+#[cfg(feature = "ipv6-iterator")]
+impl Iterator for Ipv6NetworkIterator {
+    type Item = Ipv6Addr;
+
+    fn next(&mut self) -> Option<Ipv6Addr> {
+
+        if self.next <= self.end {
+
+            let next = Ipv6Addr::from(self.next);
+            self.next += 1;
+            Some(next)
+
+        } else {
+
+            None
+
+        }
+
+    }
+}
+
 impl fmt::Display for Ipv6Network {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         write!(fmt, "{}/{}", self.ip(), self.prefix())
@@ -208,4 +258,36 @@ mod test {
         let prefix = ipv6_mask_to_prefix(mask);
         assert!(prefix.is_err());
     }
+
+    #[test]
+    #[cfg(feature = "ipv6-iterator")]
+    fn iterator_v6() {
+        let cidr: Ipv6Network = "2001:db8::/126".parse().unwrap();
+        let mut iter = cidr.iter();
+        assert_eq!(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0), iter.next().unwrap());
+        assert_eq!(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1), iter.next().unwrap());
+        assert_eq!(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 2), iter.next().unwrap());
+        assert_eq!(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 3), iter.next().unwrap());
+        assert_eq!(None, iter.next());
+    }
+
+    #[test]
+    #[cfg(feature = "ipv6-iterator")]
+    fn iterator_v6_tiny() {
+        let cidr: Ipv6Network = "2001:db8::/128".parse().unwrap();
+        let mut iter = cidr.iter();
+        assert_eq!(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0), iter.next().unwrap());
+        assert_eq!(None, iter.next());
+    }
+
+    #[test]
+    #[cfg(feature = "ipv6-iterator")]
+    fn iterator_v6_huge() {
+        let cidr: Ipv6Network = "2001:db8::/0".parse().unwrap();
+        let mut iter = cidr.iter();
+        assert_eq!(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0), iter.next().unwrap());
+        assert_eq!(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1), iter.next().unwrap());
+        assert_eq!(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 2), iter.next().unwrap());
+    }
+
 }
diff --git a/src/lib.rs b/src/lib.rs
index d5c8190..5758984 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,6 +3,7 @@
 //! is still WIP.
 #![cfg_attr(feature = "dev", feature(plugin))]
 #![cfg_attr(feature = "dev", plugin(clippy))]
+#![cfg_attr(feature = "ipv6-iterator", feature(i128_type))]
 #![crate_type = "lib"]
 
 use std::fmt;