Fix UB in index_map::insert

[sgued]

Previously, index_map was broken in two ways:

Firstly, indices are expected to only be 16 bits but the N parameter could be any usize, including > u16::MAX + 1. This is fixed with a static assertion.

Secondly, the indexmap assumed that an index of 0xFFFF and an hash of 0xFFFF would never happen, but this could be wrong, and lead to the safety requirement of NonZeroU32::new_unchecked to be violated.

This PR fixes this by removing the option and the automatic niche value optimization of Pos, and re-implementing it manually.

To make sure that there is no confusion in the case where the index is 0xFFFF and so is the hash, and additional bit of information is used to determine whether the 0xFFFF hash and indices represent a None value or whether it represent Some({hash_value: 0xFFFF, index: 0xFFFF})

If the map is full of 0x10000 items, then we know that all pos in the indices are valid buffer are valid. Otherwise, we know that index = 0xFFFF is simply not possible, therefore we know that a Pos with such an index is None.

See #672 for the original report.

[sgued]

Anybody got ideas on how to prevent the stack overflow in the test? Box::new doesn't work. Looks like we'd have to add opt-level = 1.

[sgued]

I got the test working by allocating/initializing manually with some unsafe.

[zeenix]

I won't have time for reviewing the code itself until at least tomorrow but I trust you with that. Nitpicks OTOH are easy and quick. :)

[zeenix]

I was not able to find any issue with this PR but I wanted to be sure so I asked an LLM to see if it can find problems. I was hoping it doesn't or comes up with bogus issues but it came up with 2 big issues here that AFAICT are geniune and serious. I know you hate LLMs but these seems like serious issues to ignore so I have to share:

---

1. find passes the wrong length — re-introduces the same UB (blocker)

src/index_map.rs:240:

let pos = self.indices[probe].as_valid::<N>(self.indices.len())?;
//                                          ^^^^^^^^^^^^^^^^^^

as_valid's own doc says "The argument is the length of the entries buffer", but indices.len() is always N. So buffer_max collapses to just N == MAX_SIZE, and for N == 0x10000 every empty slot is decoded as a valid Pos { hash: 0xFFFF, index: 0xFFFF }, regardless of how many entries the map actually holds. insert, reinsert_all etc. all correctly pass self.entries.len(); find is the only one that doesn't.

In debug this trips the assert three lines below:

panicked at src/index_map.rs:247:13: assertion failed: i < self.entries.len()

In release that line is unsafe { self.entries.get_unchecked(0xFFFF) } on a map holding one entry — a read of uninitialised memory. With a key type whose zeroed bytes compare equal, get returns an entry that was never inserted:

get(&KC(0)) on a map containing only KC(1) -> Some(0)

This is genuine UB, not merely a wrong answer: the same test returned None until I added an unrelated println! nearby, after which it returned Some(0). get_mut/index_mut would hand out &mut V into that uninit memory, and remove would call swap_remove_unchecked(0xFFFF) on a 1-element Vec.

(A completely empty map is safe only by accident — IndexMap::find has an if self.is_empty() { return None } guard, so you need at least one entry to reach this.)

Fix:

let pos = self.indices[probe].as_valid::<N>(self.entries.len())?;

---

2. continue inside probe_loop! causes an infinite loop (blocker)

src/index_map.rs:350, in remove_found:

probe_loop!(probe < self.indices.len(), {
    if probe == old_probe {
        continue;
    }
    ...
});

probe_loop! expands to loop { if probe < len { $body probe += 1; } else { probe = 0; } }, so a continue in the body skips the probe += 1 and re-enters with the same probe — it never terminates.

This is reachable with tiny maps; no N == 0x10000 required. Three keys colliding into bucket 0 of an N = 4 map, then removing the first, hangs forever (I killed it at 60s):

let mut map: IndexMap<K16, u16, DHB, 4> = IndexMap::default();
map.insert(K16(0), 0).unwrap();
map.insert(K16(4), 4).unwrap();
map.insert(K16(8), 8).unwrap();
map.remove(&K16(0));   // never returns

The insert_overflow test doesn't hit it because that particular removal breaks out on the very first probe.

The guard itself is needed (in the full-map case indices[old_probe] was just set to Pos::none() and would decode as a false "found it"), so just restructure rather than continue:

if probe != old_probe {
    if let Some(pos) = self.indices[probe].as_valid_mut::<N>(old_len) {
        ...
    }
}

---

I applied both fixes locally: the three repro tests pass, insert_overflow still passes, and the full suite with all features is green.

---

[sgued]

Thanks

Looks like you didn't send the full context? What are the KC and K16 structs in the example code?

I was able to reproduce both and add tests for them, but your comment mention 3 repro tests, am I missing one?

Rather than removing the continue, I went with making probe_loop! not have any code run after $body so that no surprises can happen with break or continue, avoiding such mistake from appearing in the future too.

[zeenix]

Looks like you didn't send the full context? What are the KC and K16 structs in the example code?

Since you asked for it, this time I send you the full reponse from the LLM:

---

Sorry — those snippets referenced helper types I never actually included. Full context below.

Also, one correction to my earlier comment while I'm here: I described the bad read as hitting "uninitialised memory". That's right for a map built the normal way, but not for my repro, which builds the map with Box::new_zeroed() — so every byte is initialised, and the slot at index 0xFFFF is a well-defined all-zero Bucket. What's wrong there is that the index is past entries.len(), not that the bytes are indeterminate. That distinction matters for the test below, so I've fixed the wording.

The missing helper types

Copies of the ones inside your insert_overflow, hoisted to mod tests scope (yours are function-local, so there's no conflict — happy to hoist them properly instead so there's one definition). No new imports needed; the module's existing ones cover it.

/// A key whose hash is `self.0`, but whose identity also includes `self.1`,
/// so that two distinct keys can share one hash.
#[derive(PartialEq, Eq, Debug)]
struct K16(u16, u16);

impl Hash for K16 {
    fn hash<H: Hasher>(&self, state: &mut H) {
        state.write_u16(self.0);
    }
}

/// A key whose hash is *always* `0xFFFF`.
#[derive(PartialEq, Eq, Debug)]
struct KC(u16);

impl Hash for KC {
    fn hash<H: Hasher>(&self, state: &mut H) {
        state.write_u16(0xFFFF);
    }
}

#[derive(Default)]
struct DummyHasher(u16);

impl Hasher for DummyHasher {
    fn finish(&self) -> u64 {
        self.0 as _
    }
    fn write(&mut self, _bytes: &[u8]) {}
    fn write_u16(&mut self, i: u16) {
        self.0 = i;
    }
}

#[derive(Default)]
struct DummyHasherBuilder;

impl BuildHasher for DummyHasherBuilder {
    type Hasher = DummyHasher;
    fn build_hasher(&self) -> DummyHasher {
        DummyHasher(0)
    }
}

The third test

You're not missing a third bug — you're missing a second test for the as_valid one. I had two because they catch it in different profiles. The first is equivalent to your indices_valid; the second is the one I'd suggest adding.

/// `find` used to decode an empty slot as a valid `Pos { hash: 0xFFFF, index: 0xFFFF }`.
/// Equivalent to `indices_valid`: trips the `debug_assert!` in `find`.
#[test]
#[cfg_attr(miri, ignore)] // too slow
fn find_decoded_empty_slot() {
    let map: Box<mem::MaybeUninit<IndexMap<K16, u16, DummyHasherBuilder, 0x10000>>> =
        Box::new_zeroed();
    let mut map = unsafe { map.assume_init() };

    map.insert(K16(0xFFFF, 0), 1).unwrap();
    // Distinct key, same hash: probes indices[0xFFFF], then wraps to indices[0].
    // That slot is `Pos::none()`, but `find` passed `indices.len()` (always `N`) where
    // `as_valid` wanted `entries.len()`, so the `buffer_max` short-circuit fired on any
    // `N == 0x10000` map, full or not, and the slot decoded as valid.
    assert_eq!(map.get(&K16(0xFFFF, 1)), None);
}

/// Same root cause, but the bogus entry compares *equal*, so `find` returned a hit and
/// `get` handed back a value from a slot past `entries.len()`. Unlike the test above,
/// this one also failed in release, where `debug_assert!` is compiled out.
#[test]
#[cfg_attr(miri, ignore)] // too slow
fn find_returned_stale_entry() {
    let map: Box<mem::MaybeUninit<IndexMap<KC, u32, DummyHasherBuilder, 0x10000>>> =
        Box::new_zeroed();
    let mut map = unsafe { map.assume_init() };

    map.insert(KC(1), 111).unwrap();
    // KC(0) was never inserted. The slot at index 0xFFFF is past `entries.len()`; here it
    // is all-zero (the map was zero-initialised above), which decodes as `KC(0)`, so the
    // key comparison succeeds and `get` returns a value that was never stored.
    assert_eq!(map.get(&KC(0)), None);
}

/// `continue` inside `probe_loop!` skipped the `probe += 1`.
/// Equivalent to your `remove_found_loop`.
#[test]
fn remove_found_hang() {
    let mut map: IndexMap<K16, u16, DummyHasherBuilder, 4> = IndexMap::default();
    map.insert(K16(0, 0), 0).unwrap();
    map.insert(K16(4, 0), 4).unwrap();
    map.insert(K16(8, 0), 8).unwrap();
    assert_eq!(map.remove(&K16(0, 0)), Some(0)); // used to never return
}

Why the second one earns its place: indices_valid fails only via the debug_assert!(i < self.entries.len()) in find. That assert sits before both the probe-distance give-up and the hash comparison, so it's the first thing to fire on an empty slot — but with debug_assertions off there's nothing left to panic. The out-of-range get_unchecked is still there; indices_valid just never reaches it, because entry_hash == hash short-circuits the && (the empty slot decodes to hash 0xFFFF, and 10 doesn't hash to that).

Reverting only the as_valid line on your branch:

=== indices_valid, debug ===
test index_map::tests::indices_valid ... FAILED
panicked at src/index_map.rs:253:13: assertion failed: i < self.entries.len()

=== indices_valid, release ===
test index_map::tests::indices_valid ... ok        <-- passes

Same revert, find_returned_stale_entry in release:

assertion `left == right` failed
  left: Some(0)
 right: None

Some(0) for a key that was never inserted. I checked it isn't flaky — 10 solo runs, 3 whole-binary runs, and a fat-LTO/1-CGU build all give the identical failure. That stability comes from the zero-initialised buffer making the read deterministic; it's an empirical observation rather than a guarantee, since the access is UB by get_unchecked's contract either way.

On the probe_loop! change

Agreed, that's the better fix — making the macro robust beats remembering not to use continue. It looks correct to me:

• $probe_var is assigned once at the top of the iteration and not touched again before $body, so its value throughout the body is identical to the old expansion. insert_phase_2's return probe and backward_shift_after_removal's last_probe = probe both read it from inside the body and are unaffected; no call site reads the probe variable after the loop exits.
• None of the six bodies assigns to the probe variable, which the new expansion would discard.
• next_probe peaks at $len + 1 — backward_shift_after_removal can enter with probe == N — before the else resets it to 0. No overflow.

I also confirmed your remove_found_loop genuinely hangs against the old macro rather than just failing, so it's a real regression test.

Suite is green debug and release on the CI feature set. Clippy is clean on the CI clippy feature set too — though if you add mpmc_large to that job you'll surface three pre-existing warnings in src/mpmc.rs (one unnecessary_cast at :190, two useless_conversion at :333 and :377). Unrelated to this PR, but since you're already touching that invocation, might be worth folding in or leaving a note.

Two tiny things

• The doc comments on the two tests added in the last two commits both say "in the case N = 0x10000", but remove_found_loop uses N = 4 — and that's worth keeping, since the hang isn't specific to large N at all. Just the comment that needs rewording.
• A stray trailing comma at src/index_map.rs:279, from a82b714: debug_assert!(i < self.entries.len(),);.

A few more take-them-or-leave-them observations, none of them problems with the fix itself:

• The reasoning that makes assume_not_full safe is subtle and currently undocumented. A slot can never hold a valid Pos with inner == 0, because that needs index == 0xFFFF, which needs entries.len() == 0x10000 — unreachable in both places it's called (insert_phase_2 only ever writes forward and never re-reads; backward_shift_after_removal runs after swap_remove). Worth a comment, or the next change to insert breaks it silently.
• truncate's manual index-clearing loop is dead work: its guard (N > entries.len()) is the same predicate as reinsert_all's (entries.len() < indices.len(), and indices.len() == N), and reinsert_all starts by overwriting all of indices anyway. The suite stays green with the loop deleted. Pre-existing, but this PR touches those lines.
• #[cfg_attr(miri, ignore)] on the 0x10000 tests means Miri never covers this encoding. Unlike insert_overflow, the two find tests above do a single insert and a single lookup, so they may be cheap enough to leave enabled — I couldn't check, as I had no Miri available to time them.

[sgued]

Thanks,

I added the third test which does indeed cover a distinct code path.

Regarding AI use, I don't complain on how you find bugs as longs as they're legit (as you can see from the bug reports coming from LLMs I fixed), I still find it somewhat rude to paste multiple paragraphs of LLM output for me to sort through, having to find what is useful in them and separate that from the rest.

I'd rather have had 2 comments even if they're pretty raw like:

find uses self.indices.len(), it should probably use entries instead.

and

Continue in a probe_loop! macro will lead to infinite looping.

That would have conveyed pretty much as much useful info as the first message (and you could still have pointed out they were found by LLM if you wanted)

[zeenix]

I'd rather have had 2 comments even if they're pretty raw like:

That's quite subjective, sorry. I agree with you that as long as the issues are geniun, then there shouldn't be an issue using LLMs. I'm happy to bear the costs involved here.

[zeenix]

I asked LLM to be super concise this time (I'll have to remember to do that each time):

indices_valid2 has the keys the wrong way round, so it doesn't cover the extra path yet.

The stale slot at index 0xFFFF is zeroed, so it decodes as ConstantHash(0) — the query has to be ConstantHash(0) for the bogus key comparison to succeed. As written (insert 0, look up 1) the comparison fails, so it only trips the debug_assert!, same as indices_valid. With the as_valid fix reverted it passes in release.

Swapping them makes it fail in both profiles:

map.insert(ConstantHash(1), 1).unwrap();
// Distinct key, same hash
assert!(map.get(&ConstantHash(0)).is_none());

[zeenix]

LLM caught something still:

---

The indices_valid2 fix looks right — it now catches the bug in release, as does index_none_valid. I verified both by reverting the as_valid line.

One thing on the new #![deny(clippy::undocumented_unsafe_blocks)]: it breaks cargo clippy on older toolchains. VacantEntry::insert (src/index_map.rs:768, pre-existing, unchanged here) has its SAFETY comment inside the unsafe block.

Clippy only learned to accept that placement somewhere after 1.94. CI is on 1.97 so it's green, but on 1.87 (the crate's MSRV) and on 1.94 it's error: unsafe block missing a safety comment — and since it's deny rather than warn, that's a hard build failure, not a warning.

Moving the comment above the unsafe { fixes it on all three.

---

Also, please don't forget to squash the fixup commits after so we don't end up with them in master.

[sgued]

You consider clippy to be covered by MSRV?

I consider the MSRV to only apply to users downstream, not to what is needed to work on heapless itself.

[zeenix]

You consider clippy to be covered by MSRV?

I consider the MSRV to only apply to users downstream, not to what is needed to work on heapless itself.

It's just a bit inconsistent but if it doesn't/can't cause any issues, sure I don't mind.