2 files changed, 44 insertions, 28 deletions

diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs
index 62de223..50a781c 100644
--- a/src/decoder/mod.rs
+++ b/src/decoder/mod.rs
@@ -1,5 +1,7 @@
 use std::error;
 use std::fmt;
+use std::collections::{btree_map, BTreeMap};
+use std::iter::Peekable;
 
 use Value;
 use self::DecodeErrorKind::*;
@@ -17,6 +19,8 @@ pub struct Decoder {
     /// whether fields were decoded or not.
     pub toml: Option<Value>,
     cur_field: Option<String>,
+    cur_map: Peekable<btree_map::IntoIter<String, Value>>,
+    leftover_map: ::Table,
 }
 
 /// Description for errors which can occur while decoding a type.
@@ -105,7 +109,12 @@ impl Decoder {
     /// This decoder can be passed to the `Decodable` methods or driven
     /// manually.
     pub fn new(toml: Value) -> Decoder {
-        Decoder { toml: Some(toml), cur_field: None }
+        Decoder {
+            toml: Some(toml),
+            cur_field: None,
+            leftover_map: BTreeMap::new(),
+            cur_map: BTreeMap::new().into_iter().peekable(),
+        }
     }
 
     fn sub_decoder(&self, toml: Option<Value>, field: &str) -> Decoder {
@@ -118,7 +127,9 @@ impl Decoder {
                     None => Some(format!("{}", field)),
                     Some(ref s) => Some(format!("{}.{}", s, field))
                 }
-            }
+            },
+            leftover_map: BTreeMap::new(),
+            cur_map: BTreeMap::new().into_iter().peekable(),
         }
     }
 
diff --git a/src/decoder/rustc_serialize.rs b/src/decoder/rustc_serialize.rs
index b84eeb9..9d2644e 100644
--- a/src/decoder/rustc_serialize.rs
+++ b/src/decoder/rustc_serialize.rs
@@ -1,5 +1,6 @@
 use rustc_serialize;
 use std::mem;
+use std::collections::BTreeMap;
 
 use super::{Decoder, DecodeError};
 use super::DecodeErrorKind::*;
@@ -275,46 +276,50 @@ impl rustc_serialize::Decoder for Decoder {
         -> Result<T, DecodeError>
         where F: FnOnce(&mut Decoder, usize) -> Result<T, DecodeError>
     {
-        let len = match self.toml {
-            Some(Value::Table(ref table)) => table.len(),
-            ref found => return Err(self.mismatch("table", found)),
+        let map = match self.toml.take() {
+            Some(Value::Table(table)) => table,
+            found => {
+                self.toml = found;
+                return Err(self.mismatch("table", &self.toml))
+            }
         };
-        let ret = try!(f(self, len));
-        self.toml.take();
+        let amt = map.len();
+        let prev_iter = mem::replace(&mut self.cur_map,
+                                     map.into_iter().peekable());
+        let prev_map = mem::replace(&mut self.leftover_map, BTreeMap::new());
+        let ret = try!(f(self, amt));
+        let leftover = mem::replace(&mut self.leftover_map, prev_map);
+        self.cur_map = prev_iter;
+        if leftover.len() > 0 {
+            self.toml = Some(Value::Table(leftover));
+        }
         Ok(ret)
     }
     fn read_map_elt_key<T, F>(&mut self, idx: usize, f: F)
         -> Result<T, DecodeError>
         where F: FnOnce(&mut Decoder) -> Result<T, DecodeError>
     {
-        match self.toml {
-            Some(Value::Table(ref table)) => {
-                match table.iter().skip(idx).next() {
-                    Some((key, _)) => {
-                        let val = Value::String(key.to_string());
-                        f(&mut self.sub_decoder(Some(val), key))
-                    }
-                    None => Err(self.err(ExpectedMapKey(idx))),
-                }
-            }
-            ref found => Err(self.mismatch("table", found)),
-        }
+        let key = match self.cur_map.peek().map(|p| p.0.clone()) {
+            Some(k) => k,
+            None => return Err(self.err(ExpectedMapKey(idx))),
+        };
+        let val = Value::String(key.clone());
+        f(&mut self.sub_decoder(Some(val), &key))
     }
     fn read_map_elt_val<T, F>(&mut self, idx: usize, f: F)
         -> Result<T, DecodeError>
         where F: FnOnce(&mut Decoder) -> Result<T, DecodeError>
     {
-        match self.toml {
-            Some(Value::Table(ref table)) => {
-                match table.iter().skip(idx).next() {
-                    Some((key, value)) => {
-                        // XXX: this shouldn't clone
-                        f(&mut self.sub_decoder(Some(value.clone()), key))
-                    }
-                    None => Err(self.err(ExpectedMapElement(idx))),
+        match self.cur_map.next() {
+            Some((key, value)) => {
+                let mut d = self.sub_decoder(Some(value), &key);
+                let ret = f(&mut d);
+                if let Some(toml) = d.toml.take() {
+                    self.leftover_map.insert(key, toml);
                 }
+                return ret
             }
-            ref found => Err(self.mismatch("table", found)),
+            None => return Err(self.err(ExpectedMapElement(idx))),
         }
     }