<commit_before>/* Multi-producer/multi-consumer bounded queue
 * Copyright (c) 2010-2011 Dmitry Vyukov. All rights reserved.
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *    1. Redistributions of source code must retain the above copyright notice,
 *       this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY DMITRY VYUKOV "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
 * SHALL DMITRY VYUKOV OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * The views and conclusions contained in the software and documentation are
 * those of the authors and should not be interpreted as representing official
 * policies, either expressed or implied, of Dmitry Vyukov.
 */

use unstable::sync::UnsafeArc;
use unstable::atomics::{AtomicUint,Relaxed,Release,Acquire};
use option::*;
use vec;
use clone::Clone;
use kinds::Send;
use num::{Exponential,Algebraic,Round};

struct Node<T> {
    sequence: AtomicUint,
    value: Option<T>,
}

struct State<T> {
    buffer: ~[Node<T>],
    mask: uint,
    enqueue_pos: AtomicUint,
    dequeue_pos: AtomicUint,
}

struct Queue<T> {
    priv state: UnsafeArc<State<T>>,
}

impl<T: Send> State<T> {
    fn with_capacity(capacity: uint) -> State<T> {
        let capacity = if capacity < 2 || (capacity & (capacity - 1)) != 0 {
            // use next power of 2 as capacity
            2f64.pow(&((capacity as f64).log2().floor()+1f64)) as uint
        } else {
            capacity
        };
        let buffer = do vec::from_fn(capacity) |i:uint| {
            Node{sequence:AtomicUint::new(i),value:None}
        };
        State{
            buffer: buffer,
            mask: capacity-1,
            enqueue_pos: AtomicUint::new(0),
            dequeue_pos: AtomicUint::new(0),
        }
    }

    fn push(&mut self, value: T) -> bool {
        let mask = self.mask;
        let mut pos = self.enqueue_pos.load(Relaxed);
        loop {
            let node = &mut self.buffer[pos & mask];
            let seq = node.sequence.load(Acquire);
            let diff: int = seq as int - pos as int;

            if diff == 0 {
                let enqueue_pos = self.enqueue_pos.compare_and_swap(pos, pos+1, Relaxed);
                if enqueue_pos == pos {
                    node.value = Some(value);
                    node.sequence.store(pos+1, Release);
                    break
                } else {
                    pos = enqueue_pos;
                }
            } else if (diff < 0) {
                return false
            } else {
                pos = self.enqueue_pos.load(Relaxed);
            }
        }
        true
    }

    fn pop(&mut self) -> Option<T> {
        let mask = self.mask;
        let mut pos = self.dequeue_pos.load(Relaxed);
        loop {
            let node = &mut self.buffer[pos & mask];
            let seq = node.sequence.load(Acquire);
            let diff: int = seq as int - (pos + 1) as int;
            if diff == 0 {
                let dequeue_pos = self.dequeue_pos.compare_and_swap(pos, pos+1, Relaxed);
                if dequeue_pos == pos {
                    let value = node.value.take();
                    node.sequence.store(pos + mask + 1, Release);
                    return value
                } else {
                    pos = dequeue_pos;
                }
            } else if diff < 0 {
                return None
            } else {
                pos = self.dequeue_pos.load(Relaxed);
            }
        }
    }
}

impl<T: Send> Queue<T> {
    pub fn with_capacity(capacity: uint) -> Queue<T> {
        Queue{
            state: UnsafeArc::new(State::with_capacity(capacity))
        }
    }

    pub fn push(&mut self, value: T) -> bool {
        unsafe { (*self.state.get()).push(value) }
    }

    pub fn pop(&mut self) -> Option<T> {
        unsafe { (*self.state.get()).pop() }
    }
}

impl<T: Send> Clone for Queue<T> {
    fn clone(&self) -> Queue<T> {
        Queue {
            state: self.state.clone()
        }
    }
}

#[cfg(test)]
mod tests {
    use prelude::*;
    use option::*;
    use task;
    use comm;
    use super::Queue;

    #[test]
    fn test() {
        let nthreads = 8u;
        let nmsgs = 1000u;
        let mut q = Queue::with_capacity(nthreads*nmsgs);
        assert_eq!(None, q.pop());

        for _ in range(0, nthreads) {
            let (port, chan)  = comm::stream();
            chan.send(q.clone());
            do task::spawn_sched(task::SingleThreaded) {
                let mut q = port.recv();
                for i in range(0, nmsgs) {
                    assert!(q.push(i));
                }
            }
        }

        let mut completion_ports = ~[];
        for _ in range(0, nthreads) {
            let (completion_port, completion_chan) = comm::stream();
            completion_ports.push(completion_port);
            let (port, chan)  = comm::stream();
            chan.send(q.clone());
            do task::spawn_sched(task::SingleThreaded) {
                let mut q = port.recv();
                let mut i = 0u;
                loop {
                    match q.pop() {
                        None => {},
                        Some(_) => {
                            i += 1;
                            if i == nmsgs { break }
                        }
                    }
                }
                completion_chan.send(i);
            }
        }

        for completion_port in completion_ports.iter() {
            assert_eq!(nmsgs, completion_port.recv());
        }
    }
}
<commit_msg>add padding to prevent false sharing<commit_after>/* Multi-producer/multi-consumer bounded queue
 * Copyright (c) 2010-2011 Dmitry Vyukov. All rights reserved.
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *    1. Redistributions of source code must retain the above copyright notice,
 *       this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY DMITRY VYUKOV "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
 * SHALL DMITRY VYUKOV OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * The views and conclusions contained in the software and documentation are
 * those of the authors and should not be interpreted as representing official
 * policies, either expressed or implied, of Dmitry Vyukov.
 */

use unstable::sync::UnsafeArc;
use unstable::atomics::{AtomicUint,Relaxed,Release,Acquire};
use option::*;
use vec;
use clone::Clone;
use kinds::Send;
use num::{Exponential,Algebraic,Round};

struct Node<T> {
    sequence: AtomicUint,
    value: Option<T>,
}

struct State<T> {
    pad0: [u8, ..64],
    buffer: ~[Node<T>],
    mask: uint,
    pad1: [u8, ..64],
    enqueue_pos: AtomicUint,
    pad2: [u8, ..64],
    dequeue_pos: AtomicUint,
    pad3: [u8, ..64],
}

struct Queue<T> {
    priv state: UnsafeArc<State<T>>,
}

impl<T: Send> State<T> {
    fn with_capacity(capacity: uint) -> State<T> {
        let capacity = if capacity < 2 || (capacity & (capacity - 1)) != 0 {
            // use next power of 2 as capacity
            2f64.pow(&((capacity as f64).log2().floor()+1f64)) as uint
        } else {
            capacity
        };
        let buffer = do vec::from_fn(capacity) |i:uint| {
            Node{sequence:AtomicUint::new(i),value:None}
        };
        State{
            pad0: [0, ..64],
            buffer: buffer,
            mask: capacity-1,
            pad1: [0, ..64],
            enqueue_pos: AtomicUint::new(0),
            pad2: [0, ..64],
            dequeue_pos: AtomicUint::new(0),
            pad3: [0, ..64],
        }
    }

    fn push(&mut self, value: T) -> bool {
        let mask = self.mask;
        let mut pos = self.enqueue_pos.load(Relaxed);
        loop {
            let node = &mut self.buffer[pos & mask];
            let seq = node.sequence.load(Acquire);
            let diff: int = seq as int - pos as int;

            if diff == 0 {
                let enqueue_pos = self.enqueue_pos.compare_and_swap(pos, pos+1, Relaxed);
                if enqueue_pos == pos {
                    node.value = Some(value);
                    node.sequence.store(pos+1, Release);
                    break
                } else {
                    pos = enqueue_pos;
                }
            } else if (diff < 0) {
                return false
            } else {
                pos = self.enqueue_pos.load(Relaxed);
            }
        }
        true
    }

    fn pop(&mut self) -> Option<T> {
        let mask = self.mask;
        let mut pos = self.dequeue_pos.load(Relaxed);
        loop {
            let node = &mut self.buffer[pos & mask];
            let seq = node.sequence.load(Acquire);
            let diff: int = seq as int - (pos + 1) as int;
            if diff == 0 {
                let dequeue_pos = self.dequeue_pos.compare_and_swap(pos, pos+1, Relaxed);
                if dequeue_pos == pos {
                    let value = node.value.take();
                    node.sequence.store(pos + mask + 1, Release);
                    return value
                } else {
                    pos = dequeue_pos;
                }
            } else if diff < 0 {
                return None
            } else {
                pos = self.dequeue_pos.load(Relaxed);
            }
        }
    }
}

impl<T: Send> Queue<T> {
    pub fn with_capacity(capacity: uint) -> Queue<T> {
        Queue{
            state: UnsafeArc::new(State::with_capacity(capacity))
        }
    }

    pub fn push(&mut self, value: T) -> bool {
        unsafe { (*self.state.get()).push(value) }
    }

    pub fn pop(&mut self) -> Option<T> {
        unsafe { (*self.state.get()).pop() }
    }
}

impl<T: Send> Clone for Queue<T> {
    fn clone(&self) -> Queue<T> {
        Queue {
            state: self.state.clone()
        }
    }
}

#[cfg(test)]
mod tests {
    use prelude::*;
    use option::*;
    use task;
    use comm;
    use super::Queue;

    #[test]
    fn test() {
        let nthreads = 8u;
        let nmsgs = 1000u;
        let mut q = Queue::with_capacity(nthreads*nmsgs);
        assert_eq!(None, q.pop());

        for _ in range(0, nthreads) {
            let (port, chan)  = comm::stream();
            chan.send(q.clone());
            do task::spawn_sched(task::SingleThreaded) {
                let mut q = port.recv();
                for i in range(0, nmsgs) {
                    assert!(q.push(i));
                }
            }
        }

        let mut completion_ports = ~[];
        for _ in range(0, nthreads) {
            let (completion_port, completion_chan) = comm::stream();
            completion_ports.push(completion_port);
            let (port, chan)  = comm::stream();
            chan.send(q.clone());
            do task::spawn_sched(task::SingleThreaded) {
                let mut q = port.recv();
                let mut i = 0u;
                loop {
                    match q.pop() {
                        None => {},
                        Some(_) => {
                            i += 1;
                            if i == nmsgs { break }
                        }
                    }
                }
                completion_chan.send(i);
            }
        }

        for completion_port in completion_ports.iter() {
            assert_eq!(nmsgs, completion_port.recv());
        }
    }
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2016-2017, Robert Ou <rqou@robertou.com>
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

use std::marker::PhantomData;

use analyzer::util::*;

// We need this because the Vec can be reallocated around, but we want to keep
// around some kind of reference into the storage that can keep working after
// the reallocation happens.
#[derive(Copy, Clone)]
pub struct StringPoolIndex {
    start: usize,
    end: usize,
}

pub struct StringPool {
    storage: Vec<u8>
}

impl StringPool {
    pub fn new() -> StringPool {
        StringPool {storage: Vec::new()}
    }

    // TODO: Might interning strings improve performance in the future?
    // TODO: Do we care about freeing things?
    pub fn add_latin1_str(&mut self, inp: &[u8]) -> StringPoolIndex {
        let addition_begin = self.storage.len();
        let addition_length = inp.len();

        self.storage.extend(inp.iter().cloned());

        StringPoolIndex {
            start: addition_begin,
            end: addition_begin + addition_length
        }
    }

    pub fn retrieve_latin1_str(&self, i: StringPoolIndex) -> Latin1Str {
        let the_slice = &self.storage[i.start..i.end];
        Latin1Str::new(the_slice)
    }
}

pub struct ObjPoolIndex<T> {
    i: usize,
    type_marker: PhantomData<T>
}

impl<T> Copy for ObjPoolIndex<T> { }

impl<T> Clone for ObjPoolIndex<T> {
    fn clone(&self) -> ObjPoolIndex<T> {
        *self
    }
}

pub struct ObjPool<T> {
    storage: Vec<T>
}

impl<T: Default> ObjPool<T> {
    pub fn new() -> ObjPool<T> {
        ObjPool {storage: Vec::new()}
    }

    pub fn alloc(&mut self) -> ObjPoolIndex<T> {
        let i = self.storage.len();
        let o = T::default();

        self.storage.push(o);

        ObjPoolIndex::<T> {i: i, type_marker: PhantomData}
    }

    pub fn get(&self, i: ObjPoolIndex<T>) -> &T {
        &self.storage[i.i]
    }

    pub fn get_mut(&mut self, i: ObjPoolIndex<T>) -> &mut T {
        &mut self.storage[i.i]
    }
}

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

    #[test]
    fn stringpool_basic_works() {
        let mut sp = StringPool::new();
        let x = sp.add_latin1_str(b"test1");
        let y = sp.add_latin1_str(b"test2");
        let sx = sp.retrieve_latin1_str(x);
        let sy = sp.retrieve_latin1_str(y);
        assert_eq!(sx.raw_name(), b"test1");
        assert_eq!(sy.raw_name(), b"test2");
    }

    #[test]
    fn stringpool_actually_pools() {
        let mut sp = StringPool::new();
        sp.add_latin1_str(b"test1");
        sp.add_latin1_str(b"test2");
        assert_eq!(sp.storage, b"test1test2");
    }

    #[derive(Default)]
    struct ObjPoolTestObject {
        foo: u32
    }

    #[test]
    fn objpool_basic_works() {
        let mut pool = ObjPool::<ObjPoolTestObject>::new();
        let x = pool.alloc();
        let y = pool.alloc();
        {
            let o = pool.get_mut(x);
            o.foo = 123;
        }
        {
            let o = pool.get_mut(y);
            o.foo = 456;
        }
        let ox = pool.get(x);
        let oy = pool.get(y);
        assert_eq!(ox.foo, 123);
        assert_eq!(oy.foo, 456);
    }
}
<commit_msg>Added string interning in the string pool<commit_after>/*
Copyright (c) 2016-2017, Robert Ou <rqou@robertou.com>
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

use std::collections::HashMap;
use std::marker::PhantomData;

use analyzer::util::*;

// We need this because the Vec can be reallocated around, but we want to keep
// around some kind of reference into the storage that can keep working after
// the reallocation happens.
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub struct StringPoolIndex {
    start: usize,
    end: usize,
}

pub struct StringPool {
    storage: Vec<u8>,
    hashtable: HashMap<Vec<u8>, StringPoolIndex>,
}

impl<'a> StringPool {
    pub fn new() -> StringPool {
        StringPool {
            storage: Vec::new(),
            hashtable: HashMap::new(),
        }
    }

    // TODO: Do we care about freeing things?
    pub fn add_latin1_str(&mut self, inp: &[u8]) -> StringPoolIndex {
        // We now need to intern the string because we need to ensure that
        // we always can compare string pool indices. This is needed so that
        // Identifiers can be compared without having to drag the string
        // pool around, which is in turn needed so that Identifiers can be
        // properly hashed.
        if let Some(existing_idx) = self.hashtable.get(inp) {
            return *existing_idx;
        }

        let addition_begin = self.storage.len();
        let addition_end = addition_begin + inp.len();

        self.storage.extend(inp.iter().cloned());

        let new_idx = StringPoolIndex {
            start: addition_begin,
            end: addition_end
        };

        self.hashtable.insert(inp.to_vec(), new_idx);

        new_idx
    }

    pub fn retrieve_latin1_str(&self, i: StringPoolIndex) -> Latin1Str {
        let the_slice = &self.storage[i.start..i.end];
        Latin1Str::new(the_slice)
    }
}

pub struct ObjPoolIndex<T> {
    i: usize,
    type_marker: PhantomData<T>
}

impl<T> Copy for ObjPoolIndex<T> { }

impl<T> Clone for ObjPoolIndex<T> {
    fn clone(&self) -> ObjPoolIndex<T> {
        *self
    }
}

pub struct ObjPool<T> {
    storage: Vec<T>
}

impl<T: Default> ObjPool<T> {
    pub fn new() -> ObjPool<T> {
        ObjPool {storage: Vec::new()}
    }

    pub fn alloc(&mut self) -> ObjPoolIndex<T> {
        let i = self.storage.len();
        let o = T::default();

        self.storage.push(o);

        ObjPoolIndex::<T> {i: i, type_marker: PhantomData}
    }

    pub fn get(&self, i: ObjPoolIndex<T>) -> &T {
        &self.storage[i.i]
    }

    pub fn get_mut(&mut self, i: ObjPoolIndex<T>) -> &mut T {
        &mut self.storage[i.i]
    }
}

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

    #[test]
    fn stringpool_basic_works() {
        let mut sp = StringPool::new();
        let x = sp.add_latin1_str(b"test1");
        let y = sp.add_latin1_str(b"test2");
        let sx = sp.retrieve_latin1_str(x);
        let sy = sp.retrieve_latin1_str(y);
        assert_eq!(sx.raw_name(), b"test1");
        assert_eq!(sy.raw_name(), b"test2");
    }

    #[test]
    fn stringpool_actually_pools() {
        let mut sp = StringPool::new();
        sp.add_latin1_str(b"test1");
        sp.add_latin1_str(b"test2");
        assert_eq!(sp.storage, b"test1test2");
    }

    #[test]
    fn stringpool_interns() {
        let mut sp = StringPool::new();
        let x = sp.add_latin1_str(b"test1");
        let y = sp.add_latin1_str(b"test1");
        assert_eq!(x, y);
        assert_eq!(sp.storage, b"test1");
    }

    #[derive(Default)]
    struct ObjPoolTestObject {
        foo: u32
    }

    #[test]
    fn objpool_basic_works() {
        let mut pool = ObjPool::<ObjPoolTestObject>::new();
        let x = pool.alloc();
        let y = pool.alloc();
        {
            let o = pool.get_mut(x);
            o.foo = 123;
        }
        {
            let o = pool.get_mut(y);
            o.foo = 456;
        }
        let ox = pool.get(x);
        let oy = pool.get(y);
        assert_eq!(ox.foo, 123);
        assert_eq!(oy.foo, 456);
    }
}
<|endoftext|>


<commit_before>extern crate byteorder;
extern crate futures;
extern crate linear_map;
extern crate protobuf;

extern crate librespot_core as core;
extern crate librespot_protocol as protocol;

pub mod cover;

use futures::Future;
use linear_map::LinearMap;

use core::mercury::MercuryError;
use core::session::Session;
use core::util::{SpotifyId, FileId, StrChunksExt};

pub use protocol::metadata::AudioFile_Format as FileFormat;

fn countrylist_contains(list: &str, country: &str) -> bool {
    list.chunks(2).any(|cc| cc == country)
}

fn parse_restrictions<'s, I>(restrictions: I, country: &str, catalogue: &str) -> bool
    where I: IntoIterator<Item = &'s protocol::metadata::Restriction>
{
    let mut forbidden = "".to_string();
    let mut has_forbidden = false;

    let mut allowed = "".to_string();
    let mut has_allowed = false;

    let rs = restrictions.into_iter().filter(|r|
        r.get_catalogue_str().contains(&catalogue.to_owned())
    );

    for r in rs {
        if r.has_countries_forbidden() {
            forbidden.push_str(r.get_countries_forbidden());
            has_forbidden = true;
        }

        if r.has_countries_allowed() {
            allowed.push_str(r.get_countries_allowed());
            has_allowed = true;
        }
    }

    (has_forbidden || has_allowed) &&
    (!has_forbidden || !countrylist_contains(forbidden.as_str(), country)) &&
    (!has_allowed || countrylist_contains(allowed.as_str(), country))
}

pub trait Metadata : Send + Sized + 'static {
    type Message: protobuf::MessageStatic;

    fn base_url() -> &'static str;
    fn parse(msg: &Self::Message, session: &Session) -> Self;

    fn get(session: &Session, id: SpotifyId) -> Box<Future<Item = Self, Error = MercuryError>> {
        let uri = format!("{}/{}", Self::base_url(), id.to_base16());
        let request = session.mercury().get(uri);

        let session = session.clone();
        Box::new(request.and_then(move |response| {
            let data = response.payload.first().expect("Empty payload");
            let msg: Self::Message = protobuf::parse_from_bytes(data).unwrap();

            Ok(Self::parse(&msg, &session))
        }))
    }
}

#[derive(Debug, Clone)]
pub struct Track {
    pub id: SpotifyId,
    pub name: String,
    pub album: SpotifyId,
    pub artists: Vec<SpotifyId>,
    pub files: LinearMap<FileFormat, FileId>,
    pub alternatives: Vec<SpotifyId>,
    pub available: bool,
}

#[derive(Debug, Clone)]
pub struct Album {
    pub id: SpotifyId,
    pub name: String,
    pub artists: Vec<SpotifyId>,
    pub tracks: Vec<SpotifyId>,
    pub covers: Vec<FileId>,
}

#[derive(Debug, Clone)]
pub struct Artist {
    pub id: SpotifyId,
    pub name: String,
    pub top_tracks: Vec<SpotifyId>,
}

impl Metadata for Track {
    type Message = protocol::metadata::Track;

    fn base_url() -> &'static str {
        "hm://metadata/3/track"
    }

    fn parse(msg: &Self::Message, session: &Session) -> Self {
        let country = session.country();

        let artists = msg.get_artist()
                         .iter()
                         .filter(|artist| artist.has_gid())
                         .map(|artist| SpotifyId::from_raw(artist.get_gid()))
                         .collect::<Vec<_>>();

        let files = msg.get_file()
                       .iter()
                       .filter(|file| file.has_file_id())
                       .map(|file| {
                           let mut dst = [0u8; 20];
                           dst.clone_from_slice(file.get_file_id());
                           (file.get_format(), FileId(dst))
                       })
                       .collect();

        Track {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            album: SpotifyId::from_raw(msg.get_album().get_gid()),
            artists: artists,
            files: files,
            alternatives: msg.get_alternative()
                             .iter()
                             .map(|alt| SpotifyId::from_raw(alt.get_gid()))
                             .collect(),
            available: parse_restrictions(msg.get_restriction(),
                                          &country,
                                          "premium"),
        }
    }
}

impl Metadata for Album {
    type Message = protocol::metadata::Album;

    fn base_url() -> &'static str {
        "hm://metadata/3/album"
    }

    fn parse(msg: &Self::Message, _: &Session) -> Self {
        let artists = msg.get_artist()
                         .iter()
                         .filter(|artist| artist.has_gid())
                         .map(|artist| SpotifyId::from_raw(artist.get_gid()))
                         .collect::<Vec<_>>();

        let tracks = msg.get_disc()
                        .iter()
                        .flat_map(|disc| disc.get_track())
                        .filter(|track| track.has_gid())
                        .map(|track| SpotifyId::from_raw(track.get_gid()))
                        .collect::<Vec<_>>();

        let covers = msg.get_cover_group()
                        .get_image()
                        .iter()
                        .filter(|image| image.has_file_id())
                        .map(|image| {
                            let mut dst = [0u8; 20];
                            dst.clone_from_slice(image.get_file_id());
                            FileId(dst)
                        })
                        .collect::<Vec<_>>();

        Album {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            artists: artists,
            tracks: tracks,
            covers: covers,
        }
    }
}

impl Metadata for Artist {
    type Message = protocol::metadata::Artist;

    fn base_url() -> &'static str {
        "hm://metadata/3/artist"
    }

    fn parse(msg: &Self::Message, session: &Session) -> Self {
        let country = session.country();

        let top_tracks: Vec<SpotifyId> = match msg.get_top_track()
                            .iter()
                            .find(|tt| !tt.has_country() || countrylist_contains(tt.get_country(), &country)) {
                                Some(tracks) => {
                                    tracks.get_track()
                                    .iter()
                                    .filter(|track| track.has_gid())
                                    .map(|track| SpotifyId::from_raw(track.get_gid()))
                                    .collect::<Vec<_>>()
                                },
                                None => Vec::new()
                            };

        Artist {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            top_tracks: top_tracks
        }
    }
}

<commit_msg>Add track duration to `Track` metadata<commit_after>extern crate byteorder;
extern crate futures;
extern crate linear_map;
extern crate protobuf;

extern crate librespot_core as core;
extern crate librespot_protocol as protocol;

pub mod cover;

use futures::Future;
use linear_map::LinearMap;

use core::mercury::MercuryError;
use core::session::Session;
use core::util::{SpotifyId, FileId, StrChunksExt};

pub use protocol::metadata::AudioFile_Format as FileFormat;

fn countrylist_contains(list: &str, country: &str) -> bool {
    list.chunks(2).any(|cc| cc == country)
}

fn parse_restrictions<'s, I>(restrictions: I, country: &str, catalogue: &str) -> bool
    where I: IntoIterator<Item = &'s protocol::metadata::Restriction>
{
    let mut forbidden = "".to_string();
    let mut has_forbidden = false;

    let mut allowed = "".to_string();
    let mut has_allowed = false;

    let rs = restrictions.into_iter().filter(|r|
        r.get_catalogue_str().contains(&catalogue.to_owned())
    );

    for r in rs {
        if r.has_countries_forbidden() {
            forbidden.push_str(r.get_countries_forbidden());
            has_forbidden = true;
        }

        if r.has_countries_allowed() {
            allowed.push_str(r.get_countries_allowed());
            has_allowed = true;
        }
    }

    (has_forbidden || has_allowed) &&
    (!has_forbidden || !countrylist_contains(forbidden.as_str(), country)) &&
    (!has_allowed || countrylist_contains(allowed.as_str(), country))
}

pub trait Metadata : Send + Sized + 'static {
    type Message: protobuf::MessageStatic;

    fn base_url() -> &'static str;
    fn parse(msg: &Self::Message, session: &Session) -> Self;

    fn get(session: &Session, id: SpotifyId) -> Box<Future<Item = Self, Error = MercuryError>> {
        let uri = format!("{}/{}", Self::base_url(), id.to_base16());
        let request = session.mercury().get(uri);

        let session = session.clone();
        Box::new(request.and_then(move |response| {
            let data = response.payload.first().expect("Empty payload");
            let msg: Self::Message = protobuf::parse_from_bytes(data).unwrap();

            Ok(Self::parse(&msg, &session))
        }))
    }
}

#[derive(Debug, Clone)]
pub struct Track {
    pub id: SpotifyId,
    pub name: String,
    pub duration: i32,
    pub album: SpotifyId,
    pub artists: Vec<SpotifyId>,
    pub files: LinearMap<FileFormat, FileId>,
    pub alternatives: Vec<SpotifyId>,
    pub available: bool,
}

#[derive(Debug, Clone)]
pub struct Album {
    pub id: SpotifyId,
    pub name: String,
    pub artists: Vec<SpotifyId>,
    pub tracks: Vec<SpotifyId>,
    pub covers: Vec<FileId>,
}

#[derive(Debug, Clone)]
pub struct Artist {
    pub id: SpotifyId,
    pub name: String,
    pub top_tracks: Vec<SpotifyId>,
}

impl Metadata for Track {
    type Message = protocol::metadata::Track;

    fn base_url() -> &'static str {
        "hm://metadata/3/track"
    }

    fn parse(msg: &Self::Message, session: &Session) -> Self {
        let country = session.country();

        let artists = msg.get_artist()
                         .iter()
                         .filter(|artist| artist.has_gid())
                         .map(|artist| SpotifyId::from_raw(artist.get_gid()))
                         .collect::<Vec<_>>();

        let files = msg.get_file()
                       .iter()
                       .filter(|file| file.has_file_id())
                       .map(|file| {
                           let mut dst = [0u8; 20];
                           dst.clone_from_slice(file.get_file_id());
                           (file.get_format(), FileId(dst))
                       })
                       .collect();

        Track {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            duration: msg.get_duration(),
            album: SpotifyId::from_raw(msg.get_album().get_gid()),
            artists: artists,
            files: files,
            alternatives: msg.get_alternative()
                             .iter()
                             .map(|alt| SpotifyId::from_raw(alt.get_gid()))
                             .collect(),
            available: parse_restrictions(msg.get_restriction(),
                                          &country,
                                          "premium"),
        }
    }
}

impl Metadata for Album {
    type Message = protocol::metadata::Album;

    fn base_url() -> &'static str {
        "hm://metadata/3/album"
    }

    fn parse(msg: &Self::Message, _: &Session) -> Self {
        let artists = msg.get_artist()
                         .iter()
                         .filter(|artist| artist.has_gid())
                         .map(|artist| SpotifyId::from_raw(artist.get_gid()))
                         .collect::<Vec<_>>();

        let tracks = msg.get_disc()
                        .iter()
                        .flat_map(|disc| disc.get_track())
                        .filter(|track| track.has_gid())
                        .map(|track| SpotifyId::from_raw(track.get_gid()))
                        .collect::<Vec<_>>();

        let covers = msg.get_cover_group()
                        .get_image()
                        .iter()
                        .filter(|image| image.has_file_id())
                        .map(|image| {
                            let mut dst = [0u8; 20];
                            dst.clone_from_slice(image.get_file_id());
                            FileId(dst)
                        })
                        .collect::<Vec<_>>();

        Album {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            artists: artists,
            tracks: tracks,
            covers: covers,
        }
    }
}

impl Metadata for Artist {
    type Message = protocol::metadata::Artist;

    fn base_url() -> &'static str {
        "hm://metadata/3/artist"
    }

    fn parse(msg: &Self::Message, session: &Session) -> Self {
        let country = session.country();

        let top_tracks: Vec<SpotifyId> = match msg.get_top_track()
                            .iter()
                            .find(|tt| !tt.has_country() || countrylist_contains(tt.get_country(), &country)) {
                                Some(tracks) => {
                                    tracks.get_track()
                                    .iter()
                                    .filter(|track| track.has_gid())
                                    .map(|track| SpotifyId::from_raw(track.get_gid()))
                                    .collect::<Vec<_>>()
                                },
                                None => Vec::new()
                            };

        Artist {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            top_tracks: top_tracks
        }
    }
}
<|endoftext|>


<commit_before>pub use prometheus::Error;
use prometheus::{Histogram, HistogramOpts, IntCounter, Opts, Registry};

pub struct Metrics {
    pub block_processing_requests: IntCounter,
    pub block_processing_successes: IntCounter,
    pub block_processing_times: Histogram,
    pub block_production_requests: IntCounter,
    pub block_production_successes: IntCounter,
    pub block_production_times: Histogram,
    pub attestation_production_requests: IntCounter,
    pub attestation_production_successes: IntCounter,
    pub attestation_production_times: Histogram,
    pub attestation_processing_requests: IntCounter,
    pub attestation_processing_successes: IntCounter,
    pub attestation_processing_times: Histogram,
    pub fork_choice_requests: IntCounter,
    pub fork_choice_changed_head: IntCounter,
    pub fork_choice_reorg_count: IntCounter,
    pub fork_choice_times: Histogram,
}

impl Metrics {
    pub fn new() -> Result<Self, Error> {
        Ok(Self {
            block_processing_requests: {
                let opts = Opts::new("blocks_processed", "total_blocks_processed");
                IntCounter::with_opts(opts)?
            },
            block_processing_successes: {
                let opts = Opts::new("valid_blocks_processed", "total_valid_blocks_processed");
                IntCounter::with_opts(opts)?
            },
            block_processing_times: {
                let opts = HistogramOpts::new("block_processing_times", "block_processing_time");
                Histogram::with_opts(opts)?
            },
            block_production_requests: {
                let opts = Opts::new("block_production_requests", "attempts_to_produce_new_block");
                IntCounter::with_opts(opts)?
            },
            block_production_successes: {
                let opts = Opts::new("block_production_successes", "blocks_successfully_produced");
                IntCounter::with_opts(opts)?
            },
            block_production_times: {
                let opts = HistogramOpts::new("block_production_times", "block_production_time");
                Histogram::with_opts(opts)?
            },
            attestation_production_requests: {
                let opts = Opts::new(
                    "attestation_production_requests",
                    "total_attestation_production_requests",
                );
                IntCounter::with_opts(opts)?
            },
            attestation_production_successes: {
                let opts = Opts::new(
                    "attestation_production_successes",
                    "total_attestation_production_successes",
                );
                IntCounter::with_opts(opts)?
            },
            attestation_production_times: {
                let opts = HistogramOpts::new(
                    "attestation_production_times",
                    "attestation_production_time",
                );
                Histogram::with_opts(opts)?
            },
            attestation_processing_requests: {
                let opts = Opts::new(
                    "attestation_processing_requests",
                    "total_attestation_processing_requests",
                );
                IntCounter::with_opts(opts)?
            },
            attestation_processing_successes: {
                let opts = Opts::new(
                    "attestation_processing_successes",
                    "total_attestation_processing_successes",
                );
                IntCounter::with_opts(opts)?
            },
            attestation_processing_times: {
                let opts = HistogramOpts::new(
                    "attestation_processing_times",
                    "attestation_processing_time",
                );
                Histogram::with_opts(opts)?
            },
            fork_choice_requests: {
                let opts = Opts::new("fork_choice_requests", "total_times_fork_choice_called");
                IntCounter::with_opts(opts)?
            },
            fork_choice_changed_head: {
                let opts = Opts::new(
                    "fork_choice_changed_head",
                    "total_times_fork_choice_chose_a_new_head",
                );
                IntCounter::with_opts(opts)?
            },
            fork_choice_reorg_count: {
                let opts = Opts::new("fork_choice_reorg_depth", "depth_of_reorg");
                IntCounter::with_opts(opts)?
            },
            fork_choice_times: {
                let opts = HistogramOpts::new("fork_choice_time", "total_time_to_run_fork_choice");
                Histogram::with_opts(opts)?
            },
        })
    }

    pub fn register(&self, registry: &Registry) -> Result<(), Error> {
        registry.register(Box::new(self.block_processing_requests.clone()))?;
        registry.register(Box::new(self.block_processing_successes.clone()))?;
        registry.register(Box::new(self.block_processing_times.clone()))?;
        registry.register(Box::new(self.block_production_requests.clone()))?;
        registry.register(Box::new(self.block_production_successes.clone()))?;
        registry.register(Box::new(self.block_production_times.clone()))?;
        registry.register(Box::new(self.attestation_production_requests.clone()))?;
        registry.register(Box::new(self.attestation_production_successes.clone()))?;
        registry.register(Box::new(self.attestation_production_times.clone()))?;
        registry.register(Box::new(self.attestation_processing_requests.clone()))?;
        registry.register(Box::new(self.attestation_processing_successes.clone()))?;
        registry.register(Box::new(self.attestation_processing_times.clone()))?;
        registry.register(Box::new(self.fork_choice_requests.clone()))?;
        registry.register(Box::new(self.fork_choice_changed_head.clone()))?;
        registry.register(Box::new(self.fork_choice_reorg_count.clone()))?;
        registry.register(Box::new(self.fork_choice_times.clone()))?;

        Ok(())
    }
}
<commit_msg>Fix some metrics labels<commit_after>pub use prometheus::Error;
use prometheus::{Histogram, HistogramOpts, IntCounter, Opts, Registry};

pub struct Metrics {
    pub block_processing_requests: IntCounter,
    pub block_processing_successes: IntCounter,
    pub block_processing_times: Histogram,
    pub block_production_requests: IntCounter,
    pub block_production_successes: IntCounter,
    pub block_production_times: Histogram,
    pub attestation_production_requests: IntCounter,
    pub attestation_production_successes: IntCounter,
    pub attestation_production_times: Histogram,
    pub attestation_processing_requests: IntCounter,
    pub attestation_processing_successes: IntCounter,
    pub attestation_processing_times: Histogram,
    pub fork_choice_requests: IntCounter,
    pub fork_choice_changed_head: IntCounter,
    pub fork_choice_reorg_count: IntCounter,
    pub fork_choice_times: Histogram,
}

impl Metrics {
    pub fn new() -> Result<Self, Error> {
        Ok(Self {
            block_processing_requests: {
                let opts = Opts::new("block_processing_requests", "total_blocks_processed");
                IntCounter::with_opts(opts)?
            },
            block_processing_successes: {
                let opts = Opts::new("block_processing_successes", "total_valid_blocks_processed");
                IntCounter::with_opts(opts)?
            },
            block_processing_times: {
                let opts = HistogramOpts::new("block_processing_times", "block_processing_time");
                Histogram::with_opts(opts)?
            },
            block_production_requests: {
                let opts = Opts::new("block_production_requests", "attempts_to_produce_new_block");
                IntCounter::with_opts(opts)?
            },
            block_production_successes: {
                let opts = Opts::new("block_production_successes", "blocks_successfully_produced");
                IntCounter::with_opts(opts)?
            },
            block_production_times: {
                let opts = HistogramOpts::new("block_production_times", "block_production_time");
                Histogram::with_opts(opts)?
            },
            attestation_production_requests: {
                let opts = Opts::new(
                    "attestation_production_requests",
                    "total_attestation_production_requests",
                );
                IntCounter::with_opts(opts)?
            },
            attestation_production_successes: {
                let opts = Opts::new(
                    "attestation_production_successes",
                    "total_attestation_production_successes",
                );
                IntCounter::with_opts(opts)?
            },
            attestation_production_times: {
                let opts = HistogramOpts::new(
                    "attestation_production_times",
                    "attestation_production_time",
                );
                Histogram::with_opts(opts)?
            },
            attestation_processing_requests: {
                let opts = Opts::new(
                    "attestation_processing_requests",
                    "total_attestation_processing_requests",
                );
                IntCounter::with_opts(opts)?
            },
            attestation_processing_successes: {
                let opts = Opts::new(
                    "attestation_processing_successes",
                    "total_attestation_processing_successes",
                );
                IntCounter::with_opts(opts)?
            },
            attestation_processing_times: {
                let opts = HistogramOpts::new(
                    "attestation_processing_times",
                    "attestation_processing_time",
                );
                Histogram::with_opts(opts)?
            },
            fork_choice_requests: {
                let opts = Opts::new("fork_choice_requests", "total_times_fork_choice_called");
                IntCounter::with_opts(opts)?
            },
            fork_choice_changed_head: {
                let opts = Opts::new(
                    "fork_choice_changed_head",
                    "total_times_fork_choice_chose_a_new_head",
                );
                IntCounter::with_opts(opts)?
            },
            fork_choice_reorg_count: {
                let opts = Opts::new("fork_choice_reorg_count", "number_of_reorgs");
                IntCounter::with_opts(opts)?
            },
            fork_choice_times: {
                let opts = HistogramOpts::new("fork_choice_time", "total_time_to_run_fork_choice");
                Histogram::with_opts(opts)?
            },
        })
    }

    pub fn register(&self, registry: &Registry) -> Result<(), Error> {
        registry.register(Box::new(self.block_processing_requests.clone()))?;
        registry.register(Box::new(self.block_processing_successes.clone()))?;
        registry.register(Box::new(self.block_processing_times.clone()))?;
        registry.register(Box::new(self.block_production_requests.clone()))?;
        registry.register(Box::new(self.block_production_successes.clone()))?;
        registry.register(Box::new(self.block_production_times.clone()))?;
        registry.register(Box::new(self.attestation_production_requests.clone()))?;
        registry.register(Box::new(self.attestation_production_successes.clone()))?;
        registry.register(Box::new(self.attestation_production_times.clone()))?;
        registry.register(Box::new(self.attestation_processing_requests.clone()))?;
        registry.register(Box::new(self.attestation_processing_successes.clone()))?;
        registry.register(Box::new(self.attestation_processing_times.clone()))?;
        registry.register(Box::new(self.fork_choice_requests.clone()))?;
        registry.register(Box::new(self.fork_choice_changed_head.clone()))?;
        registry.register(Box::new(self.fork_choice_reorg_count.clone()))?;
        registry.register(Box::new(self.fork_choice_times.clone()))?;

        Ok(())
    }
}
<|endoftext|>


<commit_before>#![crate_type = "dylib"]

extern crate intercom;
extern crate winapi;
use intercom::prelude::*;
use std::convert::TryInto;
use winapi::shared::{minwindef::DWORD, ntdef::ULARGE_INTEGER, windef::HBITMAP};
use winapi::um::{
    objidlbase::STATSTG,
    wingdi::{CreateBitmap, DeleteObject},
};

com_library! {
    on_load=on_load,
    class ThumbnailProvider
}

/// Called when the DLL is loaded.
///
/// Sets up logging to the Cargo.toml directory for debug purposes.
fn on_load()
{
    // Set up logging to the proejct directory.
    use log::LevelFilter;
    simple_logging::log_to_file(
        &format!("{}\\debug.log", env!("CARGO_MANIFEST_DIR")),
        LevelFilter::Trace,
    )
    .unwrap();
}

#[com_class(IInitializeWithStream, IThumbnailProvider)]
#[derive(Default)]
struct ThumbnailProvider
{
    bitmap: std::cell::Cell<Option<HBITMAP>>,
}

impl IInitializeWithStream for ThumbnailProvider
{
    fn initialize(&self, _stream: &ComItf<dyn IStream>, _mode: DWORD) -> ComResult<()>
    {
        // This sample always returns the same thumbnail so we don't need to read
        // the stream itself.
        Ok(())
    }
}

impl IThumbnailProvider for ThumbnailProvider
{
    fn get_thumbnail(&self, cx: u32) -> ComResult<(ComHBITMAP, WTS_ALPHATYPE)>
    {
        // Render a faded circle as an example thumbnail.
        let mut data = Vec::new();
        data.resize((cx * cx) as usize, 0);
        let midpoint = cx as f64 / 2.0;
        for x in 0..cx {
            for y in 0..cx {
                let x_coord = x as f64;
                let y_coord = y as f64;
                let dist_squared = (midpoint - x_coord).powf(2.0) + (midpoint - y_coord).powf(2.0);
                let mut value = (dist_squared.sqrt() / midpoint) * 255.0;
                if value > 255.0 {
                    value = 255.0
                }
                let value = value as u32;
                data[(x * cx + y) as usize] = value + (value << 8) + (value << 16) + (value << 24);
            }
        }

        // Create a bitmap from the data and return it.
        //
        // We'll store the bitmap handle in the struct so it can destroy the data when it's
        // not needed anymore.
        let icx: i32 = cx.try_into().map_err(|_| ComError::E_INVALIDARG)?;
        let bitmap = unsafe { CreateBitmap(icx, icx, 1, 32, data.as_ptr() as *const _) };
        self.bitmap.set(Some(bitmap));
        Ok((ComHBITMAP(bitmap), 0))
    }
}

impl Drop for ThumbnailProvider
{
    fn drop(&mut self)
    {
        // Delete the bitmap once it's not needed anymore.
        if let Some(bitmap) = self.bitmap.get() {
            unsafe { DeleteObject(bitmap as _) };
        }
    }
}

// New types for deriving Intercom traits.

#[derive(intercom::ExternType, intercom::ForeignType, intercom::ExternOutput)]
#[allow(non_camel_case_types)]
#[repr(transparent)]
struct ComHBITMAP(HBITMAP);

#[derive(intercom::ExternType, intercom::ForeignType, intercom::ExternOutput, intercom::ExternInput)]
#[allow(non_camel_case_types)]
#[repr(transparent)]
struct ComULARGE_INTEGER(ULARGE_INTEGER);

#[derive(intercom::ExternType, intercom::ForeignType, intercom::ExternOutput, intercom::ExternInput)]
#[allow(non_camel_case_types)]
#[repr(transparent)]
struct ComSTATSTG(STATSTG);

#[allow(non_camel_case_types)]
type WTS_ALPHATYPE = u32;

// COM interface definitions.

#[com_interface(com_iid = "e357fccd-a995-4576-b01f-234630154e96")]
trait IThumbnailProvider
{
    fn get_thumbnail(&self, cx: u32) -> ComResult<(ComHBITMAP, WTS_ALPHATYPE)>;
}

// The IStream definition is untested since this provider doesn't initialize itself
// from the stream.
#[com_interface(com_iid = "0000000c-0000-0000-C000-000000000046")]
trait IStream
{
    fn clone(&self) -> ComResult<ComRc<dyn IStream>>;
    fn commit(&self, commit_flags: DWORD) -> ComResult<()>;
    fn copy_to(
        &self,
        target: &ComItf<dyn IStream>,
        count: ComULARGE_INTEGER,
    ) -> ComResult<(ComULARGE_INTEGER, ComULARGE_INTEGER)>;
    fn lock_region(
        &self,
        lock_offset: ComULARGE_INTEGER,
        count: ComULARGE_INTEGER,
        lock_type: DWORD,
    ) -> ComResult<()>;
    fn revert(&self) -> ComResult<()>;
    fn seek(&self, move_to: ComULARGE_INTEGER, origin: DWORD) -> ComResult<ComULARGE_INTEGER>;
    fn set_size(&self, new_size: ComULARGE_INTEGER) -> ComResult<()>;
    fn stat(&self, stat: *mut ComSTATSTG, stat_flag: DWORD) -> ComResult<()>;
    fn unlock_region(
        &self,
        lock_offset: ComULARGE_INTEGER,
        count: ComULARGE_INTEGER,
        lock_type: DWORD,
    ) -> ComResult<()>;
}

#[com_interface(com_iid = "b824b49d-22ac-4161-ac8a-9916e8fa3f7f")]
trait IInitializeWithStream
{
    fn initialize(&self, stream: &ComItf<dyn IStream>, mode: DWORD) -> ComResult<()>;
}
<commit_msg>Use hardcoded CLSID in the thumbnail_provider<commit_after>#![crate_type = "dylib"]

extern crate intercom;
extern crate winapi;
use intercom::prelude::*;
use std::convert::TryInto;
use winapi::shared::{minwindef::DWORD, ntdef::ULARGE_INTEGER, windef::HBITMAP};
use winapi::um::{
    objidlbase::STATSTG,
    wingdi::{CreateBitmap, DeleteObject},
};

com_library! {
    on_load=on_load,
    class ThumbnailProvider
}

/// Called when the DLL is loaded.
///
/// Sets up logging to the Cargo.toml directory for debug purposes.
fn on_load()
{
    // Set up logging to the proejct directory.
    use log::LevelFilter;
    simple_logging::log_to_file(
        &format!("{}\\debug.log", env!("CARGO_MANIFEST_DIR")),
        LevelFilter::Trace,
    )
    .unwrap();
}

#[com_class(
    clsid = "61a6080e-0e9f-3cef-50a7-622d16145b78",
    IInitializeWithStream,
    IThumbnailProvider)]
#[derive(Default)]
struct ThumbnailProvider
{
    bitmap: std::cell::Cell<Option<HBITMAP>>,
}

impl IInitializeWithStream for ThumbnailProvider
{
    fn initialize(&self, _stream: &ComItf<dyn IStream>, _mode: DWORD) -> ComResult<()>
    {
        // This sample always returns the same thumbnail so we don't need to read
        // the stream itself.
        Ok(())
    }
}

impl IThumbnailProvider for ThumbnailProvider
{
    fn get_thumbnail(&self, cx: u32) -> ComResult<(ComHBITMAP, WTS_ALPHATYPE)>
    {
        // Render a faded circle as an example thumbnail.
        let mut data = Vec::new();
        data.resize((cx * cx) as usize, 0);
        let midpoint = cx as f64 / 2.0;
        for x in 0..cx {
            for y in 0..cx {
                let x_coord = x as f64;
                let y_coord = y as f64;
                let dist_squared = (midpoint - x_coord).powf(2.0) + (midpoint - y_coord).powf(2.0);
                let mut value = (dist_squared.sqrt() / midpoint) * 255.0;
                if value > 255.0 {
                    value = 255.0
                }
                let value = value as u32;
                data[(x * cx + y) as usize] = value + (value << 8) + (value << 16) + (value << 24);
            }
        }

        // Create a bitmap from the data and return it.
        //
        // We'll store the bitmap handle in the struct so it can destroy the data when it's
        // not needed anymore.
        let icx: i32 = cx.try_into().map_err(|_| ComError::E_INVALIDARG)?;
        let bitmap = unsafe { CreateBitmap(icx, icx, 1, 32, data.as_ptr() as *const _) };
        self.bitmap.set(Some(bitmap));
        Ok((ComHBITMAP(bitmap), 0))
    }
}

impl Drop for ThumbnailProvider
{
    fn drop(&mut self)
    {
        // Delete the bitmap once it's not needed anymore.
        if let Some(bitmap) = self.bitmap.get() {
            unsafe { DeleteObject(bitmap as _) };
        }
    }
}

// New types for deriving Intercom traits.

#[derive(intercom::ExternType, intercom::ForeignType, intercom::ExternOutput)]
#[allow(non_camel_case_types)]
#[repr(transparent)]
struct ComHBITMAP(HBITMAP);

#[derive(intercom::ExternType, intercom::ForeignType, intercom::ExternOutput, intercom::ExternInput)]
#[allow(non_camel_case_types)]
#[repr(transparent)]
struct ComULARGE_INTEGER(ULARGE_INTEGER);

#[derive(intercom::ExternType, intercom::ForeignType, intercom::ExternOutput, intercom::ExternInput)]
#[allow(non_camel_case_types)]
#[repr(transparent)]
struct ComSTATSTG(STATSTG);

#[allow(non_camel_case_types)]
type WTS_ALPHATYPE = u32;

// COM interface definitions.

#[com_interface(com_iid = "e357fccd-a995-4576-b01f-234630154e96")]
trait IThumbnailProvider
{
    fn get_thumbnail(&self, cx: u32) -> ComResult<(ComHBITMAP, WTS_ALPHATYPE)>;
}

// The IStream definition is untested since this provider doesn't initialize itself
// from the stream.
#[com_interface(com_iid = "0000000c-0000-0000-C000-000000000046")]
trait IStream
{
    fn clone(&self) -> ComResult<ComRc<dyn IStream>>;
    fn commit(&self, commit_flags: DWORD) -> ComResult<()>;
    fn copy_to(
        &self,
        target: &ComItf<dyn IStream>,
        count: ComULARGE_INTEGER,
    ) -> ComResult<(ComULARGE_INTEGER, ComULARGE_INTEGER)>;
    fn lock_region(
        &self,
        lock_offset: ComULARGE_INTEGER,
        count: ComULARGE_INTEGER,
        lock_type: DWORD,
    ) -> ComResult<()>;
    fn revert(&self) -> ComResult<()>;
    fn seek(&self, move_to: ComULARGE_INTEGER, origin: DWORD) -> ComResult<ComULARGE_INTEGER>;
    fn set_size(&self, new_size: ComULARGE_INTEGER) -> ComResult<()>;
    fn stat(&self, stat: *mut ComSTATSTG, stat_flag: DWORD) -> ComResult<()>;
    fn unlock_region(
        &self,
        lock_offset: ComULARGE_INTEGER,
        count: ComULARGE_INTEGER,
        lock_type: DWORD,
    ) -> ComResult<()>;
}

#[com_interface(com_iid = "b824b49d-22ac-4161-ac8a-9916e8fa3f7f")]
trait IInitializeWithStream
{
    fn initialize(&self, stream: &ComItf<dyn IStream>, mode: DWORD) -> ComResult<()>;
}
<|endoftext|>


<commit_before>pub fn xor_byte_streams(a: &Vec<u8>, b: &Vec<u8>) -> Vec<u8> {
	assert_eq!(a.len(), b.len());

	let mut result: Vec<u8> = vec![];
	for i in 0..a.len() {
		result.push(a[i] ^ b[i]);
	}

	result
}

pub fn xor_hex_strings(a: &str, b: &str) -> String {
	use byte_conversion::*;
	bytes_to_hex(
		&xor_byte_streams(
			&hex_to_bytes(a),
			&hex_to_bytes(b)
		)
	)
}

pub fn transpose_chunks(chunks: &Vec<Vec<u8>>) -> Vec<Vec<u8>> {
	let chunksCount = chunks.len();
	if chunksCount == 0 {
		return vec![];
	}

	let chunkSize = chunks[0].len();
	if chunkSize == 0 {
		return vec![];
	}

	let mut result: Vec<Vec<u8>> = Vec::new();
	for _ in 0..chunkSize {
		result.push(vec![]);
	}

	for x in 0..chunksCount {
		let chunk = chunks.get(x).unwrap();
		assert_eq!(chunkSize, chunk.len());

		for y in 0..chunkSize {
			result[y].push(*chunk.get(y).unwrap());
		}
	}

	result
}

//-----------------------------------------------------------------------------

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

	#[test]
	fn xor_hex_strings_i_empty() {
		assert_eq!("", xor_hex_strings("", ""));
	}

	#[test]
	fn xor_hex_strings_i_example() {
		assert_eq!(
			"746865206b696420646f6e277420706c6179",
			xor_hex_strings(
				"1c0111001f010100061a024b53535009181c",
				"686974207468652062756c6c277320657965"
			)
		);
	}

	#[test]
	#[should_panic]
	fn xor_hex_strings_i_unequal_sizes_1() {
		xor_hex_strings("08AF", "");
	}

	#[test]
	#[should_panic]
	fn xor_hex_strings_i_unequal_sizes_2() {
		xor_hex_strings("", "08AF");
	}

	#[test]
	fn transpose_chunks_i_empty() {
		let a: Vec<Vec<u8>> = vec![];
		let expected: Vec<Vec<u8>> = vec![];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_empty_chunks() {
		let a: Vec<Vec<u8>> = vec![vec![], vec![]];
		let expected: Vec<Vec<u8>> = vec![];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_single_byte() {
		// 5 -> 5

		let a = vec![vec![5]];
		let expected: Vec<Vec<u8>> = vec![vec![5]];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_double_single() {
		// 5 6  ->  5
		//          6

		let a = vec![vec![5, 6]];
		let expected: Vec<Vec<u8>> = vec![vec![5], vec![6]];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_double_double() {
		// 5 6  ->  5 7
		// 7 8  ->  6 8

		let a = vec![vec![5, 6], vec![7, 8]];
		let expected: Vec<Vec<u8>> = vec![vec![5, 7], vec![6, 8]];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_single_double() {
		// 5  ->  5 7
		// 7

		let a = vec![vec![5], vec![7]];
		let expected: Vec<Vec<u8>> = vec![vec![5, 7]];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_double_triple() {
		// 1 2 3 ->  1 4
		// 4 5 6 ->  2 5
		//       ->  3 6

		let a = vec![vec![1, 2, 3], vec![4, 5, 6]];
		let expected: Vec<Vec<u8>> = vec![vec![1, 4], vec![2, 5], vec![3, 6]];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	#[should_panic]
	fn transpose_chunks_i_mismatched_sized_chunks() {
		let a = vec![vec![5], vec![7, 8]];

		transpose_chunks(&a);
	}
}<commit_msg>xor_hex_strings reversable encryption test<commit_after>pub fn xor_byte_streams(a: &Vec<u8>, b: &Vec<u8>) -> Vec<u8> {
	assert_eq!(a.len(), b.len());

	let mut result: Vec<u8> = vec![];
	for i in 0..a.len() {
		result.push(a[i] ^ b[i]);
	}

	result
}

pub fn xor_hex_strings(a: &str, b: &str) -> String {
	use byte_conversion::*;
	bytes_to_hex(
		&xor_byte_streams(
			&hex_to_bytes(a),
			&hex_to_bytes(b)
		)
	)
}

pub fn transpose_chunks(chunks: &Vec<Vec<u8>>) -> Vec<Vec<u8>> {
	let chunksCount = chunks.len();
	if chunksCount == 0 {
		return vec![];
	}

	let chunkSize = chunks[0].len();
	if chunkSize == 0 {
		return vec![];
	}

	let mut result: Vec<Vec<u8>> = Vec::new();
	for _ in 0..chunkSize {
		result.push(vec![]);
	}

	for x in 0..chunksCount {
		let chunk = chunks.get(x).unwrap();
		assert_eq!(chunkSize, chunk.len());

		for y in 0..chunkSize {
			result[y].push(*chunk.get(y).unwrap());
		}
	}

	result
}

//-----------------------------------------------------------------------------

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

	#[test]
	fn xor_hex_strings_i_empty() {
		assert_eq!("", xor_hex_strings("", ""));
	}

	#[test]
	fn xor_hex_strings_i_example() {
		assert_eq!(
			"746865206b696420646f6e277420706c6179",
			xor_hex_strings(
				"1c0111001f010100061a024b53535009181c",
				"686974207468652062756c6c277320657965"
			)
		);
	}

	#[test]
	#[should_panic]
	fn xor_hex_strings_i_unequal_sizes_1() {
		xor_hex_strings("08AF", "");
	}

	#[test]
	#[should_panic]
	fn xor_hex_strings_i_unequal_sizes_2() {
		xor_hex_strings("", "08AF");
	}

	#[test]
	fn xor_hex_strings_i_reversable() {
		let key = "aa44";
		let word = "baba";

		let encrypted = xor_hex_strings(&key, &word);
		let decrypted = xor_hex_strings(&key, &encrypted);

		assert_eq!(
			word,
			decrypted
		);
	}

	#[test]
	fn transpose_chunks_i_empty() {
		let a: Vec<Vec<u8>> = vec![];
		let expected: Vec<Vec<u8>> = vec![];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_empty_chunks() {
		let a: Vec<Vec<u8>> = vec![vec![], vec![]];
		let expected: Vec<Vec<u8>> = vec![];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_single_byte() {
		// 5 -> 5

		let a = vec![vec![5]];
		let expected: Vec<Vec<u8>> = vec![vec![5]];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_double_single() {
		// 5 6  ->  5
		//          6

		let a = vec![vec![5, 6]];
		let expected: Vec<Vec<u8>> = vec![vec![5], vec![6]];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_double_double() {
		// 5 6  ->  5 7
		// 7 8  ->  6 8

		let a = vec![vec![5, 6], vec![7, 8]];
		let expected: Vec<Vec<u8>> = vec![vec![5, 7], vec![6, 8]];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_single_double() {
		// 5  ->  5 7
		// 7

		let a = vec![vec![5], vec![7]];
		let expected: Vec<Vec<u8>> = vec![vec![5, 7]];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	fn transpose_chunks_i_double_triple() {
		// 1 2 3 ->  1 4
		// 4 5 6 ->  2 5
		//       ->  3 6

		let a = vec![vec![1, 2, 3], vec![4, 5, 6]];
		let expected: Vec<Vec<u8>> = vec![vec![1, 4], vec![2, 5], vec![3, 6]];

		assert_eq!(expected, transpose_chunks(&a));
	}

	#[test]
	#[should_panic]
	fn transpose_chunks_i_mismatched_sized_chunks() {
		let a = vec![vec![5], vec![7, 8]];

		transpose_chunks(&a);
	}
}<|endoftext|>


<commit_before>use super::simplex_vertex::SimplexVertex;
use super::super::common::math::vec_distance;
use super::super::common::Vec2;

pub struct Simplex {
    pub m_v1: SimplexVertex,
    pub m_v2: SimplexVertex,
    pub m_v3: SimplexVertex,
    pub m_count: i32
}

impl Simplex {
    // TODO implement
    // pub fn read_cache()

    // TODO implement
    // pub fn write_cache()

    // TODO implement {
    // pub fn get_search_direction() -> Vec2

    // TODO implement
    // pub fn get_closest_point() -> Vec2

    // TODO implement
    // pub fn get_witness_points()

    pub fn get_metric(&self) -> f32 {
        match self.m_count {
            0 => {
                assert!(false);
                return 0.0;
            },
            1 => 0.0,
            2 => {
                return vec_distance(self.m_v1.w, self.m_v2.w);
            },
            3 => {
                return Vec2::cross_vec(self.m_v2.w - self.m_v1.w, self.m_v3.w - self.m_v1.w);
            },
            _ => {
                assert!(false);
                return 0.0;
            }
        }
    }

    // TODO implement
    pub fn solve_2() {

    }

    // TODO implement
    pub fn solve_3() {

    }
}
<commit_msg>Adding get_witness_points to Simplex<commit_after>use super::simplex_vertex::SimplexVertex;
use super::super::common::math::vec_distance;
use super::super::common::Vec2;

pub struct Simplex {
    pub m_v1: SimplexVertex,
    pub m_v2: SimplexVertex,
    pub m_v3: SimplexVertex,
    pub m_count: i32
}

impl Simplex {
    // TODO implement
    // pub fn read_cache()

    // TODO implement
    // pub fn write_cache()

    // TODO implement {
    // pub fn get_search_direction() -> Vec2

    // TODO implement
    // pub fn get_closest_point() -> Vec2

    pub fn get_witness_points(&self) -> (Vec2, Vec2) {
        match self.m_count {
            0 => {
                assert!(false);
                return (Vec2::zero(), Vec2::zero());
            },
            1 => {
                return (self.m_v1.w_a, self.m_v1.w_b);
            },
            2 => {
                return (self.m_v1.a * self.m_v1.w_a + self.m_v2.a * self.m_v2.w_a,
                    self.m_v1.a * self.m_v1.w_b + self.m_v2.a * self.m_v2.w_b)
            },
            3 => {
                let p_a = self.m_v1.a * self.m_v1.w_a + self.m_v2.a * self.m_v2.w_a + self.m_v3.a * self.m_v3.w_a;
                return (p_a, p_a)
            },
            _ => {
                assert!(false);
                return (Vec2::zero(), Vec2::zero());
            }
        }
    }

    pub fn get_metric(&self) -> f32 {
        match self.m_count {
            0 => {
                assert!(false);
                return 0.0;
            },
            1 => 0.0,
            2 => {
                return vec_distance(self.m_v1.w, self.m_v2.w);
            },
            3 => {
                return Vec2::cross_vec(self.m_v2.w - self.m_v1.w, self.m_v3.w - self.m_v1.w);
            },
            _ => {
                assert!(false);
                return 0.0;
            }
        }
    }

    // TODO implement
    pub fn solve_2() {

    }

    // TODO implement
    pub fn solve_3() {

    }
}
<|endoftext|>


<commit_before>use std::sync::Arc;

use rustfft::num_complex::Complex;
use rustfft::Length;

use crate::common::dct_error_inplace;
use crate::{twiddles, RequiredScratch, DctNum};
use crate::{Dct2, Dct3, Dst2, Dst3, TransformType2And3};

/// DCT2, DCT3, DST2, and DST3 implemention that recursively divides the problem in half.
///
/// The problem size must be 2^n, n > 1
///
/// ~~~
/// // Computes a DCT Type 2 of size 1024
/// use rustdct::algorithm::Type2And3SplitRadix;
/// use rustdct::Dct2;
/// use rustdct::DctPlanner;
///
/// let len = 1024;
///
/// let mut planner = DctPlanner::new();
/// let quarter_dct = planner.plan_dct2(len / 4);
/// let half_dct = planner.plan_dct2(len / 2);
///
/// let dct = Type2And3SplitRadix::new(half_dct, quarter_dct);
///
/// let mut buffer = vec![0f32; len];
/// dct.process_dct2(&mut buffer);
/// ~~~
pub struct Type2And3SplitRadix<T> {
    half_dct: Arc<dyn TransformType2And3<T>>,
    quarter_dct: Arc<dyn TransformType2And3<T>>,
    twiddles: Box<[Complex<T>]>,
}

impl<T: DctNum> Type2And3SplitRadix<T> {
    /// Creates a new DCT2, DCT3, DST2, and DST3 context that will process signals of length `half_dct.len() * 2`
    pub fn new(
        half_dct: Arc<dyn TransformType2And3<T>>,
        quarter_dct: Arc<dyn TransformType2And3<T>>,
    ) -> Self {
        let half_len = half_dct.len();
        let quarter_len = quarter_dct.len();
        let len = half_len * 2;

        assert!(
            len.is_power_of_two() && len > 2,
            "The DCT2SplitRadix algorithm requires a power-of-two input size greater than two. Got {}", len
        );
        assert_eq!(half_len, quarter_len * 2,
            "half_dct.len() must be 2 * quarter_dct.len(). Got half_dct.len()={}, quarter_dct.len()={}", half_len, quarter_len
        );

        let twiddles: Vec<Complex<T>> = (0..(len / 4))
            .map(|i| twiddles::single_twiddle(2 * i + 1, len * 4).conj())
            .collect();

        Self {
            half_dct: half_dct,
            quarter_dct: quarter_dct,
            twiddles: twiddles.into_boxed_slice(),
        }
    }
}

impl<T: DctNum> Dct2<T> for Type2And3SplitRadix<T> {
    fn process_dct2_with_scratch(&self, buffer: &mut [T], scratch: &mut [T]) {
        let scratch = validate_buffers!(buffer, scratch, self.len(), self.get_scratch_len());

        let len = self.len();
        let half_len = len / 2;
        let quarter_len = len / 4;

        //preprocess the data by splitting it up into vectors of size n/2, n/4, and n/4
        let (input_dct2, input_dct4) = scratch.split_at_mut(half_len);
        let (input_dct4_even, input_dct4_odd) = input_dct4.split_at_mut(quarter_len);

        for i in 0..quarter_len {
            unsafe {
                let input_bottom = *buffer.get_unchecked(i);
                let input_top = *buffer.get_unchecked(len - i - 1);

                let input_half_bottom = *buffer.get_unchecked(half_len - i - 1);
                let input_half_top = *buffer.get_unchecked(half_len + i);

                //prepare the inner DCT2
                *input_dct2.get_unchecked_mut(i) = input_top + input_bottom;
                *input_dct2.get_unchecked_mut(half_len - i - 1) =
                    input_half_bottom + input_half_top;

                //prepare the inner DCT4 - which consists of two DCT2s of half size
                let lower_dct4 = input_bottom - input_top;
                let upper_dct4 = input_half_bottom - input_half_top;
                let twiddle = self.twiddles.get_unchecked(i);

                let cos_input = lower_dct4 * twiddle.re + upper_dct4 * twiddle.im;
                let sin_input = upper_dct4 * twiddle.re - lower_dct4 * twiddle.im;

                *input_dct4_even.get_unchecked_mut(i) = cos_input;
                *input_dct4_odd.get_unchecked_mut(quarter_len - i - 1) =
                    if i % 2 == 0 { sin_input } else { -sin_input };
            }
        }

        // compute the recursive DCT2s, using the original buffer as scratch space
        self.half_dct.process_dct2_with_scratch(input_dct2, buffer);
        self.quarter_dct
            .process_dct2_with_scratch(input_dct4_even, buffer);
        self.quarter_dct
            .process_dct2_with_scratch(input_dct4_odd, buffer);

        unsafe {
            //post process the 3 DCT2 outputs. the first few and the last will be done outside of the loop
            *buffer.get_unchecked_mut(0) = *input_dct2.get_unchecked(0);
            *buffer.get_unchecked_mut(1) = *input_dct4_even.get_unchecked(0);
            *buffer.get_unchecked_mut(2) = *input_dct2.get_unchecked(1);

            for i in 1..quarter_len {
                let dct4_cos_output = *input_dct4_even.get_unchecked(i);
                let dct4_sin_output = if (i + quarter_len) % 2 == 0 {
                    -*input_dct4_odd.get_unchecked(quarter_len - i)
                } else {
                    *input_dct4_odd.get_unchecked(quarter_len - i)
                };

                *buffer.get_unchecked_mut(i * 4 - 1) = dct4_cos_output + dct4_sin_output;
                *buffer.get_unchecked_mut(i * 4) = *input_dct2.get_unchecked(i * 2);

                *buffer.get_unchecked_mut(i * 4 + 1) = dct4_cos_output - dct4_sin_output;
                *buffer.get_unchecked_mut(i * 4 + 2) = *input_dct2.get_unchecked(i * 2 + 1);
            }

            *buffer.get_unchecked_mut(len - 1) = -*input_dct4_odd.get_unchecked(0);
        }
    }
}
impl<T: DctNum> Dst2<T> for Type2And3SplitRadix<T> {
    fn process_dst2_with_scratch(&self, buffer: &mut [T], scratch: &mut [T]) {
        for i in 0..(self.len() / 2) {
            buffer[2 * i + 1] = buffer[2 * i + 1].neg();
        }

        self.process_dct2_with_scratch(buffer, scratch);

        buffer.reverse();
    }
}
impl<T: DctNum> Dct3<T> for Type2And3SplitRadix<T> {
    fn process_dct3_with_scratch(&self, buffer: &mut [T], scratch: &mut [T]) {
        let scratch = validate_buffers!(buffer, scratch, self.len(), self.get_scratch_len());

        let len = buffer.len();
        let half_len = len / 2;
        let quarter_len = len / 4;

        // divide the output into 3 sub-lists to use for our inner DCTs, one of size N/2 and two of size N/4
        let (recursive_input_evens, recursive_input_odds) = scratch.split_at_mut(half_len);
        let (recursive_input_n1, recursive_input_n3) =
            recursive_input_odds.split_at_mut(quarter_len);

        // do the same pre-loop setup as DCT4ViaDCT3, and since we're skipping the first iteration of the loop we
        // to also set up the corresponding evens cells
        recursive_input_evens[0] = buffer[0];
        recursive_input_evens[1] = buffer[2];
        recursive_input_n1[0] = buffer[1] * T::two();
        recursive_input_n3[0] = buffer[len - 1] * T::two();

        // populate the recursive input arrays
        for i in 1..quarter_len {
            let k = 4 * i;

            unsafe {
                // the evens are the easy ones - just copy straight over
                *recursive_input_evens.get_unchecked_mut(i * 2) = *buffer.get_unchecked(k);
                *recursive_input_evens.get_unchecked_mut(i * 2 + 1) = *buffer.get_unchecked(k + 2);

                // for the odd ones we're going to do the same addition/subtraction we do in the setup for DCT4ViaDCT3
                *recursive_input_n1.get_unchecked_mut(i) =
                    *buffer.get_unchecked(k - 1) + *buffer.get_unchecked(k + 1);
                *recursive_input_n3.get_unchecked_mut(quarter_len - i) =
                    *buffer.get_unchecked(k - 1) - *buffer.get_unchecked(k + 1);
            }
        }

        //perform our recursive DCTs, using the original buffer as scratch space
        self.half_dct
            .process_dct3_with_scratch(recursive_input_evens, buffer);
        self.quarter_dct
            .process_dct3_with_scratch(recursive_input_n1, buffer);
        self.quarter_dct
            .process_dct3_with_scratch(recursive_input_n3, buffer);

        //merge the results. we're going to combine 2 separate things:
        // - merging the two smaller DCT3 outputs into a DCT4 output
        // - marging the DCT4 outputand the larger DCT3 output into the final output
        for i in 0..quarter_len {
            let twiddle = self.twiddles[i];
            let cosine_value = recursive_input_n1[i];

            // flip the sign of every other sine value to finish the job of using a DCT3 to compute a DST3
            let sine_value = if i % 2 == 0 {
                recursive_input_n3[i]
            } else {
                -recursive_input_n3[i]
            };

            let lower_dct4 = cosine_value * twiddle.re + sine_value * twiddle.im;
            let upper_dct4 = cosine_value * twiddle.im - sine_value * twiddle.re;

            unsafe {
                let lower_dct3 = *recursive_input_evens.get_unchecked(i);
                let upper_dct3 = *recursive_input_evens.get_unchecked(half_len - i - 1);

                *buffer.get_unchecked_mut(i) = lower_dct3 + lower_dct4;
                *buffer.get_unchecked_mut(len - i - 1) = lower_dct3 - lower_dct4;

                *buffer.get_unchecked_mut(half_len - i - 1) = upper_dct3 + upper_dct4;
                *buffer.get_unchecked_mut(half_len + i) = upper_dct3 - upper_dct4;
            }
        }
    }
}
impl<T: DctNum> Dst3<T> for Type2And3SplitRadix<T> {
    fn process_dst3_with_scratch(&self, buffer: &mut [T], scratch: &mut [T]) {
        buffer.reverse();

        self.process_dct3_with_scratch(buffer, scratch);

        for i in 0..(self.len() / 2) {
            buffer[2 * i + 1] = buffer[2 * i + 1].neg();
        }
    }
}
impl<T: DctNum> TransformType2And3<T> for Type2And3SplitRadix<T> {}
impl<T> Length for Type2And3SplitRadix<T> {
    fn len(&self) -> usize {
        self.twiddles.len() * 4
    }
}
impl<T> RequiredScratch for Type2And3SplitRadix<T> {
    fn get_scratch_len(&self) -> usize {
        self.len()
    }
}

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

    use crate::test_utils::{compare_float_vectors, random_signal};

    /// Verify that our fast implementation of the DCT2 gives the same output as the slow version, for many different inputs
    #[test]
    fn test_dct2_splitradix() {
        for i in 2..8 {
            let size = 1 << i;
            println!("len: {}", size);

            let mut expected_buffer = random_signal(size);
            let mut actual_buffer = expected_buffer.clone();

            let naive_dct = Type2And3Naive::new(size);
            naive_dct.process_dct2(&mut expected_buffer);

            let quarter_dct = Arc::new(Type2And3Naive::new(size / 4));
            let half_dct = Arc::new(Type2And3Naive::new(size / 2));

            let dct = Type2And3SplitRadix::new(half_dct, quarter_dct);
            dct.process_dct2(&mut actual_buffer);

            println!("expected:    {:?}", expected_buffer);
            println!("fast output: {:?}", actual_buffer);

            assert!(
                compare_float_vectors(&actual_buffer, &expected_buffer),
                "len = {}",
                size
            );
        }
    }

    /// Verify that our fast implementation of the DCT3 gives the same output as the slow version, for many different inputs
    #[test]
    fn test_dct3_splitradix() {
        for i in 2..8 {
            let size = 1 << i;
            println!("len: {}", size);

            let mut expected_buffer = random_signal(size);
            let mut actual_buffer = expected_buffer.clone();

            let naive_dct = Type2And3Naive::new(size);
            naive_dct.process_dct3(&mut expected_buffer);

            let quarter_dct = Arc::new(Type2And3Naive::new(size / 4));
            let half_dct = Arc::new(Type2And3Naive::new(size / 2));

            let dct = Type2And3SplitRadix::new(half_dct, quarter_dct);
            dct.process_dct3(&mut actual_buffer);

            println!("expected:    {:?}", expected_buffer);
            println!("fast output: {:?}", actual_buffer);

            assert!(
                compare_float_vectors(&actual_buffer, &expected_buffer),
                "len = {}",
                size
            );
        }
    }
}
<commit_msg>Shrink unsafe block<commit_after>use std::sync::Arc;

use rustfft::num_complex::Complex;
use rustfft::Length;

use crate::common::dct_error_inplace;
use crate::{twiddles, RequiredScratch, DctNum};
use crate::{Dct2, Dct3, Dst2, Dst3, TransformType2And3};

/// DCT2, DCT3, DST2, and DST3 implemention that recursively divides the problem in half.
///
/// The problem size must be 2^n, n > 1
///
/// ~~~
/// // Computes a DCT Type 2 of size 1024
/// use rustdct::algorithm::Type2And3SplitRadix;
/// use rustdct::Dct2;
/// use rustdct::DctPlanner;
///
/// let len = 1024;
///
/// let mut planner = DctPlanner::new();
/// let quarter_dct = planner.plan_dct2(len / 4);
/// let half_dct = planner.plan_dct2(len / 2);
///
/// let dct = Type2And3SplitRadix::new(half_dct, quarter_dct);
///
/// let mut buffer = vec![0f32; len];
/// dct.process_dct2(&mut buffer);
/// ~~~
pub struct Type2And3SplitRadix<T> {
    half_dct: Arc<dyn TransformType2And3<T>>,
    quarter_dct: Arc<dyn TransformType2And3<T>>,
    twiddles: Box<[Complex<T>]>,
}

impl<T: DctNum> Type2And3SplitRadix<T> {
    /// Creates a new DCT2, DCT3, DST2, and DST3 context that will process signals of length `half_dct.len() * 2`
    pub fn new(
        half_dct: Arc<dyn TransformType2And3<T>>,
        quarter_dct: Arc<dyn TransformType2And3<T>>,
    ) -> Self {
        let half_len = half_dct.len();
        let quarter_len = quarter_dct.len();
        let len = half_len * 2;

        assert!(
            len.is_power_of_two() && len > 2,
            "The DCT2SplitRadix algorithm requires a power-of-two input size greater than two. Got {}", len
        );
        assert_eq!(half_len, quarter_len * 2,
            "half_dct.len() must be 2 * quarter_dct.len(). Got half_dct.len()={}, quarter_dct.len()={}", half_len, quarter_len
        );

        let twiddles: Vec<Complex<T>> = (0..(len / 4))
            .map(|i| twiddles::single_twiddle(2 * i + 1, len * 4).conj())
            .collect();

        Self {
            half_dct: half_dct,
            quarter_dct: quarter_dct,
            twiddles: twiddles.into_boxed_slice(),
        }
    }
}

impl<T: DctNum> Dct2<T> for Type2And3SplitRadix<T> {
    fn process_dct2_with_scratch(&self, buffer: &mut [T], scratch: &mut [T]) {
        let scratch = validate_buffers!(buffer, scratch, self.len(), self.get_scratch_len());

        let len = self.len();
        let half_len = len / 2;
        let quarter_len = len / 4;

        //preprocess the data by splitting it up into vectors of size n/2, n/4, and n/4
        let (input_dct2, input_dct4) = scratch.split_at_mut(half_len);
        let (input_dct4_even, input_dct4_odd) = input_dct4.split_at_mut(quarter_len);

        for i in 0..quarter_len {

                let input_bottom = unsafe { *buffer.get_unchecked(i) };
                let input_top = unsafe { *buffer.get_unchecked(len - i - 1) };

                let input_half_bottom = unsafe { *buffer.get_unchecked(half_len - i - 1) };
                let input_half_top = unsafe { *buffer.get_unchecked(half_len + i) };

                //prepare the inner DCT2
                unsafe {*input_dct2.get_unchecked_mut(i)  = input_top + input_bottom };
                unsafe {*input_dct2.get_unchecked_mut(half_len - i - 1)  =
                    input_half_bottom + input_half_top };

                //prepare the inner DCT4 - which consists of two DCT2s of half size
                let lower_dct4 = input_bottom - input_top;
                let upper_dct4 = input_half_bottom - input_half_top;
                let twiddle = unsafe { self.twiddles.get_unchecked(i) };

                let cos_input = lower_dct4 * twiddle.re + upper_dct4 * twiddle.im;
                let sin_input = upper_dct4 * twiddle.re - lower_dct4 * twiddle.im;

                unsafe {*input_dct4_even.get_unchecked_mut(i) = cos_input };
                unsafe {*input_dct4_odd.get_unchecked_mut(quarter_len - i - 1) =
                    if i % 2 == 0 { sin_input } else { -sin_input } };

        }

        // compute the recursive DCT2s, using the original buffer as scratch space
        self.half_dct.process_dct2_with_scratch(input_dct2, buffer);
        self.quarter_dct
            .process_dct2_with_scratch(input_dct4_even, buffer);
        self.quarter_dct
            .process_dct2_with_scratch(input_dct4_odd, buffer);

        unsafe {
            //post process the 3 DCT2 outputs. the first few and the last will be done outside of the loop
            *buffer.get_unchecked_mut(0) = *input_dct2.get_unchecked(0);
            *buffer.get_unchecked_mut(1) = *input_dct4_even.get_unchecked(0);
            *buffer.get_unchecked_mut(2) = *input_dct2.get_unchecked(1);

            for i in 1..quarter_len {
                let dct4_cos_output = *input_dct4_even.get_unchecked(i);
                let dct4_sin_output = if (i + quarter_len) % 2 == 0 {
                    -*input_dct4_odd.get_unchecked(quarter_len - i)
                } else {
                    *input_dct4_odd.get_unchecked(quarter_len - i)
                };

                *buffer.get_unchecked_mut(i * 4 - 1) = dct4_cos_output + dct4_sin_output;
                *buffer.get_unchecked_mut(i * 4) = *input_dct2.get_unchecked(i * 2);

                *buffer.get_unchecked_mut(i * 4 + 1) = dct4_cos_output - dct4_sin_output;
                *buffer.get_unchecked_mut(i * 4 + 2) = *input_dct2.get_unchecked(i * 2 + 1);
            }

            *buffer.get_unchecked_mut(len - 1) = -*input_dct4_odd.get_unchecked(0);
        }
    }
}
impl<T: DctNum> Dst2<T> for Type2And3SplitRadix<T> {
    fn process_dst2_with_scratch(&self, buffer: &mut [T], scratch: &mut [T]) {
        for i in 0..(self.len() / 2) {
            buffer[2 * i + 1] = buffer[2 * i + 1].neg();
        }

        self.process_dct2_with_scratch(buffer, scratch);

        buffer.reverse();
    }
}
impl<T: DctNum> Dct3<T> for Type2And3SplitRadix<T> {
    fn process_dct3_with_scratch(&self, buffer: &mut [T], scratch: &mut [T]) {
        let scratch = validate_buffers!(buffer, scratch, self.len(), self.get_scratch_len());

        let len = buffer.len();
        let half_len = len / 2;
        let quarter_len = len / 4;

        // divide the output into 3 sub-lists to use for our inner DCTs, one of size N/2 and two of size N/4
        let (recursive_input_evens, recursive_input_odds) = scratch.split_at_mut(half_len);
        let (recursive_input_n1, recursive_input_n3) =
            recursive_input_odds.split_at_mut(quarter_len);

        // do the same pre-loop setup as DCT4ViaDCT3, and since we're skipping the first iteration of the loop we
        // to also set up the corresponding evens cells
        recursive_input_evens[0] = buffer[0];
        recursive_input_evens[1] = buffer[2];
        recursive_input_n1[0] = buffer[1] * T::two();
        recursive_input_n3[0] = buffer[len - 1] * T::two();

        // populate the recursive input arrays
        for i in 1..quarter_len {
            let k = 4 * i;

            unsafe {
                // the evens are the easy ones - just copy straight over
                *recursive_input_evens.get_unchecked_mut(i * 2) = *buffer.get_unchecked(k);
                *recursive_input_evens.get_unchecked_mut(i * 2 + 1) = *buffer.get_unchecked(k + 2);

                // for the odd ones we're going to do the same addition/subtraction we do in the setup for DCT4ViaDCT3
                *recursive_input_n1.get_unchecked_mut(i) =
                    *buffer.get_unchecked(k - 1) + *buffer.get_unchecked(k + 1);
                *recursive_input_n3.get_unchecked_mut(quarter_len - i) =
                    *buffer.get_unchecked(k - 1) - *buffer.get_unchecked(k + 1);
            }
        }

        //perform our recursive DCTs, using the original buffer as scratch space
        self.half_dct
            .process_dct3_with_scratch(recursive_input_evens, buffer);
        self.quarter_dct
            .process_dct3_with_scratch(recursive_input_n1, buffer);
        self.quarter_dct
            .process_dct3_with_scratch(recursive_input_n3, buffer);

        //merge the results. we're going to combine 2 separate things:
        // - merging the two smaller DCT3 outputs into a DCT4 output
        // - marging the DCT4 outputand the larger DCT3 output into the final output
        for i in 0..quarter_len {
            let twiddle = self.twiddles[i];
            let cosine_value = recursive_input_n1[i];

            // flip the sign of every other sine value to finish the job of using a DCT3 to compute a DST3
            let sine_value = if i % 2 == 0 {
                recursive_input_n3[i]
            } else {
                -recursive_input_n3[i]
            };

            let lower_dct4 = cosine_value * twiddle.re + sine_value * twiddle.im;
            let upper_dct4 = cosine_value * twiddle.im - sine_value * twiddle.re;

            unsafe {
                let lower_dct3 = *recursive_input_evens.get_unchecked(i);
                let upper_dct3 = *recursive_input_evens.get_unchecked(half_len - i - 1);

                *buffer.get_unchecked_mut(i) = lower_dct3 + lower_dct4;
                *buffer.get_unchecked_mut(len - i - 1) = lower_dct3 - lower_dct4;

                *buffer.get_unchecked_mut(half_len - i - 1) = upper_dct3 + upper_dct4;
                *buffer.get_unchecked_mut(half_len + i) = upper_dct3 - upper_dct4;
            }
        }
    }
}
impl<T: DctNum> Dst3<T> for Type2And3SplitRadix<T> {
    fn process_dst3_with_scratch(&self, buffer: &mut [T], scratch: &mut [T]) {
        buffer.reverse();

        self.process_dct3_with_scratch(buffer, scratch);

        for i in 0..(self.len() / 2) {
            buffer[2 * i + 1] = buffer[2 * i + 1].neg();
        }
    }
}
impl<T: DctNum> TransformType2And3<T> for Type2And3SplitRadix<T> {}
impl<T> Length for Type2And3SplitRadix<T> {
    fn len(&self) -> usize {
        self.twiddles.len() * 4
    }
}
impl<T> RequiredScratch for Type2And3SplitRadix<T> {
    fn get_scratch_len(&self) -> usize {
        self.len()
    }
}

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

    use crate::test_utils::{compare_float_vectors, random_signal};

    /// Verify that our fast implementation of the DCT2 gives the same output as the slow version, for many different inputs
    #[test]
    fn test_dct2_splitradix() {
        for i in 2..8 {
            let size = 1 << i;
            println!("len: {}", size);

            let mut expected_buffer = random_signal(size);
            let mut actual_buffer = expected_buffer.clone();

            let naive_dct = Type2And3Naive::new(size);
            naive_dct.process_dct2(&mut expected_buffer);

            let quarter_dct = Arc::new(Type2And3Naive::new(size / 4));
            let half_dct = Arc::new(Type2And3Naive::new(size / 2));

            let dct = Type2And3SplitRadix::new(half_dct, quarter_dct);
            dct.process_dct2(&mut actual_buffer);

            println!("expected:    {:?}", expected_buffer);
            println!("fast output: {:?}", actual_buffer);

            assert!(
                compare_float_vectors(&actual_buffer, &expected_buffer),
                "len = {}",
                size
            );
        }
    }

    /// Verify that our fast implementation of the DCT3 gives the same output as the slow version, for many different inputs
    #[test]
    fn test_dct3_splitradix() {
        for i in 2..8 {
            let size = 1 << i;
            println!("len: {}", size);

            let mut expected_buffer = random_signal(size);
            let mut actual_buffer = expected_buffer.clone();

            let naive_dct = Type2And3Naive::new(size);
            naive_dct.process_dct3(&mut expected_buffer);

            let quarter_dct = Arc::new(Type2And3Naive::new(size / 4));
            let half_dct = Arc::new(Type2And3Naive::new(size / 2));

            let dct = Type2And3SplitRadix::new(half_dct, quarter_dct);
            dct.process_dct3(&mut actual_buffer);

            println!("expected:    {:?}", expected_buffer);
            println!("fast output: {:?}", actual_buffer);

            assert!(
                compare_float_vectors(&actual_buffer, &expected_buffer),
                "len = {}",
                size
            );
        }
    }
}
<|endoftext|>