<commit_before>use std::io;
use std::io::Write;
use std::sync::Arc;
use std::fs::OpenOptions;
use std::path::Path;

use storage::{Storage, Result, ByteEncodable, ByteDecodable};
use chunked_storage::indexing_chunk::HotIndexingChunk;
pub use chunked_storage::indexing_chunk::IndexingChunk;

mod indexing_chunk;

pub const SIZE: usize = 104;
const HOTCHUNKS_FILENAME: &'static str = "hot_chunks.bin";
const ASSOCIATED_FILES: &'static [&'static str; 1] = &[HOTCHUNKS_FILENAME];

pub struct MutChunkRef<'a> {
    chunk: &'a mut HotIndexingChunk,
    archive: &'a mut Box<Storage<IndexingChunk>>,
}

pub struct ChunkRef<'a> {
    read_ptr: usize,
    chunk: &'a HotIndexingChunk,
    archive: &'a Box<Storage<IndexingChunk>>,
}

// The idea here is the abstract the inner workings of chunked storage and indexing chunk from the index
// To do this, we implement Read and Write
impl<'a> io::Write for MutChunkRef<'a> {
    // Fill the HotIndexingChunk
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let bytes_written = try!((&mut self.chunk.data[SIZE - self.chunk.capacity as usize..]).write(buf));
        self.chunk.capacity -= bytes_written as u16;
        if self.chunk.capacity == 0 {
            let id = self.archive.len();
            // TODO: Errorhandling
            self.archive.store(id as u64, self.chunk.archive(id as u32)).unwrap();
        }
        Ok(bytes_written)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl<'a> io::Read for ChunkRef<'a> {
    // BULLSHIT. RETHINK THIS!
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut bytes_read = 0;
        loop {
            let chunk_id = self.read_ptr / SIZE;
            if chunk_id < self.chunk.archived_chunks().len() {
                let read = (&self.archive.get(chunk_id as u64).unwrap().get_bytes()[self.read_ptr % SIZE..])
                    .read(&mut buf[bytes_read..])
                    .unwrap();
                if read == 0 {
                    return Ok(bytes_read);
                }
                bytes_read += read;
                self.read_ptr += read;
            } else {
                break;
            }
        }
        let read = (&self.chunk.get_bytes()[self.read_ptr % SIZE..]).read(&mut buf[bytes_read..]).unwrap();
        bytes_read += read;
        self.read_ptr += read;
        Ok(bytes_read)
    }
}

impl<'a> MutChunkRef<'a> {
    #[inline]
    pub fn get_last_doc_id(&self) -> u64 {
        self.chunk.get_last_doc_id()
    }

    pub fn set_last_doc_id(&mut self, new_id: u64) {
        self.chunk.set_last_doc_id(new_id);
    }
}

// TODO: Think about implementing `Persistent` and `Volatile`
pub struct ChunkedStorage {
    hot_chunks: Vec<HotIndexingChunk>, // Size of vocabulary
    archive: Box<Storage<IndexingChunk>>,
}

impl ChunkedStorage {
    pub fn new(capacity: usize, archive: Box<Storage<IndexingChunk>>) -> Self {
        ChunkedStorage {
            hot_chunks: Vec::with_capacity(capacity),
            archive: archive,
        }
    }

    /// Persists hot_chunks to a file.
    /// We currently only need to persist hot_chunks
    /// Archive takes care of the rest
    pub fn persist(&self, target: &Path) -> Result<()> {
        let mut file =
            try!(OpenOptions::new().write(true).create(true).truncate(true).open(target.join(HOTCHUNKS_FILENAME)));
        for chunk in &self.hot_chunks {
            let bytes = chunk.encode();
            try!(file.write(&bytes));
        }
        Ok(())
    }

    pub fn load(source: &Path, archive: Box<Storage<IndexingChunk>>) -> Result<Self> {
        let mut file = try!(OpenOptions::new().read(true).open(source.join(HOTCHUNKS_FILENAME)));
        let mut hot_chunks = Vec::new();
        while let Ok(decoded_chunk) = HotIndexingChunk::decode(&mut file) {
            hot_chunks.push(decoded_chunk);
        }
        Ok(ChunkedStorage {
            hot_chunks: hot_chunks,
            archive: archive,
        })
    }

    fn chunk_ref(&mut self, term_id: u64) -> MutChunkRef {
        MutChunkRef {
            chunk: &mut self.hot_chunks[term_id as usize],
            archive: &mut self.archive,
        }
    }

    pub fn new_chunk(&mut self, term_id: u64) -> MutChunkRef {
        // The following code-uglyness is due to the fact that in indexing one sort thread can
        // overtake the other. That means: ids are not just comming in an incremental order but can jump

        // If the id is larger than the len (e.g. can not be just pushed)
        // Push uninitialized chunks until the desired chunk can be created and pushed
        let diff = term_id as usize - self.hot_chunks.len();
        for _ in 0..diff + 1 {
            self.hot_chunks.push(HotIndexingChunk::new());
        }

        self.hot_chunks[term_id as usize] = HotIndexingChunk::new();
        MutChunkRef {
            chunk: &mut self.hot_chunks[term_id as usize],
            archive: &mut self.archive,
        }
    }

    pub fn next_chunk(&mut self, id: u64) -> Result<&mut HotIndexingChunk> {
        let to_archive = self.hot_chunks[id as usize].archive(self.archive.len() as u32);
        // TODO: Needs to go
        let new_id = self.archive.len() as u64;
        // That's more fun than I thought
        try!(self.archive.store(new_id, to_archive));
        Ok(&mut self.hot_chunks[id as usize])
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.hot_chunks.len()
    }

    #[inline]
    pub fn get_current(&self, id: u64) -> ChunkRef {
        ChunkRef {
            read_ptr: 0,
            chunk: &self.hot_chunks[id as usize],
            archive: &self.archive
        }
    }

    #[inline]
    pub fn get_current_mut(&mut self, id: u64) -> MutChunkRef {
        MutChunkRef {
            chunk: &mut self.hot_chunks[id as usize],
            archive: &mut self.archive,
        }
    }

    #[inline]
    pub fn get_archived(&self, pos: usize) -> Arc<IndexingChunk> {
        self.archive.get(pos as u64).unwrap()
    }

    pub fn associated_files() -> &'static [&'static str] {
        ASSOCIATED_FILES
    }
}

//TODO: Fix Tests

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use utils::persistence::Volatile;
    use storage::RamStorage;

    // #[test]
    // fn basic() {
    //     let mut store = ChunkedStorage::new(10, Box::new(RamStorage::new()));
    //     {
    //         let chunk = store.new_chunk(0);
    //         let listing = vec![(0, vec![0, 10, 20]), (20, vec![24, 25, 289]), (204, vec![209, 2456])];
    //         chunk.append(&listing).unwrap();
    //         let data = chunk.get_bytes();
    //         assert_eq!(data.len(), 18);
    //     }

    //     let chunk = store.get_current(0);
    //     let data = chunk.get_bytes();
    //     assert_eq!(data.len(), 18);
    // }

    // #[test]
    // fn continued() {
    //     let mut store = ChunkedStorage::new(10, Box::new(RamStorage::new()));
    //     let listing = vec![(0, vec![0, 10, 20]), (20, vec![24, 25, 289]), (204, vec![209, 2456])];
    //     let next_listing = vec![(205, vec![0, 10, 20]), (225, vec![24, 25, 289]), (424, vec![209, 2456])];
    //     {
    //         let chunk = store.new_chunk(0);
    //         chunk.append(&listing).unwrap();
    //         assert_eq!(chunk.get_bytes().len(), 18);
    //     }
    //     {
    //         let new_chunk = store.next_chunk(0).unwrap();
    //         new_chunk.append(&next_listing).unwrap();
    //         assert_eq!(new_chunk.get_bytes().len(), 18);
    //     }
    //     let chunk = store.get_current(0);
    //     let old_chunk = store.get_archived((chunk.archived_chunks()[0]) as usize);
    //     assert_eq!(old_chunk.get_bytes().len(), 18);
    //     assert_eq!(chunk.get_bytes().len(), 18);
    // }

    // #[test]
    // fn overflowing_chunk_ref() {
    //     let mut store = ChunkedStorage::new(10, Box::new(RamStorage::new()));
    //     store.new_chunk(0);
    //     let data = (0..255u8).collect::<Vec<_>>();
    //     {
    //         let mut chunk_ref = store.chunk_ref(0);
    //         chunk_ref.write_all(&data);
    //     }
    //     {
    //         let chunk_ref = store.get_current(0);
    //         let first_chunk = store.get_archived((chunk.archived_chunks()[0]) as usize);
    //         assert_eq!(first_chunk.get_bytes(), &data[0..SIZE]);
    //     }
    //     let mut read_data = Vec::new();
    //     {
    //         let mut chunk_ref = store.chunk_ref(0);
    //         chunk_ref.read_to_end(&mut read_data);
    //     }
    //     assert_eq!(data, read_data);
    // }

    // #[test]
    // fn basic_chunk_ref() {
    //     let mut store = ChunkedStorage::new(10, Box::new(RamStorage::new()));
    //     store.new_chunk(0);
    //     let data = (0..20u8).collect::<Vec<_>>();
    //     {
    //         let mut chunk_ref = store.chunk_ref(0);
    //         chunk_ref.write_all(&data);
    //     }
    //     let mut read_data = Vec::new();
    //     {
    //         let mut chunk_ref = store.chunk_ref(0);
    //         chunk_ref.read_to_end(&mut read_data);
    //     }
    //     assert_eq!(data, read_data);
    // }

    // #[test]
    // fn repeated_writes() {
    //     let mut store = ChunkedStorage::new(10, Box::new(RamStorage::new()));
    //     store.new_chunk(0);
    //     let data = (0..20u8).collect::<Vec<_>>();
    //     {
    //         let mut chunk_ref = store.chunk_ref(0);
    //         chunk_ref.write_all(&data[0..10]);
    //         chunk_ref.write_all(&data[10..20]);
    //     }
    //     let mut read_data = Vec::new();
    //     {
    //         let mut chunk_ref = store.chunk_ref(0);
    //         chunk_ref.read_to_end(&mut read_data);
    //     }
    //     assert_eq!(data, read_data);
    // }

    // #[test]
    // fn repeated_writes_overflowing() {
    //     let mut store = ChunkedStorage::new(10, Box::new(RamStorage::new()));
    //     store.new_chunk(0);
    //     let data = (0..255u8).collect::<Vec<_>>();
    //     {
    //         let mut chunk_ref = store.chunk_ref(0);
    //         chunk_ref.write_all(&data);
    //         chunk_ref.write_all(&data);
    //         chunk_ref.write_all(&data);
    //     }
    //     let mut read_data = Vec::new();
    //     {
    //         let mut chunk_ref = store.chunk_ref(0);
    //         chunk_ref.read_to_end(&mut read_data);
    //     }
    //     assert_eq!(data, &read_data[0..255]);
    //     assert_eq!(data, &read_data[255..510]);
    //     assert_eq!(data, &read_data[510..765]);
    // }
}
<commit_msg>fix chunked_storage tests<commit_after>use std::io;
use std::io::Write;
use std::sync::Arc;
use std::fs::OpenOptions;
use std::path::Path;

use storage::{Storage, Result, ByteEncodable, ByteDecodable};
use chunked_storage::indexing_chunk::HotIndexingChunk;
pub use chunked_storage::indexing_chunk::IndexingChunk;

mod indexing_chunk;

pub const SIZE: usize = 104;
const HOTCHUNKS_FILENAME: &'static str = "hot_chunks.bin";
const ASSOCIATED_FILES: &'static [&'static str; 1] = &[HOTCHUNKS_FILENAME];

pub struct MutChunkRef<'a> {
    chunk: &'a mut HotIndexingChunk,
    archive: &'a mut Box<Storage<IndexingChunk>>,
}

pub struct ChunkRef<'a> {
    read_ptr: usize,
    chunk: &'a HotIndexingChunk,
    archive: &'a Box<Storage<IndexingChunk>>,
}

// The idea here is the abstract the inner workings of chunked storage and indexing chunk from the index
// To do this, we implement Read and Write
impl<'a> io::Write for MutChunkRef<'a> {
    // Fill the HotIndexingChunk
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let bytes_written = try!((&mut self.chunk.data[SIZE - self.chunk.capacity as usize..]).write(buf));
        self.chunk.capacity -= bytes_written as u16;
        if self.chunk.capacity == 0 {
            let id = self.archive.len();
            // TODO: Errorhandling
            self.archive.store(id as u64, self.chunk.archive(id as u32)).unwrap();
        }
        Ok(bytes_written)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl<'a> io::Read for ChunkRef<'a> {
    // BULLSHIT. RETHINK THIS!
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut bytes_read = 0;
        loop {
            let chunk_id = self.read_ptr / SIZE;
            if chunk_id < self.chunk.archived_chunks().len() {
                let read = (&self.archive.get(chunk_id as u64).unwrap().get_bytes()[self.read_ptr % SIZE..])
                    .read(&mut buf[bytes_read..])
                    .unwrap();
                if read == 0 {
                    return Ok(bytes_read);
                }
                bytes_read += read;
                self.read_ptr += read;
            } else {
                break;
            }
        }
        let read = (&self.chunk.get_bytes()[self.read_ptr % SIZE..]).read(&mut buf[bytes_read..]).unwrap();
        bytes_read += read;
        self.read_ptr += read;
        Ok(bytes_read)
    }
}

impl<'a> MutChunkRef<'a> {
    #[inline]
    pub fn get_last_doc_id(&self) -> u64 {
        self.chunk.get_last_doc_id()
    }

    pub fn set_last_doc_id(&mut self, new_id: u64) {
        self.chunk.set_last_doc_id(new_id);
    }
}

// TODO: Think about implementing `Persistent` and `Volatile`
pub struct ChunkedStorage {
    hot_chunks: Vec<HotIndexingChunk>, // Size of vocabulary
    archive: Box<Storage<IndexingChunk>>,
}

impl ChunkedStorage {
    pub fn new(capacity: usize, archive: Box<Storage<IndexingChunk>>) -> Self {
        ChunkedStorage {
            hot_chunks: Vec::with_capacity(capacity),
            archive: archive,
        }
    }

    /// Persists hot_chunks to a file.
    /// We currently only need to persist hot_chunks
    /// Archive takes care of the rest
    pub fn persist(&self, target: &Path) -> Result<()> {
        let mut file =
            try!(OpenOptions::new().write(true).create(true).truncate(true).open(target.join(HOTCHUNKS_FILENAME)));
        for chunk in &self.hot_chunks {
            let bytes = chunk.encode();
            try!(file.write(&bytes));
        }
        Ok(())
    }

    pub fn load(source: &Path, archive: Box<Storage<IndexingChunk>>) -> Result<Self> {
        let mut file = try!(OpenOptions::new().read(true).open(source.join(HOTCHUNKS_FILENAME)));
        let mut hot_chunks = Vec::new();
        while let Ok(decoded_chunk) = HotIndexingChunk::decode(&mut file) {
            hot_chunks.push(decoded_chunk);
        }
        Ok(ChunkedStorage {
            hot_chunks: hot_chunks,
            archive: archive,
        })
    }

    pub fn new_chunk(&mut self, term_id: u64) -> MutChunkRef {
        // The following code-uglyness is due to the fact that in indexing one sort thread can
        // overtake the other. That means: ids are not just comming in an incremental order but can jump

        // If the id is larger than the len (e.g. can not be just pushed)
        // Push uninitialized chunks until the desired chunk can be created and pushed
        let diff = term_id as usize - self.hot_chunks.len();
        for _ in 0..diff + 1 {
            self.hot_chunks.push(HotIndexingChunk::new());
        }

        self.hot_chunks[term_id as usize] = HotIndexingChunk::new();
        MutChunkRef {
            chunk: &mut self.hot_chunks[term_id as usize],
            archive: &mut self.archive,
        }
    }

    pub fn next_chunk(&mut self, id: u64) -> Result<&mut HotIndexingChunk> {
        let to_archive = self.hot_chunks[id as usize].archive(self.archive.len() as u32);
        // TODO: Needs to go
        let new_id = self.archive.len() as u64;
        // That's more fun than I thought
        try!(self.archive.store(new_id, to_archive));
        Ok(&mut self.hot_chunks[id as usize])
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.hot_chunks.len()
    }

    #[inline]
    pub fn get_current(&self, id: u64) -> ChunkRef {
        ChunkRef {
            read_ptr: 0,
            chunk: &self.hot_chunks[id as usize],
            archive: &self.archive
        }
    }

    #[inline]
    pub fn get_current_mut(&mut self, id: u64) -> MutChunkRef {
        MutChunkRef {
            chunk: &mut self.hot_chunks[id as usize],
            archive: &mut self.archive,
        }
    }

    #[inline]
    pub fn get_archived(&self, pos: usize) -> Arc<IndexingChunk> {
        self.archive.get(pos as u64).unwrap()
    }

    pub fn associated_files() -> &'static [&'static str] {
        ASSOCIATED_FILES
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use utils::persistence::Volatile;
    use storage::RamStorage;

    #[test]
    fn overflowing_chunk_ref() {
        let mut store = ChunkedStorage::new(10, Box::new(RamStorage::new()));
        store.new_chunk(0);
        let data = (0..255u8).collect::<Vec<_>>();
        {
            let mut chunk_ref = store.get_current_mut(0);
            chunk_ref.write_all(&data).unwrap();
        }
        let mut read_data = Vec::new();
        {
            let mut chunk_ref = store.get_current(0);
            chunk_ref.read_to_end(&mut read_data).unwrap();
        }
        assert_eq!(data, read_data);
    }

    #[test]
    fn basic_chunk_ref() {
        let mut store = ChunkedStorage::new(10, Box::new(RamStorage::new()));
        store.new_chunk(0);
        let data = (0..20u8).collect::<Vec<_>>();
        {
            let mut chunk_ref = store.get_current_mut(0);
            chunk_ref.write_all(&data).unwrap();
        }
        let mut read_data = Vec::new();
        {
            let mut chunk_ref = store.get_current(0);
            chunk_ref.read_to_end(&mut read_data).unwrap();
        }
        assert_eq!(data, read_data);
    }

    #[test]
    fn repeated_writes() {
        let mut store = ChunkedStorage::new(10, Box::new(RamStorage::new()));
        store.new_chunk(0);
        let data = (0..20u8).collect::<Vec<_>>();
        {
            let mut chunk_ref = store.get_current_mut(0);
            chunk_ref.write_all(&data[0..10]).unwrap();
            chunk_ref.write_all(&data[10..20]).unwrap();
        }
        let mut read_data = Vec::new();
        {
            let mut chunk_ref = store.get_current(0);
            chunk_ref.read_to_end(&mut read_data).unwrap();
        }
        assert_eq!(data, read_data);
    }

    #[test]
    fn repeated_writes_overflowing() {
        let mut store = ChunkedStorage::new(10, Box::new(RamStorage::new()));
        store.new_chunk(0);
        let data = (0..255u8).collect::<Vec<_>>();
        {
            let mut chunk_ref = store.get_current_mut(0);
            chunk_ref.write_all(&data).unwrap();
            chunk_ref.write_all(&data).unwrap();
            chunk_ref.write_all(&data).unwrap();
        }
        let mut read_data = Vec::new();
        {
            let mut chunk_ref = store.get_current(0);
            chunk_ref.read_to_end(&mut read_data).unwrap();
        }
        assert_eq!(data, &read_data[0..255]);
        assert_eq!(data, &read_data[255..510]);
        assert_eq!(data, &read_data[510..765]);
    }
}
<|endoftext|>


<commit_before>pub mod base;
mod flag;

use self::base::{BaseRand, RAND_MAX};
use self::flag::Flag;
use std::f64::consts::PI;

pub enum Style {
    Normal,
    Pythagoras,
    Gauss(u32),
    BMgauss(f64, f64, bool),
    Marsaglia(f64, f64, f64, bool),
}

pub enum JumpStyle {
    Static,
    Next,
    Jump,
    DoubleJump,
}

/// Solver is random number solver
/// Please do not assume that the fields are any type
pub struct Solver {
    base: Vec<BaseRand>,
    handle: Flag,
    style: Style,
    jump: JumpStyle,
}

impl Solver {
    /// A lazy way to get a random solver
    pub fn new() -> Self {
        Solver {
            base: vec![BaseRand::new()],
            handle: Flag::On(0),
            style: Style::Normal,
            jump: JumpStyle::Static,
        }
    }

    /// Insert a new base
    pub fn new_multibase(mul: usize) -> Self {
        let mut _base = Solver::new();
        let mut _seed: Vec<usize> = Vec::new();
        for _ in 0..mul {
            _seed.push(_base.base());
        }
        for _ in 1..mul {
            _base.new_base();
        }
        _base.multisrand(_seed);
        _base.jump_style(JumpStyle::DoubleJump);
        _base
    }

    /// Set random solver style
    pub fn set_style(&mut self, style: Style) {
        self.style = style;
    }

    /// Get handle base
    pub fn get_base(&self, handle: usize) -> Option<BaseRand> {
        if handle<self.base.len() {
            Some(self.base[handle])
        } else {
            None
        }
    }

    pub fn set_base(&mut self, handle: usize, base: BaseRand) -> bool {
        if handle<self.base.len() {
            self.base[handle] = base;
            return true;
        }
        false
    }

    /// Set the random seed
    pub fn multisrand(&mut self, seed: Vec<usize>) {
        for (i, item) in seed.iter().enumerate() {
            let mut _base = BaseRand::new();
            _base.srand(*item);
            self.set_base(i, _base);
        }
    }

    /// Lazy way to set the random seed
    pub fn lazy_srand(&mut self) {
        for i in 0..self.base.len() {
            self.base[i].lazy_srand();
        }
    }

    /// Get a random number
    pub fn rand(&mut self) -> f64 {
        match self.style {
            Style::Normal => self.base() as f64,
            Style::Pythagoras => pythagoras(self),
            Style::Gauss(nsum) => gauss(self, nsum),
            Style::BMgauss(u, v, phase) => {
                let mut _u = u;
                let mut _v = v;
                let mut _phase = phase;
                let _return = bmgauss(self, &mut _u, &mut _v, &mut _phase);
                self.style = Style::BMgauss(_u, _v, _phase);
                _return
            }
            Style::Marsaglia(v1, v2, s, phase) => {
                let mut _v1 = v1;
                let mut _v2 = v2;
                let mut _s = s;
                let mut _phase = phase;
                let _return = marsaglia(self, &mut _v1, &mut _v2, &mut _s, &mut _phase);
                self.style = Style::Marsaglia(_v1, _v2, _s, _phase);
                _return
            }
        }
    }

    /// Get a random number for multibase
    pub fn multirand(&mut self) -> f64 {
        let _return = self.rand();
        self.jump();
        _return
    }

    pub fn base(&mut self) -> usize {
        match self.handle {
            Flag::On(e) => self.base[e].rand(),
            _ => {
                let mut _base = 0;
                for i in 0..self.base.len() {
                    _base += self.base[i].rand() / self.base.len();
                }
                _base
            }
        }
    }

    pub fn new_base(&mut self) {
        self.base.push(BaseRand::new());
    }

    pub fn del_base(&mut self) -> Option<BaseRand> {
        self.base.pop()
    }

    pub fn add_base(&mut self, base: BaseRand) {
        self.base.push(base);
    }

    pub fn jump_style(&mut self, style: JumpStyle) {
        self.jump = style;
    }

    pub fn jump(&mut self) {
        if let Flag::On(e) = self.handle {
            match self.jump {
                JumpStyle::Next => self.handle = Flag::On((e + 1) % self.base.len()),
                JumpStyle::Jump => {
                    let _jump = self.base();
                    self.handle = Flag::On(_jump % self.base.len());
                }
                JumpStyle::DoubleJump => {
                    let mut _gap = Flag::new();
                    _gap.on(Some(self.base()));
                    while _gap.is_on() {
                        let _jump = self.base();
                        self.handle = Flag::On(_jump % self.base.len());
                        _gap.down();
                    }
                }
                _ => {}
            }
        }
    }
}

impl Default for Solver {
    fn default() -> Self {
        Self::new()
    }
}

fn base64(base: &mut Solver, base_offset: usize, max_offset: usize) -> f64 {
    (base.base() + base_offset) as f64 / (RAND_MAX + base_offset + max_offset) as f64
}

fn gauss(base: &mut Solver, nsum: u32) -> f64 {
    let mut x = 0.0;
    for _ in 0..nsum {
        x += base64(base, 0, 0);
    }
    x -= nsum as f64 / 2.0;
    x /= (nsum as f64 / 12.0).sqrt();
    x
}

fn bmgauss(base: &mut Solver, u: &mut f64, v: &mut f64, phase: &mut bool) -> f64 {
    if *phase {
        *phase = false;
        *u = base64(base, 1, 1);
        *v = base64(base, 0, 1);
        (-2.0 * (*u).log10()).sqrt() * (2.0 * PI * (*v)).sin()
    } else {
        *phase = true;
        (-2.0 * (*u).log10()).sqrt() * (2.0 * PI * (*v)).cos()
    }
}

fn marsaglia(base: &mut Solver, v1: &mut f64, v2: &mut f64, s: &mut f64, phase: &mut bool) -> f64 {
    if *phase {
        *phase = false;
        let u1 = base64(base, 0, 0);
        let u2 = base64(base, 0, 0);
        *v1 = 2.0 * u1 - 1.0;
        *v2 = 2.0 * u2 - 1.0;
        *v1 * (-2.0 * (*s).log10() / (*s)).sqrt()
    } else {
        *phase = true;
        *v2 * (-2.0 * (*s).log10() / (*s)).sqrt()
    }
}

fn pythagoras(base: &mut Solver) -> f64 {
    let a = base64(base, 0, 1);
    let b = base64(base, 0, 1);
    let c_sqr = a * a + b * b;
    c_sqr.sqrt()
}
<commit_msg>Update<commit_after>pub mod base;
mod flag;

use self::base::{BaseRand, RAND_MAX};
use self::flag::Flag;
use std::f64::consts::PI;

pub enum Style {
    Normal,
    Pythagoras,
    Gauss(u32),
    BMgauss(f64, f64, bool),
    Marsaglia(f64, f64, f64, bool),
}

pub enum JumpStyle {
    Static,
    Next,
    Jump,
    DoubleJump,
}

/// Solver is random number solver
/// Please do not assume that the fields are any type
pub struct Solver {
    base: Vec<BaseRand>,
    handle: Flag,
    style: Style,
    jump: JumpStyle,
}

impl Solver {
    /// A lazy way to get a random solver
    pub fn new() -> Self {
        Solver {
            base: vec![BaseRand::new()],
            handle: Flag::On(0),
            style: Style::Normal,
            jump: JumpStyle::Static,
        }
    }

    /// Insert a new base
    pub fn new_multibase(mul: usize) -> Self {
        let mut _base = Solver::new();
        let mut _seed: Vec<usize> = Vec::new();
        for _ in 0..mul {
            _seed.push(_base.base());
        }
        for _ in 1..mul {
            _base.new_base();
        }
        _base.multisrand(_seed);
        _base.jump_style(JumpStyle::DoubleJump);
        _base
    }

    /// Set random solver style
    pub fn set_style(&mut self, style: Style) {
        self.style = style;
    }

    /// Get handle base
    pub fn get_base(&self, handle: usize) -> Option<BaseRand> {
        if handle<self.base.len() {
            Some(self.base[handle])
        } else {
            None
        }
    }

    pub fn set_base(&mut self, handle: usize, base: BaseRand) -> bool {
        if handle<self.base.len() {
            self.base[handle] = base;
            return true;
        }
        false
    }

    /// Set the random seed
    pub fn multisrand(&mut self, seed: Vec<usize>) {
        for (i, item) in seed.iter().enumerate() {
            let mut _base = BaseRand::new();
            _base.srand(*item);
            self.set_base(i, _base);
        }
    }

    /// Lazy way to set the random seed
    pub fn lazy_srand(&mut self) {
        for i in 0..self.base.len() {
            self.base[i].lazy_srand();
        }
    }

    /// Get a random number
    pub fn rand(&mut self) -> f64 {
        match self.style {
            Style::Normal => self.base() as f64,
            Style::Pythagoras => pythagoras(self),
            Style::Gauss(nsum) => gauss(self, nsum),
            Style::BMgauss(u, v, phase) => {
                let mut _u = u;
                let mut _v = v;
                let mut _phase = phase;
                let _return = bmgauss(self,
                                      &mut _u,
                                      &mut _v,
                                      &mut _phase);
                self.style = Style::BMgauss(_u, _v, _phase);
                _return
            }
            Style::Marsaglia(v1, v2, s, phase) => {
                let mut _v1 = v1;
                let mut _v2 = v2;
                let mut _s = s;
                let mut _phase = phase;
                let _return = marsaglia(self,
                                        &mut _v1,
                                        &mut _v2,
                                        &mut _s,
                                        &mut _phase);
                self.style = Style::Marsaglia(_v1,
                                              _v2,
                                              _s,
                                              _phase);
                _return
            }
        }
    }

    /// Get a random number for multibase
    pub fn multirand(&mut self) -> f64 {
        let _return = self.rand();
        self.jump();
        _return
    }

    pub fn base(&mut self) -> usize {
        match self.handle {
            Flag::On(e) => self.base[e].rand(),
            _ => {
                let mut _base = 0;
                for i in 0..self.base.len() {
                    _base += self.base[i].rand() / self.base.len();
                }
                _base
            }
        }
    }

    pub fn new_base(&mut self) {
        self.base.push(BaseRand::new());
    }

    pub fn del_base(&mut self) -> Option<BaseRand> {
        self.base.pop()
    }

    pub fn add_base(&mut self, base: BaseRand) {
        self.base.push(base);
    }

    pub fn jump_style(&mut self, style: JumpStyle) {
        self.jump = style;
    }

    pub fn jump(&mut self) {
        if let Flag::On(e) = self.handle {
            match self.jump {
                JumpStyle::Next => self.handle = Flag::On((e + 1) % self.base.len()),
                JumpStyle::Jump => {
                    let _jump = self.base();
                    self.handle = Flag::On(_jump % self.base.len());
                }
                JumpStyle::DoubleJump => {
                    let mut _gap = Flag::new();
                    _gap.on(Some(self.base()));
                    while _gap.is_on() {
                        let _jump = self.base();
                        self.handle = Flag::On(_jump % self.base.len());
                        _gap.down();
                    }
                }
                _ => {}
            }
        }
    }
}

impl Default for Solver {
    fn default() -> Self {
        Self::new()
    }
}

fn base64(base: &mut Solver,
          base_offset: usize,
          max_offset: usize) -> f64 {
    (base.base() + base_offset) as f64 / (RAND_MAX + base_offset + max_offset) as f64
}

fn gauss(base: &mut Solver, nsum: u32) -> f64 {
    let mut x = 0.0;
    for _ in 0..nsum {
        x += base64(base, 0, 0);
    }
    x -= nsum as f64 / 2.0;
    x /= (nsum as f64 / 12.0).sqrt();
    x
}

fn bmgauss(base: &mut Solver, u: &mut f64, v: &mut f64, phase: &mut bool) -> f64 {
    if *phase {
        *phase = false;
        *u = base64(base, 1, 1);
        *v = base64(base, 0, 1);
        (-2.0 * (*u).log10()).sqrt() * (2.0 * PI * (*v)).sin()
    } else {
        *phase = true;
        (-2.0 * (*u).log10()).sqrt() * (2.0 * PI * (*v)).cos()
    }
}

fn marsaglia(base: &mut Solver,
             v1: &mut f64,
             v2: &mut f64,
             s: &mut f64,
             phase: &mut bool) -> f64 {
    if *phase {
        *phase = false;
        let u1 = base64(base, 0, 0);
        let u2 = base64(base, 0, 0);
        *v1 = 2.0 * u1 - 1.0;
        *v2 = 2.0 * u2 - 1.0;
        *v1 * (-2.0 * (*s).log10() / (*s)).sqrt()
    } else {
        *phase = true;
        *v2 * (-2.0 * (*s).log10() / (*s)).sqrt()
    }
}

fn pythagoras(base: &mut Solver) -> f64 {
    let a = base64(base, 0, 1);
    let b = base64(base, 0, 1);
    let c_sqr = a * a + b * b;
    c_sqr.sqrt()
}
<|endoftext|>


<commit_before>//! This example will showcase the beauty of collectors.
//! They allow to await messages or reactions from a user in the middle
//! of a control flow, one being a command.
use std::collections::HashSet;
use std::env;
use std::time::Duration;

use serenity::async_trait;
use serenity::collector::{EventCollectorBuilder, MessageCollectorBuilder};
use serenity::framework::standard::macros::{command, group, help};
use serenity::framework::standard::{
    help_commands,
    Args,
    CommandGroup,
    CommandResult,
    HelpOptions,
    StandardFramework,
};
// Collectors are streams, that means we can use `StreamExt` and `TryStreamExt`.
use serenity::futures::stream::StreamExt;
use serenity::http::Http;
use serenity::model::prelude::*;
use serenity::prelude::*;

#[group("collector")]
#[commands(challenge)]
struct Collector;

#[help]
async fn my_help(
    context: &Context,
    msg: &Message,
    args: Args,
    help_options: &'static HelpOptions,
    groups: &[&'static CommandGroup],
    owners: HashSet<UserId>,
) -> CommandResult {
    let _ = help_commands::with_embeds(context, msg, args, help_options, groups, owners).await;
    Ok(())
}

struct Handler;

#[async_trait]
impl EventHandler for Handler {
    async fn ready(&self, _: Context, ready: Ready) {
        println!("{} is connected!", ready.user.name);
    }
}

#[tokio::main]
async fn main() {
    // Configure the client with your Discord bot token in the environment.
    let token = env::var("DISCORD_TOKEN").expect("Expected a token in the environment");

    let http = Http::new(&token);

    // We will fetch your bot's id.
    let bot_id = match http.get_current_user().await {
        Ok(info) => info.id,
        Err(why) => panic!("Could not access user info: {:?}", why),
    };

    let framework = StandardFramework::new()
        .configure(|c| {
            c.with_whitespace(true).on_mention(Some(bot_id)).prefix("~").delimiters(vec![", ", ","])
        })
        .help(&MY_HELP)
        .group(&COLLECTOR_GROUP);

    let intents = GatewayIntents::GUILD_MESSAGES
        | GatewayIntents::DIRECT_MESSAGES
        | GatewayIntents::MESSAGE_CONTENT;
    let mut client = Client::builder(&token, intents)
        .event_handler(Handler)
        .framework(framework)
        .await
        .expect("Err creating client");

    if let Err(why) = client.start().await {
        println!("Client error: {:?}", why);
    }
}

#[command]
async fn challenge(ctx: &Context, msg: &Message, _: Args) -> CommandResult {
    let mut score = 0u32;
    let _ = msg.reply(ctx, "How was that crusty crab called again? 10 seconds time!").await;

    // There are methods implemented for some models to conveniently collect replies.
    // This one returns a future that will await a single message only.
    // The other method for messages is called `await_replies` and returns a future
    // which builds a stream to easily handle them.
    if let Some(answer) = &msg.author.await_reply(&ctx).timeout(Duration::from_secs(10)).await {
        if answer.content.to_lowercase() == "ferris" {
            let _ = answer.reply(ctx, "That's correct!").await;
            score += 1;
        } else {
            let _ = answer.reply(ctx, "Wrong, it's Ferris!").await;
        }
    } else {
        let _ = msg.reply(ctx, "No answer within 10 seconds.").await;
    };

    let react_msg = msg
        .reply(ctx, "React with the reaction representing 1, you got 10 seconds!")
        .await
        .unwrap();

    // The message model has a way to collect reactions on it.
    // Other methods are `await_n_reactions` and `await_all_reactions`.
    // Same goes for messages!
    if let Some(reaction) = &react_msg
        .await_reaction(&ctx)
        .timeout(Duration::from_secs(10))
        .author_id(msg.author.id)
        .await
    {
        // By default, the collector will collect only added reactions.
        // We could also pattern-match the reaction in case we want
        // to handle added or removed reactions.
        // In this case we will just get the inner reaction.
        let emoji = &reaction.as_inner_ref().emoji;

        let _ = match emoji.as_data().as_str() {
            "1⃣" => {
                score += 1;
                msg.reply(ctx, "That's correct!").await
            },
            _ => msg.reply(ctx, "Wrong!").await,
        };
    } else {
        let _ = msg.reply(ctx, "No reaction within 10 seconds.").await;
    };

    let _ = msg.reply(ctx, "Write 5 messages in 10 seconds").await;

    // We can create a collector from scratch too using this builder future.
    let collector = MessageCollectorBuilder::new(&ctx)
    // Only collect messages by this user.
        .author_id(msg.author.id)
        .channel_id(msg.channel_id)
        .collect_limit(5u32)
        .timeout(Duration::from_secs(10))
    // Build the collector.
        .build();

    // Let's acquire borrow HTTP to send a message inside the `async move`.
    let http = &ctx.http;

    // We want to process each message and get the length.
    // There are a couple of ways to do this. Folding the stream with `fold`
    // is one way.
    // Using `then` to first reply and then create a new stream with all
    // messages is another way to do it, which can be nice if you want
    // to further process the messages.
    // If you don't want to collect the stream, `for_each` may be sufficient.
    let collected: Vec<_> = collector
        .then(|msg| async move {
            let _ = msg.reply(http, format!("I repeat: {}", msg.content)).await;

            msg
        })
        .collect()
        .await;

    if collected.len() >= 5 {
        score += 1;
    }

    // We can also collect arbitrary events using the generic EventCollector. For example, here we
    // collect updates to the messages that the user sent above and check for them updating all 5 of
    // them.
    let builder = EventCollectorBuilder::new(&ctx)
        .add_event_type(EventType::MessageUpdate)
        .timeout(Duration::from_secs(20));
    // Only collect MessageUpdate events for the 5 MessageIds we're interested in.
    let mut collector =
        collected.iter().fold(builder, |b, msg| b.add_message_id(msg.id)).build()?;

    let _ = msg.reply(ctx, "Edit each of those 5 messages in 20 seconds").await;
    let mut edited = HashSet::new();
    while let Some(event) = collector.next().await {
        match event.as_ref() {
            Event::MessageUpdate(e) => {
                edited.insert(e.id);
            },
            e => panic!("Unexpected event type received: {:?}", e.event_type()),
        }
        if edited.len() >= 5 {
            break;
        }
    }

    if edited.len() >= 5 {
        score += 1;
        let _ = msg.reply(ctx, "Great! You edited 5 out of 5").await;
    } else {
        let _ = msg.reply(ctx, &format!("You only edited {} out of 5", edited.len())).await;
    }

    let _ = msg
        .reply(ctx, &format!("TIME'S UP! You completed {} out of 4 tasks correctly!", score))
        .await;

    Ok(())
}
<commit_msg>Add missing reactions intent to example 10<commit_after>//! This example will showcase the beauty of collectors.
//! They allow to await messages or reactions from a user in the middle
//! of a control flow, one being a command.
use std::collections::HashSet;
use std::env;
use std::time::Duration;

use serenity::async_trait;
use serenity::collector::{EventCollectorBuilder, MessageCollectorBuilder};
use serenity::framework::standard::macros::{command, group, help};
use serenity::framework::standard::{
    help_commands,
    Args,
    CommandGroup,
    CommandResult,
    HelpOptions,
    StandardFramework,
};
// Collectors are streams, that means we can use `StreamExt` and `TryStreamExt`.
use serenity::futures::stream::StreamExt;
use serenity::http::Http;
use serenity::model::prelude::*;
use serenity::prelude::*;

#[group("collector")]
#[commands(challenge)]
struct Collector;

#[help]
async fn my_help(
    context: &Context,
    msg: &Message,
    args: Args,
    help_options: &'static HelpOptions,
    groups: &[&'static CommandGroup],
    owners: HashSet<UserId>,
) -> CommandResult {
    let _ = help_commands::with_embeds(context, msg, args, help_options, groups, owners).await;
    Ok(())
}

struct Handler;

#[async_trait]
impl EventHandler for Handler {
    async fn ready(&self, _: Context, ready: Ready) {
        println!("{} is connected!", ready.user.name);
    }
}

#[tokio::main]
async fn main() {
    // Configure the client with your Discord bot token in the environment.
    let token = env::var("DISCORD_TOKEN").expect("Expected a token in the environment");

    let http = Http::new(&token);

    // We will fetch your bot's id.
    let bot_id = match http.get_current_user().await {
        Ok(info) => info.id,
        Err(why) => panic!("Could not access user info: {:?}", why),
    };

    let framework = StandardFramework::new()
        .configure(|c| {
            c.with_whitespace(true).on_mention(Some(bot_id)).prefix("~").delimiters(vec![", ", ","])
        })
        .help(&MY_HELP)
        .group(&COLLECTOR_GROUP);

    let intents = GatewayIntents::GUILD_MESSAGES
        | GatewayIntents::DIRECT_MESSAGES
        | GatewayIntents::MESSAGE_CONTENT
        | GatewayIntents::GUILD_MESSAGE_REACTIONS;

    let mut client = Client::builder(&token, intents)
        .event_handler(Handler)
        .framework(framework)
        .await
        .expect("Err creating client");

    if let Err(why) = client.start().await {
        println!("Client error: {:?}", why);
    }
}

#[command]
async fn challenge(ctx: &Context, msg: &Message, _: Args) -> CommandResult {
    let mut score = 0u32;
    let _ = msg.reply(ctx, "How was that crusty crab called again? 10 seconds time!").await;

    // There are methods implemented for some models to conveniently collect replies.
    // This one returns a future that will await a single message only.
    // The other method for messages is called `await_replies` and returns a future
    // which builds a stream to easily handle them.
    if let Some(answer) = &msg.author.await_reply(&ctx).timeout(Duration::from_secs(10)).await {
        if answer.content.to_lowercase() == "ferris" {
            let _ = answer.reply(ctx, "That's correct!").await;
            score += 1;
        } else {
            let _ = answer.reply(ctx, "Wrong, it's Ferris!").await;
        }
    } else {
        let _ = msg.reply(ctx, "No answer within 10 seconds.").await;
    };

    let react_msg = msg
        .reply(ctx, "React with the reaction representing 1, you got 10 seconds!")
        .await
        .unwrap();

    // The message model has a way to collect reactions on it.
    // Other methods are `await_n_reactions` and `await_all_reactions`.
    // Same goes for messages!
    if let Some(reaction) = &react_msg
        .await_reaction(&ctx)
        .timeout(Duration::from_secs(10))
        .author_id(msg.author.id)
        .await
    {
        // By default, the collector will collect only added reactions.
        // We could also pattern-match the reaction in case we want
        // to handle added or removed reactions.
        // In this case we will just get the inner reaction.
        let emoji = &reaction.as_inner_ref().emoji;

        let _ = match emoji.as_data().as_str() {
            "1⃣" => {
                score += 1;
                msg.reply(ctx, "That's correct!").await
            },
            _ => msg.reply(ctx, "Wrong!").await,
        };
    } else {
        let _ = msg.reply(ctx, "No reaction within 10 seconds.").await;
    };

    let _ = msg.reply(ctx, "Write 5 messages in 10 seconds").await;

    // We can create a collector from scratch too using this builder future.
    let collector = MessageCollectorBuilder::new(&ctx)
    // Only collect messages by this user.
        .author_id(msg.author.id)
        .channel_id(msg.channel_id)
        .collect_limit(5u32)
        .timeout(Duration::from_secs(10))
    // Build the collector.
        .build();

    // Let's acquire borrow HTTP to send a message inside the `async move`.
    let http = &ctx.http;

    // We want to process each message and get the length.
    // There are a couple of ways to do this. Folding the stream with `fold`
    // is one way.
    // Using `then` to first reply and then create a new stream with all
    // messages is another way to do it, which can be nice if you want
    // to further process the messages.
    // If you don't want to collect the stream, `for_each` may be sufficient.
    let collected: Vec<_> = collector
        .then(|msg| async move {
            let _ = msg.reply(http, format!("I repeat: {}", msg.content)).await;

            msg
        })
        .collect()
        .await;

    if collected.len() >= 5 {
        score += 1;
    }

    // We can also collect arbitrary events using the generic EventCollector. For example, here we
    // collect updates to the messages that the user sent above and check for them updating all 5 of
    // them.
    let builder = EventCollectorBuilder::new(&ctx)
        .add_event_type(EventType::MessageUpdate)
        .timeout(Duration::from_secs(20));
    // Only collect MessageUpdate events for the 5 MessageIds we're interested in.
    let mut collector =
        collected.iter().fold(builder, |b, msg| b.add_message_id(msg.id)).build()?;

    let _ = msg.reply(ctx, "Edit each of those 5 messages in 20 seconds").await;
    let mut edited = HashSet::new();
    while let Some(event) = collector.next().await {
        match event.as_ref() {
            Event::MessageUpdate(e) => {
                edited.insert(e.id);
            },
            e => panic!("Unexpected event type received: {:?}", e.event_type()),
        }
        if edited.len() >= 5 {
            break;
        }
    }

    if edited.len() >= 5 {
        score += 1;
        let _ = msg.reply(ctx, "Great! You edited 5 out of 5").await;
    } else {
        let _ = msg.reply(ctx, &format!("You only edited {} out of 5", edited.len())).await;
    }

    let _ = msg
        .reply(ctx, &format!("TIME'S UP! You completed {} out of 4 tasks correctly!", score))
        .await;

    Ok(())
}
<|endoftext|>


<commit_before>pub mod error;
pub mod helpers;
#[cfg(test)]
mod test;

use self::error::Error;

use std::result::Result as StdResult;

use serde::de::{self, Deserialize, IntoDeserializer};
use serde::de::value::SeqDeserializer;

macro_rules! forward_parsed_values {
  ($($ty:ident => $method:ident,)*) => {
    $(
      fn $method<V>(self, visitor: V) -> Result<V::Value>
        where V: de::Visitor<'de>
      {
        match self.parts.remove(0).parse::<$ty>() {
          Ok(val) => val.into_deserializer().$method(visitor),
          Err(e) => Err(de::Error::custom(e))
        }
      }
    )*
  }
}

pub type Result<T> = StdResult<T, Error>;

struct Deserializer<'de> {
  parts: Vec<&'de str>,
  count: usize
}

pub fn from_str<'a, T>(input: &'a str) -> Result<T>
  where T: Deserialize<'a>
{
  let parts: Vec<&str> = input.split_whitespace().collect();
  let mut deserializer = Deserializer { parts: parts, count: 0 };
  let t = T::deserialize(&mut deserializer)?;
  if deserializer.parts.is_empty() {
    Ok(t)
  } else {
    Err(Error::Custom("trailing characters".to_owned()))
  }
}

impl<'de, 'a> de::Deserializer<'de> for &'a mut Deserializer<'de> {
  type Error = Error;

  fn deserialize_struct<V>(self, _name: &'static str, fields: &'static [&'static str], visitor: V) -> Result<V::Value>
    where V: de::Visitor<'de>
  {
    self.deserialize_tuple(fields.len(), visitor)
  }

  fn deserialize_tuple<V>(self, _len: usize, visitor: V) -> Result<V::Value>
    where V: de::Visitor<'de>
  {
    visitor.visit_seq(self)
  }

  fn deserialize_any<V>(self, visitor: V) -> Result<V::Value>
    where V: de::Visitor<'de>
  {
    if self.parts.is_empty() {
      return Err(Error::MissingParams);
    }
    let first = self.parts.remove(0);
    first.into_deserializer().deserialize_any(visitor)
  }

  fn deserialize_option<V>(self, visitor: V) -> Result<V::Value>
    where V: de::Visitor<'de>
  {
    if self.parts.is_empty() {
      visitor.visit_none()
    } else {
      visitor.visit_some(self)
    }
  }

  fn deserialize_seq<V>(self, visitor: V) -> Result<V::Value>
    where V: de::Visitor<'de>
  {
    if self.parts.is_empty() {
      return Err(Error::MissingParams);
    }

    SeqDeserializer::new(self.parts.drain(..)).deserialize_seq(visitor)
  }

  forward_parsed_values! {
    bool => deserialize_bool,
    u8 => deserialize_u8,
    u16 => deserialize_u16,
    u32 => deserialize_u32,
    u64 => deserialize_u64,
    i8 => deserialize_i8,
    i16 => deserialize_i16,
    i32 => deserialize_i32,
    i64 => deserialize_i64,
    f32 => deserialize_f32,
    f64 => deserialize_f64,
  }

  forward_to_deserialize_any! {
    char str string unit
    bytes byte_buf map unit_struct tuple_struct
    identifier ignored_any newtype_struct enum
  }
}

impl<'de> de::SeqAccess<'de> for Deserializer<'de> {
  type Error = Error;

  fn next_element_seed<V>(&mut self, seed: V) -> Result<Option<V::Value>>
    where V: de::DeserializeSeed<'de>
  {
    self.count += 1;
    seed.deserialize(self).map(Some)
  }

  fn size_hint(&self) -> Option<usize> {
    match self.parts.iter().size_hint() {
      (lower, Some(upper)) if lower == upper => Some(upper),
      _ => None,
    }
  }
}
<commit_msg>docs: add docs to from_str parser<commit_after>pub mod error;
pub mod helpers;
#[cfg(test)]
mod test;

use self::error::Error;

use std::result::Result as StdResult;

use serde::de::{self, Deserialize, IntoDeserializer};
use serde::de::value::SeqDeserializer;

macro_rules! forward_parsed_values {
  ($($ty:ident => $method:ident,)*) => {
    $(
      fn $method<V>(self, visitor: V) -> Result<V::Value>
        where V: de::Visitor<'de>
      {
        match self.parts.remove(0).parse::<$ty>() {
          Ok(val) => val.into_deserializer().$method(visitor),
          Err(e) => Err(de::Error::custom(e))
        }
      }
    )*
  }
}

pub type Result<T> = StdResult<T, Error>;

struct Deserializer<'de> {
  parts: Vec<&'de str>,
  count: usize
}

/// Parse a command-line parameter string into a Rust type, most often a struct.
///
/// The order that fields are declared in the struct is the order in which they are parsed.
///
/// Note that this handles certain types in radically different ways than what you would expect.
///
/// `Option<T>` will only ever result in `None` if it is at the end of a string of parameters.
/// `Options`s can be chained at the end of a struct, but as soon as a `None` is hit, the rest are
/// guaranteed to be `None`, as well.
///
/// `Vec<T>` will consume every parameter except the ones before it. If you want a list of a
/// specific number of parts, use a fixed-length slice, e.g. `[String; 2]`.
pub fn from_str<'a, T>(input: &'a str) -> Result<T>
  where T: Deserialize<'a>
{
  let parts: Vec<&str> = input.split_whitespace().collect();
  let mut deserializer = Deserializer { parts: parts, count: 0 };
  let t = T::deserialize(&mut deserializer)?;
  if deserializer.parts.is_empty() {
    Ok(t)
  } else {
    Err(Error::Custom("trailing characters".to_owned()))
  }
}

impl<'de, 'a> de::Deserializer<'de> for &'a mut Deserializer<'de> {
  type Error = Error;

  fn deserialize_struct<V>(self, _name: &'static str, fields: &'static [&'static str], visitor: V) -> Result<V::Value>
    where V: de::Visitor<'de>
  {
    self.deserialize_tuple(fields.len(), visitor)
  }

  fn deserialize_tuple<V>(self, _len: usize, visitor: V) -> Result<V::Value>
    where V: de::Visitor<'de>
  {
    visitor.visit_seq(self)
  }

  fn deserialize_any<V>(self, visitor: V) -> Result<V::Value>
    where V: de::Visitor<'de>
  {
    if self.parts.is_empty() {
      return Err(Error::MissingParams);
    }
    let first = self.parts.remove(0);
    first.into_deserializer().deserialize_any(visitor)
  }

  fn deserialize_option<V>(self, visitor: V) -> Result<V::Value>
    where V: de::Visitor<'de>
  {
    if self.parts.is_empty() {
      visitor.visit_none()
    } else {
      visitor.visit_some(self)
    }
  }

  fn deserialize_seq<V>(self, visitor: V) -> Result<V::Value>
    where V: de::Visitor<'de>
  {
    if self.parts.is_empty() {
      return Err(Error::MissingParams);
    }

    SeqDeserializer::new(self.parts.drain(..)).deserialize_seq(visitor)
  }

  forward_parsed_values! {
    bool => deserialize_bool,
    u8 => deserialize_u8,
    u16 => deserialize_u16,
    u32 => deserialize_u32,
    u64 => deserialize_u64,
    i8 => deserialize_i8,
    i16 => deserialize_i16,
    i32 => deserialize_i32,
    i64 => deserialize_i64,
    f32 => deserialize_f32,
    f64 => deserialize_f64,
  }

  forward_to_deserialize_any! {
    char str string unit
    bytes byte_buf map unit_struct tuple_struct
    identifier ignored_any newtype_struct enum
  }
}

impl<'de> de::SeqAccess<'de> for Deserializer<'de> {
  type Error = Error;

  fn next_element_seed<V>(&mut self, seed: V) -> Result<Option<V::Value>>
    where V: de::DeserializeSeed<'de>
  {
    self.count += 1;
    seed.deserialize(self).map(Some)
  }

  fn size_hint(&self) -> Option<usize> {
    match self.parts.iter().size_hint() {
      (lower, Some(upper)) if lower == upper => Some(upper),
      _ => None,
    }
  }
}
<|endoftext|>


<commit_before>#![allow(non_snake_case, non_upper_case_globals)]
// #![cfg_attr(test, feature(test))]
extern crate time;
extern crate num_cpus;
extern crate scoped_threadpool;

use time::precise_time_s;
use scoped_threadpool::Pool;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const aalpha:f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const bbeta:f64 = 0.95;           // Discount factor

// Productivity values
const vProductivity:[f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
static mTransition:[[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                    [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                    [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                    [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                    [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const nGridCapital: usize = 17820;
const nGridProductivity: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let thread_count = num_cpus::get();
    let mut pool = Pool::new(thread_count as u32);

    let capitalSteadyState: f64 = (aalpha * bbeta).powf(1_f64 / (1_f64 - aalpha));
    let outputSteadyState: f64 = capitalSteadyState.powf(aalpha);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let mut vGridCapital = [0f64; nGridCapital];

    for (i, val) in vGridCapital.iter_mut().enumerate() {
        *val = 0.5 * capitalSteadyState + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut mValueFunction: [Vec<f64>; nGridProductivity] = [
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
    ];

    let mut mPolicyFunction: [Vec<f64>; nGridProductivity] = [
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
    ];

    let mut expectedValueFunction: [Vec<f64>; nGridProductivity] = [
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
    ];

    // 4. We pre-build output for each point in the grid
    let mOutput : Vec<[f64; nGridProductivity]> = (0..nGridCapital).map(|nCapital| {
        let mut arr = [0.0; nGridProductivity];

        for (nProductivity, slot) in arr.iter_mut().enumerate() {
            *slot = vProductivity[nProductivity] * vGridCapital[nCapital].powf(aalpha)
        }

        arr
    }).collect();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const tolerance: f64 = 0.0000001; // compiler warn: variable does not need to be mutable

    let mut iteration = 0;

    // small array to split where the writes are going, so we don't need a mutex
    let mut differences = [-100000.0; nGridProductivity];

    while maxDifference > tolerance {
        pool.scoped(|scoped| {
            for (expected_values, transitions) in expectedValueFunction.iter_mut()
                                                                       .zip(mTransition.iter()) {
                // Only capture refs
                let mValueFunction = &mValueFunction;

                scoped.execute(move || {
                    for (idx, expected_value) in expected_values.iter_mut().enumerate() {
                        *expected_value = transitions.iter()
                                                     .zip(mValueFunction.iter())
                                                     .fold(0.0f64, |acc, (transition, value_fns)| {
                            acc + (transition * value_fns[idx])
                        });
                    }
                });
            }
        });

        pool.scoped(|scoped| {
            for (nProductivity, (((policies, value_fns), maxDifference), expected_values)) in mPolicyFunction.iter_mut()
                                                                         .zip(mValueFunction.iter_mut())
                                                                         .zip(differences.iter_mut())
                                                                         .zip(expectedValueFunction.iter())
                                                                         .enumerate() {
                *maxDifference = -100000.0;

                // Only capture refs
                let mOutput = &mOutput;

                scoped.execute(move || {
                    // We start from previous choice (monotonicity of policy function)
                    let mut gridCapitalNextPeriod = 0;

                    for ((policy, output), value_fn) in policies.iter_mut()
                                                                .zip(mOutput.iter())
                                                                .zip(value_fns.iter_mut()) {
                        let mut valueHighSoFar = -100000.0;
                        let mut capitalChoice = vGridCapital[0];
                        let mOutput_cache = &output[nProductivity];

                        for nCapitalNextPeriod in gridCapitalNextPeriod..nGridCapital {
                            let consumption = mOutput_cache - &vGridCapital[nCapitalNextPeriod];
                            let valueProvisional = (1_f64 - bbeta) * (consumption.ln()) +
                                                    bbeta *
                                                    expected_values[nCapitalNextPeriod];

                            if valueProvisional > valueHighSoFar {
                                valueHighSoFar = valueProvisional;
                                capitalChoice = vGridCapital[nCapitalNextPeriod];
                                gridCapitalNextPeriod = nCapitalNextPeriod;
                            } else {
                                break; // We break when we have achieved the max
                            }
                        }

                        let old = mem::replace(value_fn, valueHighSoFar);
                        let diff = (old - valueHighSoFar).abs();
                        if diff > *maxDifference {
                            *maxDifference = diff
                        }

                        *policy = capitalChoice;
                    }
                })
            }
        });

        maxDifference = -100000.0;
        for &diff in &differences {
            if diff > maxDifference {
                maxDifference = diff
            }
        }

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[2][999]
}

fn main() {
    use std::env::args;

    let sample = args().any(|arg| arg == "--sample");

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<commit_msg>Rename globals to rust style<commit_after>#![allow(non_snake_case)]
// #![cfg_attr(test, feature(test))]
extern crate time;
extern crate num_cpus;
extern crate scoped_threadpool;

use time::precise_time_s;
use scoped_threadpool::Pool;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const ALPHA: f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const BETA: f64 = 0.95;           // Discount factor

// Productivity values
const PRODUCTIVITY: [f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
static TRANSITIONS: [[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                     [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                     [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                     [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                     [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const GRID_CAPITAL: usize = 17820;
const GRID_PRODUCTIVITY: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let thread_count = num_cpus::get();
    let mut pool = Pool::new(thread_count as u32);

    let capitalSteadyState: f64 = (ALPHA * BETA).powf(1_f64 / (1_f64 - ALPHA));
    let outputSteadyState: f64 = capitalSteadyState.powf(ALPHA);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let mut vGridCapital = [0f64; GRID_CAPITAL];

    for (i, val) in vGridCapital.iter_mut().enumerate() {
        *val = 0.5 * capitalSteadyState + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut mValueFunction: [Vec<f64>; GRID_PRODUCTIVITY] = [
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
    ];

    let mut mPolicyFunction: [Vec<f64>; GRID_PRODUCTIVITY] = [
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
    ];

    let mut expectedValueFunction: [Vec<f64>; GRID_PRODUCTIVITY] = [
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
    ];

    // 4. We pre-build output for each point in the grid
    let mOutput : Vec<[f64; GRID_PRODUCTIVITY]> = (0..GRID_CAPITAL).map(|nCapital| {
        let mut arr = [0.0; GRID_PRODUCTIVITY];

        for (nProductivity, slot) in arr.iter_mut().enumerate() {
            *slot = PRODUCTIVITY[nProductivity] * vGridCapital[nCapital].powf(ALPHA)
        }

        arr
    }).collect();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const TOLERANCE: f64 = 0.0000001; // compiler warn: variable does not need to be mutable

    let mut iteration = 0;

    // small array to split where the writes are going, so we don't need a mutex
    let mut differences = [-100000.0; GRID_PRODUCTIVITY];

    while maxDifference > TOLERANCE {
        pool.scoped(|scoped| {
            for (expected_values, transitions) in expectedValueFunction.iter_mut()
                                                                       .zip(TRANSITIONS.iter()) {
                // Only capture refs
                let mValueFunction = &mValueFunction;

                scoped.execute(move || {
                    for (idx, expected_value) in expected_values.iter_mut().enumerate() {
                        *expected_value = transitions.iter()
                                                     .zip(mValueFunction.iter())
                                                     .fold(0.0f64, |acc, (transition, value_fns)| {
                            acc + (transition * value_fns[idx])
                        });
                    }
                });
            }
        });

        pool.scoped(|scoped| {
            for (nProductivity, (((policies, value_fns), maxDifference), expected_values)) in mPolicyFunction.iter_mut()
                                                                         .zip(mValueFunction.iter_mut())
                                                                         .zip(differences.iter_mut())
                                                                         .zip(expectedValueFunction.iter())
                                                                         .enumerate() {
                *maxDifference = -100000.0;

                // Only capture refs
                let mOutput = &mOutput;

                scoped.execute(move || {
                    // We start from previous choice (monotonicity of policy function)
                    let mut gridCapitalNextPeriod = 0;

                    for ((policy, output), value_fn) in policies.iter_mut()
                                                                .zip(mOutput.iter())
                                                                .zip(value_fns.iter_mut()) {
                        let mut valueHighSoFar = -100000.0;
                        let mut capitalChoice = vGridCapital[0];
                        let mOutput_cache = &output[nProductivity];

                        for nCapitalNextPeriod in gridCapitalNextPeriod..GRID_CAPITAL {
                            let consumption = mOutput_cache - &vGridCapital[nCapitalNextPeriod];
                            let valueProvisional = (1_f64 - BETA) * (consumption.ln()) +
                                                    BETA *
                                                    expected_values[nCapitalNextPeriod];

                            if valueProvisional > valueHighSoFar {
                                valueHighSoFar = valueProvisional;
                                capitalChoice = vGridCapital[nCapitalNextPeriod];
                                gridCapitalNextPeriod = nCapitalNextPeriod;
                            } else {
                                break; // We break when we have achieved the max
                            }
                        }

                        let old = mem::replace(value_fn, valueHighSoFar);
                        let diff = (old - valueHighSoFar).abs();
                        if diff > *maxDifference {
                            *maxDifference = diff
                        }

                        *policy = capitalChoice;
                    }
                })
            }
        });

        maxDifference = -100000.0;
        for &diff in &differences {
            if diff > maxDifference {
                maxDifference = diff
            }
        }

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[2][999]
}

fn main() {
    use std::env::args;

    let sample = args().any(|arg| arg == "--sample");

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<|endoftext|>


<commit_before>#![allow(non_snake_case, non_upper_case_globals)]
// #![cfg_attr(test, feature(test))]
extern crate time;
extern crate num_cpus;
extern crate scoped_threadpool;

use time::precise_time_s;
use scoped_threadpool::Pool;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const aalpha:f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const bbeta:f64 = 0.95;           // Discount factor

// Productivity values
const vProductivity:[f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
static mTransition:[[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                    [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                    [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                    [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                    [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const nGridCapital: usize = 17820;
const nGridProductivity: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let thread_count = num_cpus::get();
    let mut pool = Pool::new(thread_count as u32);

    let capitalSteadyState: f64 = (aalpha * bbeta).powf(1_f64 / (1_f64 - aalpha));
    let outputSteadyState: f64 = capitalSteadyState.powf(aalpha);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let mut vGridCapital = [0f64; nGridCapital];

    for (i, val) in vGridCapital.iter_mut().enumerate() {
        *val = 0.5 * capitalSteadyState + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut mValueFunction: [Vec<f64>; nGridProductivity] = [
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
    ];

    let mut mPolicyFunction: [Vec<f64>; nGridProductivity] = [
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
    ];

    let mut expectedValueFunction: [Vec<f64>; nGridProductivity] = [
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
    ];

    // 4. We pre-build output for each point in the grid
    let mOutput : Vec<[f64; nGridProductivity]> = (0..nGridCapital).map(|nCapital| {
        let mut arr = [0.0; nGridProductivity];

        for (nProductivity, slot) in arr.iter_mut().enumerate() {
            *slot = vProductivity[nProductivity] * vGridCapital[nCapital].powf(aalpha)
        }

        arr
    }).collect();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const tolerance: f64 = 0.0000001; // compiler warn: variable does not need to be mutable

    let mut iteration = 0;

    while maxDifference > tolerance {
        // small array to split where the writes are going, so we don't need a mutex
        let mut differences = [-100000.0; nGridProductivity];

        pool.scoped(|scoped| {
            for (expected_values, transitions) in expectedValueFunction.iter_mut()
                                                                       .zip(mTransition.iter()) {
                // Only capture refs
                let mValueFunction = &mValueFunction;

                scoped.execute(move || {
                    for (idx, expected_value) in expected_values.iter_mut().enumerate() {
                        *expected_value = transitions.iter()
                                                     .zip(mValueFunction.iter())
                                                     .fold(0.0f64, |acc, (transition, value_fns)| {
                            acc + (transition * value_fns[idx])
                        });
                    }
                });
            }
        });

        pool.scoped(|scoped| {
            for (nProductivity, (((policies, value_fns), maxDifference), expected_values)) in mPolicyFunction.iter_mut()
                                                                         .zip(mValueFunction.iter_mut())
                                                                         .zip(differences.iter_mut())
                                                                         .zip(expectedValueFunction.iter())
                                                                         .enumerate() {
                // Only capture refs
                let mOutput = &mOutput;

                scoped.execute(move || {
                    // We start from previous choice (monotonicity of policy function)
                    let mut gridCapitalNextPeriod = 0;

                    for ((policy, output), value_fn) in policies.iter_mut()
                                                                .zip(mOutput.iter())
                                                                .zip(value_fns.iter_mut()) {
                        let mut valueHighSoFar = -100000.0;
                        let mut capitalChoice = vGridCapital[0];
                        let mOutput_cache = &output[nProductivity];

                        for nCapitalNextPeriod in gridCapitalNextPeriod..nGridCapital {
                            let consumption = mOutput_cache - &vGridCapital[nCapitalNextPeriod];
                            let valueProvisional = (1_f64 - bbeta) * (consumption.ln()) +
                                                    bbeta *
                                                    expected_values[nCapitalNextPeriod];

                            if valueProvisional > valueHighSoFar {
                                valueHighSoFar = valueProvisional;
                                capitalChoice = vGridCapital[nCapitalNextPeriod];
                                gridCapitalNextPeriod = nCapitalNextPeriod;
                            } else {
                                break; // We break when we have achieved the max
                            }
                        }

                        let old = mem::replace(value_fn, valueHighSoFar);
                        let diff = (old - valueHighSoFar).abs();
                        if diff > *maxDifference {
                            *maxDifference = diff
                        }

                        *policy = capitalChoice;
                    }
                })
            }
        });

        maxDifference = -100000.0;
        for &diff in &differences {
            if diff > maxDifference {
                maxDifference = diff
            }
        }

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[2][999]
}

fn main() {
    let sample = false;

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<commit_msg>Move differences array initialization<commit_after>#![allow(non_snake_case, non_upper_case_globals)]
// #![cfg_attr(test, feature(test))]
extern crate time;
extern crate num_cpus;
extern crate scoped_threadpool;

use time::precise_time_s;
use scoped_threadpool::Pool;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const aalpha:f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const bbeta:f64 = 0.95;           // Discount factor

// Productivity values
const vProductivity:[f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
static mTransition:[[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                    [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                    [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                    [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                    [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const nGridCapital: usize = 17820;
const nGridProductivity: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let thread_count = num_cpus::get();
    let mut pool = Pool::new(thread_count as u32);

    let capitalSteadyState: f64 = (aalpha * bbeta).powf(1_f64 / (1_f64 - aalpha));
    let outputSteadyState: f64 = capitalSteadyState.powf(aalpha);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let mut vGridCapital = [0f64; nGridCapital];

    for (i, val) in vGridCapital.iter_mut().enumerate() {
        *val = 0.5 * capitalSteadyState + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut mValueFunction: [Vec<f64>; nGridProductivity] = [
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
    ];

    let mut mPolicyFunction: [Vec<f64>; nGridProductivity] = [
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
    ];

    let mut expectedValueFunction: [Vec<f64>; nGridProductivity] = [
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
        vec![0f64; nGridCapital],
    ];

    // 4. We pre-build output for each point in the grid
    let mOutput : Vec<[f64; nGridProductivity]> = (0..nGridCapital).map(|nCapital| {
        let mut arr = [0.0; nGridProductivity];

        for (nProductivity, slot) in arr.iter_mut().enumerate() {
            *slot = vProductivity[nProductivity] * vGridCapital[nCapital].powf(aalpha)
        }

        arr
    }).collect();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const tolerance: f64 = 0.0000001; // compiler warn: variable does not need to be mutable

    let mut iteration = 0;

    // small array to split where the writes are going, so we don't need a mutex
    let mut differences = [-100000.0; nGridProductivity];

    while maxDifference > tolerance {
        pool.scoped(|scoped| {
            for (expected_values, transitions) in expectedValueFunction.iter_mut()
                                                                       .zip(mTransition.iter()) {
                // Only capture refs
                let mValueFunction = &mValueFunction;

                scoped.execute(move || {
                    for (idx, expected_value) in expected_values.iter_mut().enumerate() {
                        *expected_value = transitions.iter()
                                                     .zip(mValueFunction.iter())
                                                     .fold(0.0f64, |acc, (transition, value_fns)| {
                            acc + (transition * value_fns[idx])
                        });
                    }
                });
            }
        });

        pool.scoped(|scoped| {
            for (nProductivity, (((policies, value_fns), maxDifference), expected_values)) in mPolicyFunction.iter_mut()
                                                                         .zip(mValueFunction.iter_mut())
                                                                         .zip(differences.iter_mut())
                                                                         .zip(expectedValueFunction.iter())
                                                                         .enumerate() {
                *maxDifference = -100000.0;

                // Only capture refs
                let mOutput = &mOutput;

                scoped.execute(move || {
                    // We start from previous choice (monotonicity of policy function)
                    let mut gridCapitalNextPeriod = 0;

                    for ((policy, output), value_fn) in policies.iter_mut()
                                                                .zip(mOutput.iter())
                                                                .zip(value_fns.iter_mut()) {
                        let mut valueHighSoFar = -100000.0;
                        let mut capitalChoice = vGridCapital[0];
                        let mOutput_cache = &output[nProductivity];

                        for nCapitalNextPeriod in gridCapitalNextPeriod..nGridCapital {
                            let consumption = mOutput_cache - &vGridCapital[nCapitalNextPeriod];
                            let valueProvisional = (1_f64 - bbeta) * (consumption.ln()) +
                                                    bbeta *
                                                    expected_values[nCapitalNextPeriod];

                            if valueProvisional > valueHighSoFar {
                                valueHighSoFar = valueProvisional;
                                capitalChoice = vGridCapital[nCapitalNextPeriod];
                                gridCapitalNextPeriod = nCapitalNextPeriod;
                            } else {
                                break; // We break when we have achieved the max
                            }
                        }

                        let old = mem::replace(value_fn, valueHighSoFar);
                        let diff = (old - valueHighSoFar).abs();
                        if diff > *maxDifference {
                            *maxDifference = diff
                        }

                        *policy = capitalChoice;
                    }
                })
            }
        });

        maxDifference = -100000.0;
        for &diff in &differences {
            if diff > maxDifference {
                maxDifference = diff
            }
        }

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[2][999]
}

fn main() {
    let sample = false;

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<|endoftext|>