feat!: Provide a Map::entry API for tables (#154)

To make inserting dotted keys into an `InlineTable` easier, I wanted
`Table::entry` but translating it directly over didn't work because it
relied on `Item::None`, so instead switching the API to what you would
see for a stdlib map.

BREAKING CHANGE: `entry`s return type changed

src/inline_table.rs

@@ -60,6 +60,56 @@ impl InlineTable {
         self.items.clear()
     }
 
+    /// Gets the given key's corresponding entry in the Table for in-place manipulation.
+    pub fn entry<'a>(&'a mut self, key: &str) -> InlineEntry<'a> {
+        // Accept a `&str` rather than an owned type to keep `InternalString`, well, internal
+        match self.items.entry(key.to_owned()) {
+            linked_hash_map::Entry::Occupied(mut entry) => {
+                // Ensure it is a `Value` to simplify `InlineOccupiedEntry`'s code.
+                let mut scratch = Item::None;
+                std::mem::swap(&mut scratch, &mut entry.get_mut().value);
+                let mut scratch = Item::Value(
+                    scratch
+                        .into_value()
+                        // HACK: `Item::None` is a corner case of a corner case, let's just pick a
+                        // "safe" value
+                        .unwrap_or_else(|_| Value::InlineTable(Default::default())),
+                );
+                std::mem::swap(&mut scratch, &mut entry.get_mut().value);
+
+                InlineEntry::Occupied(InlineOccupiedEntry { entry })
+            }
+            linked_hash_map::Entry::Vacant(entry) => {
+                InlineEntry::Vacant(InlineVacantEntry { entry, key: None })
+            }
+        }
+    }
+
+    /// Gets the given key's corresponding entry in the Table for in-place manipulation.
+    pub fn entry_format<'a>(&'a mut self, key: &'a Key) -> InlineEntry<'a> {
+        // Accept a `&Key` to be consistent with `entry`
+        match self.items.entry(key.get().to_owned()) {
+            linked_hash_map::Entry::Occupied(mut entry) => {
+                // Ensure it is a `Value` to simplify `InlineOccupiedEntry`'s code.
+                let mut scratch = Item::None;
+                std::mem::swap(&mut scratch, &mut entry.get_mut().value);
+                let mut scratch = Item::Value(
+                    scratch
+                        .into_value()
+                        // HACK: `Item::None` is a corner case of a corner case, let's just pick a
+                        // "safe" value
+                        .unwrap_or_else(|_| Value::InlineTable(Default::default())),
+                );
+                std::mem::swap(&mut scratch, &mut entry.get_mut().value);
+
+                InlineEntry::Occupied(InlineOccupiedEntry { entry })
+            }
+            linked_hash_map::Entry::Vacant(entry) => InlineEntry::Vacant(InlineVacantEntry {
+                entry,
+                key: Some(key),
+            }),
+        }
+    }
     /// Return an optional reference to the value at the given the key.
     pub fn get(&self, key: &str) -> Option<&Value> {
         self.items.get(key).and_then(|kv| kv.value.as_value())
@@ -103,8 +153,18 @@ impl InlineTable {
     }
 
     /// Removes an item given the key.
-    pub fn remove(&mut self, key: &str) -> Option<Item> {
-        self.items.remove(key).map(|kv| kv.value)
+    pub fn remove(&mut self, key: &str) -> Option<Value> {
+        self.items
+            .remove(key)
+            .and_then(|kv| kv.value.into_value().ok())
+    }
+
+    /// Removes a key from the map, returning the stored key and value if the key was previously in the map.
+    pub fn remove_entry(&mut self, key: &str) -> Option<(Key, Value)> {
+        self.items.remove(key).and_then(|kv| {
+            let key = kv.key;
+            kv.value.into_value().ok().map(|value| (key, value))
+        })
     }
 }
 
@@ -213,3 +273,137 @@ impl TableLike for InlineTable {
 
 // `{ key1 = value1, ... }`
 pub(crate) const DEFAULT_INLINE_KEY_DECOR: (&str, &str) = (" ", " ");
+
+/// A view into a single location in a map, which may be vacant or occupied.
+pub enum InlineEntry<'a> {
+    /// An occupied Entry.
+    Occupied(InlineOccupiedEntry<'a>),
+    /// A vacant Entry.
+    Vacant(InlineVacantEntry<'a>),
+}
+
+impl<'a> InlineEntry<'a> {
+    /// Returns the entry key
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use toml_edit::Table;
+    ///
+    /// let mut map = Table::new();
+    ///
+    /// assert_eq!("hello", map.entry("hello").key());
+    /// ```
+    pub fn key(&self) -> &str {
+        match self {
+            InlineEntry::Occupied(e) => e.key(),
+            InlineEntry::Vacant(e) => e.key(),
+        }
+    }
+
+    /// Ensures a value is in the entry by inserting the default if empty, and returns
+    /// a mutable reference to the value in the entry.
+    pub fn or_insert(self, default: Value) -> &'a mut Value {
+        match self {
+            InlineEntry::Occupied(entry) => entry.into_mut(),
+            InlineEntry::Vacant(entry) => entry.insert(default),
+        }
+    }
+
+    /// Ensures a value is in the entry by inserting the result of the default function if empty,
+    /// and returns a mutable reference to the value in the entry.
+    pub fn or_insert_with<F: FnOnce() -> Value>(self, default: F) -> &'a mut Value {
+        match self {
+            InlineEntry::Occupied(entry) => entry.into_mut(),
+            InlineEntry::Vacant(entry) => entry.insert(default()),
+        }
+    }
+}
+
+/// A view into a single occupied location in a `LinkedHashMap`.
+pub struct InlineOccupiedEntry<'a> {
+    entry: linked_hash_map::OccupiedEntry<'a, InternalString, TableKeyValue>,
+}
+
+impl<'a> InlineOccupiedEntry<'a> {
+    /// Gets a reference to the entry key
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use toml_edit::Table;
+    ///
+    /// let mut map = Table::new();
+    ///
+    /// assert_eq!("foo", map.entry("foo").key());
+    /// ```
+    pub fn key(&self) -> &str {
+        self.entry.key().as_str()
+    }
+
+    /// Gets a reference to the value in the entry.
+    pub fn get(&self) -> &Value {
+        self.entry.get().value.as_value().unwrap()
+    }
+
+    /// Gets a mutable reference to the value in the entry.
+    pub fn get_mut(&mut self) -> &mut Value {
+        self.entry.get_mut().value.as_value_mut().unwrap()
+    }
+
+    /// Converts the OccupiedEntry into a mutable reference to the value in the entry
+    /// with a lifetime bound to the map itself
+    pub fn into_mut(self) -> &'a mut Value {
+        self.entry.into_mut().value.as_value_mut().unwrap()
+    }
+
+    /// Sets the value of the entry, and returns the entry's old value
+    pub fn insert(&mut self, value: Value) -> Value {
+        let mut value = Item::Value(value);
+        std::mem::swap(&mut value, &mut self.entry.get_mut().value);
+        value.into_value().unwrap()
+    }
+
+    /// Takes the value out of the entry, and returns it
+    pub fn remove(self) -> Value {
+        self.entry.remove().value.into_value().unwrap()
+    }
+}
+
+/// A view into a single empty location in a `LinkedHashMap`.
+pub struct InlineVacantEntry<'a> {
+    entry: linked_hash_map::VacantEntry<'a, InternalString, TableKeyValue>,
+    key: Option<&'a Key>,
+}
+
+impl<'a> InlineVacantEntry<'a> {
+    /// Gets a reference to the entry key
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use toml_edit::Table;
+    ///
+    /// let mut map = Table::new();
+    ///
+    /// assert_eq!("foo", map.entry("foo").key());
+    /// ```
+    pub fn key(&self) -> &str {
+        self.entry.key().as_str()
+    }
+
+    /// Sets the value of the entry with the VacantEntry's key,
+    /// and returns a mutable reference to it
+    pub fn insert(self, value: Value) -> &'a mut Value {
+        let key = self
+            .key
+            .cloned()
+            .unwrap_or_else(|| Key::with_key(self.key()));
+        let value = Item::Value(value);
+        self.entry
+            .insert(TableKeyValue::new(key, value))
+            .value
+            .as_value_mut()
+            .unwrap()
+    }
+}

src/lib.rs

@@ -86,11 +86,14 @@ pub use crate::array_of_tables::ArrayOfTables;
 pub use crate::datetime::*;
 pub use crate::document::Document;
 pub use crate::inline_table::{
-    InlineTable, InlineTableIntoIter, InlineTableIter, InlineTableIterMut,
+    InlineEntry, InlineOccupiedEntry, InlineTable, InlineTableIntoIter, InlineTableIter,
+    InlineTableIterMut, InlineVacantEntry,
 };
 pub use crate::item::{array, table, value, Item};
 pub use crate::key::Key;
 pub use crate::parser::TomlError;
 pub use crate::repr::{Decor, Repr};
-pub use crate::table::{IntoIter, Iter, IterMut, Table, TableLike};
+pub use crate::table::{
+    Entry, IntoIter, Iter, IterMut, OccupiedEntry, Table, TableLike, VacantEntry,
+};
 pub use crate::value::Value;

src/parser/table.rs

@@ -5,7 +5,7 @@ use crate::parser::key::key;
 use crate::parser::trivia::line_trailing;
 use crate::parser::TomlParser;
 use crate::repr::Decor;
-use crate::{Item, Table};
+use crate::{Entry, Item, Table};
 use combine::parser::char::char;
 use combine::parser::range::range;
 use combine::stream::RangeStream;
@@ -81,17 +81,16 @@ pub(crate) fn duplicate_key(path: &[Key], i: usize) -> CustomError {
 impl TomlParser {
     pub(crate) fn descend_path<'a>(
         table: &'a mut Table,
-        path: &[Key],
+        path: &'a [Key],
         i: usize,
     ) -> Result<&'a mut Table, CustomError> {
         if let Some(key) = path.get(i) {
-            let entry = table.entry_format(key);
-            if entry.is_none() {
+            let entry = table.entry_format(key).or_insert_with(|| {
                 let mut new_table = Table::new();
                 new_table.set_implicit(true);
 
-                *entry = Item::Table(new_table);
-            }
+                Item::Table(new_table)
+            });
             match *entry {
                 Item::Value(..) => Err(duplicate_key(path, i)),
                 Item::ArrayOfTables(ref mut array) => {
@@ -119,23 +118,15 @@ impl TomlParser {
         let table = self.document.as_table_mut();
         self.current_table_position += 1;
 
-        let table = Self::descend_path(table, &path[..path.len() - 1], 0);
+        let table = Self::descend_path(table, &path[..path.len() - 1], 0)?;
         let key = &path[path.len() - 1];
 
-        match table {
-            Ok(table) => {
-                let decor = Decor::new(leading, trailing);
+        let decor = Decor::new(leading, trailing);
 
-                let entry = table.entry_format(key);
-                if entry.is_none() {
-                    *entry = Item::Table(Table::with_decor_and_pos(
-                        decor,
-                        Some(self.current_table_position),
-                    ));
-                    self.current_table_path = path.to_vec();
-                    return Ok(());
-                }
-                match *entry {
+        let entry = table.entry_format(key);
+        match entry {
+            Entry::Occupied(entry) => {
+                match entry.into_mut() {
                     // if [a.b.c] header preceded [a.b]
                     Item::Table(ref mut t) if t.implicit => {
                         debug_assert!(t.values_len() == 0);
@@ -145,13 +136,20 @@ impl TomlParser {
                         t.set_implicit(false);
 
                         self.current_table_path = path.to_vec();
-                        return Ok(());
+                        Ok(())
                     }
-                    _ => {}
+                    _ => Err(duplicate_key(path, path.len() - 1)),
                 }
-                Err(duplicate_key(path, path.len() - 1))
             }
-            Err(e) => Err(e),
+            Entry::Vacant(entry) => {
+                let item = Item::Table(Table::with_decor_and_pos(
+                    decor,
+                    Some(self.current_table_position),
+                ));
+                entry.insert(item);
+                self.current_table_path = path.to_vec();
+                Ok(())
+            }
         }
     }