1 files changed, 95 insertions, 31 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 21cdf3e..d81b9cd 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -44,7 +44,6 @@
 
 use std::collections::BTreeMap;
 use std::str::FromStr;
-use std::str::Split;
 
 pub use parser::{Parser, ParserError};
 
@@ -208,28 +207,6 @@ impl Value {
         Some(cur_value)
     }
 
-    fn lookup_mut_recurse<'a>(&'a mut self, matches: &mut Split<'a, char>) -> Option<&'a mut Value> {
-        if let Some(key) = matches.next() {
-            match *self {
-                Value::Table(ref mut hm) => {
-                    match hm.get_mut(key) {
-                      Some(v) => return v.lookup_mut_recurse(matches),
-                      None => return None,
-                  }
-              },
-              Value::Array(ref mut v) => {
-                  match key.parse::<usize>().ok() {
-                      Some(idx) if idx < v.len()
-                          => return (&mut v[idx]).lookup_mut_recurse(matches),
-                        _ => return None,
-                    }
-                },
-            _ => return None
-            }
-        }
-        Some(self)
-    }
-
     /// Lookups for mutable value at specified path.
     ///
     /// Uses '.' as a path separator.
@@ -259,15 +236,32 @@ impl Value {
     /// let result = value.lookup_mut("test.foo").unwrap();
     /// assert_eq!(result.as_str().unwrap(), "foo");
     /// ```
-    pub fn lookup_mut<'a>(&'a mut self, path: &'a str) -> Option<&'a mut Value> {
-      if path.len() == 0 {
-          return Some(self)
-      }
-
-      let mut matches = path.split('.');
-      self.lookup_mut_recurse(&mut matches)
-  }
+    pub fn lookup_mut(&mut self, path: &str) -> Option<&mut Value> {
+        let mut cur = self;
+        if path.len() == 0 {
+            return Some(cur)
+        }
 
+        for key in path.split('.') {
+            let tmp = cur;
+            match *tmp {
+                Value::Table(ref mut hm) => {
+                    match hm.get_mut(key) {
+                        Some(v) => cur = v,
+                        None => return None
+                    }
+                }
+                Value::Array(ref mut v) => {
+                    match key.parse::<usize>().ok() {
+                        Some(idx) if idx < v.len() => cur = &mut v[idx],
+                        _ => return None
+                    }
+                }
+                _ => return None
+           }
+        }
+        Some(cur)
+    }
 }
 
 impl FromStr for Value {
@@ -286,6 +280,76 @@ mod tests {
     use super::Value;
 
     #[test]
+    fn lookup_mut_valid() {
+        let toml = r#"
+              [test]
+              foo = "bar"
+
+              [[values]]
+              foo = "baz"
+
+              [[values]]
+              foo = "qux"
+        "#;
+
+        let mut value: Value = toml.parse().unwrap();
+
+        {
+            let test_foo = value.lookup_mut("test.foo").unwrap();
+            assert_eq!(test_foo.as_str().unwrap(), "bar");
+        }
+
+        {
+            let foo1 = value.lookup_mut("values.1.foo").unwrap();
+            assert_eq!(foo1.as_str().unwrap(), "qux");
+        }
+
+        assert!(value.lookup_mut("test.bar").is_none());
+        assert!(value.lookup_mut("test.foo.bar").is_none());
+    }
+
+    #[test]
+    fn lookup_mut_invalid_index() {
+        let toml = r#"
+            [[values]]
+            foo = "baz"
+        "#;
+
+        let mut value: Value = toml.parse().unwrap();
+
+        {
+            let foo = value.lookup_mut("test.foo");
+            assert!(foo.is_none());
+        }
+
+        {
+            let foo = value.lookup_mut("values.100.foo");
+            assert!(foo.is_none());
+        }
+
+        {
+            let foo = value.lookup_mut("values.str.foo");
+            assert!(foo.is_none());
+        }
+    }
+
+    #[test]
+    fn lookup_mut_self() {
+        let mut value: Value = r#"foo = "bar""#.parse().unwrap();
+
+        {
+            let foo = value.lookup_mut("foo").unwrap();
+            assert_eq!(foo.as_str().unwrap(), "bar");
+        }
+
+        let foo = value.lookup_mut("").unwrap();
+        assert!(foo.as_table().is_some());
+
+        let baz = foo.lookup_mut("foo").unwrap();
+        assert_eq!(baz.as_str().unwrap(), "bar");
+    }
+
+    #[test]
     fn lookup_valid() {
         let toml = r#"
               [test]