Add a count method to Source (#329)

src/kv/source.rs

@@ -5,13 +5,45 @@ use kv::{Error, Key, ToKey, Value, ToValue};
 /// A source of key-value pairs.
 /// 
 /// The source may be a single pair, a set of pairs, or a filter over a set of pairs.
-/// Use the [`Visitor`](struct.Visitor.html) trait to inspect the structured data
+/// Use the [`Visitor`](trait.Visitor.html) trait to inspect the structured data
 /// in a source.
 pub trait Source {
     /// Visit key-value pairs.
     /// 
-    /// A source doesn't have to guarantee any ordering or uniqueness of pairs.
+    /// A source doesn't have to guarantee any ordering or uniqueness of key-value pairs.
+    /// If the given visitor returns an error then the source may early-return with it,
+    /// even if there are more key-value pairs.
+    /// 
+    /// # Implementation notes
+    /// 
+    /// A source should yield the same key-value pairs to a subsequent visitor unless
+    /// that visitor itself fails.
     fn visit<'kvs>(&'kvs self, visitor: &mut Visitor<'kvs>) -> Result<(), Error>;
+
+    /// Count the number of key-value pairs that can be visited.
+    /// 
+    /// # Implementation notes
+    /// 
+    /// A source that knows the number of key-value pairs upfront may provide a more
+    /// efficient implementation.
+    /// 
+    /// A subsequent call to `visit` should yield the same number of key-value pairs
+    /// to the visitor, unless that visitor fails part way through.
+    fn count(&self) -> usize {
+        struct Count(usize);
+
+        impl<'kvs> Visitor<'kvs> for Count {
+            fn visit_pair(&mut self, _: Key<'kvs>, _: Value<'kvs>) -> Result<(), Error> {
+                self.0 += 1;
+
+                Ok(())
+            }
+        }
+
+        let mut count = Count(0);
+        let _ = self.visit(&mut count);
+        count.0
+    }
 }
 
 impl<'a, T> Source for &'a T
@@ -21,6 +53,10 @@ where
     fn visit<'kvs>(&'kvs self, visitor: &mut Visitor<'kvs>) -> Result<(), Error> {
         (**self).visit(visitor)
     }
+
+    fn count(&self) -> usize {
+        (**self).count()
+    }
 }
 
 impl<K, V> Source for (K, V)
@@ -31,6 +67,10 @@ where
     fn visit<'kvs>(&'kvs self, visitor: &mut Visitor<'kvs>) -> Result<(), Error> {
         visitor.visit_pair(self.0.to_key(), self.1.to_value())
     }
+
+    fn count(&self) -> usize {
+        1
+    }
 }
 
 impl<S> Source for [S]
@@ -44,6 +84,10 @@ where
 
         Ok(())
     }
+
+    fn count(&self) -> usize {
+        self.len()
+    }
 }
 
 impl<S> Source for Option<S>
@@ -57,6 +101,10 @@ where
 
         Ok(())
     }
+
+    fn count(&self) -> usize {
+        self.as_ref().map(Source::count).unwrap_or(0)
+    }
 }
 
 /// A visitor for the key-value pairs in a [`Source`](trait.Source.html).
@@ -85,6 +133,10 @@ mod std_support {
         fn visit<'kvs>(&'kvs self, visitor: &mut Visitor<'kvs>) -> Result<(), Error> {
             (**self).visit(visitor)
         }
+
+        fn count(&self) -> usize {
+            (**self).count()
+        }
     }
 
     impl<S> Source for Vec<S>
@@ -94,6 +146,10 @@ mod std_support {
         fn visit<'kvs>(&'kvs self, visitor: &mut Visitor<'kvs>) -> Result<(), Error> {
             (**self).visit(visitor)
         }
+
+        fn count(&self) -> usize {
+            (**self).count()
+        }
     }
 
     impl<'kvs, V> Visitor<'kvs> for Box<V>
@@ -104,6 +160,17 @@ mod std_support {
             (**self).visit_pair(key, value)
         }
     }
+
+    #[cfg(test)]
+    mod tests {
+        use super::*;
+
+        #[test]
+        fn count() {
+            assert_eq!(1, Source::count(&Box::new(("a", 1))));
+            assert_eq!(2, Source::count(&vec![("a", 1), ("b", 2)]));
+        }
+    }
 }
 
 #[cfg(test)]
@@ -119,4 +186,23 @@ mod tests {
     fn visitor_is_object_safe() {
         fn _check(_: &Visitor) {}
     }
+
+    #[test]
+    fn count() {
+        struct OnePair {
+            key: &'static str,
+            value: i32,
+        }
+
+        impl Source for OnePair {
+            fn visit<'kvs>(&'kvs self, visitor: &mut Visitor<'kvs>) -> Result<(), Error> {
+                visitor.visit_pair(self.key.to_key(), self.value.to_value())
+            }
+        }
+
+        assert_eq!(1, Source::count(&("a", 1)));
+        assert_eq!(2, Source::count(&[("a", 1), ("b", 2)] as &[_]));
+        assert_eq!(0, Source::count(&Option::None::<(&str, i32)>));
+        assert_eq!(1, Source::count(&OnePair { key: "a", value: 1 }));
+    }
 }