Cap calls to with_capacity with sane default

src/collection_impls.rs

@@ -12,7 +12,7 @@
 
 use std::hash::Hash;
 
-use {Decodable, Encodable, Decoder, Encoder};
+use {Decodable, Encodable, Decoder, Encoder, cap_capacity};
 use std::collections::{LinkedList, VecDeque, BTreeMap, BTreeSet, HashMap, HashSet};
 
 impl<
@@ -149,7 +149,7 @@ impl<K, V> Decodable for HashMap<K, V>
 {
     fn decode<D: Decoder>(d: &mut D) -> Result<HashMap<K, V>, D::Error> {
         d.read_map(|d, len| {
-            let mut map = HashMap::with_capacity(len);
+            let mut map = HashMap::with_capacity(cap_capacity::<(K, V)>(len));
             for i in 0..len {
                 let key = try!(d.read_map_elt_key(i, |d| Decodable::decode(d)));
                 let val = try!(d.read_map_elt_val(i, |d| Decodable::decode(d)));
@@ -176,7 +176,7 @@ impl<T> Encodable for HashSet<T> where T: Encodable + Hash + Eq {
 impl<T> Decodable for HashSet<T> where T: Decodable + Hash + Eq, {
     fn decode<D: Decoder>(d: &mut D) -> Result<HashSet<T>, D::Error> {
         d.read_seq(|d, len| {
-            let mut set = HashSet::with_capacity(len);
+            let mut set = HashSet::with_capacity(cap_capacity::<T>(len));
             for i in 0..len {
                 set.insert(try!(d.read_seq_elt(i, |d| Decodable::decode(d))));
             }

src/lib.rs

@@ -43,6 +43,20 @@
 pub use self::serialize::{Decoder, Encoder, Decodable, Encodable,
                           DecoderHelpers, EncoderHelpers};
 
+
+// Limit collections from allocating more than
+// 1 MB for calls to `with_capacity`.
+fn cap_capacity<T>(given_len: usize) -> usize {
+    use std::cmp::min;
+    use std::mem::size_of;
+    const PRE_ALLOCATE_CAP: usize = 0x100000;
+
+    match size_of::<T>() {
+        0 => min(given_len, PRE_ALLOCATE_CAP),
+        n => min(given_len, PRE_ALLOCATE_CAP / n)
+    }
+}
+
 mod serialize;
 mod collection_impls;
 

src/serialize.rs

@@ -22,6 +22,8 @@ use std::sync::Arc;
 use std::marker::PhantomData;
 use std::borrow::Cow;
 
+use cap_capacity;
+
 pub trait Encoder {
     type Error;
 
@@ -479,7 +481,7 @@ impl<T:Encodable> Encodable for Vec<T> {
 impl<T:Decodable> Decodable for Vec<T> {
     fn decode<D: Decoder>(d: &mut D) -> Result<Vec<T>, D::Error> {
         d.read_seq(|d, len| {
-            let mut v = Vec::with_capacity(len);
+            let mut v = Vec::with_capacity(cap_capacity::<T>(len));
             for i in 0..len {
                 v.push(try!(d.read_seq_elt(i, |d| Decodable::decode(d))));
             }
@@ -722,11 +724,120 @@ impl<D: Decoder> DecoderHelpers for D {
         FnMut(&mut D) -> Result<T, D::Error>,
     {
         self.read_seq(|this, len| {
-            let mut v = Vec::with_capacity(len);
+            let mut v = Vec::with_capacity(cap_capacity::<T>(len));
             for i in 0..len {
                 v.push(try!(this.read_seq_elt(i, |this| f(this))));
             }
             Ok(v)
         })
     }
 }
+
+#[test]
+#[allow(unused_variables)]
+fn capacity_rules() {
+    use std::usize::MAX;
+    use std::collections::{HashMap, HashSet};
+
+    struct MyDecoder;
+    impl Decoder for MyDecoder {
+        type Error = ();
+
+        // Primitive types:
+        fn read_nil(&mut self) -> Result<(), Self::Error> { Err(()) }
+        fn read_usize(&mut self) -> Result<usize, Self::Error> { Err(()) }
+        fn read_u64(&mut self) -> Result<u64, Self::Error> { Err(()) }
+        fn read_u32(&mut self) -> Result<u32, Self::Error> { Err(()) }
+        fn read_u16(&mut self) -> Result<u16, Self::Error> { Err(()) }
+        fn read_u8(&mut self) -> Result<u8, Self::Error> { Err(()) }
+        fn read_isize(&mut self) -> Result<isize, Self::Error> { Err(()) }
+        fn read_i64(&mut self) -> Result<i64, Self::Error> { Err(()) }
+        fn read_i32(&mut self) -> Result<i32, Self::Error> { Err(()) }
+        fn read_i16(&mut self) -> Result<i16, Self::Error> { Err(()) }
+        fn read_i8(&mut self) -> Result<i8, Self::Error> { Err(()) }
+        fn read_bool(&mut self) -> Result<bool, Self::Error> { Err(()) }
+        fn read_f64(&mut self) -> Result<f64, Self::Error> { Err(()) }
+        fn read_f32(&mut self) -> Result<f32, Self::Error> { Err(()) }
+        fn read_char(&mut self) -> Result<char, Self::Error> { Err(()) }
+        fn read_str(&mut self) -> Result<String, Self::Error> { Err(()) }
+
+        // Compound types:
+        fn read_enum<T, F>(&mut self, name: &str, f: F) -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+
+        fn read_enum_variant<T, F>(&mut self, names: &[&str], f: F)
+                                   -> Result<T, Self::Error>
+            where F: FnMut(&mut Self, usize) -> Result<T, Self::Error> { Err(()) }
+        fn read_enum_variant_arg<T, F>(&mut self, a_idx: usize, f: F)
+                                       -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+
+        fn read_enum_struct_variant<T, F>(&mut self, names: &[&str], f: F)
+                                          -> Result<T, Self::Error>
+            where F: FnMut(&mut Self, usize) -> Result<T, Self::Error> { Err(()) }
+        fn read_enum_struct_variant_field<T, F>(&mut self,
+                                                f_name: &str,
+                                                f_idx: usize,
+                                                f: F)
+                                                -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+
+        fn read_struct<T, F>(&mut self, s_name: &str, len: usize, f: F)
+                             -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+        fn read_struct_field<T, F>(&mut self,
+                                   f_name: &str,
+                                   f_idx: usize,
+                                   f: F)
+                                   -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+
+        fn read_tuple<T, F>(&mut self, len: usize, f: F) -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+        fn read_tuple_arg<T, F>(&mut self, a_idx: usize, f: F)
+                                -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+
+        fn read_tuple_struct<T, F>(&mut self, s_name: &str, len: usize, f: F)
+                                   -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+        fn read_tuple_struct_arg<T, F>(&mut self, a_idx: usize, f: F)
+                                       -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+
+        // Specialized types:
+        fn read_option<T, F>(&mut self, f: F) -> Result<T, Self::Error>
+            where F: FnMut(&mut Self, bool) -> Result<T, Self::Error> { Err(()) }
+
+        fn read_seq<T, F>(&mut self, f: F) -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self, usize) -> Result<T, Self::Error> {
+                f(self, MAX)
+            }
+        fn read_seq_elt<T, F>(&mut self, idx: usize, f: F) -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+
+        fn read_map<T, F>(&mut self, f: F) -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self, usize) -> Result<T, Self::Error> {
+            f(self, MAX)
+        }
+        fn read_map_elt_key<T, F>(&mut self, idx: usize, f: F)
+                                  -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+        fn read_map_elt_val<T, F>(&mut self, idx: usize, f: F)
+                                  -> Result<T, Self::Error>
+            where F: FnOnce(&mut Self) -> Result<T, Self::Error> { Err(()) }
+
+        // Failure
+        fn error(&mut self, err: &str) -> Self::Error { () }
+    }
+
+    let mut dummy = MyDecoder;
+    let vec_result: Result<Vec<u8>, ()> = Decodable::decode(&mut dummy);
+    assert!(vec_result.is_err());
+
+    let map_result: Result<HashMap<u8, u8>, ()> = Decodable::decode(&mut dummy);
+    assert!(map_result.is_err());
+
+    let set_result: Result<HashSet<u8>, ()> = Decodable::decode(&mut dummy);
+    assert!(set_result.is_err());
+}