From 6d1eb90946d7805e84e4ae501e3fef2a2eb81652 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Grosso?= <81065257+atgrosso@users.noreply.github.com> Date: Fri, 11 Oct 2024 16:04:32 -0300 Subject: [PATCH] Add array sorting iterator (#102) --- stwo_cairo_verifier/src/lib.cairo | 1 + stwo_cairo_verifier/src/sort.cairo | 196 +++++++++++++++++++++++++++++ 2 files changed, 197 insertions(+) create mode 100644 stwo_cairo_verifier/src/sort.cairo diff --git a/stwo_cairo_verifier/src/lib.cairo b/stwo_cairo_verifier/src/lib.cairo index 8331a87c..9dbe333c 100644 --- a/stwo_cairo_verifier/src/lib.cairo +++ b/stwo_cairo_verifier/src/lib.cairo @@ -5,6 +5,7 @@ mod poly; mod utils; mod vcs; mod fri; +mod sort; mod queries; pub use fields::{BaseField, SecureField}; diff --git a/stwo_cairo_verifier/src/sort.cairo b/stwo_cairo_verifier/src/sort.cairo new file mode 100644 index 00000000..0e873fc5 --- /dev/null +++ b/stwo_cairo_verifier/src/sort.cairo @@ -0,0 +1,196 @@ +use core::array::ToSpanTrait; +use core::array::ArrayTrait; +use core::option::OptionTrait; + +trait Compare { + fn compare(self: @C, a: T, b: T) -> bool; +} + +#[derive(Drop, Copy)] +pub struct LowerThan {} + +impl LowerThanCompare> of Compare { + fn compare(self: @LowerThan, a: T, b: T) -> bool { + return a < b; + } +} + +#[derive(Drop, Copy)] +pub struct GreaterThan {} + +impl GreaterThanCompare, +Copy, +Drop> of Compare { + fn compare(self: @GreaterThan, a: T, b: T) -> bool { + return a > b; + } +} + +#[derive(Drop)] +pub struct SortedIterator { + comparer: C, + array: Span, + last_index: Option, +} + +trait SortedIteratorTrait< + T, C, +PartialOrd, +PartialEq, +Copy, +Drop, +Compare, +Drop, +Copy +> { + fn iterate(array_to_iterate: Span) -> SortedIterator; + + fn next_deduplicated( + ref self: SortedIterator + ) -> Option<(u32, T)> { + next_deduplicated::(ref self) + } + + fn next( + ref self: SortedIterator + ) -> Option< + (u32, T) + > { + if self.last_index.is_some() { + let last_index = self.last_index.unwrap(); + let last_value = *self.array[last_index]; + let mut is_repeated = false; + + let mut i = last_index + 1; + while i < self.array.len() { + if *self.array[i] == last_value { + is_repeated = true; + self.last_index = Option::Some(i); + break; + } + i += 1; + }; + + if is_repeated { + return Option::Some((self.last_index.unwrap(), last_value)); + } + } + next_deduplicated::(ref self) + } +} + +fn next_deduplicated< + T, C, +PartialOrd, +PartialEq, +Copy, +Drop, +Compare, +Drop, +Copy +>( + ref self: SortedIterator +) -> Option<(u32, T)> { + let mut candidate_index = Option::None; + let mut candidate_value = Option::None; + + let last_value = if let Option::Some(last_index) = self.last_index { + Option::Some(*self.array[last_index]) + } else { + Option::None + }; + + let mut i = 0; + while i < self.array.len() { + let is_better_than_last = if let Option::Some(last_value) = last_value { + self.comparer.compare(last_value, *self.array[i]) + } else { + true + }; + let is_nearer_than_candidate = if let Option::Some(candidate_value) = candidate_value { + self.comparer.compare(*self.array[i], candidate_value) + } else { + true + }; + if is_better_than_last && is_nearer_than_candidate { + candidate_index = Option::Some(i); + candidate_value = Option::Some(*self.array[i]); + } + i += 1; + }; + + if candidate_value.is_none() { + Option::None + } else { + self.last_index = candidate_index; + Option::Some((candidate_index.unwrap(), candidate_value.unwrap())) + } +} + +pub impl MaximumToMinimumSortedIterator< + T, +PartialOrd, +PartialEq, +Copy, +Drop +> of SortedIteratorTrait { + fn iterate(array_to_iterate: Span) -> SortedIterator { + SortedIterator { + comparer: GreaterThan {}, array: array_to_iterate, last_index: Option::None + } + } +} + +pub impl MinimumToMaximumSortedIterator< + T, +PartialOrd, +PartialEq, +Copy, +Drop +> of SortedIteratorTrait { + fn iterate(array_to_iterate: Span) -> SortedIterator { + SortedIterator { comparer: LowerThan {}, array: array_to_iterate, last_index: Option::None } + } +} + + +#[test] +fn test_sort_lowest_to_greatest() { + let my_array: Array = array![3, 5, 2, 4]; + let expected_array: Array = array![2, 3, 4, 5]; + + let mut sorted_array = array![]; + + let mut iterator = MinimumToMaximumSortedIterator::iterate(my_array.span()); + while let Option::Some((_index, value)) = iterator.next_deduplicated() { + sorted_array.append(value); + }; + + assert_eq!(expected_array, sorted_array); +} + +#[test] +fn test_sort_greatest_to_lowest() { + let my_array: Array = array![3, 5, 2, 4]; + let expected_array: Array = array![5, 4, 3, 2]; + + let mut sorted_array = array![]; + + let mut iterator = MaximumToMinimumSortedIterator::iterate(my_array.span()); + while let Option::Some((_index, value)) = iterator.next_deduplicated() { + sorted_array.append(value); + }; + + assert_eq!(expected_array, sorted_array); +} + +#[test] +fn test_sort_indexes_are_correct() { + let my_array: Array = array![3, 5, 2, 4]; + let expected_indexes: Array = array![2, 0, 3, 1]; + + let mut sorted_indexes = array![]; + + let mut iterator = MinimumToMaximumSortedIterator::iterate(my_array.span()); + while let Option::Some((index, _value)) = iterator.next_deduplicated() { + sorted_indexes.append(index); + }; + + assert_eq!(expected_indexes, sorted_indexes); +} + +#[test] +fn test_sort_with_duplicates() { + let my_array: Array = array![3, 5, 2, 3, 4, 3, 4]; + let expected_indexes: Array = array![2, 0, 3, 5, 4, 6, 1]; + let expected_array: Array = array![2, 3, 3, 3, 4, 4, 5]; + + let mut sorted_indexes = array![]; + let mut sorted_array = array![]; + + let mut iterator = MinimumToMaximumSortedIterator::iterate(my_array.span()); + while let Option::Some((index, value)) = iterator.next() { + sorted_array.append(value); + sorted_indexes.append(index); + }; + + assert_eq!(expected_indexes, sorted_indexes); + assert_eq!(expected_array, sorted_array); +} +