[FEAT]add is_stable

[deadlywing]

Pull Request

What problem does this PR solve?

Issue Number: Fixed #

Possible side effects?

• Performance:

• Backward compatibility:

[gemini-code-assist]

Summary of Changes

Hello @deadlywing, 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 enhances the sorting capabilities by introducing an is_stable parameter across the sorting function hierarchy. This allows callers to explicitly request a stable sort, which is then handled by leveraging the radix_sort algorithm for secret-shared data, particularly within the Cheetah protocol. The change ensures that elements with identical keys retain their original relative order after sorting, a crucial feature for certain data processing scenarios. Comprehensive tests have been added to validate this new behavior.

Highlights

• New is_stable parameter for sorting functions: Introduced a boolean parameter is_stable to various sorting functions (fallback_sort1d, simple_sort1d, SimpleSort) to control whether the sort operation should be stable.
• Stable Sort Implementation using Radix Sort: When is_stable is set to true, the simple_sort1d function now explicitly dispatches to internal::radix_sort, which is noted to be a stable sorting method, particularly for secret-shared data in the Cheetah protocol.
• Comprehensive Unit Tests for Stable Sorting: Added a new test suite, CheetahStableSortTest, to verify the correctness of stable sorting in both ascending and descending orders, ensuring that elements with equal keys maintain their original relative order. It also includes a test to confirm that unstable sorts do not guarantee this property.

🧠 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.

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. ↩

[gemini-code-assist]

Code Review

This pull request introduces a is_stable flag to the sorting functions, enabling stable sort capabilities. The changes primarily involve propagating this flag through the function call stack and adding new tests to verify the stable sort behavior. My review includes three points: one critical issue in a test file that would break the build, one high-severity issue regarding a missing check that could cause a runtime error, and one medium-severity comment about an outdated code comment.

[gemini-code-assist]

The <random> header has been removed, but it is still used by the LargeNumel test case (lines 192-194) for std::random_device and std::mt19937. Removing this header will cause a compilation failure. Please add it back.

Suggested change

	#include <memory>
	#include <memory>
	#include <random>

[gemini-code-assist]

When is_stable is true, internal::radix_sort is called directly. However, radix_sort might not be supported by all protocols, which could lead to a runtime failure if a protocol without the necessary kernels is used. It's important to add a check to ensure that radix sort is supported before calling it, similar to the checks performed later in this function for the default sort method.

    SPU_ENFORCE(internal::_check_method_require(ctx, RuntimeConfig::SORT_RADIX),
                "Stable sort is requested, but radix sort is not supported by "
                "the current protocol.");
    return internal::radix_sort(ctx, inputs, direction, num_keys, valid_bits);

[gemini-code-assist]

This comment is now misleading. The sort1d function does support stable sorting when the keys are public by using std::stable_sort. Since this fallback_sort1d function is called from simple_sort1d with public keys, stable sorting is a valid path. Please update the comment to accurately reflect the current implementation.

Suggested change

	// currently, general sort1d only supports odd-even sorting network which is
	// an unstable sort method.
	// For public inputs, `sort1d` can perform a stable sort. For secret inputs,
	// it falls back to an unstable odd-even sorting network.