From 0532ea10c27f3898ca0a6517e6d35a1925d0efba Mon Sep 17 00:00:00 2001
From: Josh Stone <cuviper@gmail.com>
Date: Thu, 11 Jun 2020 12:50:27 -0700
Subject: [PATCH] Implement direct usize indexing

```rust
where
    IndexMap<K, V, S>: IndexMut<usize, Output = V>,
    IndexSet<T, S>: Index<usize, Output = T>,
```

This allows `map[i]` and `set[i]` indexing to access values directly,
panicking if the index is out of bounds, similar to slices.

On maps, this somewhat overlaps with `Index<&Q> + IndexMut<&Q>` where
`Q: Equivalent<K>`. The reference makes this indexing unambiguous, but
it could be confusing to users if the key type is also an integer.
---
 src/map.rs     | 20 ++++++++++++++++++++
 src/set.rs     | 11 ++++++++++-
 tests/quick.rs | 22 +++++++++++++++++++++-
 3 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/src/map.rs b/src/map.rs
index 1f026fd6..45778d54 100644
--- a/src/map.rs
+++ b/src/map.rs
@@ -980,6 +980,26 @@ where
     }
 }
 
+impl<K, V, S> Index<usize> for IndexMap<K, V, S> {
+    type Output = V;
+
+    /// ***Panics*** if `index` is out of bounds.
+    fn index(&self, index: usize) -> &V {
+        self.get_index(index).expect("IndexMap: index out of bounds").1
+    }
+}
+
+/// Mutable indexing allows changing / updating indexed values
+/// that are already present.
+///
+/// You can **not** insert new values with index syntax, use `.insert()`.
+impl<K, V, S> IndexMut<usize> for IndexMap<K, V, S> {
+    /// ***Panics*** if `index` is out of bounds.
+    fn index_mut(&mut self, index: usize) -> &mut V {
+        self.get_index_mut(index).expect("IndexMap: index out of bounds").1
+    }
+}
+
 impl<K, V, S> FromIterator<(K, V)> for IndexMap<K, V, S>
 where
     K: Hash + Eq,
diff --git a/src/set.rs b/src/set.rs
index 1b90fb01..d9b5c4f5 100644
--- a/src/set.rs
+++ b/src/set.rs
@@ -14,7 +14,7 @@ use std::fmt;
 use std::hash::{BuildHasher, Hash};
 use std::iter::{Chain, FromIterator};
 use std::ops::RangeFull;
-use std::ops::{BitAnd, BitOr, BitXor, Sub};
+use std::ops::{BitAnd, BitOr, BitXor, Index, Sub};
 use std::slice;
 use std::vec;
 
@@ -594,6 +594,15 @@ impl<T, S> IndexSet<T, S> {
     }
 }
 
+impl<T, S> Index<usize> for IndexSet<T, S> {
+    type Output = T;
+
+    /// ***Panics*** if `index` is out of bounds.
+    fn index(&self, index: usize) -> &T {
+        self.get_index(index).expect("IndexSet: index out of bounds")
+    }
+}
+
 /// An owning iterator over the items of a `IndexSet`.
 ///
 /// This `struct` is created by the [`into_iter`] method on [`IndexSet`]
diff --git a/tests/quick.rs b/tests/quick.rs
index 1399d0a8..f0b17d03 100644
--- a/tests/quick.rs
+++ b/tests/quick.rs
@@ -6,7 +6,7 @@ extern crate rand;
 
 extern crate fnv;
 
-use indexmap::IndexMap;
+use indexmap::{IndexMap, IndexSet};
 use itertools::Itertools;
 
 use quickcheck::Arbitrary;
@@ -142,6 +142,26 @@ quickcheck! {
             elements.iter().all(|k| map.get(k).is_some())
     }
 
+    fn indexing(insert: Vec<u8>) -> bool {
+        let mut map: IndexMap<_, _> = insert.into_iter().map(|x| (x, x)).collect();
+        let set: IndexSet<_> = map.keys().cloned().collect();
+        assert_eq!(map.len(), set.len());
+
+        for (i, &key) in set.iter().enumerate() {
+            assert_eq!(map.get_index(i), Some((&key, &key)));
+            assert_eq!(set.get_index(i), Some(&key));
+            assert_eq!(map[i], key);
+            assert_eq!(set[i], key);
+
+            *map.get_index_mut(i).unwrap().1 >>= 1;
+            map[i] <<= 1;
+        }
+
+        set.iter().enumerate().all(|(i, &key)| {
+            let value = key & !1;
+            map[&key] == value && map[i] == value
+        })
+    }
 }
 
 use Op::*;