Update location to conform to spec

1 files changed, 23 insertions, 7 deletions

diff --git a/src/lib.rs b/src/lib.rs
index dc29621..4a12e5b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -663,22 +663,27 @@ impl Configuration {
 	}
 
 	/// Returns the index of the first key which is greater than `key` + "."
-	fn first_subkey_index(&self, key: &str) -> usize {
+	fn subkey_start_idx(&self, key: &str) -> usize {
 		let key_dot = format!("{key}.");
 		self.binary_search_for(&key_dot)
 			.expect_err("items should not contain a key ending in .")
 	}
 
+	/// Returns the index of the first key which is greater than `key` + "." + `x` for all strings `x`.
+	fn subkey_end_idx(&self, key: &str) -> usize {
+		// NB: / is the next ASCII character after .
+		let key_slash = format!("{key}/");
+		self.binary_search_for(&key_slash).unwrap_or_else(|x| x)
+	}
+
 	/// Extract a section out of a configuration.
 	///
 	/// More specifically, this will give you the configuration consisting of all
 	/// keys starting with `key.` in `self`, together with their values.
 	#[must_use]
 	pub fn section(&self, key: &str) -> Configuration {
-		let start_idx = self.first_subkey_index(key);
-		// NB: / is the next ASCII character after .
-		let key_slash = format!("{key}/");
-		let end_idx = self.binary_search_for(&key_slash).unwrap_or_else(|i| i);
+		let start_idx = self.subkey_start_idx(key);
+		let end_idx = self.subkey_end_idx(key);
 		Configuration {
 			items: self.items[start_idx..end_idx].to_owned(),
 		}
@@ -721,7 +726,7 @@ impl Configuration {
 	/// in future versions without notice.
 	pub fn subkeys(&self, key: &str) -> impl '_ + Iterator<Item = &str> {
 		let key_dot = format!("{key}.");
-		let start_idx = self.first_subkey_index(key);
+		let start_idx = self.subkey_start_idx(key);
 		(start_idx..self.items.len()).map_while(move |i| {
 			let this_key = &self.items[i].0;
 			let suffix = this_key.strip_prefix(&key_dot)?;
@@ -743,7 +748,18 @@ impl Configuration {
 	/// Get location in the configuration file where `key` is defined, if any.
 	#[must_use]
 	pub fn location(&self, key: &str) -> Option<Location> {
-		Some(self.get_val(key)?.defined_at.clone())
+		if let Some(val) = self.get_val(key) {
+			Some(val.defined_at.clone())
+		} else {
+			// Check if `key` has any defined subkeys
+			let start_idx = self.subkey_start_idx(key);
+			let (subkey, subval) = self.items.get(start_idx)?;
+			if subkey.starts_with(key) && subkey[key.len()..].starts_with('.') {
+				Some(subval.defined_at.clone())
+			} else {
+				None
+			}
+		}
 	}
 
 	/// Returns `true` if `key` is defined in this configuration.