Fix inline insert spill and Miri UB

[ihciah]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @ihciah, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request primarily focuses on enhancing the robustness and correctness of the SmallMap implementation by addressing potential Undefined Behavior (UB) under Miri and optimizing its behavior during insertions. It introduces a new feature for explicit control byte initialization to prevent uninitialized SIMD reads and refines the insertion logic to avoid unnecessary heap allocations when updating existing entries in a full inline map. Additionally, internal data access mechanisms were simplified by removing a Bucket abstraction, and serialization logic was updated to reflect generic parameter changes.

Highlights

• Miri UB Fix & Control Byte Initialization: A new init-ctrl feature and Miri-specific logic have been introduced to zero-fill control bytes and add padding during AlignedGroups initialization. This prevents uninitialized SIMD reads, addressing potential Undefined Behavior (UB) when running with Miri.
• Inline Insert Spill Prevention: The SmallMap::insert logic has been modified to prevent unnecessary heap spills when inserting an existing key into a full inline map. The map now updates the existing entry in-place, preserving its inline state.
• Bucket Abstraction Removal: The internal Bucket struct and associated raw pointer operations have been removed, simplifying data access within RawInline by directly using indices into the data array.
• Serialization Parameter Update: The serde implementations for SmallMap were updated to align with changes in its generic parameters, specifically to support distinct hashers for keys and values.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• Cargo.toml
  • Added init-ctrl feature flag to enable zero-filling of control bytes and padding.
• README.md
  • Documented the new init-ctrl feature, explaining its purpose for Miri/UB checking and performance implications.
• src/inline.rs
  • Removed Bucket and raw pointer related imports (transmute, NonNull).
  • Refactored AlignedGroups initialization to use a safe new() constructor, with conditional zero-filling for init-ctrl and Miri.
  • Replaced ctrl method with ctrl_read and ctrl_write using offset_of! for safer pointer access.
  • Updated find, insert_in_slot, and remove methods to operate directly with indices instead of Bucket objects.
  • Removed bucket and bucket_index helper methods.
• src/lib.rs
  • Implemented a check in SmallMap::insert to update existing keys in-place when the inline map is full, preventing a spill to the heap.
  • Reduced fuzzing test iterations for Miri builds to improve performance.
  • Added a new test case (insert_existing_when_full_keeps_inline) to verify the inline insert spill fix.
• src/raw/util.rs
  • Removed Bucket struct definition and related utility functions (invalid_mut, IS_ZERO_SIZED).
• src/serde.rs
  • Updated Serialize and Deserialize implementations for SmallMap to include SH, SI, and LINEAR_THRESHOLD generic parameters.
  • Changed SmallMap::with_capacity_and_hasher to SmallMap::with_capacity_and_hashers in deserialization to accommodate two hashers.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[Copilot]

Pull request overview

This PR fixes two issues: (1) a bug where inserting an existing key when inline storage is full would incorrectly spill to heap instead of updating in-place, and (2) Miri undefined behavior caused by SIMD operations reading uninitialized control bytes beyond the valid array bounds.

Changes:

• Added early check in SmallMap::insert to update existing keys in-place when inline storage is full, preventing unnecessary spill to heap
• Removed Bucket abstraction and replaced bucket-based operations with direct index-based access
• Introduced init-ctrl feature flag and conditional padding to AlignedGroups to eliminate UB from SIMD reads
• Updated generic parameters throughout to accommodate the new LINEAR_THRESHOLD const generic and separate inline/heap hashers

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
src/lib.rs	Adds check for existing key before spilling to heap when inline storage is full; reduces fuzzing iterations under Miri; adds test for the fix
src/inline.rs	Removes Bucket abstraction in favor of direct indexing; adds AlignedGroups::new() with conditional initialization; refactors ctrl byte access with separate read/write methods; updates all operations to use indices instead of Buckets
src/raw/util.rs	Removes Bucket struct and related functions (from_base_index, to_base_index, etc.); removes IS_ZERO_SIZED constant and invalid_mut pointer functions
src/serde.rs	Updates Serialize and Deserialize implementations to include new generic parameters (SH, SI, LINEAR_THRESHOLD) and uses with_capacity_and_hashers
README.md	Documents the new init-ctrl feature flag and its trade-offs
Cargo.toml	Adds init-ctrl feature flag definition

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The function documentation states it "returns its raw bucket" but the function signature shows it returns nothing (void). This documentation should be updated to remove the mention of returning a bucket, as the Bucket type has been removed in this refactoring.

Suggested change

	/// Inserts a new element into the table in the given slot, and returns its
	/// raw bucket.
	/// Inserts a new element into the table in the given slot.

[Copilot]

This safety comment references "RawTableInner::set_ctrl_h2" which appears to be copied from hashbrown's codebase but doesn't apply to this implementation. The comment should either be removed or updated to describe the actual safety requirements for this function (e.g., that slot.index must be less than N and the slot must not already be occupied).

Suggested change

	// SAFETY: The caller must uphold the safety rules for the [`RawTableInner::set_ctrl_h2`]
	// SAFETY: The caller must ensure that `slot.index` is within bounds of the
	// control bytes and `data` arrays, that the slot is currently vacant
	// (i.e., does not contain an initialized element), and that the table
	// has remaining capacity so that incrementing `self.len` is consistent.

[Copilot]

Grammar error: "an can be inserted" should be "an element can be inserted" or simply "a value can be inserted".

[gemini-code-assist]

Code Review

This pull request introduces significant improvements by fixing a logic bug and addressing undefined behavior found by Miri. The fix for the inline insert spill prevents unnecessary heap allocations when updating an existing key in a full inline map, which is a great correctness improvement. The changes to address Miri UB are comprehensive, including adding padding to prevent out-of-bounds SIMD reads, introducing the init-ctrl feature for safe initialization, and refactoring pointer handling to fix aliasing issues. The removal of the Bucket abstraction is a particularly noteworthy improvement, as it greatly simplifies the code, reduces unsafe blocks, and enhances overall soundness. The code quality is excellent, and the addition of a specific test case for the spill bug is appreciated. Overall, this is a very solid and well-executed pull request.

[hsqStephenZhang]

to understand it fully, the original code is semantically correct because of the mask, right?

[hsqStephenZhang]

maybe add miri to CI?