diff --git a/src/lib.rs b/src/lib.rs index c665a4f..ba2165c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -47,7 +47,7 @@ use std::borrow::{Cow, ToOwned}; use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, LinkedList, VecDeque}; use std::ffi::{CString, OsString}; use std::hash::BuildHasher; -use std::net::{Ipv4Addr, Ipv6Addr}; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; use std::ops::Bound; use std::path::PathBuf; use std::rc::Rc; @@ -1140,6 +1140,23 @@ impl<'a> Arbitrary<'a> for Ipv6Addr { } } +impl<'a> Arbitrary<'a> for IpAddr { + fn arbitrary(u: &mut Unstructured<'a>) -> Result { + if u.arbitrary()? { + Ok(IpAddr::V4(u.arbitrary()?)) + } else { + Ok(IpAddr::V6(u.arbitrary()?)) + } + } + + fn size_hint(depth: usize) -> (usize, Option) { + size_hint::and( + bool::size_hint(depth), + size_hint::or(Ipv4Addr::size_hint(depth), Ipv6Addr::size_hint(depth)), + ) + } +} + #[cfg(test)] mod test { use super::*;