Skip to content

Commit

Permalink
fix swag
Browse files Browse the repository at this point in the history
  • Loading branch information
ngtkana committed Apr 5, 2024
1 parent dea743e commit 624ecef
Show file tree
Hide file tree
Showing 2 changed files with 120 additions and 23 deletions.
1 change: 0 additions & 1 deletion libs/swag/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,4 @@ edition = "2018"
[dependencies]

[dev-dependencies]
itertools = { workspace = true }
rand = { workspace = true }
142 changes: 120 additions & 22 deletions libs/swag/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,8 @@
//! * [`from_iter`](DequeueSwag::from_iter): [`IntoIterator`] -> [`DequeueSwag`].
//! * [`clone_from_slice`](DequeueSwag::clone_from_slice), [`copy_from_slice`](DequeueSwag::copy_from_slice): [`&[T]`] -> [`DequeueSwag`].
use std::{iter::FromIterator, ops::Index};
use std::iter::FromIterator;
use std::ops::Index;

/// Operations
pub trait Op {
Expand Down Expand Up @@ -67,10 +68,11 @@ impl<O: Op> DequeueSwag<O> {
/// use swag::DequeueSwag;
/// enum O {}
/// impl swag::Op for O {
/// type Value = i32;
/// fn op(a: &Self::Value, b: &Self::Value) -> Self::Value {
/// a + b
/// }
/// type Value = i32;
///
/// fn op(a: &Self::Value, b: &Self::Value) -> Self::Value {
/// a + b
/// }
/// }
/// let mut swag = DequeueSwag::<O>::copy_from_slice(&[2, 3, 4]);
/// swag.push_front(1);
Expand All @@ -95,10 +97,11 @@ impl<O: Op> DequeueSwag<O> {
/// use swag::DequeueSwag;
/// enum O {}
/// impl swag::Op for O {
/// type Value = i32;
/// fn op(a: &Self::Value, b: &Self::Value) -> Self::Value {
/// a + b
/// }
/// type Value = i32;
///
/// fn op(a: &Self::Value, b: &Self::Value) -> Self::Value {
/// a + b
/// }
/// }
/// let mut swag = DequeueSwag::<O>::copy_from_slice(&[1, 2, 3]);
/// swag.push_back(4);
Expand Down Expand Up @@ -147,7 +150,6 @@ impl<O: Op> DequeueSwag<O> {
}
*self = swp;
}
eprintln!("{} {}", self.front.len(), self.back.len());
let _ = self.front_sum.pop();
self.front.pop()
}
Expand All @@ -159,10 +161,11 @@ impl<O: Op> DequeueSwag<O> {
/// use swag::DequeueSwag;
/// enum O {}
/// impl swag::Op for O {
/// type Value = i32;
/// fn op(a: &Self::Value, b: &Self::Value) -> Self::Value {
/// a + b
/// }
/// type Value = i32;
///
/// fn op(a: &Self::Value, b: &Self::Value) -> Self::Value {
/// a + b
/// }
/// }
/// let mut swag = DequeueSwag::<O>::copy_from_slice(&[1, 2, 3]);
/// assert_eq!(swag.pop_back(), Some(3));
Expand All @@ -175,10 +178,10 @@ impl<O: Op> DequeueSwag<O> {
if self.back.is_empty() {
let n = self.front.len();
let mut swp = Self::new();
for x in self.front.drain(..n / 2).rev() {
for x in self.front.drain((n + 1) / 2..) {
swp.push_front(x);
}
for x in self.front.drain(..) {
for x in self.front.drain(..).rev() {
swp.push_back(x);
}
*self = swp;
Expand All @@ -194,10 +197,11 @@ impl<O: Op> DequeueSwag<O> {
/// use swag::DequeueSwag;
/// enum O {}
/// impl swag::Op for O {
/// type Value = i32;
/// fn op(a: &Self::Value, b: &Self::Value) -> Self::Value {
/// a + b
/// }
/// type Value = i32;
///
/// fn op(a: &Self::Value, b: &Self::Value) -> Self::Value {
/// a + b
/// }
/// }
/// let mut swag = DequeueSwag::<O>::copy_from_slice(&[1, 2, 3]);
/// assert_eq!(swag.fold(), Some(6));
Expand Down Expand Up @@ -306,23 +310,23 @@ where
}

impl<O: Op> IntoIterator for DequeueSwag<O> {
type Item = O::Value;
type IntoIter = std::iter::Chain<
std::iter::Rev<std::vec::IntoIter<O::Value>>,
std::vec::IntoIter<O::Value>,
>;
type Item = O::Value;

fn into_iter(self) -> Self::IntoIter {
self.front.into_iter().rev().chain(self.back.into_iter())
}
}

impl<'a, O: Op> IntoIterator for &'a DequeueSwag<O> {
type Item = &'a O::Value;
type IntoIter = std::iter::Chain<
std::iter::Rev<std::slice::Iter<'a, O::Value>>,
std::slice::Iter<'a, O::Value>,
>;
type Item = &'a O::Value;

fn into_iter(self) -> Self::IntoIter {
self.front.iter().rev().chain(self.back.iter())
Expand All @@ -348,3 +352,97 @@ where
}
}
}

#[cfg(test)]
mod tests {
use super::*;
use rand::rngs::StdRng;
use rand::Rng;
use rand::SeedableRng;
use std::collections::VecDeque;
use std::fmt::Debug;
use std::ops::Add;
use std::ops::Mul;

const P: u64 = 998244353;

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct Fp(u64);
impl Debug for Fp {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl Add for Fp {
type Output = Self;

fn add(self, rhs: Self) -> Self::Output {
Fp((self.0 + rhs.0) % P)
}
}

impl Mul for Fp {
type Output = Self;

fn mul(self, rhs: Self) -> Self::Output {
Fp(self.0 * rhs.0 % P)
}
}

enum Affine {}
impl Op for Affine {
type Value = (Fp, Fp);

fn op(&(a, b): &Self::Value, &(c, d): &Self::Value) -> Self::Value {
(a * c, d + c * b)
}
}

#[test]
fn test_swag() {
let mut rng = StdRng::seed_from_u64(42);
for _ in 0..1000 {
let q = 10;
let mut swag = DequeueSwag::<Affine>::new();
let mut deque = VecDeque::new();
for _ in 0..q {
match rng.gen_range(0..4) {
0 => {
let a = Fp(rng.gen_range(0..4));
let b = Fp(rng.gen_range(0..4));
eprintln!("push_front {:?}", (a, b));
swag.push_front((a, b));
deque.push_front((a, b));
}
1 => {
let a = Fp(rng.gen_range(0..4));
let b = Fp(rng.gen_range(0..4));
eprintln!("push_back {:?}", (a, b));
swag.push_back((a, b));
deque.push_back((a, b));
}
2 => {
eprintln!("pop_front");
assert_eq!(swag.pop_front(), deque.pop_front());
}
3 => {
eprintln!("pop_back");
assert_eq!(swag.pop_back(), deque.pop_back());
}
_ => unreachable!(),
}
eprintln!("{:?}", deque);
eprintln!("{:?}", swag);
let result = swag.fold().unwrap_or((Fp(1), Fp(0)));
let expected = deque
.iter()
.fold((Fp(1), Fp(0)), |acc, x| Affine::op(&acc, x));
assert_eq!(result, expected);
let result = swag.collect_vec();
let expected = deque.iter().copied().collect::<Vec<_>>();
assert_eq!(result, expected);
eprintln!("---");
}
}
}
}

0 comments on commit 624ecef

Please sign in to comment.