<commit_before>use super::internal::*;
use super::*;

/// Unzips the items of a parallel iterator into a pair of arbitrary
/// `ParallelExtend` containers.
pub fn unzip<I, A, B, FromA, FromB>(pi: I) -> (FromA, FromB)
    where I: ParallelIterator<Item = (A, B)>,
          FromA: Default + ParallelExtend<A>,
          FromB: Default + ParallelExtend<B>,
          A: Send,
          B: Send
{
    let mut a = FromA::default();
    let mut b = FromB::default();
    {
        // We have no idea what the consumers will look like for these
        // collections' `par_extend`, but we can intercept them in our own
        // `drive_unindexed`.  Start with the left side, type `A`:
        let iter = UnzipA {
            base: pi,
            b: &mut b,
        };
        a.par_extend(iter);
    }
    (a, b)
}

/// Unzip an `IndexedParallelIterator` into two arbitrary `Consumer`s.
pub fn unzip_indexed<I, A, B, CA, CB>(pi: I, left: CA, right: CB) -> (CA::Result, CB::Result)
    where I: IndexedParallelIterator<Item = (A, B)>,
          CA: Consumer<A>,
          CB: Consumer<B>,
          A: Send,
          B: Send
{
    let consumer = UnzipConsumer {
        left: left,
        right: right,
    };
    pi.drive(consumer)
}

/// Unzip a `ParallelIterator` into two arbitrary `UnindexedConsumer`s.
pub fn unzip_unindexed<I, A, B, CA, CB>(pi: I, left: CA, right: CB) -> (CA::Result, CB::Result)
    where I: ParallelIterator<Item = (A, B)>,
          CA: UnindexedConsumer<A>,
          CB: UnindexedConsumer<B>,
          A: Send,
          B: Send
{
    let consumer = UnzipConsumer {
        left: left,
        right: right,
    };
    pi.drive_unindexed(consumer)
}

/// A fake iterator to intercept the `Consumer` for type `A`.
struct UnzipA<'b, I, FromB: 'b> {
    base: I,
    b: &'b mut FromB,
}

impl<'b, I, A, B, FromB> ParallelIterator for UnzipA<'b, I, FromB>
    where I: ParallelIterator<Item = (A, B)>,
          FromB: Default + ParallelExtend<B>,
          A: Send,
          B: Send
{
    type Item = A;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        let mut result = None;
        {
            // Now it's time to find the consumer for type `B`
            let iter = UnzipB {
                base: self.base,
                left_consumer: consumer,
                left_result: &mut result,
            };
            self.b.par_extend(iter);
        }
        result.unwrap()
    }

    fn opt_len(&mut self) -> Option<usize> {
        self.base.opt_len()
    }
}

/// A fake iterator to intercept the `Consumer` for type `B`.
struct UnzipB<'r, I, A, CA>
    where CA: UnindexedConsumer<A>,
          CA::Result: 'r
{
    base: I,
    left_consumer: CA,
    left_result: &'r mut Option<CA::Result>,
}

impl<'r, I, A, B, CA> ParallelIterator for UnzipB<'r, I, A, CA>
    where I: ParallelIterator<Item = (A, B)>,
          CA: UnindexedConsumer<A>,
          A: Send,
          B: Send
{
    type Item = B;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        // Now that we have two consumers, we can unzip the real iterator.
        let result = unzip_unindexed(self.base, self.left_consumer, consumer);
        *self.left_result = Some(result.0);
        result.1
    }

    fn opt_len(&mut self) -> Option<usize> {
        self.base.opt_len()
    }
}

/// `Consumer` that unzips into two other `Consumer`s
struct UnzipConsumer<CA, CB> {
    left: CA,
    right: CB,
}

impl<A, B, CA, CB> Consumer<(A, B)> for UnzipConsumer<CA, CB>
    where CA: Consumer<A>,
          CB: Consumer<B>
{
    type Folder = UnzipFolder<CA::Folder, CB::Folder>;
    type Reducer = UnzipReducer<CA::Reducer, CB::Reducer>;
    type Result = (CA::Result, CB::Result);

    fn split_at(self, index: usize) -> (Self, Self, Self::Reducer) {
        let (left1, left2, left_reducer) = self.left.split_at(index);
        let (right1, right2, right_reducer) = self.right.split_at(index);

        (UnzipConsumer {
             left: left1,
             right: right1,
         },
         UnzipConsumer {
             left: left2,
             right: right2,
         },
         UnzipReducer {
             left: left_reducer,
             right: right_reducer,
         })
    }

    fn into_folder(self) -> Self::Folder {
        UnzipFolder {
            left: self.left.into_folder(),
            right: self.right.into_folder(),
        }
    }

    fn full(&self) -> bool {
        // don't stop until everyone is full
        self.left.full() && self.right.full()
    }
}

impl<A, B, CA, CB> UnindexedConsumer<(A, B)> for UnzipConsumer<CA, CB>
    where CA: UnindexedConsumer<A>,
          CB: UnindexedConsumer<B>
{
    fn split_off_left(&self) -> Self {
        UnzipConsumer {
            left: self.left.split_off_left(),
            right: self.right.split_off_left(),
        }
    }

    fn to_reducer(&self) -> Self::Reducer {
        UnzipReducer {
            left: self.left.to_reducer(),
            right: self.right.to_reducer(),
        }
    }
}

/// `Folder` that unzips into two other `Folder`s
struct UnzipFolder<FA, FB> {
    left: FA,
    right: FB,
}

impl<A, B, FA, FB> Folder<(A, B)> for UnzipFolder<FA, FB>
    where FA: Folder<A>,
          FB: Folder<B>
{
    type Result = (FA::Result, FB::Result);

    fn consume(self, (left, right): (A, B)) -> Self {
        UnzipFolder {
            left: self.left.consume(left),
            right: self.right.consume(right),
        }
    }

    fn complete(self) -> Self::Result {
        (self.left.complete(), self.right.complete())
    }

    fn full(&self) -> bool {
        // don't stop until everyone is full
        self.left.full() && self.right.full()
    }
}

/// `Reducer` that unzips into two other `Reducer`s
struct UnzipReducer<RA, RB> {
    left: RA,
    right: RB,
}

impl<A, B, RA, RB> Reducer<(A, B)> for UnzipReducer<RA, RB>
    where RA: Reducer<A>,
          RB: Reducer<B>
{
    fn reduce(self, left: (A, B), right: (A, B)) -> (A, B) {
        (self.left.reduce(left.0, right.0), self.right.reduce(left.1, right.1))
    }
}
<commit_msg>Abstract the unzip implementation<commit_after>use super::internal::*;
use super::*;

trait UnzipOp<T>: Sync {
    type Left: Send;
    type Right: Send;

    fn consume<FA, FB>(&self, item: T, left: FA, right: FB) -> (FA, FB)
        where FA: Folder<Self::Left>,
              FB: Folder<Self::Right>;

    fn indexable() -> bool {
        false
    }
}

fn execute<I, OP, FromA, FromB>(pi: I, op: OP) -> (FromA, FromB)
    where I: ParallelIterator,
          OP: UnzipOp<I::Item>,
          FromA: Default + ParallelExtend<OP::Left>,
          FromB: Default + ParallelExtend<OP::Right>
{
    let mut a = FromA::default();
    let mut b = FromB::default();
    {
        // We have no idea what the consumers will look like for these
        // collections' `par_extend`, but we can intercept them in our own
        // `drive_unindexed`.  Start with the left side, type `A`:
        let iter = UnzipA {
            base: pi,
            op: op,
            b: &mut b,
        };
        a.par_extend(iter);
    }
    (a, b)
}

/// Unzips the items of a parallel iterator into a pair of arbitrary
/// `ParallelExtend` containers.
pub fn unzip<I, A, B, FromA, FromB>(pi: I) -> (FromA, FromB)
    where I: ParallelIterator<Item = (A, B)>,
          FromA: Default + ParallelExtend<A>,
          FromB: Default + ParallelExtend<B>,
          A: Send,
          B: Send
{
    execute(pi, Unzip)
}

/// Unzip an `IndexedParallelIterator` into two arbitrary `Consumer`s.
pub fn unzip_indexed<I, A, B, CA, CB>(pi: I, left: CA, right: CB) -> (CA::Result, CB::Result)
    where I: IndexedParallelIterator<Item = (A, B)>,
          CA: Consumer<A>,
          CB: Consumer<B>,
          A: Send,
          B: Send
{
    let consumer = UnzipConsumer {
        op: &Unzip,
        left: left,
        right: right,
    };
    pi.drive(consumer)
}

struct Unzip;

impl<A: Send, B: Send> UnzipOp<(A, B)> for Unzip {
    type Left = A;
    type Right = B;

    fn consume<FA, FB>(&self, item: (A, B), left: FA, right: FB) -> (FA, FB)
        where FA: Folder<A>,
              FB: Folder<B>
    {
        (left.consume(item.0), right.consume(item.1))
    }

    fn indexable() -> bool {
        true
    }
}

/// A fake iterator to intercept the `Consumer` for type `A`.
struct UnzipA<'b, I, OP, FromB: 'b> {
    base: I,
    op: OP,
    b: &'b mut FromB,
}

impl<'b, I, OP, FromB> ParallelIterator for UnzipA<'b, I, OP, FromB>
    where I: ParallelIterator,
          OP: UnzipOp<I::Item>,
          FromB: Default + ParallelExtend<OP::Right>
{
    type Item = OP::Left;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        let mut result = None;
        {
            // Now it's time to find the consumer for type `B`
            let iter = UnzipB {
                base: self.base,
                op: self.op,
                left_consumer: consumer,
                left_result: &mut result,
            };
            self.b.par_extend(iter);
        }
        result.unwrap()
    }

    fn opt_len(&mut self) -> Option<usize> {
        if OP::indexable() {
            self.base.opt_len()
        } else {
            None
        }
    }
}

/// A fake iterator to intercept the `Consumer` for type `B`.
struct UnzipB<'r, I, OP, CA>
    where I: ParallelIterator,
          OP: UnzipOp<I::Item>,
          CA: UnindexedConsumer<OP::Left>,
          CA::Result: 'r
{
    base: I,
    op: OP,
    left_consumer: CA,
    left_result: &'r mut Option<CA::Result>,
}

impl<'r, I, OP, CA> ParallelIterator for UnzipB<'r, I, OP, CA>
    where I: ParallelIterator,
          OP: UnzipOp<I::Item>,
          CA: UnindexedConsumer<OP::Left>
{
    type Item = OP::Right;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        // Now that we have two consumers, we can unzip the real iterator.
        let consumer = UnzipConsumer {
            op: &self.op,
            left: self.left_consumer,
            right: consumer,
        };

        let result = self.base.drive_unindexed(consumer);
        *self.left_result = Some(result.0);
        result.1
    }

    fn opt_len(&mut self) -> Option<usize> {
        if OP::indexable() {
            self.base.opt_len()
        } else {
            None
        }
    }
}

/// `Consumer` that unzips into two other `Consumer`s
struct UnzipConsumer<'a, OP: 'a, CA, CB> {
    op: &'a OP,
    left: CA,
    right: CB,
}

impl<'a, T, OP, CA, CB> Consumer<T> for UnzipConsumer<'a, OP, CA, CB>
    where OP: UnzipOp<T>,
          CA: Consumer<OP::Left>,
          CB: Consumer<OP::Right>
{
    type Folder = UnzipFolder<'a, OP, CA::Folder, CB::Folder>;
    type Reducer = UnzipReducer<CA::Reducer, CB::Reducer>;
    type Result = (CA::Result, CB::Result);

    fn split_at(self, index: usize) -> (Self, Self, Self::Reducer) {
        let (left1, left2, left_reducer) = self.left.split_at(index);
        let (right1, right2, right_reducer) = self.right.split_at(index);

        (UnzipConsumer {
             op: self.op,
             left: left1,
             right: right1,
         },
         UnzipConsumer {
             op: self.op,
             left: left2,
             right: right2,
         },
         UnzipReducer {
             left: left_reducer,
             right: right_reducer,
         })
    }

    fn into_folder(self) -> Self::Folder {
        UnzipFolder {
            op: self.op,
            left: self.left.into_folder(),
            right: self.right.into_folder(),
        }
    }

    fn full(&self) -> bool {
        // don't stop until everyone is full
        self.left.full() && self.right.full()
    }
}

impl<'a, T, OP, CA, CB> UnindexedConsumer<T> for UnzipConsumer<'a, OP, CA, CB>
    where OP: UnzipOp<T>,
          CA: UnindexedConsumer<OP::Left>,
          CB: UnindexedConsumer<OP::Right>
{
    fn split_off_left(&self) -> Self {
        UnzipConsumer {
            op: self.op,
            left: self.left.split_off_left(),
            right: self.right.split_off_left(),
        }
    }

    fn to_reducer(&self) -> Self::Reducer {
        UnzipReducer {
            left: self.left.to_reducer(),
            right: self.right.to_reducer(),
        }
    }
}

/// `Folder` that unzips into two other `Folder`s
struct UnzipFolder<'a, OP: 'a, FA, FB> {
    op: &'a OP,
    left: FA,
    right: FB,
}

impl<'a, T, OP, FA, FB> Folder<T> for UnzipFolder<'a, OP, FA, FB>
    where OP: UnzipOp<T>,
          FA: Folder<OP::Left>,
          FB: Folder<OP::Right>
{
    type Result = (FA::Result, FB::Result);

    fn consume(self, item: T) -> Self {
        let (left, right) = self.op.consume(item, self.left, self.right);
        UnzipFolder {
            op: self.op,
            left: left,
            right: right,
        }
    }

    fn complete(self) -> Self::Result {
        (self.left.complete(), self.right.complete())
    }

    fn full(&self) -> bool {
        // don't stop until everyone is full
        self.left.full() && self.right.full()
    }
}

/// `Reducer` that unzips into two other `Reducer`s
struct UnzipReducer<RA, RB> {
    left: RA,
    right: RB,
}

impl<A, B, RA, RB> Reducer<(A, B)> for UnzipReducer<RA, RB>
    where RA: Reducer<A>,
          RB: Reducer<B>
{
    fn reduce(self, left: (A, B), right: (A, B)) -> (A, B) {
        (self.left.reduce(left.0, right.0), self.right.reduce(left.1, right.1))
    }
}
<|endoftext|>


<commit_before>use std::{ fs, io };
use std::path::Path;
use bytes::{ BytesMut, Bytes };
use once_cell::sync::OnceCell;
use ritsu::actions;

pub struct File {
    inner: ArcFile,
    buf: Option<BytesMut>,
    pos: u64,
    len: u64
}

impl File {
    pub async fn open(path: &Path) -> io::Result<File> {
        let handle = GLOBAL_HANDLE.get_or_init(init_ritsu_runtime);

        let fd = actions::fs::open(handle, path).await?;

        let len = fd.metadata()?.len();
        Ok(File {
            inner: ArcFile(Arc::new(fd)),
            buf: Some(BytesMut::with_capacity(1 << 64)),
            pos: 0,
            len
        })
    }

    pub async fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        fn add(x: u64, y: i64) -> u64 {
            if y >= 0 {
                x + (y as u64)
            } else {
                x - (-y as u64)
            }
        }

        match pos {
            io::SeekFrom::Start(offset) => self.pos = offset,
            io::SeekFrom::End(offset) => self.pos = add(self.len, offset),
            io::SeekFrom::Current(offset) => self.pos = add(self.pos, offset)
        }

        Ok(self.pos)
    }

    pub async fn next_chunk(&mut self) -> io::Result<Option<Bytes>> {
        let handle = GLOBAL_HANDLE.get_or_init(init_ritsu_runtime);

        let (_, mut buf) = actions::io::read_buf(
            handle,
            &mut Some(self.inner.clone()),
            self.buf.take().unwrap_or_else(|| BytesMut::with_capacity(1 << 16)),
            Some(self.pos as _)
        ).await?;

        Ok(if buf.is_empty() {
            None
        } else {
            self.pos += buf.len() as u64;
            let buf2 = buf.split();

            if buf.capacity() > 64 {
                self.buf = Some(buf);
            }

            Some(buf2.freeze())
        })
    }
}

use std::thread;
use tokio::sync::mpsc::{ UnboundedSender, unbounded_channel };
use io_uring::squeue;
use ritsu::{ Proactor, Handle };

static GLOBAL_HANDLE: OnceCell<RemoteHandle> = OnceCell::new();

struct RemoteHandle(UnboundedSender<squeue::Entry>);

impl Handle for RemoteHandle {
    unsafe fn push(&self, entry: squeue::Entry) {
        let _ = self.0.send(entry);
    }
}

fn init_ritsu_runtime() -> RemoteHandle {
    let (tx, mut rx) = unbounded_channel();

    thread::spawn(move || {
        let mut proactor = Proactor::new().unwrap();
        let handle = proactor.handle();

        proactor.block_on(async move {
            while let Some(entry) = rx.recv().await {
                unsafe {
                    handle.push(entry);
                }
            }
        }).unwrap();
    });

    RemoteHandle(tx)
}

use std::sync::Arc;
use std::os::unix::io::{ AsRawFd, RawFd };
use ritsu::actions::io::TrustedAsRawFd;

#[derive(Clone)]
struct ArcFile(Arc<fs::File>);

impl AsRawFd for ArcFile {
    #[inline]
    fn as_raw_fd(&self) -> RawFd {
        self.0.as_raw_fd()
    }
}

unsafe impl TrustedAsRawFd for ArcFile {}
<commit_msg>make it panic<commit_after>use std::{ fs, io };
use std::path::Path;
use bytes::{ BytesMut, Bytes };
use once_cell::sync::OnceCell;
use ritsu::actions;

pub struct File {
    inner: ArcFile,
    buf: Option<BytesMut>,
    pos: u64,
    len: u64
}

impl File {
    pub async fn open(path: &Path) -> io::Result<File> {
        let handle = GLOBAL_HANDLE.get_or_init(init_ritsu_runtime);

        let fd = actions::fs::open(handle, path).await?;

        let len = fd.metadata()?.len();
        Ok(File {
            inner: ArcFile(Arc::new(fd)),
            buf: Some(BytesMut::with_capacity(1 << 64)),
            pos: 0,
            len
        })
    }

    pub async fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        fn add(x: u64, y: i64) -> u64 {
            if y >= 0 {
                x + (y as u64)
            } else {
                x - (-y as u64)
            }
        }

        match pos {
            io::SeekFrom::Start(offset) => self.pos = offset,
            io::SeekFrom::End(offset) => self.pos = add(self.len, offset),
            io::SeekFrom::Current(offset) => self.pos = add(self.pos, offset)
        }

        Ok(self.pos)
    }

    pub async fn next_chunk(&mut self) -> io::Result<Option<Bytes>> {
        let handle = GLOBAL_HANDLE.get_or_init(init_ritsu_runtime);

        let (_, mut buf) = actions::io::read_buf(
            handle,
            &mut Some(self.inner.clone()),
            self.buf.take().unwrap_or_else(|| BytesMut::with_capacity(1 << 16)),
            Some(self.pos as _)
        ).await?;

        Ok(if buf.is_empty() {
            None
        } else {
            self.pos += buf.len() as u64;
            let buf2 = buf.split();

            if buf.capacity() > 64 {
                self.buf = Some(buf);
            }

            Some(buf2.freeze())
        })
    }
}

use std::thread;
use tokio::sync::mpsc::{ UnboundedSender, unbounded_channel };
use io_uring::squeue;
use ritsu::{ Proactor, Handle };

static GLOBAL_HANDLE: OnceCell<RemoteHandle> = OnceCell::new();

struct RemoteHandle(UnboundedSender<squeue::Entry>);

impl Handle for RemoteHandle {
    unsafe fn push(&self, entry: squeue::Entry) {
        self.0.send(entry).ok().expect("ritsu runtime not found");
    }
}

fn init_ritsu_runtime() -> RemoteHandle {
    let (tx, mut rx) = unbounded_channel();

    thread::spawn(move || {
        let mut proactor = Proactor::new().unwrap();
        let handle = proactor.handle();

        proactor.block_on(async move {
            while let Some(entry) = rx.recv().await {
                unsafe {
                    handle.push(entry);
                }
            }
        }).unwrap();
    });

    RemoteHandle(tx)
}

use std::sync::Arc;
use std::os::unix::io::{ AsRawFd, RawFd };
use ritsu::actions::io::TrustedAsRawFd;

#[derive(Clone)]
struct ArcFile(Arc<fs::File>);

impl AsRawFd for ArcFile {
    #[inline]
    fn as_raw_fd(&self) -> RawFd {
        self.0.as_raw_fd()
    }
}

unsafe impl TrustedAsRawFd for ArcFile {}
<|endoftext|>


<commit_before>use std::{iter, local_data};
use std::iter::MultiplicativeIterator;
use std::local_data::Key;

use arith;
use data::monoid::{Sum, MergeMonoidIterator, MergeMultiMonoidIterator, Wrap};

static PRIMES_BELOW100: &'static [uint] = &[
    2,  3,  5,  7, 11, 13, 17, 19, 23, 29, 31, 37, 41,
    43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97
];

static TASK_PRIME_KEY: Key<~[uint]> = &Key;

#[inline]
fn with_task_prime<T>(f: proc(&mut ~[uint]) -> T) -> T {
    let mut nums = match local_data::pop(TASK_PRIME_KEY) {
        Some(x) => x,
        None    => PRIMES_BELOW100.to_owned()
    };
    let result = f(&mut nums);
    local_data::set(TASK_PRIME_KEY, nums);
    result
}

#[inline(always)]
fn is_coprime(nums: &[uint], n: uint) -> bool {
    return nums.iter()
        .take_while(|& &p| p * p <= n)
        .all(|&p| n % p != 0);
}

#[inline(always)]
fn grow(nums: &mut ~[uint], len: uint) {
    if nums.len() >= len { return; }

    let mut it = iter::count(nums.last() + 2, 2)
        .filter(|&n| is_coprime(*nums, n));
    for n in it {
        nums.push(n);
        if nums.len() >= len { return; }
    }
}

#[inline(always)]
pub fn iter() -> PrimeIterator { PrimeIterator::new() }

#[inline(always)]
pub fn contains(n: uint) -> bool {
    do with_task_prime |nums| {
        let len = nums.len();
        let last = nums[len - 1];
        if n < last {
            nums.bsearch_elem(&n).is_some()
        } else {
            iter::count(0u, 1)
                .map(|i| { grow(nums, i + 1); nums[i] })
                .take_while(|&p| p * p <= n)
                .all(|p| n % p != 0)
        }
    }
}

#[inline(always)]
pub fn nth(i: uint) -> uint {
    do with_task_prime |nums| {
        grow(nums, i + 1);
        nums[i]
    }
}

pub struct PrimeIterator {
    priv idx: uint
}

impl PrimeIterator {
    #[inline(always)]
    pub fn new() -> PrimeIterator { PrimeIterator { idx: 0 } }
}

impl Iterator<uint> for PrimeIterator {
    #[inline(always)]
    fn next(&mut self) -> Option<uint> {
        let p = nth(self.idx);
        self.idx += 1;
        return Some(p);
    }
}

#[inline(always)]
pub fn factorize(n: uint) -> FactorIterator { FactorIterator::new(n) }

pub type Factor = (uint, int);

struct FactorIterator {
    priv num: uint,
    priv prime_iter: PrimeIterator
}

impl FactorIterator {
    #[inline(always)]
    pub fn new(num: uint) -> FactorIterator {
        FactorIterator { num: num, prime_iter: PrimeIterator::new() }
    }
}

impl Iterator<Factor> for FactorIterator {
    fn next(&mut self) -> Option<Factor> {
        if self.num == 0 || self.num == 1 { return None; }

        while self.num > 1 {
            let p = self.prime_iter.next().unwrap();

            if p * p > self.num {
                let n = self.num;
                self.num = 1;
                return Some((n, 1));
            }

            let mut exp = 0;
            while self.num % p == 0 {
                exp += 1;
                self.num /= p;
            }
            if exp > 0 { return Some((p, exp)); }
        }

        return None;
    }
}

#[inline(always)]
pub fn factors_to_uint<IA: Iterator<Factor>>(mut fs: IA) -> uint {
    let mut result = 1;
    for (base, exp) in fs {
        if exp > 0 {
            result *= arith::pow(base, exp as uint);
        } else {
            result /= arith::pow(base, (-exp) as uint);
        }
    }
    return result;
}

#[inline(always)]
pub fn comb(n: uint, r: uint) -> uint {
    let ns = range(r + 1, n + 1)
        .map(factorize)
        .map(|fs| fs.map(|(base, exp)| (base, Sum(exp))))
        .to_owned_vec();
    let numer = MergeMultiMonoidIterator::new(ns);

    let ds = range(1, n - r + 1)
        .map(factorize)
        .map(|fs| fs.map(|(base, exp)| (base, Sum(-exp))))
        .to_owned_vec();
    let denom = MergeMultiMonoidIterator::new(ds);

    return factors_to_uint(
        MergeMonoidIterator::new(numer, denom).map(|(a, m): (uint, Sum<int>)| (a, m.unwrap()))
    );
}

#[inline(always)]
pub fn num_of_divisors(n: uint) -> uint {
    if n == 0 { return 0; }
    return factorize(n)
        .map(|(_base, exp)| (exp as uint) + 1)
        .product();
}

#[inline(always)]
pub fn num_of_proper_divisors(n: uint) -> uint {
    num_of_proper_divisors(n) - 1
}

#[inline(always)]
pub fn sum_of_divisors(n: uint) -> uint {
    if n == 0 { return 0; }
    return factorize(n)
        .map(|(base, exp)| {
            (arith::pow(base, (exp as uint) + 1) - 1) / (base - 1)
        }).product();
}

#[inline(always)]
pub fn sum_of_proper_divisors(n: uint) -> uint {
    sum_of_divisors(n) - n
}

#[cfg(test)]
mod test {
    use super::*;

    static PRIMES_BELOW200: &'static [uint] = &[
        2,  3,  5,  7, 11, 13, 17, 19, 23, 29, 31, 37, 41,
        43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97,
        101, 103, 107, 109, 113, 127, 131, 137, 139, 149,
        151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199
    ];

    #[test]
    fn test_prime_contains() {
        assert!(!contains(0));
        assert!(!contains(1));
        assert!(contains(2));
        assert!(contains(3));
        assert!(!contains(4));
        assert!(contains(5));
        assert!(!contains(6));
        assert!(contains(7));
        assert!(!contains(100));
    }

    #[test]
    fn test_prime_index() {
        // Generated primes
        for (i, &p) in PRIMES_BELOW200.iter().enumerate() { assert_eq!(nth(i), p); }
        // Memoized primes
        for (i, &p) in PRIMES_BELOW200.iter().enumerate() { assert_eq!(nth(i), p); }
    }

    #[test]
    fn test_iter() {
        let mut it = iter();

        let mut i = 0;
        let ys = &[ 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41 ];
        for x in it {
            if i >= ys.len() { break; }
            assert_eq!(x, ys[i]);
            i += 1;
        }
    }

    #[test]
    fn test_factorize() {
        fn check(n: uint, fs: &[Factor]) {
            let mut v = ~[];
            let mut it = factorize(n);
            for f in it { v.push(f); }

            assert_eq!(v.initn(0), fs);
        }

        check(0, &[]);
        check(1, &[]);
        check(2, &[(2, 1)]);
        check(3, &[(3, 1)]);
        check(4, &[(2, 2)]);
        check(5, &[(5, 1)]);
        check(6, &[(2, 1), (3, 1)]);
        check(7, &[(7, 1)]);
        check(8, &[(2, 3)]);
        check(9, &[(3, 2)]);
        check(10, &[(2, 1), (5, 1)]);

        check(8 * 27, &[(2, 3), (3, 3)]);
        check(97, &[(97, 1)]);
        check(97 * 41, &[(41, 1), (97, 1)]);
    }

    #[test]
    fn test_factors_to_uint() {
        fn check(n: uint, fs: &[Factor]) {
            assert_eq!(factors_to_uint(fs.iter().map(|n| *n)), n);
        }

        check(1, &[]);
        check(2, &[(2, 1)]);
        check(3, &[(3, 1)]);
        check(4, &[(2, 2)]);
        check(5, &[(5, 1)]);
        check(6, &[(2, 1), (3, 1)]);
        check(7, &[(7, 1)]);
        check(8, &[(2, 3)]);
        check(9, &[(3, 2)]);
        check(10, &[(2, 1), (5, 1)]);

        check(8 * 27, &[(2, 3), (3, 3)]);
        check(97, &[(97, 1)]);
        check(97 * 41, &[(41, 1), (97, 1)]);

        check(1, &[(1, 1)]);
    }

    #[test]
    fn test_comb() {
        assert_eq!(comb(2, 2), 1);
        assert_eq!(comb(3, 2), 3);
        assert_eq!(comb(4, 2), 6);
        assert_eq!(comb(5, 2), 10);

        assert_eq!(comb(40, 20), 137846528820);
    }

    #[test]
    fn test_num_of_divisors() {
        assert_eq!(num_of_divisors(1), 1);
        assert_eq!(num_of_divisors(2), 2);
        assert_eq!(num_of_divisors(3), 2);
        assert_eq!(num_of_divisors(4), 3);
        assert_eq!(num_of_divisors(5), 2);
        assert_eq!(num_of_divisors(6), 4);
        assert_eq!(num_of_divisors(7), 2);
        assert_eq!(num_of_divisors(8), 4);
        assert_eq!(num_of_divisors(9), 3);
        assert_eq!(num_of_divisors(10), 4);
        assert_eq!(num_of_divisors(11), 2);
        assert_eq!(num_of_divisors(12), 6);

        assert_eq!(num_of_divisors(24), 8);
        assert_eq!(num_of_divisors(36), 9);
        assert_eq!(num_of_divisors(48), 10);
        assert_eq!(num_of_divisors(60), 12);

        assert_eq!(num_of_divisors(50), 6);
    }

    #[test]
    fn test_sum_of_divisors() {
        assert_eq!(sum_of_divisors(1), 1);
        assert_eq!(sum_of_divisors(2), 3);
        assert_eq!(sum_of_divisors(3), 4);
        assert_eq!(sum_of_divisors(4), 7);
        assert_eq!(sum_of_divisors(5), 6);
        assert_eq!(sum_of_divisors(6), 12);
        assert_eq!(sum_of_divisors(7), 8);
        assert_eq!(sum_of_divisors(8), 15);
        assert_eq!(sum_of_divisors(9), 13);
        assert_eq!(sum_of_divisors(10), 18);
        assert_eq!(sum_of_divisors(11), 12);
        assert_eq!(sum_of_divisors(12), 28);

        assert_eq!(sum_of_divisors(24), 60);
        assert_eq!(sum_of_divisors(36), 91);
        assert_eq!(sum_of_divisors(48), 124);
        assert_eq!(sum_of_divisors(60), 168);

        assert_eq!(sum_of_divisors(50), 93);
    }
}

#[cfg(test)]
mod bench {
    use super::{TASK_PRIME_KEY, PRIMES_BELOW100};
    use std::local_data;
    use extra::test::BenchHarness;

    fn reset_task_prime() {
        local_data::set(TASK_PRIME_KEY, PRIMES_BELOW100.to_owned());
    }

    #[bench]
    fn get_5000th(bh: &mut BenchHarness) {
        bh.iter(|| { super::nth(5000); });
    }

    #[bench]
    fn get_5000th_nocache(bh: &mut BenchHarness) {
        bh.iter(|| { reset_task_prime(); super::nth(5000); });
    }

    #[bench]
    fn factorial_600851475143(bh: &mut BenchHarness) {
        bh.iter(|| { super::factorize(600851475143).fold(0, |a, (b, _)| a + b); } )
    }

    #[bench]
    fn factorial_600851475143_nocache(bh: &mut BenchHarness) {
        bh.iter(|| { reset_task_prime(); super::factorize(600851475143).fold(0, |a, (b, _)| a + b); } )
    }
}
<commit_msg>Add more prime benchmarks<commit_after>use std::{iter, local_data};
use std::iter::MultiplicativeIterator;
use std::local_data::Key;

use arith;
use data::monoid::{Sum, MergeMonoidIterator, MergeMultiMonoidIterator, Wrap};

static PRIMES_BELOW100: &'static [uint] = &[
    2,  3,  5,  7, 11, 13, 17, 19, 23, 29, 31, 37, 41,
    43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97
];

static TASK_PRIME_KEY: Key<~[uint]> = &Key;

#[inline]
fn with_task_prime<T>(f: proc(&mut ~[uint]) -> T) -> T {
    let mut nums = match local_data::pop(TASK_PRIME_KEY) {
        Some(x) => x,
        None    => PRIMES_BELOW100.to_owned()
    };
    let result = f(&mut nums);
    local_data::set(TASK_PRIME_KEY, nums);
    result
}

#[inline(always)]
fn is_coprime(nums: &[uint], n: uint) -> bool {
    return nums.iter()
        .take_while(|& &p| p * p <= n)
        .all(|&p| n % p != 0);
}

#[inline(always)]
fn grow(nums: &mut ~[uint], len: uint) {
    if nums.len() >= len { return; }

    let mut it = iter::count(nums.last() + 2, 2)
        .filter(|&n| is_coprime(*nums, n));
    for n in it {
        nums.push(n);
        if nums.len() >= len { return; }
    }
}

#[inline(always)]
pub fn iter() -> PrimeIterator { PrimeIterator::new() }

#[inline(always)]
pub fn contains(n: uint) -> bool {
    do with_task_prime |nums| {
        let len = nums.len();
        let last = nums[len - 1];
        if n < last {
            nums.bsearch_elem(&n).is_some()
        } else {
            iter::count(0u, 1)
                .map(|i| { grow(nums, i + 1); nums[i] })
                .take_while(|&p| p * p <= n)
                .all(|p| n % p != 0)
        }
    }
}

#[inline(always)]
pub fn nth(i: uint) -> uint {
    do with_task_prime |nums| {
        grow(nums, i + 1);
        nums[i]
    }
}

pub struct PrimeIterator {
    priv idx: uint
}

impl PrimeIterator {
    #[inline(always)]
    pub fn new() -> PrimeIterator { PrimeIterator { idx: 0 } }
}

impl Iterator<uint> for PrimeIterator {
    #[inline(always)]
    fn next(&mut self) -> Option<uint> {
        let p = nth(self.idx);
        self.idx += 1;
        return Some(p);
    }
}

#[inline(always)]
pub fn factorize(n: uint) -> FactorIterator { FactorIterator::new(n) }

pub type Factor = (uint, int);

struct FactorIterator {
    priv num: uint,
    priv prime_iter: PrimeIterator
}

impl FactorIterator {
    #[inline(always)]
    pub fn new(num: uint) -> FactorIterator {
        FactorIterator { num: num, prime_iter: PrimeIterator::new() }
    }
}

impl Iterator<Factor> for FactorIterator {
    fn next(&mut self) -> Option<Factor> {
        if self.num == 0 || self.num == 1 { return None; }

        while self.num > 1 {
            let p = self.prime_iter.next().unwrap();

            if p * p > self.num {
                let n = self.num;
                self.num = 1;
                return Some((n, 1));
            }

            let mut exp = 0;
            while self.num % p == 0 {
                exp += 1;
                self.num /= p;
            }
            if exp > 0 { return Some((p, exp)); }
        }

        return None;
    }
}

#[inline(always)]
pub fn factors_to_uint<IA: Iterator<Factor>>(mut fs: IA) -> uint {
    let mut result = 1;
    for (base, exp) in fs {
        if exp > 0 {
            result *= arith::pow(base, exp as uint);
        } else {
            result /= arith::pow(base, (-exp) as uint);
        }
    }
    return result;
}

#[inline(always)]
pub fn comb(n: uint, r: uint) -> uint {
    let ns = range(r + 1, n + 1)
        .map(factorize)
        .map(|fs| fs.map(|(base, exp)| (base, Sum(exp))))
        .to_owned_vec();
    let numer = MergeMultiMonoidIterator::new(ns);

    let ds = range(1, n - r + 1)
        .map(factorize)
        .map(|fs| fs.map(|(base, exp)| (base, Sum(-exp))))
        .to_owned_vec();
    let denom = MergeMultiMonoidIterator::new(ds);

    return factors_to_uint(
        MergeMonoidIterator::new(numer, denom).map(|(a, m): (uint, Sum<int>)| (a, m.unwrap()))
    );
}

#[inline(always)]
pub fn num_of_divisors(n: uint) -> uint {
    if n == 0 { return 0; }
    return factorize(n)
        .map(|(_base, exp)| (exp as uint) + 1)
        .product();
}

#[inline(always)]
pub fn num_of_proper_divisors(n: uint) -> uint {
    num_of_proper_divisors(n) - 1
}

#[inline(always)]
pub fn sum_of_divisors(n: uint) -> uint {
    if n == 0 { return 0; }
    return factorize(n)
        .map(|(base, exp)| {
            (arith::pow(base, (exp as uint) + 1) - 1) / (base - 1)
        }).product();
}

#[inline(always)]
pub fn sum_of_proper_divisors(n: uint) -> uint {
    sum_of_divisors(n) - n
}

#[cfg(test)]
mod test {
    use super::*;

    static PRIMES_BELOW200: &'static [uint] = &[
        2,  3,  5,  7, 11, 13, 17, 19, 23, 29, 31, 37, 41,
        43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97,
        101, 103, 107, 109, 113, 127, 131, 137, 139, 149,
        151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199
    ];

    #[test]
    fn test_prime_contains() {
        assert!(!contains(0));
        assert!(!contains(1));
        assert!(contains(2));
        assert!(contains(3));
        assert!(!contains(4));
        assert!(contains(5));
        assert!(!contains(6));
        assert!(contains(7));
        assert!(!contains(100));
    }

    #[test]
    fn test_prime_index() {
        // Generated primes
        for (i, &p) in PRIMES_BELOW200.iter().enumerate() { assert_eq!(nth(i), p); }
        // Memoized primes
        for (i, &p) in PRIMES_BELOW200.iter().enumerate() { assert_eq!(nth(i), p); }
    }

    #[test]
    fn test_iter() {
        let mut it = iter();

        let mut i = 0;
        let ys = &[ 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41 ];
        for x in it {
            if i >= ys.len() { break; }
            assert_eq!(x, ys[i]);
            i += 1;
        }
    }

    #[test]
    fn test_factorize() {
        fn check(n: uint, fs: &[Factor]) {
            let mut v = ~[];
            let mut it = factorize(n);
            for f in it { v.push(f); }

            assert_eq!(v.initn(0), fs);
        }

        check(0, &[]);
        check(1, &[]);
        check(2, &[(2, 1)]);
        check(3, &[(3, 1)]);
        check(4, &[(2, 2)]);
        check(5, &[(5, 1)]);
        check(6, &[(2, 1), (3, 1)]);
        check(7, &[(7, 1)]);
        check(8, &[(2, 3)]);
        check(9, &[(3, 2)]);
        check(10, &[(2, 1), (5, 1)]);

        check(8 * 27, &[(2, 3), (3, 3)]);
        check(97, &[(97, 1)]);
        check(97 * 41, &[(41, 1), (97, 1)]);
    }

    #[test]
    fn test_factors_to_uint() {
        fn check(n: uint, fs: &[Factor]) {
            assert_eq!(factors_to_uint(fs.iter().map(|n| *n)), n);
        }

        check(1, &[]);
        check(2, &[(2, 1)]);
        check(3, &[(3, 1)]);
        check(4, &[(2, 2)]);
        check(5, &[(5, 1)]);
        check(6, &[(2, 1), (3, 1)]);
        check(7, &[(7, 1)]);
        check(8, &[(2, 3)]);
        check(9, &[(3, 2)]);
        check(10, &[(2, 1), (5, 1)]);

        check(8 * 27, &[(2, 3), (3, 3)]);
        check(97, &[(97, 1)]);
        check(97 * 41, &[(41, 1), (97, 1)]);

        check(1, &[(1, 1)]);
    }

    #[test]
    fn test_comb() {
        assert_eq!(comb(2, 2), 1);
        assert_eq!(comb(3, 2), 3);
        assert_eq!(comb(4, 2), 6);
        assert_eq!(comb(5, 2), 10);

        assert_eq!(comb(40, 20), 137846528820);
    }

    #[test]
    fn test_num_of_divisors() {
        assert_eq!(num_of_divisors(1), 1);
        assert_eq!(num_of_divisors(2), 2);
        assert_eq!(num_of_divisors(3), 2);
        assert_eq!(num_of_divisors(4), 3);
        assert_eq!(num_of_divisors(5), 2);
        assert_eq!(num_of_divisors(6), 4);
        assert_eq!(num_of_divisors(7), 2);
        assert_eq!(num_of_divisors(8), 4);
        assert_eq!(num_of_divisors(9), 3);
        assert_eq!(num_of_divisors(10), 4);
        assert_eq!(num_of_divisors(11), 2);
        assert_eq!(num_of_divisors(12), 6);

        assert_eq!(num_of_divisors(24), 8);
        assert_eq!(num_of_divisors(36), 9);
        assert_eq!(num_of_divisors(48), 10);
        assert_eq!(num_of_divisors(60), 12);

        assert_eq!(num_of_divisors(50), 6);
    }

    #[test]
    fn test_sum_of_divisors() {
        assert_eq!(sum_of_divisors(1), 1);
        assert_eq!(sum_of_divisors(2), 3);
        assert_eq!(sum_of_divisors(3), 4);
        assert_eq!(sum_of_divisors(4), 7);
        assert_eq!(sum_of_divisors(5), 6);
        assert_eq!(sum_of_divisors(6), 12);
        assert_eq!(sum_of_divisors(7), 8);
        assert_eq!(sum_of_divisors(8), 15);
        assert_eq!(sum_of_divisors(9), 13);
        assert_eq!(sum_of_divisors(10), 18);
        assert_eq!(sum_of_divisors(11), 12);
        assert_eq!(sum_of_divisors(12), 28);

        assert_eq!(sum_of_divisors(24), 60);
        assert_eq!(sum_of_divisors(36), 91);
        assert_eq!(sum_of_divisors(48), 124);
        assert_eq!(sum_of_divisors(60), 168);

        assert_eq!(sum_of_divisors(50), 93);
    }
}

#[cfg(test)]
mod bench {
    use super::{TASK_PRIME_KEY, PRIMES_BELOW100};
    use std::local_data;
    use extra::test::BenchHarness;

    fn reset_task_prime() {
        local_data::set(TASK_PRIME_KEY, PRIMES_BELOW100.to_owned());
    }

    #[bench]
    fn get_5000th(bh: &mut BenchHarness) {
        bh.iter(|| { super::nth(5000); });
    }

    #[bench]
    fn get_5000th_nocache(bh: &mut BenchHarness) {
        bh.iter(|| { reset_task_prime(); super::nth(5000); });
    }

    #[bench]
    fn get_below_5000th(bh: &mut BenchHarness) {
        bh.iter(|| { for _p in super::iter().take(5000) {} });
    }

    #[bench]
    fn get_below_5000th_nocache(bh: &mut BenchHarness) {
        bh.iter(|| { reset_task_prime(); for _p in super::iter().take(5000) {} });
    }

    #[bench]
    fn factorial_600851475143(bh: &mut BenchHarness) {
        bh.iter(|| { super::factorize(600851475143).fold(0, |a, (b, _)| a + b); } )
    }

    #[bench]
    fn factorial_600851475143_nocache(bh: &mut BenchHarness) {
        bh.iter(|| { reset_task_prime(); super::factorize(600851475143).fold(0, |a, (b, _)| a + b); } )
    }
}
<|endoftext|>


<commit_before>use std::{fmt, str};

use common::{as_alias, column_identifier_no_alias, integer_literal, opt_multispace,
             type_identifier, Literal, SqlType};
use column::Column;

#[derive(Debug, Clone, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub enum ArithmeticOperator {
    Add,
    Subtract,
    Multiply,
    Divide,
}

#[derive(Debug, Clone, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub enum ArithmeticBase {
    Column(Column),
    Scalar(Literal),
}

#[derive(Debug, Clone, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct ArithmeticExpression {
    pub op: ArithmeticOperator,
    pub left: ArithmeticBase,
    pub right: ArithmeticBase,
    pub alias: Option<String>,
}

impl ArithmeticExpression {
    pub fn new(
        op: ArithmeticOperator,
        left: ArithmeticBase,
        right: ArithmeticBase,
        alias: Option<String>,
    ) -> Self {
        Self {
            op: op,
            left: left,
            right: right,
            alias: alias,
        }
    }
}

impl fmt::Display for ArithmeticOperator {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ArithmeticOperator::Add => write!(f, "+"),
            ArithmeticOperator::Subtract => write!(f, "-"),
            ArithmeticOperator::Multiply => write!(f, "*"),
            ArithmeticOperator::Divide => write!(f, "/"),
        }
    }
}

impl fmt::Display for ArithmeticBase {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ArithmeticBase::Column(ref col) => write!(f, "{}", col),
            ArithmeticBase::Scalar(ref lit) => write!(f, "{}", lit.to_string()),
        }
    }
}

impl fmt::Display for ArithmeticExpression {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.alias {
            Some(ref alias) => write!(f, "{} {} {} AS {}", self.left, self.op, self.right, alias),
            None => write!(f, "{} {} {}", self.left, self.op, self.right),
        }
    }
}

named!(pub arithmetic_cast<&[u8], (ArithmeticBase, Option<SqlType>)>,
    dbg_dmp!(alt_complete!(
        do_parse!(
            tag_no_case!("cast") >>
            opt_multispace >>
            tag!("(") >>
            opt_multispace >>
            // TODO(malte): should be arbitrary expr
            v: arithmetic_base >>
            opt_multispace >>
            tag_no_case!("as") >>
            opt_multispace >>
            _sign: opt!(terminated!(tag_no_case!("signed"), opt_multispace)) >>
            typ: type_identifier >>
            opt_multispace >>
            tag!(")") >>
            (v, Some(typ))
        ) |
        do_parse!(
            v: arithmetic_base >>
            (v, None)
        )
    ))
);

/// Parse standard math operators.
/// TODO(malte): this doesn't currently observe operator precedence.
named!(pub arithmetic_operator<&[u8], ArithmeticOperator>,
    alt_complete!(
          map!(tag!("+"), |_| ArithmeticOperator::Add)
        | map!(tag!("-"), |_| ArithmeticOperator::Subtract)
        | map!(tag!("*"), |_| ArithmeticOperator::Multiply)
        | map!(tag!("/"), |_| ArithmeticOperator::Divide)
    )
);

/// Base case for nested arithmetic expressions: column name or literal.
named!(pub arithmetic_base<&[u8], ArithmeticBase>,
    alt_complete!(
          map!(integer_literal, |il| ArithmeticBase::Scalar(il))
        | map!(column_identifier_no_alias, |ci| ArithmeticBase::Column(ci))
    )
);

/// Parse simple arithmetic expressions combining literals, and columns and literals.
/// TODO(malte): this doesn't currently support nested expressions.
named!(pub arithmetic_expression<&[u8], ArithmeticExpression>,
    complete!(do_parse!(
        left: arithmetic_cast >>
        opt_multispace >>
        op: arithmetic_operator >>
        opt_multispace >>
        right: arithmetic_cast >>
        alias: opt!(as_alias) >>
        (ArithmeticExpression {
            op: op,
            // TODO(malte): discards casts
            left: left.0,
            right: right.0,
            alias: match alias {
                None => None,
                Some(a) => Some(String::from(a)),
            },
        })
    ))
);

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn it_parses_arithmetic_expressions() {
        use column::FunctionExpression;
        use super::ArithmeticOperator::*;
        use super::ArithmeticBase::Scalar;
        use super::ArithmeticBase::Column as ABColumn;

        let lit_ae = [
            "5 + 42",
            "5+42",
            "5 * 42",
            "5 - 42",
            "5 / 42",
            "2 * 10 AS twenty ",
        ];

        // N.B. trailing space in "5 + foo " is required because `sql_identifier`'s keyword
        // detection requires a follow-up character (in practice, there always is one because we
        // use semicolon-terminated queries).
        let col_lit_ae = [
            "foo+5",
            "foo + 5",
            "5 + foo ",
            "foo * bar AS foobar",
            "MAX(foo)-3333",
        ];

        let expected_lit_ae = [
            ArithmeticExpression::new(Add, Scalar(5.into()), Scalar(42.into()), None),
            ArithmeticExpression::new(Add, Scalar(5.into()), Scalar(42.into()), None),
            ArithmeticExpression::new(Multiply, Scalar(5.into()), Scalar(42.into()), None),
            ArithmeticExpression::new(Subtract, Scalar(5.into()), Scalar(42.into()), None),
            ArithmeticExpression::new(Divide, Scalar(5.into()), Scalar(42.into()), None),
            ArithmeticExpression::new(
                Multiply,
                Scalar(2.into()),
                Scalar(10.into()),
                Some(String::from("twenty")),
            ),
        ];
        let expected_col_lit_ae = [
            ArithmeticExpression::new(Add, ABColumn("foo".into()), Scalar(5.into()), None),
            ArithmeticExpression::new(Add, ABColumn("foo".into()), Scalar(5.into()), None),
            ArithmeticExpression::new(Add, Scalar(5.into()), ABColumn("foo".into()), None),
            ArithmeticExpression::new(
                Multiply,
                ABColumn("foo".into()),
                ABColumn("bar".into()),
                Some(String::from("foobar")),
            ),
            ArithmeticExpression::new(
                Subtract,
                ABColumn(Column {
                    name: String::from("max(foo)"),
                    alias: None,
                    table: None,
                    function: Some(Box::new(FunctionExpression::Max("foo".into()))),
                }),
                Scalar(3333.into()),
                None,
            ),
        ];

        for (i, e) in lit_ae.iter().enumerate() {
            let res = arithmetic_expression(e.as_bytes());
            assert!(res.is_done());
            assert_eq!(res.unwrap().1, expected_lit_ae[i]);
        }

        for (i, e) in col_lit_ae.iter().enumerate() {
            let res = arithmetic_expression(e.as_bytes());
            assert!(res.is_done());
            assert_eq!(res.unwrap().1, expected_col_lit_ae[i]);
        }
    }

    #[test]
    fn it_displays_arithmetic_expressions() {
        use super::ArithmeticOperator::*;
        use super::ArithmeticBase::Scalar;
        use super::ArithmeticBase::Column as ABColumn;

        let expressions = [
            ArithmeticExpression::new(Add, ABColumn("foo".into()), Scalar(5.into()), None),
            ArithmeticExpression::new(Subtract, Scalar(5.into()), ABColumn("foo".into()), None),
            ArithmeticExpression::new(
                Multiply,
                ABColumn("foo".into()),
                ABColumn("bar".into()),
                None,
            ),
            ArithmeticExpression::new(Divide, Scalar(10.into()), Scalar(2.into()), None),
            ArithmeticExpression::new(
                Add,
                Scalar(10.into()),
                Scalar(2.into()),
                Some(String::from("bob")),
            ),
        ];

        let expected_strings = ["foo + 5", "5 - foo", "foo * bar", "10 / 2", "10 + 2 AS bob"];
        for (i, e) in expressions.iter().enumerate() {
            assert_eq!(expected_strings[i], format!("{}", e));
        }
    }

    #[test]
    fn it_parses_arithmetic_casts() {
        use super::ArithmeticOperator::*;
        use super::ArithmeticBase::Scalar;
        use super::ArithmeticBase::Column as ABColumn;

        let exprs = [
            "CAST(foo AS signed int) + CAST(bar AS signed int) ",
            "CAST(5 AS bigint) - foo ",
        ];

        // XXX(malte): currently discards the cast and type information!
        let expected = [
            ArithmeticExpression::new(Add, ABColumn("foo".into()), ABColumn("bar".into()), None),
            ArithmeticExpression::new(Subtract, Scalar(5.into()), ABColumn("foo".into()), None),
        ];

        for (i, e) in exprs.iter().enumerate() {
            let res = arithmetic_expression(e.as_bytes());
            assert!(res.is_done(), "{} failed to parse", e);
            assert_eq!(res.unwrap().1, expected[i]);
        }
    }

}
<commit_msg>Remove debug rule<commit_after>use std::{fmt, str};

use common::{as_alias, column_identifier_no_alias, integer_literal, opt_multispace,
             type_identifier, Literal, SqlType};
use column::Column;

#[derive(Debug, Clone, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub enum ArithmeticOperator {
    Add,
    Subtract,
    Multiply,
    Divide,
}

#[derive(Debug, Clone, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub enum ArithmeticBase {
    Column(Column),
    Scalar(Literal),
}

#[derive(Debug, Clone, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct ArithmeticExpression {
    pub op: ArithmeticOperator,
    pub left: ArithmeticBase,
    pub right: ArithmeticBase,
    pub alias: Option<String>,
}

impl ArithmeticExpression {
    pub fn new(
        op: ArithmeticOperator,
        left: ArithmeticBase,
        right: ArithmeticBase,
        alias: Option<String>,
    ) -> Self {
        Self {
            op: op,
            left: left,
            right: right,
            alias: alias,
        }
    }
}

impl fmt::Display for ArithmeticOperator {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ArithmeticOperator::Add => write!(f, "+"),
            ArithmeticOperator::Subtract => write!(f, "-"),
            ArithmeticOperator::Multiply => write!(f, "*"),
            ArithmeticOperator::Divide => write!(f, "/"),
        }
    }
}

impl fmt::Display for ArithmeticBase {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ArithmeticBase::Column(ref col) => write!(f, "{}", col),
            ArithmeticBase::Scalar(ref lit) => write!(f, "{}", lit.to_string()),
        }
    }
}

impl fmt::Display for ArithmeticExpression {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.alias {
            Some(ref alias) => write!(f, "{} {} {} AS {}", self.left, self.op, self.right, alias),
            None => write!(f, "{} {} {}", self.left, self.op, self.right),
        }
    }
}

named!(pub arithmetic_cast<&[u8], (ArithmeticBase, Option<SqlType>)>,
    alt_complete!(
        do_parse!(
            tag_no_case!("cast") >>
            opt_multispace >>
            tag!("(") >>
            opt_multispace >>
            // TODO(malte): should be arbitrary expr
            v: arithmetic_base >>
            opt_multispace >>
            tag_no_case!("as") >>
            opt_multispace >>
            _sign: opt!(terminated!(tag_no_case!("signed"), opt_multispace)) >>
            typ: type_identifier >>
            opt_multispace >>
            tag!(")") >>
            (v, Some(typ))
        ) |
        do_parse!(
            v: arithmetic_base >>
            (v, None)
        )
    )
);

/// Parse standard math operators.
/// TODO(malte): this doesn't currently observe operator precedence.
named!(pub arithmetic_operator<&[u8], ArithmeticOperator>,
    alt_complete!(
          map!(tag!("+"), |_| ArithmeticOperator::Add)
        | map!(tag!("-"), |_| ArithmeticOperator::Subtract)
        | map!(tag!("*"), |_| ArithmeticOperator::Multiply)
        | map!(tag!("/"), |_| ArithmeticOperator::Divide)
    )
);

/// Base case for nested arithmetic expressions: column name or literal.
named!(pub arithmetic_base<&[u8], ArithmeticBase>,
    alt_complete!(
          map!(integer_literal, |il| ArithmeticBase::Scalar(il))
        | map!(column_identifier_no_alias, |ci| ArithmeticBase::Column(ci))
    )
);

/// Parse simple arithmetic expressions combining literals, and columns and literals.
/// TODO(malte): this doesn't currently support nested expressions.
named!(pub arithmetic_expression<&[u8], ArithmeticExpression>,
    complete!(do_parse!(
        left: arithmetic_cast >>
        opt_multispace >>
        op: arithmetic_operator >>
        opt_multispace >>
        right: arithmetic_cast >>
        alias: opt!(as_alias) >>
        (ArithmeticExpression {
            op: op,
            // TODO(malte): discards casts
            left: left.0,
            right: right.0,
            alias: match alias {
                None => None,
                Some(a) => Some(String::from(a)),
            },
        })
    ))
);

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn it_parses_arithmetic_expressions() {
        use column::FunctionExpression;
        use super::ArithmeticOperator::*;
        use super::ArithmeticBase::Scalar;
        use super::ArithmeticBase::Column as ABColumn;

        let lit_ae = [
            "5 + 42",
            "5+42",
            "5 * 42",
            "5 - 42",
            "5 / 42",
            "2 * 10 AS twenty ",
        ];

        // N.B. trailing space in "5 + foo " is required because `sql_identifier`'s keyword
        // detection requires a follow-up character (in practice, there always is one because we
        // use semicolon-terminated queries).
        let col_lit_ae = [
            "foo+5",
            "foo + 5",
            "5 + foo ",
            "foo * bar AS foobar",
            "MAX(foo)-3333",
        ];

        let expected_lit_ae = [
            ArithmeticExpression::new(Add, Scalar(5.into()), Scalar(42.into()), None),
            ArithmeticExpression::new(Add, Scalar(5.into()), Scalar(42.into()), None),
            ArithmeticExpression::new(Multiply, Scalar(5.into()), Scalar(42.into()), None),
            ArithmeticExpression::new(Subtract, Scalar(5.into()), Scalar(42.into()), None),
            ArithmeticExpression::new(Divide, Scalar(5.into()), Scalar(42.into()), None),
            ArithmeticExpression::new(
                Multiply,
                Scalar(2.into()),
                Scalar(10.into()),
                Some(String::from("twenty")),
            ),
        ];
        let expected_col_lit_ae = [
            ArithmeticExpression::new(Add, ABColumn("foo".into()), Scalar(5.into()), None),
            ArithmeticExpression::new(Add, ABColumn("foo".into()), Scalar(5.into()), None),
            ArithmeticExpression::new(Add, Scalar(5.into()), ABColumn("foo".into()), None),
            ArithmeticExpression::new(
                Multiply,
                ABColumn("foo".into()),
                ABColumn("bar".into()),
                Some(String::from("foobar")),
            ),
            ArithmeticExpression::new(
                Subtract,
                ABColumn(Column {
                    name: String::from("max(foo)"),
                    alias: None,
                    table: None,
                    function: Some(Box::new(FunctionExpression::Max("foo".into()))),
                }),
                Scalar(3333.into()),
                None,
            ),
        ];

        for (i, e) in lit_ae.iter().enumerate() {
            let res = arithmetic_expression(e.as_bytes());
            assert!(res.is_done());
            assert_eq!(res.unwrap().1, expected_lit_ae[i]);
        }

        for (i, e) in col_lit_ae.iter().enumerate() {
            let res = arithmetic_expression(e.as_bytes());
            assert!(res.is_done());
            assert_eq!(res.unwrap().1, expected_col_lit_ae[i]);
        }
    }

    #[test]
    fn it_displays_arithmetic_expressions() {
        use super::ArithmeticOperator::*;
        use super::ArithmeticBase::Scalar;
        use super::ArithmeticBase::Column as ABColumn;

        let expressions = [
            ArithmeticExpression::new(Add, ABColumn("foo".into()), Scalar(5.into()), None),
            ArithmeticExpression::new(Subtract, Scalar(5.into()), ABColumn("foo".into()), None),
            ArithmeticExpression::new(
                Multiply,
                ABColumn("foo".into()),
                ABColumn("bar".into()),
                None,
            ),
            ArithmeticExpression::new(Divide, Scalar(10.into()), Scalar(2.into()), None),
            ArithmeticExpression::new(
                Add,
                Scalar(10.into()),
                Scalar(2.into()),
                Some(String::from("bob")),
            ),
        ];

        let expected_strings = ["foo + 5", "5 - foo", "foo * bar", "10 / 2", "10 + 2 AS bob"];
        for (i, e) in expressions.iter().enumerate() {
            assert_eq!(expected_strings[i], format!("{}", e));
        }
    }

    #[test]
    fn it_parses_arithmetic_casts() {
        use super::ArithmeticOperator::*;
        use super::ArithmeticBase::Scalar;
        use super::ArithmeticBase::Column as ABColumn;

        let exprs = [
            "CAST(foo AS signed int) + CAST(bar AS signed int) ",
            "CAST(5 AS bigint) - foo ",
        ];

        // XXX(malte): currently discards the cast and type information!
        let expected = [
            ArithmeticExpression::new(Add, ABColumn("foo".into()), ABColumn("bar".into()), None),
            ArithmeticExpression::new(Subtract, Scalar(5.into()), ABColumn("foo".into()), None),
        ];

        for (i, e) in exprs.iter().enumerate() {
            let res = arithmetic_expression(e.as_bytes());
            assert!(res.is_done(), "{} failed to parse", e);
            assert_eq!(res.unwrap().1, expected[i]);
        }
    }

}
<|endoftext|>


<commit_before>#[macro_use] extern crate log;

extern crate aqua;
extern crate diesel;
extern crate dotenv;
extern crate env_logger;
extern crate rusqlite;

use aqua::models::{self, Entry, NewEntry, NewEntryTag, NewTag};
use aqua::schema;
use diesel::Connection as DieselConnection;
use diesel::prelude::*;
use diesel::pg::PgConnection;
use std::collections::HashMap;
use std::env;
use rusqlite::Connection;

static CLIENT_DB_NAME:  &'static str = "C:\\Hydrus Network\\db\\client.db";
static MAPPING_DB_NAME: &'static str = "C:\\Hydrus Network\\db\\client.mappings.db";
static MASTER_DB_NAME:  &'static str = "C:\\Hydrus Network\\db\\client.master.db";

#[derive(Clone, Debug)]
struct ClientHash {
    pub hash_id: i64,
    pub md5:     Vec<u8>,
    pub sha1:    Vec<u8>,
    pub sha256:  Option<Vec<u8>>,
    pub sha512:  Vec<u8>,
}

#[derive(Clone, Debug)]
struct Mapping {
    pub ns_id:   i64,
    pub tag_id:  i64,
    pub hash_id: i64,
}

#[derive(Clone, Debug)]
struct Tag {
    pub tag_id: i64,
    pub ns_id:  i64,
    pub schema: String,
    pub name: String,
}

pub fn establish_connection() -> PgConnection {
    let database_url = env::var("DATABASE_URL")
        .expect("DATABASE_URL not set");

    PgConnection::establish(&database_url)
        .expect(&format!("Error connecting to {}", database_url))
}

fn main() {
    dotenv::dotenv().expect("must provide .env file, see README (TODO: haha jk)");
    env_logger::init().expect("could not initialize console logging");

    let client_db  = Connection::open(CLIENT_DB_NAME).unwrap();
    let mapping_db = Connection::open(MAPPING_DB_NAME).unwrap();
    let master_db  = Connection::open(MASTER_DB_NAME).unwrap();

    info!("connected to hydrus database");

    let mut select_mappings = mapping_db
        .prepare("SELECT namespace_id, tag_id, hash_id FROM current_mappings_4")
        .unwrap();

    let mappings = select_mappings.query_map(&[], |row| {
        Mapping {
            ns_id:   row.get(0),
            tag_id:  row.get(1),
            hash_id: row.get(2),
        }
    });

    let mappings = mappings.unwrap()
      .map(|mb_row| mb_row.unwrap())
      .collect::<Vec<Mapping>>();

    info!("got {} mappings", mappings.len());

    let mut select_tags = master_db
        .prepare("SELECT tag_id, tag FROM tags WHERE tag_id = ?")
        .unwrap();

    let mut select_ns = master_db
        .prepare("SELECT namespace_id, namespace FROM namespaces WHERE namespace_id = ?")
        .unwrap();

    let mut tags = HashMap::new();
    for mapping in &mappings {
        let ns: String = select_ns.query_row(&[&mapping.ns_id], |row| {
            row.get(1)
        }).unwrap();

        let tag = select_tags.query_row(&[&mapping.tag_id], |row| {
            Tag { tag_id: row.get(0), ns_id: mapping.ns_id, schema: ns.clone(), name: row.get(1) }
        }).unwrap();

        tags.insert((tag.tag_id, ns.clone()), tag);
    }

    info!("got {} tags", tags.len());
    // let mut tags = vec![];

    // load the client hashes
    let mut select_hashes = client_db
        .prepare("SELECT hash_id, md5, sha1, sha512 FROM local_hashes")
        .unwrap();

    let hashes = select_hashes.query_map(&[], |row| {
        ClientHash {
            hash_id: row.get(0),
            md5:     row.get(1),
            sha1:    row.get(2),
            sha512:  row.get(3),
            sha256:  None, // filled in by master db
        }
    });

    let hash_records = hashes
        .unwrap()
        .map(|mb_row| mb_row.unwrap())
        .collect::<Vec<ClientHash>>();

    let mut hashes = HashMap::new();
    for record in &hash_records {
        hashes.insert(record.hash_id, record.clone());
    }

    info!("loaded {} client files", hashes.len());

    // load remote hashes
    let mut select_hash = master_db
        .prepare("SELECT hash_id, hash FROM hashes WHERE hash_id = ?")
        .unwrap();

    for record in &hash_records {
        use std::collections::hash_map::Entry;

        let rhash = select_hash.query_row(&[&record.hash_id], |row| {
            row.get(1)
        }).unwrap();

        if let Entry::Occupied(mut o) = hashes.entry(record.hash_id) {
            o.get_mut().sha256 = Some(rhash);
        } else { panic!("unknown hash"); }
    }

    info!("loaded {} hashes", hashes.len());

    // begin population of our database
    let pg_conn = establish_connection();

    // map hydrus IDs => our IDs for later mapping restoration
    let mut aqua_entry_ids = HashMap::new();
    let mut aqua_tag_ids = HashMap::new();

    // load the entries
    for (hash_id, hash) in &hashes {
        let hash: Vec<String> = hash.sha256
            .as_ref().unwrap()
            .iter()
            .map(|b| format!("{:02x}", b))
            .collect();

        let hash: String = hash.join("");

        let entry = NewEntry { hash: &hash, mime: None };
        let entry: Entry = diesel::insert(&entry)
            .into(schema::entries::table)
            .get_result(&pg_conn).ok().unwrap();
        aqua_entry_ids.insert(hash_id, entry.id);
    }

    // load tags
    for (_tag_id, otag) in &tags {
        // info!("tag({:?}) => {:?}", tag_id, tag);
        let new_tag = NewTag { name: &otag.name, schema: Some(&otag.schema) };
        let tag: models::Tag = diesel::insert(&new_tag)
            .into(schema::tags::table)
            .get_result(&pg_conn).ok().unwrap();
        aqua_tag_ids.insert((otag.ns_id, otag.tag_id), tag.id);
    }

    // load mappings
    for mapping in &mappings {
        let nentry = aqua_entry_ids.get(&mapping.hash_id).unwrap();
        let ntag   = aqua_tag_ids.get(&(mapping.ns_id, mapping.tag_id)).unwrap();
        let link   = NewEntryTag { tag_id: *ntag, entry_id: *nentry };

        let _entry: models::EntryTag = diesel::insert(&link)
            .into(schema::entries_tags::table)
            .get_result(&pg_conn).ok().unwrap();
    }
}
<commit_msg>(lint) run cargo clippy<commit_after>#[macro_use] extern crate log;

extern crate aqua;
extern crate diesel;
extern crate dotenv;
extern crate env_logger;
extern crate rusqlite;

use aqua::models::{self, Entry, NewEntry, NewEntryTag, NewTag};
use aqua::schema;
use diesel::Connection as DieselConnection;
use diesel::prelude::*;
use diesel::pg::PgConnection;
use std::collections::HashMap;
use std::env;
use rusqlite::Connection;

static CLIENT_DB_NAME:  &'static str = "C:\\Hydrus Network\\db\\client.db";
static MAPPING_DB_NAME: &'static str = "C:\\Hydrus Network\\db\\client.mappings.db";
static MASTER_DB_NAME:  &'static str = "C:\\Hydrus Network\\db\\client.master.db";

#[derive(Clone, Debug)]
struct ClientHash {
    pub hash_id: i64,
    pub md5:     Vec<u8>,
    pub sha1:    Vec<u8>,
    pub sha256:  Option<Vec<u8>>,
    pub sha512:  Vec<u8>,
}

#[derive(Clone, Debug)]
struct Mapping {
    pub ns_id:   i64,
    pub tag_id:  i64,
    pub hash_id: i64,
}

#[derive(Clone, Debug)]
struct Tag {
    pub tag_id: i64,
    pub ns_id:  i64,
    pub schema: String,
    pub name: String,
}

pub fn establish_connection() -> PgConnection {
    let database_url = env::var("DATABASE_URL")
        .expect("DATABASE_URL not set");

    PgConnection::establish(&database_url)
        .expect(&format!("Error connecting to {}", database_url))
}

#[allow(doc_markdown)] // nuu~ clippy, SQLite is not a type v(T.T)v
/// This application is used to import tags from SQLite databases generated
/// by the [Hydrus Network][hgit] image client.
///
/// This application puts most data in separate databases, hindering our
/// ability to have SQLite perform joins. So most of the time is spent
/// building application side indexes.
///
/// First we read the `current_mappings` table to get a list of mappings the
/// end-user established. Then we look up the tags & digests for each mapping.
/// One important transformation is that we consider tags in different namespaces
/// to be *distinct entries*, instead of having a separate table for namespaces.
///
/// After this is done we go through the mappings and create appropriate entries
/// into the `aqua` database; at the moment this panicks if it inserts a duplicate.
///
/// [hgit]: https://hydrusnetwork.github.io/hydrus/
fn main() {
    dotenv::dotenv().expect("must provide .env file, see README (TODO: haha jk)");
    env_logger::init().expect("could not initialize console logging");

    let client_db  = Connection::open(CLIENT_DB_NAME).unwrap();
    let mapping_db = Connection::open(MAPPING_DB_NAME).unwrap();
    let master_db  = Connection::open(MASTER_DB_NAME).unwrap();

    info!("connected to hydrus database");

    let mut select_mappings = mapping_db
        .prepare("SELECT namespace_id, tag_id, hash_id FROM current_mappings_4")
        .unwrap();

    let mappings = select_mappings.query_map(&[], |row| {
        Mapping {
            ns_id:   row.get(0),
            tag_id:  row.get(1),
            hash_id: row.get(2),
        }
    });

    let mappings = mappings.unwrap()
      .map(|mb_row| mb_row.unwrap())
      .collect::<Vec<Mapping>>();

    info!("got {} mappings", mappings.len());

    let mut select_tags = master_db
        .prepare("SELECT tag_id, tag FROM tags WHERE tag_id = ?")
        .unwrap();

    let mut select_ns = master_db
        .prepare("SELECT namespace_id, namespace FROM namespaces WHERE namespace_id = ?")
        .unwrap();

    let mut tags = HashMap::new();
    for mapping in &mappings {
        let ns: String = select_ns.query_row(&[&mapping.ns_id], |row| {
            row.get(1)
        }).unwrap();

        let tag = select_tags.query_row(&[&mapping.tag_id], |row| {
            Tag { tag_id: row.get(0), ns_id: mapping.ns_id, schema: ns.clone(), name: row.get(1) }
        }).unwrap();

        tags.insert((tag.tag_id, ns.clone()), tag);
    }

    info!("got {} tags", tags.len());
    // let mut tags = vec![];

    // load the client hashes
    let mut select_hashes = client_db
        .prepare("SELECT hash_id, md5, sha1, sha512 FROM local_hashes")
        .unwrap();

    let hashes = select_hashes.query_map(&[], |row| {
        ClientHash {
            hash_id: row.get(0),
            md5:     row.get(1),
            sha1:    row.get(2),
            sha512:  row.get(3),
            sha256:  None, // filled in by master db
        }
    });

    let hash_records = hashes
        .unwrap()
        .map(|mb_row| mb_row.unwrap())
        .collect::<Vec<ClientHash>>();

    let mut hashes = HashMap::new();
    for record in &hash_records {
        hashes.insert(record.hash_id, record.clone());
    }

    info!("loaded {} client files", hashes.len());

    // load remote hashes
    let mut select_hash = master_db
        .prepare("SELECT hash_id, hash FROM hashes WHERE hash_id = ?")
        .unwrap();

    for record in &hash_records {
        use std::collections::hash_map::Entry;

        let rhash = select_hash.query_row(&[&record.hash_id], |row| {
            row.get(1)
        }).unwrap();

        if let Entry::Occupied(mut o) = hashes.entry(record.hash_id) {
            o.get_mut().sha256 = Some(rhash);
        } else { panic!("unknown hash"); }
    }

    info!("loaded {} hashes", hashes.len());

    // begin population of our database
    let pg_conn = establish_connection();

    // map hydrus IDs => our IDs for later mapping restoration
    let mut aqua_entry_ids = HashMap::new();
    let mut aqua_tag_ids = HashMap::new();

    // load the entries
    for (hash_id, hash) in &hashes {
        let hash: Vec<String> = hash.sha256
            .as_ref().unwrap()
            .iter()
            .map(|b| format!("{:02x}", b))
            .collect();

        let hash: String = hash.join("");

        let entry = NewEntry { hash: &hash, mime: None };
        let entry: Entry = diesel::insert(&entry)
            .into(schema::entries::table)
            .get_result(&pg_conn).ok().unwrap();
        aqua_entry_ids.insert(hash_id, entry.id);
    }

    // load tags
    for otag in tags.values() {
        // info!("tag({:?}) => {:?}", tag_id, tag);
        let new_tag = NewTag { name: &otag.name, schema: Some(&otag.schema) };
        let tag: models::Tag = diesel::insert(&new_tag)
            .into(schema::tags::table)
            .get_result(&pg_conn).ok().unwrap();
        aqua_tag_ids.insert((otag.ns_id, otag.tag_id), tag.id);
    }

    // load mappings
    for mapping in &mappings {
        let nentry = aqua_entry_ids[&mapping.hash_id];
        let ntag   = aqua_tag_ids[&(mapping.ns_id, mapping.tag_id)];
        let link   = NewEntryTag { tag_id: ntag, entry_id: nentry };

        let _entry: models::EntryTag = diesel::insert(&link)
            .into(schema::entries_tags::table)
            .get_result(&pg_conn).ok().unwrap();
    }
}
<|endoftext|>


<commit_before>
use std::cell::Cell;
use std::mem;
use std::ptr;
use std::sync::Arc;
use std::sync::atomic::{AtomicPtr, AtomicUsize, fence};
use std::sync::atomic::Ordering::{Relaxed, Acquire, Release, AcqRel};

use alloc;
use countedindex::{CountedIndex, get_valid_wrap, INITIAL_QUEUE_FLAG};
use maybe_acquire::{maybe_acquire_fence, MAYBE_ACQUIRE};

use read_cursor::{ReadCursor, Reader};

#[derive(Clone, Copy)]
enum QueueState {
    Single,
    Multi,
}

struct QueueEntry<T> {
    val: T,
    wraps: AtomicUsize,
}

/// A bounded queue that supports multiple reader and writers
/// and supports effecient methods for single consumers and producers
#[repr(C)]
struct MultiQueue<T> {
    d1: [u8; 64],

    // Writer data
    head: CountedIndex,
    tail_cache: AtomicUsize,
    writers: AtomicUsize,
    d2: [u8; 64],

    // Shared Data
    // The data and the wraps flag are in the same location
    // to reduce the # of distinct cache lines read when getting an item
    // The tail itself is rarely modified, making it a suitable candidate
    // to be in the shared space
    tail: ReadCursor,
    data: *mut QueueEntry<T>,
    capacity: isize,
    d3: [u8; 64],
}

pub struct MultiWriter<T> {
    queue: Arc<MultiQueue<T>>,
    state: Cell<QueueState>,
}

pub struct MultiReader<T> {
    queue: Arc<MultiQueue<T>>,
    reader: AtomicPtr<Reader>,
}

impl<T> MultiQueue<T> {
    pub fn new(_capacity: u32) -> (MultiWriter<T>, MultiReader<T>) {
        let capacity = get_valid_wrap(_capacity);
        let queuedat = alloc::allocate(capacity as usize);
        unsafe {
            for i in 0..capacity as isize {
                let elem: &QueueEntry<T> = &*queuedat.offset(i);
                elem.wraps.store(INITIAL_QUEUE_FLAG, Relaxed);
            }
        }

        let (cursor, reader) = ReadCursor::new(capacity);

        let queue = MultiQueue {
            d1: unsafe { mem::uninitialized() },

            head: CountedIndex::new(capacity),
            tail_cache: AtomicUsize::new(0),
            writers: AtomicUsize::new(1),
            d2: unsafe { mem::uninitialized() },

            tail: cursor,
            data: queuedat,
            capacity: capacity as isize,

            d3: unsafe { mem::uninitialized() },
        };

        let qarc = Arc::new(queue);

        let mwriter = MultiWriter {
            queue: qarc.clone(),
            state: Cell::new(QueueState::Single),
        };

        let mreader = MultiReader {
            queue: qarc,
            reader: reader,
        };

        (mwriter, mreader)
    }

    pub fn push_multi(&self, val: T) -> Result<(), T> {
        let mut transaction = self.head.load_transaction(Relaxed);

        // This tries to ensure the tail fetch metadata is always in the cache
        // The effect of this is that whenever one has to find the minimum tail,
        // the data about the loop is in-cache so that whole loop executes deep in
        // an out-of-order engine while the branch predictor predicts there is more space
        // and continues on pushing
        self.tail.prefetch_metadata();
        unsafe {
            loop {
                let tail_cache = self.tail_cache.load(Acquire);
                let (chead, wrap_valid_tag) = transaction.get();
                let write_cell = &mut *self.data.offset(chead);
                if transaction.matches_previous(tail_cache) {
                    let new_tail = self.reload_tail_multi(tail_cache, wrap_valid_tag);
                    if transaction.matches_previous(new_tail) {
                        return Err(val);
                    }
                }
                // This isize conversion is valid since indices are 30 bits
                match transaction.commit(1, Relaxed) {
                    Some(new_transaction) => transaction = new_transaction,
                    None => {
                        ptr::write(&mut write_cell.val, val);
                        write_cell.wraps.store(wrap_valid_tag, Release);
                        return Ok(());
                    }
                }
            }
        }
    }

    pub fn push_single(&self, val: T) -> Result<(), T> {
        let transaction = self.head.load_transaction(Relaxed);
        let (chead, wrap_valid_tag) = transaction.get();
        self.tail.prefetch_metadata(); // See push_multi on this
        unsafe {
        let write_cell = &mut *self.data.offset(chead);
            let tail_cache = self.tail_cache.load(Relaxed);
            if transaction.matches_previous(tail_cache) {
                let new_tail = self.reload_tail_single(wrap_valid_tag);
                if transaction.matches_previous(new_tail) {
                    return Err(val);
                }
            }
            ptr::write(&mut write_cell.val, val);
            write_cell.wraps.store(wrap_valid_tag, Release);
            transaction.commit_direct(1, Relaxed);
            Ok(())
        }

        // Might consider letting the queue update the tail cache here preemptively
        // so it doesn't waste time before sending a message to do so
    }

    pub fn pop(&self, reader: &Reader) -> Option<T> {
        let mut ctail_attempt = reader.load_attempt(Relaxed);
        unsafe {
            loop {
                let (ctail, wrap_valid_tag) = ctail_attempt.get();
                let read_cell = &*self.data.offset(ctail);
                if read_cell.wraps.load(MAYBE_ACQUIRE) != wrap_valid_tag {
                    return None;
                }
                maybe_acquire_fence();
                let rval = ptr::read(&read_cell.val);
                match ctail_attempt.commit_attempt(1, Release) {
                    Some(new_attempt) => ctail_attempt = new_attempt,
                    None => return Some(rval),
                }
            }
        }
    }

    fn reload_tail_multi(&self, count: usize, tail_cache: usize) -> usize {
        // This shows how far behind from head the reader is
        if let Some(max_diff_from_head) = self.tail.get_max_diff(count) {
            let current_tail = CountedIndex::get_previous(count, max_diff_from_head);
            match self.tail_cache.compare_exchange(tail_cache, current_tail, AcqRel, Relaxed) {
                Ok(val) => val,
                Err(val) => val,
            }
        } else {
            self.tail_cache.load(Acquire)
        }
    }

    fn reload_tail_single(&self, count: usize) -> usize {
        let max_diff_from_head = self.tail
            .get_max_diff(count)
            .expect("The write head got ran over by consumers in single writer mode. This \
                     process is borked!");
        let current_tail = CountedIndex::get_previous(count, max_diff_from_head);
        self.tail_cache.store(current_tail, Relaxed);
        current_tail
    }
}

impl<T> MultiWriter<T> {
    #[inline(always)]
    pub fn push(&self, val: T) -> Result<(), T> {
        match self.state.get() {
            QueueState::Single => self.queue.push_single(val),
            QueueState::Multi => {
                if self.queue.writers.load(Relaxed) == 1 {
                    fence(Acquire);
                    self.state.set(QueueState::Single);
                    self.queue.push_single(val)
                } else {
                    self.queue.push_multi(val)
                }
            }
        }
    }
}

impl<T> MultiReader<T> {
    #[inline(always)]
    pub fn pop(&self) -> Option<T> {
        unsafe { self.queue.pop(&*self.reader.load(Relaxed)) }
    }

    pub fn add_reader(&self) -> MultiReader<T> {
        MultiReader {
            queue: self.queue.clone(),
            reader: unsafe { self.queue.tail.add_reader(&*self.reader.load(Relaxed)) },
        }
    }
}

impl<T> Clone for MultiWriter<T> {
    fn clone(&self) -> MultiWriter<T> {
        self.state.set(QueueState::Multi);
        let rval = MultiWriter {
            queue: self.queue.clone(),
            state: Cell::new(QueueState::Multi),
        };
        self.queue.writers.fetch_add(1, Release);
        rval
    }
}

impl<T> Clone for MultiReader<T> {
    fn clone(&self) -> MultiReader<T> {
        let reader = self.reader.load(Relaxed);
        let rval = MultiReader {
            queue: self.queue.clone(),
            reader: AtomicPtr::new(reader),
        };
        unsafe {
            (*reader).dup_consumer();
        }
        rval
    }
}

impl<T> Drop for MultiWriter<T> {
    fn drop(&mut self) {
        self.queue.writers.fetch_sub(1, Release);
    }
}

impl<T> Drop for MultiReader<T> {
    fn drop(&mut self) {
        unsafe { (*self.reader.load(Relaxed)).remove_consumer() }
    }
}

unsafe impl<T> Sync for MultiQueue<T> {}
unsafe impl<T> Send for MultiQueue<T> {}
unsafe impl<T> Send for MultiWriter<T> {}
unsafe impl<T> Send for MultiReader<T> {}

pub fn multiqueue<T>(capacity: u32) -> (MultiWriter<T>, MultiReader<T>) {
    MultiQueue::new(capacity)
}

#[cfg(test)]
mod test {

    use super::*;
    use super::MultiQueue;

    extern crate crossbeam;
    use self::crossbeam::scope;

    use std::sync::atomic::Ordering::*;

    use std::sync::Barrier;

    #[test]
    fn build_queue() {
        let _ = MultiQueue::<usize>::new(10);
    }

    #[test]
    fn push_pop_test() {
        let (writer, reader) = MultiQueue::<usize>::new(1);
        for _ in 0..100 {
            assert!(reader.pop().is_none());
            writer.push(1 as usize).expect("Push should succeed");
            assert!(writer.push(1).is_err());
            assert_eq!(1, reader.pop().unwrap());
        }
    }

    fn spsc_broadcast(receivers: usize) {
        let (writer, reader) = MultiQueue::<usize>::new(10);
        let myb = Barrier::new(receivers + 1);
        let bref = &myb;
        let num_loop = 1000000;
        scope(|scope| {
            let cur_writer = writer.clone();
            scope.spawn(move || {
                bref.wait();
                for i in 0..num_loop {
                    loop {
                        if cur_writer.push(i).is_ok() {
                            break;
                        }
                    }
                }
            });
            for i in 0..(receivers - 1) {
                let this_reader = reader.add_reader();
                scope.spawn(move || {
                    bref.wait();
                    for i in 0..num_loop {
                        loop {
                            if let Some(val) = this_reader.pop() {
                                assert_eq!(i, val);
                                break;
                            }
                        }
                    }
                });
            }
            bref.wait();
            'outer: for i in 0..num_loop {
                loop {
                    if let Some(val) = reader.pop() {
                        assert_eq!(i, val);
                        break;
                    }
                }
            }
        });
        assert!(reader.pop().is_none());
    }

    #[test]
    fn test_spsc() {
        spsc_broadcast(1);
    }

    #[test]
    fn test_spsc_broadcast() {
        spsc_broadcast(3);
    }

}
<commit_msg>Saving<commit_after>
use std::cell::Cell;
use std::mem;
use std::ptr;
use std::sync::Arc;
use std::sync::atomic::{AtomicPtr, AtomicUsize, fence};
use std::sync::atomic::Ordering::{Relaxed, Acquire, Release, AcqRel};

use alloc;
use countedindex::{CountedIndex, get_valid_wrap, INITIAL_QUEUE_FLAG};
use maybe_acquire::{maybe_acquire_fence, MAYBE_ACQUIRE};

use read_cursor::{ReadCursor, Reader};

#[derive(Clone, Copy)]
enum QueueState {
    Single,
    Multi,
}

struct QueueEntry<T> {
    val: T,
    wraps: AtomicUsize,
}

/// A bounded queue that supports multiple reader and writers
/// and supports effecient methods for single consumers and producers
#[repr(C)]
struct MultiQueue<T> {
    d1: [u8; 64],

    // Writer data
    head: CountedIndex,
    tail_cache: AtomicUsize,
    writers: AtomicUsize,
    d2: [u8; 64],

    // Shared Data
    // The data and the wraps flag are in the same location
    // to reduce the # of distinct cache lines read when getting an item
    // The tail itself is rarely modified, making it a suitable candidate
    // to be in the shared space
    tail: ReadCursor,
    data: *mut QueueEntry<T>,
    capacity: isize,
    d3: [u8; 64],
}

pub struct MultiWriter<T> {
    queue: Arc<MultiQueue<T>>,
    state: Cell<QueueState>,
}

pub struct MultiReader<T> {
    queue: Arc<MultiQueue<T>>,
    reader: AtomicPtr<Reader>,
}

impl<T> MultiQueue<T> {
    pub fn new(_capacity: u32) -> (MultiWriter<T>, MultiReader<T>) {
        let capacity = get_valid_wrap(_capacity);
        let queuedat = alloc::allocate(capacity as usize);
        unsafe {
            for i in 0..capacity as isize {
                let elem: &QueueEntry<T> = &*queuedat.offset(i);
                elem.wraps.store(INITIAL_QUEUE_FLAG, Relaxed);
            }
        }

        let (cursor, reader) = ReadCursor::new(capacity);

        let queue = MultiQueue {
            d1: unsafe { mem::uninitialized() },

            head: CountedIndex::new(capacity),
            tail_cache: AtomicUsize::new(0),
            writers: AtomicUsize::new(1),
            d2: unsafe { mem::uninitialized() },

            tail: cursor,
            data: queuedat,
            capacity: capacity as isize,

            d3: unsafe { mem::uninitialized() },
        };

        let qarc = Arc::new(queue);

        let mwriter = MultiWriter {
            queue: qarc.clone(),
            state: Cell::new(QueueState::Single),
        };

        let mreader = MultiReader {
            queue: qarc,
            reader: reader,
        };

        (mwriter, mreader)
    }

    pub fn push_multi(&self, val: T) -> Result<(), T> {
        let mut transaction = self.head.load_transaction(Relaxed);

        // This tries to ensure the tail fetch metadata is always in the cache
        // The effect of this is that whenever one has to find the minimum tail,
        // the data about the loop is in-cache so that whole loop executes deep in
        // an out-of-order engine while the branch predictor predicts there is more space
        // and continues on pushing
        self.tail.prefetch_metadata();
        unsafe {
            loop {
                let tail_cache = self.tail_cache.load(Acquire);
                let (chead, wrap_valid_tag) = transaction.get();
                let write_cell = &mut *self.data.offset(chead);
                if transaction.matches_previous(tail_cache) {
                    let new_tail = self.reload_tail_multi(tail_cache, wrap_valid_tag);
                    if transaction.matches_previous(new_tail) {
                        return Err(val);
                    }
                }
                // This isize conversion is valid since indices are 30 bits
                match transaction.commit(1, Relaxed) {
                    Some(new_transaction) => transaction = new_transaction,
                    None => {
                        ptr::write(&mut write_cell.val, val);
                        write_cell.wraps.store(wrap_valid_tag, Release);
                        return Ok(());
                    }
                }
            }
        }
    }

    pub fn push_single(&self, val: T) -> Result<(), T> {
        let transaction = self.head.load_transaction(Relaxed);
        let (chead, wrap_valid_tag) = transaction.get();
        self.tail.prefetch_metadata(); // See push_multi on this
        unsafe {
        let write_cell = &mut *self.data.offset(chead);
            let tail_cache = self.tail_cache.load(Relaxed);
            if transaction.matches_previous(tail_cache) {
                let new_tail = self.reload_tail_single(wrap_valid_tag);
                if transaction.matches_previous(new_tail) {
                    return Err(val);
                }
            }
            ptr::write(&mut write_cell.val, val);
            write_cell.wraps.store(wrap_valid_tag, Release);
            transaction.commit_direct(1, Relaxed);
            Ok(())
        }

        // Might consider letting the queue update the tail cache here preemptively
        // so it doesn't waste time before sending a message to do so
    }

    pub fn pop(&self, reader: &Reader) -> Option<T> {
        let mut ctail_attempt = reader.load_attempt(Relaxed);
        unsafe {
            loop {
                let (ctail, wrap_valid_tag) = ctail_attempt.get();
                let read_cell = &*self.data.offset(ctail);
                if read_cell.wraps.load(MAYBE_ACQUIRE) != wrap_valid_tag {
                    return None;
                }
                maybe_acquire_fence();
                let rval = ptr::read(&read_cell.val);
                match ctail_attempt.commit_attempt(1, Release) {
                    Some(new_attempt) => ctail_attempt = new_attempt,
                    None => return Some(rval),
                }
            }
        }
    }

    fn reload_tail_multi(&self, count: usize, tail_cache: usize) -> usize {
        // This shows how far behind from head the reader is
        if let Some(max_diff_from_head) = self.tail.get_max_diff(count) {
            let current_tail = CountedIndex::get_previous(count, max_diff_from_head);
            match self.tail_cache.compare_exchange(tail_cache, current_tail, AcqRel, Relaxed) {
                Ok(val) => val,
                Err(val) => val,
            }
        } else {
            self.tail_cache.load(Acquire)
        }
    }

    fn reload_tail_single(&self, count: usize) -> usize {
        let max_diff_from_head = self.tail
            .get_max_diff(count)
            .expect("The write head got ran over by consumers in single writer mode. This \
                     process is borked!");
        let current_tail = CountedIndex::get_previous(count, max_diff_from_head);
        self.tail_cache.store(current_tail, Relaxed);
        current_tail
    }
}

impl<T> MultiWriter<T> {
    #[inline(always)]
    pub fn push(&self, val: T) -> Result<(), T> {
        match self.state.get() {
            QueueState::Single => self.queue.push_single(val),
            QueueState::Multi => {
                if self.queue.writers.load(Relaxed) == 1 {
                    fence(Acquire);
                    self.state.set(QueueState::Single);
                    self.queue.push_single(val)
                } else {
                    self.queue.push_multi(val)
                }
            }
        }
    }
}

impl<T> MultiReader<T> {
    #[inline(always)]
    pub fn pop(&self) -> Option<T> {
        unsafe { self.queue.pop(&*self.reader.load(Relaxed)) }
    }

    pub fn add_reader(&self) -> MultiReader<T> {
        MultiReader {
            queue: self.queue.clone(),
            reader: unsafe { self.queue.tail.add_reader(&*self.reader.load(Relaxed)) },
        }
    }
}

impl<T> Clone for MultiWriter<T> {
    fn clone(&self) -> MultiWriter<T> {
        self.state.set(QueueState::Multi);
        let rval = MultiWriter {
            queue: self.queue.clone(),
            state: Cell::new(QueueState::Multi),
        };
        self.queue.writers.fetch_add(1, Release);
        rval
    }
}

impl<T> Clone for MultiReader<T> {
    fn clone(&self) -> MultiReader<T> {
        let reader = self.reader.load(Relaxed);
        let rval = MultiReader {
            queue: self.queue.clone(),
            reader: AtomicPtr::new(reader),
        };
        unsafe {
            (*reader).dup_consumer();
        }
        rval
    }
}

impl<T> Drop for MultiWriter<T> {
    fn drop(&mut self) {
        self.queue.writers.fetch_sub(1, Release);
    }
}

impl<T> Drop for MultiReader<T> {
    fn drop(&mut self) {
        unsafe { (*self.reader.load(Relaxed)).remove_consumer() }
    }
}

unsafe impl<T> Sync for MultiQueue<T> {}
unsafe impl<T> Send for MultiQueue<T> {}
unsafe impl<T> Send for MultiWriter<T> {}
unsafe impl<T> Send for MultiReader<T> {}

pub fn multiqueue<T>(capacity: u32) -> (MultiWriter<T>, MultiReader<T>) {
    MultiQueue::new(capacity)
}

#[cfg(test)]
mod test {

    use super::*;
    use super::MultiQueue;

    extern crate crossbeam;
    use self::crossbeam::scope;

    use std::sync::atomic::Ordering::*;

    use std::sync::Barrier;

    #[test]
    fn build_queue() {
        let _ = MultiQueue::<usize>::new(10);
    }

    #[test]
    fn push_pop_test() {
        let (writer, reader) = MultiQueue::<usize>::new(1);
        for _ in 0..100 {
            assert!(reader.pop().is_none());
            writer.push(1 as usize).expect("Push should succeed");
            assert!(writer.push(1).is_err());
            assert_eq!(1, reader.pop().unwrap());
        }
    }

    fn spsc_broadcast(receivers: usize) {
        let (writer, reader) = MultiQueue::<usize>::new(10);
        let myb = Barrier::new(receivers + 1);
        let bref = &myb;
        let num_loop = 1000000;
        scope(|scope| {
            scope.spawn(move || {
                bref.wait();
                for i in 0..num_loop {
                    loop {
                        if writer.push(i).is_ok() {
                            break;
                        }
                    }
                }
            });
            for i in 0..(receivers - 1) {
                let this_reader = reader.add_reader();
                scope.spawn(move || {
                    bref.wait();
                    for i in 0..num_loop {
                        loop {
                            if let Some(val) = this_reader.pop() {
                                assert_eq!(i, val);
                                break;
                            }
                        }
                    }
                });
            }
            bref.wait();
            'outer: for i in 0..num_loop {
                loop {
                    if let Some(val) = reader.pop() {
                        assert_eq!(i, val);
                        break;
                    }
                }
            }
        });
        assert!(reader.pop().is_none());
    }

    #[test]
    fn test_spsc() {
        spsc_broadcast(1);
    }

    #[test]
    fn test_spsc_broadcast() {
        spsc_broadcast(3);
    }

}
<|endoftext|>


<commit_before>#[macro_use]
extern crate gfx;
extern crate gfx_window_glutin;
extern crate glutin;

use gfx::traits::{Stream, ToIndexSlice, ToSlice, FactoryExt};

gfx_vertex!( Vertex {
    a_Pos@ pos: [f32; 3],
    a_Color@ color: [f32; 3],
});

pub fn main() {
    let (mut stream, mut device, mut factory) = gfx_window_glutin::init(
        glutin::Window::new().unwrap());

    stream.out.window.set_title("GFX App");

    'main: loop {
        // Quit when esc key is pressed
        for event in stream.out.window.poll_events() {
            match event {
                glutin::Event::KeyboardInput(_, _, Some(glutin::VirtualKeyCode::Escape)) => break 'main,
                glutin::Event::Closed => break 'main,
                _ => {},
            }
        }

        stream.clear(gfx::ClearData {
            color: [0.5, 0.5, 0.5, 1.0],
            depth: 1.0,
            stencil: 0
        });

        stream.present(&mut device);
    }
}
<commit_msg>Finished test implemntatino of gfx_app<commit_after>#[macro_use]
extern crate gfx;
extern crate gfx_window_glutin;
extern crate glutin;

use std::ffi::CString;
use gfx::traits::{Stream, ToIndexSlice, ToSlice, FactoryExt};

gfx_vertex!( Vertex {
    position@ pos: [f32; 3],
    color@ color: [f32; 3],
});

// Shader sources
static VERTEX_SOURCE: &'static str =
   "#version 130\n\

    in vec3 position;\n\
    in vec3 color;\n\
    out vec4 v_color;\n\

    void main() {\n\
       v_color = vec4(color, 1.0);\n\
       gl_Position = vec4(position, 1.0);\n\
    }";

static FRAGMENT_SOURCE: &'static str =
   "#version 130\n\

    in vec4 v_color;\n\
    out vec4 color;\n\

    void main() {\n\
       color = v_color;\n\
    }";

pub fn main() {
    let (mut stream, mut device, mut factory) = gfx_window_glutin::init(
        glutin::Window::new().unwrap());

    stream.out.window.set_title("GFX App");

    let vertex_data = [
        Vertex { pos: [ -0.5, -0.5, 0.0 ], color: [1.0, 0.0, 0.0] },
        Vertex { pos: [  0.5, -0.5, 0.0 ], color: [0.0, 1.0, 0.0] },
        Vertex { pos: [  0.0,  0.5, 0.0 ], color: [0.0, 0.0, 1.0] },
    ];
    let mesh = factory.create_mesh(&vertex_data);
    let slice = mesh.to_slice(gfx::PrimitiveType::TriangleList);

    let program = {
        let vs = gfx::ShaderSource {
            glsl_120: Some(VERTEX_SOURCE.as_bytes()),
            glsl_150: Some(VERTEX_SOURCE.as_bytes()),
            .. gfx::ShaderSource::empty()
        };
        let fs = gfx::ShaderSource {
            glsl_120: Some(FRAGMENT_SOURCE.as_bytes()),
            glsl_150: Some(FRAGMENT_SOURCE.as_bytes()),
            .. gfx::ShaderSource::empty()
        };
        factory.link_program_source(vs, fs).unwrap()
    };
    let state = gfx::DrawState::new();

    'main: loop {
        // Quit when esc key is pressed
        for event in stream.out.window.poll_events() {
            match event {
                glutin::Event::KeyboardInput(_, _, Some(glutin::VirtualKeyCode::Escape)) => break 'main,
                glutin::Event::Closed => break 'main,
                _ => {},
            }
        }

        stream.clear(gfx::ClearData {
            color: [0.5, 0.5, 0.5, 1.0],
            depth: 1.0,
            stencil: 0
        });

        stream.draw(&gfx::batch::bind(&state, &mesh, slice.clone(), &program, &None))
              .unwrap();

        stream.present(&mut device);
    }
}
<|endoftext|>