Repository: tnballo/scapegoat (language: rust)

Issue:
Add map `entry` APIs
To mirror `BTreeMap`'s [`entry`](https://doc.rust-lang.org/std/collections/struct.BTreeMap.html#method.entry), [`first_entry`](https://doc.rust-lang.org/std/collections/struct.BTreeMap.html#method.first_entry), and [`last_entry`](https://doc.rust-lang.org/std/collections/struct.BTreeMap.html#method.last_entry).


Discussion / hints:
Hi! I'd like to work on this.
Hey! Awesome, thank you - I'll assign the issue to you. To help you navigate the code:

* `src/tree/tree.rs` is where the majority of the logic is implemented, both sets and maps are mostly public wrappers around this tree.
    * Since v2.0, this module is internal-only (not part of crate's public API). So you're free to make any changes necessary.
    * This where you'll find `self.min_idx` and `self.max_idx`, array indexes (e.x. `&self.arena[self.min_idx]`) that the help access the first/minimum and last/maximum items (useful for `first_entry` and `last_entry`).

* `src/map_types.rs` is where `SgMap`'s return types are. This is a good place to define a public [`Entry` struct](https://doc.rust-lang.org/std/collections/btree_map/enum.Entry.html), even if you decide to put most of the logic elsewhere.

Two asides:

* One of the main ways this crate is tested is by fuzzing against the standard library. I'd be happy to accept a PR even if it doesn't update the fuzz harness, so you don't need to worry about this. But, in case you find fuzzing helpful in testing the APIs you add: [this section](https://rust-fuzz.github.io/book/cargo-fuzz/structure-aware-fuzzing.html) of the "Rust Fuzz" book explains structure-aware fuzzing, which is used in `fuzz/fuzz_targets/sg_map.rs`.

* `first_entry` and `last_entry` are currently nightly-only for `BTreeMap`, so if you can implement those APIs in this crate (stable 1.55+) I think that's pretty exciting :sunglasses:

Feel free to take you time on a PR since this isn't a bug/blocking. If you have any questions, I'll do my best to answer them on this thread but may not always reply same-day. Thanks for your time/effort!

Fix patch:
diff --git a/src/entry.rs b/src/entry.rs
new file mode 100644
index 0000000..1ebd0d4
--- /dev/null
+++ b/src/entry.rs
@@ -0,0 +1,232 @@
+use crate::tree::{Idx, SmallNode};
+use crate::SgMap;
+
+/// A view into an occupied entry in a `SgMap`.
+/// It is part of the [`Entry`] enum.
+pub struct OccupiedEntry<'a, K: Ord + Default, V: Default, const N: usize> {
+    pub(super) node_idx: usize,
+    pub(super) table: &'a mut SgMap<K, V, N>,
+}
+
+/// A view into a vacant entry in a `SgMap`.
+/// It is part of the [`Entry`] enum.
+pub struct VacantEntry<'a, K: Ord + Default, V: Default, const N: usize> {
+    pub(super) key: K,
+    pub(super) table: &'a mut SgMap<K, V, N>,
+}
+
+impl<'a, K: Ord + Default, V: Default, const N: usize> VacantEntry<'a, K, V, N> {
+    /// Gets a reference to the key that would be used when inserting a value
+    /// through the VacantEntry.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map = SgMap::<&str, usize, 2>::new();
+    /// assert_eq!(map.entry("poneyland").key(), &"poneyland");
+    /// ```
+    pub fn key(&self) -> &K {
+        &self.key
+    }
+
+    /// Take ownership of the key.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    /// use scapegoat::map_types::Entry;
+    ///
+    /// let mut map = SgMap::<&str, usize, 2>::new();
+    ///
+    /// if let Entry::Vacant(v) = map.entry("poneyland") {
+    ///     v.into_key();
+    /// }
+    /// ```
+    pub fn into_key(self) -> K {
+        self.key
+    }
+
+    /// Sets the value of the entry with the `VacantEntry`'s key,
+    /// and returns a mutable reference to it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    /// use scapegoat::map_types::Entry;
+    ///
+    /// let mut map = SgMap::<&str, u32, 2>::new();
+    ///
+    /// if let Entry::Vacant(o) = map.entry("poneyland") {
+    ///     o.insert(37);
+    /// }
+    /// assert_eq!(map["poneyland"], 37);
+    /// ```
+    pub fn insert(self, value: V) -> &'a mut V {
+        let (_, new_node_idx) = self.table.bst.priv_balancing_insert::<Idx>(self.key, value);
+
+        self.table.bst.arena[new_node_idx].get_mut().1
+    }
+}
+
+impl<'a, K: Ord + Default, V: Default, const N: usize> OccupiedEntry<'a, K, V, N> {
+    /// Gets a reference to the key in the entry.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map = SgMap::<&str, usize, 2>::new();
+    /// map.entry("poneyland").or_insert(12);
+    /// assert_eq!(map.entry("poneyland").key(), &"poneyland");
+    /// ```
+    pub fn key(&self) -> &K {
+        self.table.bst.arena[self.node_idx].key()
+    }
+
+    /// Gets a reference to the value in the entry.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    /// use scapegoat::map_types::Entry;
+    ///
+    /// let mut map = SgMap::<&str, usize, 2>::new();
+    /// map.entry("poneyland").or_insert(12);
+    ///
+    /// if let Entry::Occupied(o) = map.entry("poneyland") {
+    ///     assert_eq!(o.get(), &12);
+    /// }
+    /// ```
+    pub fn get(&self) -> &V {
+        self.table.bst.arena[self.node_idx].val()
+    }
+
+    /// Gets a mutable reference to the value in the entry.
+    ///
+    /// If you need a reference to the `OccupiedEntry` that may outlive the
+    /// destruction of the `Entry` value, see [`into_mut`].
+    ///
+    /// [`into_mut`]: OccupiedEntry::into_mut
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    /// use scapegoat::map_types::Entry;
+    ///
+    /// let mut map = SgMap::<&str, usize, 2>::new();
+    /// map.entry("poneyland").or_insert(12);
+    ///
+    /// assert_eq!(map["poneyland"], 12);
+    /// if let Entry::Occupied(mut o) = map.entry("poneyland") {
+    ///     *o.get_mut() += 10;
+    ///     assert_eq!(*o.get(), 22);
+    ///
+    ///     // We can use the same Entry multiple times.
+    ///     *o.get_mut() += 2;
+    /// }
+    /// assert_eq!(map["poneyland"], 24);
+    /// ```
+    pub fn get_mut(&mut self) -> &mut V {
+        self.table.bst.arena[self.node_idx].get_mut().1
+    }
+
+    /// Converts the entry into a mutable reference to its value.
+    ///
+    /// If you need multiple references to the `OccupiedEntry`, see [`get_mut`].
+    ///
+    /// [`get_mut`]: OccupiedEntry::get_mut
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    /// use scapegoat::map_types::Entry;
+    ///
+    /// let mut map = SgMap::<&str, usize, 2>::new();
+    /// map.entry("poneyland").or_insert(12);
+    ///
+    /// assert_eq!(map["poneyland"], 12);
+    /// if let Entry::Occupied(o) = map.entry("poneyland") {
+    ///     *o.into_mut() += 10;
+    /// }
+    /// assert_eq!(map["poneyland"], 22);
+    /// ```
+    pub fn into_mut(self) -> &'a mut V {
+        self.table.bst.arena[self.node_idx].get_mut().1
+    }
+
+    /// Sets the value of the entry with the `OccupiedEntry`'s key,
+    /// and returns the entry's old value.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    /// use scapegoat::map_types::Entry;
+    ///
+    /// let mut map = SgMap::<&str, usize, 2>::new();
+    /// map.entry("poneyland").or_insert(12);
+    ///
+    /// if let Entry::Occupied(mut o) = map.entry("poneyland") {
+    ///     assert_eq!(o.insert(15), 12);
+    /// }
+    /// assert_eq!(map["poneyland"], 15);
+    /// ```
+    pub fn insert(&mut self, value: V) -> V {
+        core::mem::replace(self.get_mut(), value)
+    }
+
+    /// Take ownership of the key and value from the map.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    /// use scapegoat::map_types::Entry;
+    ///
+    /// let mut map = SgMap::<&str, usize, 2>::new();
+    /// map.entry("poneyland").or_insert(12);
+    ///
+    /// if let Entry::Occupied(o) = map.entry("poneyland") {
+    ///     // We delete the entry from the map.
+    ///     o.remove_entry();
+    /// }
+    ///
+    /// // If now try to get the value, it will panic:
+    /// // println!("{}", map["poneyland"]);
+    /// ```
+    pub fn remove_entry(self) -> (K, V) {
+        self.table
+            .bst
+            .priv_remove_by_idx(self.node_idx)
+            .expect("Must be occupied")
+    }
+
+    /// Takes the value of the entry out of the map, and returns it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    /// use scapegoat::map_types::Entry;
+    ///
+    /// let mut map = SgMap::<&str, usize, 2>::new();
+    /// map.entry("poneyland").or_insert(12);
+    ///
+    /// if let Entry::Occupied(o) = map.entry("poneyland") {
+    ///     assert_eq!(o.remove(), 12);
+    /// }
+    /// // If we try to get "poneyland"'s value, it'll panic:
+    /// // println!("{}", map["poneyland"]);
+    /// ```
+    pub fn remove(self) -> V {
+        self.remove_entry().1
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 89a206e..5171c34 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -208,6 +208,9 @@ pub use crate::tree::{Arena, Node, NodeGetHelper, NodeRebuildHelper};
 mod tree;
 pub use crate::tree::SgError;
 
+mod entry;
+pub use crate::entry::{OccupiedEntry, VacantEntry};
+
 mod map;
 pub use crate::map::SgMap;
 
diff --git a/src/map.rs b/src/map.rs
index defca5e..584d971 100644
--- a/src/map.rs
+++ b/src/map.rs
@@ -3,7 +3,10 @@ use core::fmt::{self, Debug};
 use core::iter::FromIterator;
 use core::ops::Index;
 
-use crate::map_types::{IntoIter, IntoKeys, IntoValues, Iter, IterMut, Keys, Values, ValuesMut};
+use crate::entry::{OccupiedEntry, VacantEntry};
+use crate::map_types::{
+    Entry, IntoIter, IntoKeys, IntoValues, Iter, IterMut, Keys, Values, ValuesMut,
+};
 use crate::tree::{SgError, SgTree};
 
 /// Safe, fallible, embedded-friendly ordered map.
@@ -876,6 +879,100 @@ impl<K: Ord + Default, V: Default, const N: usize> SgMap<K, V, N> {
     pub fn len(&self) -> usize {
         self.bst.len()
     }
+
+    /// Gets the given key's corresponding entry in the map for in-place manipulation.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut count = SgMap::<&str, usize, 10>::new();
+    ///
+    /// // count the number of occurrences of letters in the vec
+    /// for x in vec!["a", "b", "a", "c", "a", "b"] {
+    ///     *count.entry(x).or_insert(0) += 1;
+    /// }
+    ///
+    /// assert_eq!(count["a"], 3);
+    /// ```
+    pub fn entry(&mut self, key: K) -> Entry<'_, K, V, N> {
+        use crate::tree::node::NodeGetHelper;
+        use crate::tree::Idx;
+
+        let ngh: NodeGetHelper<Idx> = self.bst.priv_get(None, &key);
+        match ngh.node_idx() {
+            Some(node_idx) => Entry::Occupied(OccupiedEntry {
+                node_idx,
+                table: self,
+            }),
+            None => Entry::Vacant(VacantEntry { key, table: self }),
+        }
+    }
+
+    /// Returns the first entry in the map for in-place manipulation.
+    /// The key of this entry is the minimum key in the map.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map = SgMap::<_, _, 10>::new();
+    /// map.insert(1, "a");
+    /// map.insert(2, "b");
+    /// if let Some(mut entry) = map.first_entry() {
+    ///     if *entry.key() > 0 {
+    ///         entry.insert("first");
+    ///     }
+    /// }
+    /// assert_eq!(*map.get(&1).unwrap(), "first");
+    /// assert_eq!(*map.get(&2).unwrap(), "b");
+    /// ```
+    pub fn first_entry(&mut self) -> Option<OccupiedEntry<'_, K, V, N>> {
+        if self.is_empty() {
+            return None;
+        }
+
+        let node_idx = self.bst.min_idx;
+        Some(OccupiedEntry {
+            node_idx,
+            table: self,
+        })
+    }
+
+    /// Returns the last entry in the map for in-place manipulation.
+    /// The key of this entry is the maximum key in the map.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map = SgMap::<_, _, 10>::new();
+    /// map.insert(1, "a");
+    /// map.insert(2, "b");
+    /// if let Some(mut entry) = map.last_entry() {
+    ///     if *entry.key() > 0 {
+    ///         entry.insert("last");
+    ///     }
+    /// }
+    /// assert_eq!(*map.get(&1).unwrap(), "a");
+    /// assert_eq!(*map.get(&2).unwrap(), "last");
+    /// ```
+    pub fn last_entry(&mut self) -> Option<OccupiedEntry<'_, K, V, N>> {
+        if self.is_empty() {
+            return None;
+        }
+
+        let node_idx = self.bst.max_idx;
+        Some(OccupiedEntry {
+            node_idx,
+            table: self,
+        })
+    }
 }
 
 // Convenience Traits --------------------------------------------------------------------------------------------------
diff --git a/src/map_types.rs b/src/map_types.rs
index 11742e1..90ebcaa 100644
--- a/src/map_types.rs
+++ b/src/map_types.rs
@@ -1,3 +1,4 @@
+use crate::entry::{OccupiedEntry, VacantEntry};
 use crate::map::SgMap;
 use crate::tree::{IntoIter as TreeIntoIter, Iter as TreeIter, IterMut as TreeIterMut};
 
@@ -216,3 +217,155 @@ impl<'a, K: Ord + Default, V: Default, const N: usize> ExactSizeIterator
         self.inner.len()
     }
 }
+
+// Entry API -----------------------------------------------------------------------------------------------------
+
+/// A view into a single entry in a map, which may either be vacant or occupied.
+///
+/// This `enum` is constructed from the [`entry`] method on [`SgMap`].
+pub enum Entry<'a, K: Ord + Default, V: Default, const N: usize> {
+    /// A vacant entry.
+    Vacant(VacantEntry<'a, K, V, N>),
+    /// An occupied entry.
+    Occupied(OccupiedEntry<'a, K, V, N>),
+}
+
+use Entry::*;
+
+impl<'a, K: Ord + Default, V: Default, const N: usize> Entry<'a, K, V, N> {
+    /// Ensures a value is in the entry by inserting the default if empty, and returns a mutable
+    /// reference to the value in the entry.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map = SgMap::<&str, usize, 10>::new();
+    /// map.entry("poneyland").or_insert(12);
+    ///
+    /// assert_eq!(map["poneyland"], 12);
+    /// ```
+    pub fn or_insert(self, default: V) -> &'a mut V {
+        match self {
+            Occupied(entry) => entry.into_mut(),
+            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.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map = SgMap::<&str, usize, 10>::new();
+    /// let x = 42;
+    /// map.entry("poneyland").or_insert_with(|| x);
+    ///
+    /// assert_eq!(map["poneyland"], 42);
+    /// ```
+    pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> &'a mut V {
+        match self {
+            Occupied(entry) => entry.into_mut(),
+            Vacant(entry) => entry.insert(default()),
+        }
+    }
+
+    /// Ensures a value is in the entry by inserting, if empty, the result of the default function.
+    /// This method allows for generating key-derived values for insertion by providing the default
+    /// function a reference to the key that was moved during the `.entry(key)` method call.
+    ///
+    /// The reference to the moved key is provided so that cloning or copying the key is
+    /// unnecessary, unlike with `.or_insert_with(|| ... )`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map = SgMap::<&str, usize, 10>::new();
+    ///
+    /// map.entry("poneyland").or_insert_with_key(|key| key.chars().count());
+    ///
+    /// assert_eq!(map["poneyland"], 9);
+    /// ```
+    pub fn or_insert_with_key<F: FnOnce(&K) -> V>(self, default: F) -> &'a mut V {
+        match self {
+            Occupied(entry) => entry.into_mut(),
+            Vacant(entry) => {
+                let value = default(entry.key());
+                entry.insert(value)
+            }
+        }
+    }
+
+    /// Returns a reference to this entry's key.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map = SgMap::<&str, usize, 10>::new();
+    /// assert_eq!(map.entry("poneyland").key(), &"poneyland");
+    /// ```
+    pub fn key(&self) -> &K {
+        match self {
+            Occupied(entry) => entry.key(),
+            Vacant(entry) => entry.key(),
+        }
+    }
+
+    /// Provides in-place mutable access to an occupied entry before any
+    /// potential inserts into the map.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map = SgMap::<&str, usize, 10>::new();
+    ///
+    /// map.entry("poneyland")
+    ///    .and_modify(|e| { *e += 1 })
+    ///    .or_insert(42);
+    /// assert_eq!(map["poneyland"], 42);
+    ///
+    /// map.entry("poneyland")
+    ///    .and_modify(|e| { *e += 1 })
+    ///    .or_insert(42);
+    /// assert_eq!(map["poneyland"], 43);
+    /// ```
+    pub fn and_modify<F: FnOnce(&mut V)>(self, f: F) -> Self {
+        match self {
+            Occupied(mut entry) => {
+                f(entry.get_mut());
+                Occupied(entry)
+            }
+            Vacant(entry) => Vacant(entry),
+        }
+    }
+
+    /// Ensures a value is in the entry by inserting the default value if empty,
+    /// and returns a mutable reference to the value in the entry.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use scapegoat::SgMap;
+    ///
+    /// let mut map = SgMap::<&str, Option<usize>, 10>::new();
+    /// map.entry("poneyland").or_default();
+    ///
+    /// assert_eq!(map["poneyland"], None);
+    /// ```
+    pub fn or_default(self) -> &'a mut V {
+        match self {
+            Occupied(entry) => entry.into_mut(),
+            Vacant(entry) => entry.insert(Default::default()),
+        }
+    }
+}
diff --git a/src/tree/mod.rs b/src/tree/mod.rs
index 6c1ae4e..989e829 100644
--- a/src/tree/mod.rs
+++ b/src/tree/mod.rs
@@ -1,4 +1,5 @@
 mod node_dispatch;
+pub use node_dispatch::SmallNode;
 
 #[cfg(test)]
 mod test;
@@ -7,7 +8,7 @@ mod arena;
 #[cfg(fuzzing)]
 pub use arena::Arena;
 
-mod node;
+pub(super) mod node;
 #[cfg(fuzzing)]
 pub use node::{Node, NodeGetHelper, NodeRebuildHelper};
 
diff --git a/src/tree/tree.rs b/src/tree/tree.rs
index 1fa6fd5..1e7bb24 100644
--- a/src/tree/tree.rs
+++ b/src/tree/tree.rs
@@ -32,8 +32,8 @@ pub struct SgTree<K: Default, V: Default, const N: usize> {
     pub(crate) opt_root_idx: Option<usize>,
 
     // Query cache
-    max_idx: usize,
-    min_idx: usize,
+    pub(crate) max_idx: usize,
+    pub(crate) min_idx: usize,
     curr_size: usize,
 
     // Balance control
@@ -167,7 +167,7 @@ impl<K: Ord + Default, V: Default, const N: usize> SgTree<K, V, N> {
     where
         K: Ord,
     {
-        self.priv_balancing_insert::<Idx>(key, val)
+        self.priv_balancing_insert::<Idx>(key, val).0
     }
 
     /// Insert a key-value pair into the tree.
@@ -181,7 +181,7 @@ impl<K: Ord + Default, V: Default, const N: usize> SgTree<K, V, N> {
     {
         // Replace current slot or safely fill a new one
         match self.contains_key(&key) || (self.capacity() > self.len()) {
-            true => Ok(self.priv_balancing_insert::<Idx>(key, val)),
+            true => Ok(self.priv_balancing_insert::<Idx>(key, val).0),
             false => Err(SgError::StackCapacityExceeded),
         }
     }
@@ -533,7 +533,7 @@ impl<K: Ord + Default, V: Default, const N: usize> SgTree<K, V, N> {
 
     // Iterative search. If key found, returns node idx, parent idx, and a bool indicating if node is right child
     // `opt_path` is only populated if `Some` and key is found.
-    fn priv_get<Q, U: SmallUnsigned + Default + Copy>(
+    pub(crate) fn priv_get<Q, U: SmallUnsigned + Default + Copy>(
         &self,
         mut opt_path: Option<&mut ArrayVec<[U; N]>>,
         key: &Q,
@@ -603,13 +603,15 @@ impl<K: Ord + Default, V: Default, const N: usize> SgTree<K, V, N> {
 
     // Sorted insert of node into the tree (outer).
     // Re-balances the tree if necessary.
-    fn priv_balancing_insert<U: Default + Copy + Ord + Sub + SmallUnsigned>(
+    //
+    // Returns the old value, if any, and the index of the new node in the arena.
+    pub(crate) fn priv_balancing_insert<U: Default + Copy + Ord + Sub + SmallUnsigned>(
         &mut self,
         key: K,
         val: V,
-    ) -> Option<V> {
+    ) -> (Option<V>, usize) {
         let mut path: ArrayVec<[U; N]> = Arena::<K, V, U, N>::new_idx_vec();
-        let opt_val = self.priv_insert(&mut path, key, val);
+        let (opt_val, ngh) = self.priv_insert(&mut path, key, val);
 
         #[cfg(feature = "fast_rebalance")]
         {
@@ -627,18 +629,21 @@ impl<K: Ord + Default, V: Default, const N: usize> SgTree<K, V, N> {
             }
         }
 
-        opt_val
+        let new_node_idx = ngh.node_idx().expect("Must be `Some`");
+        (opt_val, new_node_idx)
     }
 
     // Sorted insert of node into the tree (inner).
     // Maintains a traversal path to avoid nodes needing to maintain a parent index.
+    // Returns a tuple of the old value, if any, and the `NodeGetHelper` of the new node.
+    //
     // If a node with the same key existed, overwrites both that nodes key and value with the new one's and returns the old value.
     fn priv_insert<U: SmallUnsigned + Default + Copy>(
         &mut self,
         path: &mut ArrayVec<[U; N]>,
         key: K,
         val: V,
-    ) -> Option<V> {
+    ) -> (Option<V>, NodeGetHelper<U>) {
         match self.opt_root_idx {
             // Sorted insert
             Some(idx) => {
@@ -688,7 +693,7 @@ impl<K: Ord + Default, V: Default, const N: usize> SgTree<K, V, N> {
                             curr_node.set_val(val);
 
                             // Key/val updated "in-place": no need to update `curr_node`'s parent or children
-                            ngh = NodeGetHelper::new(None, None, false);
+                            ngh = NodeGetHelper::new(Some(curr_idx), None, false);
                             break;
                         }
                         Ordering::Greater => {
@@ -736,7 +741,7 @@ impl<K: Ord + Default, V: Default, const N: usize> SgTree<K, V, N> {
                 }
 
                 // Return old value if overwritten
-                opt_val
+                (opt_val, ngh)
             }
 
             // Empty tree
@@ -750,7 +755,8 @@ impl<K: Ord + Default, V: Default, const N: usize> SgTree<K, V, N> {
                 self.max_idx = root_idx;
                 self.min_idx = root_idx;
 
-                None
+                let ngh = NodeGetHelper::new(Some(root_idx), None, false);
+                (None, ngh)
             }
         }
     }


Test patch: