From 94c7624ac402eebe21a6490ce1dc621c3808d731 Mon Sep 17 00:00:00 2001
From: Tolik Zinovyev <aszinovyev@gmail.com>
Date: Wed, 8 Jan 2025 09:05:12 -0500
Subject: [PATCH] Check elements are distinct in simple lottery. (#111)

## Content

Found another issue with the simple lottery. We need to check that
elements in the proof are distinct. I guess the simple lottery is not so
simple. :)

## Pre-submit checklist

- Branch
    - [x] Tests are provided (if possible)
    - [x] Commit sequence broadly makes sense
    - [x] Key commits have useful messages
- PR
    - [x] No clippy warnings in the CI
    - [x] Self-reviewed the diff
    - [x] Useful pull request description
    - [x] Reviewer requested
- Documentation
    - [x] Update README file (if relevant)
    - [x] Update documentation website (if relevant)
---
 src/centralized_telescope/algorithm.rs | 2 ++
 src/simple_lottery/algorithm.rs        | 4 ++++
 src/utils/misc.rs                      | 8 ++++++++
 src/utils/mod.rs                       | 2 ++
 4 files changed, 16 insertions(+)
 create mode 100644 src/utils/misc.rs

diff --git a/src/centralized_telescope/algorithm.rs b/src/centralized_telescope/algorithm.rs
index 499d8e83..ca3600b1 100644
--- a/src/centralized_telescope/algorithm.rs
+++ b/src/centralized_telescope/algorithm.rs
@@ -12,6 +12,8 @@ use blake2::{Blake2s256, Digest};
 /// Calls up to setup.max_retries times the prove_index function and returns an empty
 /// proof if no suitable candidate is found.
 pub(super) fn prove(setup: &Setup, prover_set: &[Element]) -> Option<Proof> {
+    debug_assert!(crate::utils::misc::check_distinct(prover_set));
+
     // Run prove_index up to max_retries times
     (0..setup.max_retries).find_map(|retry_counter| prove_index(setup, prover_set, retry_counter).1)
 }
diff --git a/src/simple_lottery/algorithm.rs b/src/simple_lottery/algorithm.rs
index cbe5c95f..cf4335fb 100644
--- a/src/simple_lottery/algorithm.rs
+++ b/src/simple_lottery/algorithm.rs
@@ -8,12 +8,15 @@ use crate::utils::types::Element;
 use blake2::{Blake2s256, Digest};
 
 pub(super) fn prove(setup: &Setup, prover_set: &[Element]) -> Option<Proof> {
+    debug_assert!(crate::utils::misc::check_distinct(prover_set));
+
     let mut element_sequence = Vec::with_capacity(setup.proof_size as usize);
     for &element in prover_set {
         if lottery_hash(setup, element) {
             element_sequence.push(element);
         }
         if prover_set.len() as u64 >= setup.proof_size {
+            element_sequence.sort_unstable();
             return Some(Proof { element_sequence });
         }
     }
@@ -23,6 +26,7 @@ pub(super) fn prove(setup: &Setup, prover_set: &[Element]) -> Option<Proof> {
 
 pub(super) fn verify(setup: &Setup, proof: &Proof) -> bool {
     (proof.element_sequence.len() as u64 == setup.proof_size)
+        && proof.element_sequence.is_sorted_by(|a, b| a < b)
         && proof
             .element_sequence
             .iter()
diff --git a/src/utils/misc.rs b/src/utils/misc.rs
new file mode 100644
index 00000000..16148a81
--- /dev/null
+++ b/src/utils/misc.rs
@@ -0,0 +1,8 @@
+use super::types::Element;
+
+/// Returns true iff all elements in the slice are distinct.
+pub(crate) fn check_distinct(elements: &[Element]) -> bool {
+    let mut elements = elements.to_vec();
+    elements.sort_unstable();
+    elements.is_sorted_by(|a, b| a < b)
+}
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 1c13829d..bc01c2bc 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -1,3 +1,5 @@
 pub(crate) mod sample;
 
 pub(crate) mod types;
+
+pub(crate) mod misc;