<commit_before>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

use std::num::Int;
use std::collections;
use std;

use alphabets::{Alphabet, RankTransform};
use utils;

/// A classical, flexible, q-gram index implementation.
pub struct QGramIndex {
    q: u32,
    address: Vec<usize>,
    pos: Vec<usize>,
    ranks: RankTransform,
}

impl QGramIndex {
    /// Create a new q-gram index.
    /// The q has to be smaller than b / log2(|A|) with |A| being the alphabet size and b the number
    /// bits with the `usize` data type.
    pub fn new(q: u32, text: &[u8], alphabet: &Alphabet) -> Self {
        QGramIndex::with_max_count(q, text, alphabet, std::usize::MAX)
    }

    /// Create a new q-gram index, only considering q-grams that occur at most `max_count` times.
    /// The q has to be smaller than b / log2(|A|) with |A| being the alphabet size and b the number
    /// bits with the `usize` data type.
    pub fn with_max_count(q: u32, text: &[u8], alphabet: &Alphabet, max_count: usize) -> Self {
        let ranks = RankTransform::new(alphabet);

        let qgram_count = alphabet.len().pow(q as u32);
        let mut address = vec![0; qgram_count + 1];
        let mut pos = vec![0; text.len()];

        for qgram in ranks.qgrams(q, text) {
            address[qgram] += 1;
        }

        for g in 1..address.len() {
            if address[g] > max_count {
                // mask qgram
                address[g] = 0;
            }
        }

        utils::prescan(&mut address, 0, |a, b| a + b);

        {
            let mut offset = vec![0; qgram_count];
            for (i, qgram) in ranks.qgrams(q, text).enumerate() {
                let a = address[qgram as usize];
                if address[qgram as usize + 1] - a != 0 {
                    // if not masked, insert positions
                    pos[a + offset[qgram as usize]] = i;
                    offset[qgram as usize] += 1;
                }
            }
        }

        QGramIndex { q: q, address: address, pos: pos, ranks: ranks }
    }

    /// Return text positions with matching q-gram.
    pub fn qgram_matches(&self, qgram: usize) -> &[usize] {
        &self.pos[self.address[qgram]..self.address[qgram + 1]]
    }

    /// Return matches of the given pattern.
    pub fn matches(&self, pattern: &[u8], min_count: usize) -> Vec<Match> {
        let q = self.q as usize;
        let mut diagonals = collections::HashMap::new();
        for (i, qgram) in self.ranks.qgrams(self.q, pattern).enumerate() {
            for &p in self.qgram_matches(qgram) {
                let diagonal = p - i;
                if !diagonals.contains_key(&diagonal) {
                    //diagonals.insert(diagonal, 1);
                    diagonals.insert(diagonal, Match {
                        pattern: Interval { start: i, stop: i + q },
                        text: Interval { start: p, stop: p + q },
                        count: 1,
                    });
                }
                else {
                    let m = diagonals.get_mut(&diagonal).unwrap();
                    m.pattern.stop = i + q;
                    m.text.stop = p + q;
                    m.count += 1;
                }
            }
        }
        diagonals.into_iter().filter_map(
            |(_, m)| if m.count >= min_count { Some(m) } else { None }
        ).collect()
    }

    /// Return exact matches of the given pattern.
    pub fn exact_matches(&self, pattern: &[u8]) -> Vec<ExactMatch> {
        let q = self.q as usize;
        let mut diagonals: collections::HashMap<usize, ExactMatch> = collections::HashMap::new();
        let mut intervals = Vec::new();
        for (i, qgram) in self.ranks.qgrams(self.q, pattern).enumerate() {
            for &p in self.qgram_matches(qgram) {
                let diagonal = p - i;
                if !diagonals.contains_key(&diagonal) {
                    // nothing yet, start new match
                    diagonals.insert(diagonal, ExactMatch {
                            pattern: Interval{ start: i, stop: i + q  },
                            text: Interval { start: p, stop: p + q },
                    });
                }
                else {
                    let interval = diagonals.get_mut(&diagonal).unwrap();
                    if interval.pattern.stop - q + 1 == i {
                        // extend exact match
                        interval.pattern.stop = i + q;
                        interval.text.stop = p + q;
                    }
                    else {
                        // report previous match
                        intervals.push(interval.clone());
                        // mismatch or indel, start new match
                        interval.pattern.start = i;
                        interval.pattern.stop = i + q;
                        interval.text.start = p;
                        interval.text.stop = p + q;
                    }

                }
            }
        }
        // report remaining intervals
        for (_, interval) in diagonals.into_iter() {
            intervals.push(interval);
        }
        intervals
    }
}

#[derive(PartialEq)]
#[derive(Debug)]
pub struct Match {
    pub pattern: Interval,
    pub text: Interval,
    pub count: usize,
}

/// An interval, consisting of start and stop position (the latter exclusive).
#[derive(Clone)]
#[derive(PartialEq)]
#[derive(Debug)]
pub struct Interval {
    pub start: usize,
    pub stop: usize
}

impl Interval {
    pub fn get<'a>(&self, text: &'a [u8]) -> &'a [u8] {
        &text[self.start..self.stop]
    }
}

#[derive(Clone)]
#[derive(PartialEq)]
#[derive(Debug)]
pub struct ExactMatch {
    pub pattern: Interval,
    pub text: Interval,
}

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

    fn setup() -> (&'static [u8], alphabets::Alphabet) {
        let text = b"ACGGCTGAGATGAT";
        let alphabet = alphabets::dna::alphabet();

        (text, alphabet)
    }

    #[test]
    fn test_qgram_matches() {
        let (text, alphabet) = setup();
        let q = 3;
        let qgram_index = QGramIndex::new(q, text, &alphabet);

        let ranks = alphabets::RankTransform::new(&alphabet);

        let qgram = ranks.qgrams(q, b"TGA").next().unwrap();

        let matches = qgram_index.qgram_matches(qgram);
        assert_eq!(matches, [5, 10]);
    }

    #[test]
    fn test_matches() {
        let (text, alphabet) = setup();
        let q = 3;
        let qgram_index = QGramIndex::new(q, text, &alphabet);

        let pattern = b"GCTG";
        let matches = qgram_index.matches(pattern, 1);
        assert_eq!(matches, [
            Match {
                pattern: Interval { start: 0, stop: 4 },
                text: Interval { start: 3, stop: 7 },
                count: 2
            }
        ]);
    }

    #[test]
    fn test_exact_matches() {
        let (text, alphabet) = setup();
        let q = 3;
        let qgram_index = QGramIndex::new(q, text, &alphabet);

        let pattern = b"GCTGA";
        let exact_matches = qgram_index.exact_matches(pattern);
        assert!(exact_matches.len() == 2);
        for m in exact_matches {
            assert_eq!(m.pattern.get(pattern), m.text.get(text));
        }
    }
}
<commit_msg>Docs, fixes and improvements.<commit_after>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

/// A classical, flexible, q-gram index implementation.
///
/// # Example
///
/// ```
/// use bio::data_structures::qgram_index;
/// use bio::alphabets;
///
/// let text = b"ACGGCTGAGATGAT";
/// let alphabet = alphabets::dna::alphabet();
/// let q = 3;
/// let qgram_index = qgram_index::QGramIndex::new(q, text, &alphabet);
///
/// let pattern = b"GCTG";
/// let matches = qgram_index.matches(pattern, 1);
/// assert_eq!(matches, [
///     qgram_index::Match {
///         pattern: qgram_index::Interval { start: 0, stop: 4 },
///         text: qgram_index::Interval { start: 3, stop: 7 },
///         count: 2
///     }
/// ]);
/// ```

use std::num::Int;
use std::collections;
use std;

use alphabets::{Alphabet, RankTransform};
use utils;

/// A classical, flexible, q-gram index implementation.
pub struct QGramIndex {
    q: u32,
    address: Vec<usize>,
    pos: Vec<usize>,
    ranks: RankTransform,
}

impl QGramIndex {
    /// Create a new q-gram index.
    /// The q has to be smaller than b / log2(|A|) with |A| being the alphabet size and b the number
    /// bits with the `usize` data type.
    pub fn new(q: u32, text: &[u8], alphabet: &Alphabet) -> Self {
        QGramIndex::with_max_count(q, text, alphabet, std::usize::MAX)
    }

    /// Create a new q-gram index, only considering q-grams that occur at most `max_count` times.
    /// The q has to be smaller than b / log2(|A|) with |A| being the alphabet size and b the number
    /// bits with the `usize` data type.
    pub fn with_max_count(q: u32, text: &[u8], alphabet: &Alphabet, max_count: usize) -> Self {
        let ranks = RankTransform::new(alphabet);

        let qgram_count = alphabet.len().pow(q as u32);
        let mut address = vec![0; qgram_count + 1];
        let mut pos = vec![0; text.len()];

        for qgram in ranks.qgrams(q, text) {
            address[qgram] += 1;
        }

        for g in 1..address.len() {
            if address[g] > max_count {
                // mask qgram
                address[g] = 0;
            }
        }

        utils::prescan(&mut address, 0, |a, b| a + b);

        {
            let mut offset = vec![0; qgram_count];
            for (i, qgram) in ranks.qgrams(q, text).enumerate() {
                let a = address[qgram as usize];
                if address[qgram as usize + 1] - a != 0 {
                    // if not masked, insert positions
                    pos[a + offset[qgram as usize]] = i;
                    offset[qgram as usize] += 1;
                }
            }
        }

        QGramIndex { q: q, address: address, pos: pos, ranks: ranks }
    }

    /// Return text positions with matching q-gram.
    pub fn qgram_matches(&self, qgram: usize) -> &[usize] {
        &self.pos[self.address[qgram]..self.address[qgram + 1]]
    }

    /// Return matches of the given pattern.
    pub fn matches(&self, pattern: &[u8], min_count: usize) -> Vec<Match> {
        let q = self.q as usize;
        let mut diagonals = collections::HashMap::new();
        for (i, qgram) in self.ranks.qgrams(self.q, pattern).enumerate() {
            for &p in self.qgram_matches(qgram) {
                let diagonal = p - i;
                if !diagonals.contains_key(&diagonal) {
                    diagonals.insert(diagonal, Match {
                        pattern: Interval { start: i, stop: i + q },
                        text: Interval { start: p, stop: p + q },
                        count: 1,
                    });
                }
                else {
                    let m = diagonals.get_mut(&diagonal).unwrap();
                    m.pattern.stop = i + q;
                    m.text.stop = p + q;
                    m.count += 1;
                }
            }
        }
        diagonals.into_iter().filter_map(
            |(_, m)| if m.count >= min_count { Some(m) } else { None }
        ).collect()
    }

    /// Return exact matches of the given pattern.
    pub fn exact_matches(&self, pattern: &[u8]) -> Vec<ExactMatch> {
        let q = self.q as usize;
        let mut diagonals = collections::HashMap::new();
        let mut matches = Vec::new();

        for (i, qgram) in self.ranks.qgrams(self.q, pattern).enumerate() {
            for &p in self.qgram_matches(qgram) {
                let diagonal = p - i;
                if !diagonals.contains_key(&diagonal) {
                    diagonals.insert(diagonal, ExactMatch {
                        pattern: Interval { start: i, stop: i + q },
                        text: Interval { start: p, stop: p + q },
                    });
                }
                else {
                    let m = diagonals.get_mut(&diagonal).unwrap();
                    if m.pattern.stop - q + 1 != i {
                        // discontinue match
                        matches.push(*m);
                        // start new match
                        m.pattern.start = i;
                        m.pattern.stop = i + q;
                        m.text.start = p;
                        m.text.stop = p + q;
                    }
                    else {
                        m.pattern.stop = i + q;
                        m.text.stop = p + q;
                    }
                }
            }
        }
        for (_, m) in diagonals.into_iter() {
            matches.push(m);
        }

        matches
    }
}

/// An interval, consisting of start and stop position (the latter exclusive).
#[derive(PartialEq)]
#[derive(Debug)]
#[derive(Copy)]
pub struct Interval {
    pub start: usize,
    pub stop: usize
}

impl Interval {
    pub fn get<'a>(&self, text: &'a [u8]) -> &'a [u8] {
        &text[self.start..self.stop]
    }
}

#[derive(PartialEq)]
#[derive(Debug)]
#[derive(Copy)]
pub struct Match {
    pub pattern: Interval,
    pub text: Interval,
    pub count: usize,
}

#[derive(PartialEq)]
#[derive(Debug)]
#[derive(Copy)]
pub struct ExactMatch {
    pub pattern: Interval,
    pub text: Interval,
}

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

    fn setup() -> (&'static [u8], alphabets::Alphabet) {
        let text = b"ACGGCTGAGATGAT";
        let alphabet = alphabets::dna::alphabet();

        (text, alphabet)
    }

    #[test]
    fn test_qgram_matches() {
        let (text, alphabet) = setup();
        let q = 3;
        let qgram_index = QGramIndex::new(q, text, &alphabet);

        let ranks = alphabets::RankTransform::new(&alphabet);

        let qgram = ranks.qgrams(q, b"TGA").next().unwrap();

        let matches = qgram_index.qgram_matches(qgram);
        assert_eq!(matches, [5, 10]);
    }

    #[test]
    fn test_matches() {
        let (text, alphabet) = setup();
        let q = 3;
        let qgram_index = QGramIndex::new(q, text, &alphabet);

        let pattern = b"GCTG";
        let matches = qgram_index.matches(pattern, 1);
        assert_eq!(matches, [
            Match {
                pattern: Interval { start: 0, stop: 4 },
                text: Interval { start: 3, stop: 7 },
                count: 2
            }
        ]);
    }

    #[test]
    fn test_exact_matches() {
        let (text, alphabet) = setup();
        let q = 3;
        let qgram_index = QGramIndex::new(q, text, &alphabet);

        let pattern = b"GCTGA";
        let exact_matches = qgram_index.exact_matches(pattern);
        assert!(exact_matches.len() == 2);
        for m in exact_matches {
            assert_eq!(m.pattern.get(pattern), m.text.get(text));
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::iter;
use std::sync::Arc;

use buffer::TrackedBuffer;
use command_buffer::DynamicState;
use command_buffer::pool::CommandPool;
use command_buffer::states_manager::StatesManager;
use command_buffer::submit::Submit;
use command_buffer::submit::SubmitInfo;
use command_buffer::sys::PipelineBarrierBuilder;
use command_buffer::sys::UnsafeCommandBufferBuilder;
use descriptor::PipelineLayout;
use descriptor::descriptor_set::collection::TrackedDescriptorSetsCollection;
use device::Device;
use device::Queue;
use framebuffer::traits::TrackedFramebuffer;
use framebuffer::RenderPass;
use framebuffer::RenderPassClearValues;
use instance::QueueFamily;
use pipeline::ComputePipeline;
use pipeline::GraphicsPipeline;
use pipeline::vertex::Source;
use sync::Fence;
use vk;

pub use self::empty::PrimaryCb;
pub use self::empty::PrimaryCbBuilder;

pub mod dispatch;
pub mod draw;
pub mod empty;
pub mod execute;
pub mod fill_buffer;
pub mod render_pass;
pub mod update_buffer;

/// A list of commands that can be turned into a command buffer.
pub unsafe trait CommandsList {
    /// Adds a command that writes the content of a buffer.
    ///
    /// After this command is executed, the content of `buffer` will become `data`.
    #[inline]
    fn update_buffer<'a, B, D: ?Sized>(self, buffer: B, data: &'a D)
                                       -> update_buffer::UpdateCommand<'a, Self, B, D>
        where Self: Sized + CommandsListPossibleOutsideRenderPass, B: TrackedBuffer, D: Copy + 'static
    {
        update_buffer::UpdateCommand::new(self, buffer, data)
    }

    /// Adds a command that writes the content of a buffer.
    #[inline]
    fn fill_buffer<B>(self, buffer: B, data: u32) -> fill_buffer::FillCommand<Self, B>
        where Self: Sized + CommandsListPossibleOutsideRenderPass, B: TrackedBuffer
    {
        fill_buffer::FillCommand::new(self, buffer, data)
    }

    /// Adds a command that executes a secondary command buffer.
    ///
    /// When you create a command buffer, you have the possibility to create either a primary
    /// command buffer or a secondary command buffer. Secondary command buffers can't be executed
    /// directly, but can be executed from a primary command buffer.
    ///
    /// A secondary command buffer can't execute another secondary command buffer. The only way
    /// you can use `execute` is to make a primary command buffer call a secondary command buffer.
    #[inline]
    fn execute<Cb>(self, command_buffer: Cb) -> execute::ExecuteCommand<Cb, Self>
        where Self: Sized, Cb: CommandsListOutput       /* FIXME: */
    {
        execute::ExecuteCommand::new(self, command_buffer)
    }

    /// Adds a command that executes a compute shader.
    ///
    /// The `dimensions` are the number of working groups to start. The GPU will execute the
    /// compute shader `dimensions[0] * dimensions[1] * dimensions[2]` times.
    ///
    /// The `pipeline` is the compute pipeline that will be executed, and the sets and push
    /// constants will be accessible to all the invocations.
    #[inline]
    fn dispatch<'a, Pl, S, Pc>(self, pipeline: Arc<ComputePipeline<Pl>>, sets: S,
                               dimensions: [u32; 3], push_constants: &'a Pc)
                               -> dispatch::DispatchCommand<'a, Self, Pl, S, Pc>
        where Self: Sized + CommandsList + CommandsListPossibleOutsideRenderPass, Pl: PipelineLayout,
              S: TrackedDescriptorSetsCollection, Pc: 'a
    {
        dispatch::DispatchCommand::new(self, pipeline, sets, dimensions, push_constants)
    }

    /// Adds a command that starts a render pass.
    ///
    /// If `secondary` is true, then you will only be able to add secondary command buffers while
    /// you're inside the first subpass on the render pass. If `secondary` is false, you will only
    /// be able to add inline draw commands and not secondary command buffers.
    ///
    /// You must call this before you can add draw commands.
    #[inline]
    fn begin_render_pass<F, C>(self, framebuffer: F, secondary: bool, clear_values: C)
                               -> render_pass::BeginRenderPassCommand<Self, F::RenderPass, F>
        where Self: Sized + CommandsListPossibleOutsideRenderPass,
              F: TrackedFramebuffer, F::RenderPass: RenderPass + RenderPassClearValues<C>
    {
        render_pass::BeginRenderPassCommand::new(self, framebuffer, secondary, clear_values)
    }

    /// Adds a command that jumps to the next subpass of the current render pass.
    #[inline]
    fn next_subpass(self, secondary: bool) -> render_pass::NextSubpassCommand<Self>
        where Self: Sized + CommandsListPossibleInsideRenderPass
    {
        render_pass::NextSubpassCommand::new(self, secondary)
    }

    /// Adds a command that ends the current render pass.
    ///
    /// This must be called after you went through all the subpasses and before you can build
    /// the command buffer or add further commands.
    #[inline]
    fn end_render_pass(self) -> render_pass::EndRenderPassCommand<Self>
        where Self: Sized + CommandsListPossibleInsideRenderPass
    {
        render_pass::EndRenderPassCommand::new(self)
    }

    /// Adds a command that draws.
    ///
    /// Can only be used from inside a render pass.
    #[inline]
    fn draw<'a, Pv, Pl, Prp, S, Pc, V>(self, pipeline: Arc<GraphicsPipeline<Pv, Pl, Prp>>,
                                       dynamic: &DynamicState, vertices: V, sets: S,
                                       push_constants: &'a Pc)
                                       -> draw::DrawCommand<'a, Self, Pv, Pl, Prp, S, Pc>
        where Self: Sized + CommandsList + CommandsListPossibleInsideRenderPass, Pl: PipelineLayout,
              S: TrackedDescriptorSetsCollection, Pc: 'a, Pv: Source<V>
    {
        draw::DrawCommand::regular(self, pipeline, dynamic, vertices, sets, push_constants)
    }

    /// Turns the commands list into a command buffer that can be submitted.
    #[inline]
    fn build(self) -> CommandBuffer<Self::Output> where Self: Sized + CommandsListConcrete {
        CommandsListConcrete::build(self)
    }

    /// Builds a boxed command buffer object.
    #[inline]
    fn abstract_build(self) -> CommandBuffer where Self: Sized + CommandsListAbstract {
        CommandsListAbstract::abstract_build(self)
    }

    /// Returns true if the command buffer can be built. This function should always return true,
    /// except when we're building a primary command buffer that is inside a render pass.
    fn buildable_state(&self) -> bool;

    /// Returns the number of commands in the commands list.
    ///
    /// Note that multiple actual commands may count for just 1.
    fn num_commands(&self) -> usize;

    /// Checks whether the command can be executed on the given queue family.
    // TODO: error type?
    fn check_queue_validity(&self, queue: QueueFamily) -> Result<(), ()>;

    /// Extracts the object that contains the states of all the resources of the commands list.
    ///
    /// Panics if the states were already extracted.
    fn extract_states(&mut self) -> StatesManager;

    /// Returns true if the given compute pipeline is currently binded in the commands list.
    // TODO: better API?
    fn is_compute_pipeline_bound(&self, pipeline: vk::Pipeline) -> bool;

    /// Returns true if the given graphics pipeline is currently binded in the commands list.
    // TODO: better API?
    fn is_graphics_pipeline_bound(&self, pipeline: vk::Pipeline) -> bool;
}

pub unsafe trait CommandsListConcrete: CommandsList {
    type Pool: CommandPool;
    /// The type of the command buffer that will be generated.
    type Output: CommandsListOutput;

    /// Turns the commands list into a command buffer.
    ///
    /// This function accepts additional arguments that will customize the output:
    ///
    /// - `additional_elements` is a closure that must be called on the command buffer builder
    ///   after it has finished building and before `final_barrier` are added.
    /// - `barriers` is a list of pipeline barriers accompanied by a command number. The
    ///   pipeline barrier must happen after the given command number. Usually you want all the
    ///   the command numbers to be inferior to `num_commands`.
    /// - `final_barrier` is a pipeline barrier that must be added at the end of the
    ///   command buffer builder.
    ///
    /// This function doesn't check that `buildable_state` returns true.
    unsafe fn raw_build<I, F>(self, in_s: &mut StatesManager, out: &mut StatesManager,
                              additional_elements: F, barriers: I,
                              final_barrier: PipelineBarrierBuilder) -> Self::Output
        where F: FnOnce(&mut UnsafeCommandBufferBuilder<Self::Pool>),
              I: Iterator<Item = (usize, PipelineBarrierBuilder)>;

    /// Turns the commands list into a command buffer that can be submitted.
    // This function isn't inline because `raw_build` implementations usually are inline.
    fn build(mut self) -> CommandBuffer<Self::Output> where Self: Sized {
        assert!(self.buildable_state(), "Tried to build a command buffer still inside a \
                                         render pass");

        let mut states_in = self.extract_states();
        let mut states_out = StatesManager::new();

        let output = unsafe {
            self.raw_build(&mut states_in, &mut states_out, |_| {},
                           iter::empty(), PipelineBarrierBuilder::new())
        };

        CommandBuffer {
            states: states_out,
            commands: output,
        }
    }
}

pub unsafe trait CommandsListAbstract: CommandsList {
    /// Turns the commands list into a command buffer that can be submitted.
    fn abstract_build(self) -> CommandBuffer where Self: Sized;
}

unsafe impl<C> CommandsListAbstract for C
    where C: CommandsListConcrete, <C as CommandsListConcrete>::Output: 'static
{
    #[inline]
    fn abstract_build(self) -> CommandBuffer {
        let tmp = CommandsListConcrete::build(self);

        CommandBuffer {
            states: tmp.states,
            commands: Box::new(tmp.commands) as Box<_>,
        }
    }
}

/// Extension trait for both `CommandsList` and `CommandsListOutput` that indicates that we're
/// possibly outside a render pass.
///
/// In other words, if this trait is *not* implemented then we're guaranteed *not* to be outside
/// of a render pass. If it is implemented, then we maybe are but that's not sure.
pub unsafe trait CommandsListPossibleOutsideRenderPass {
    /// Returns `true` if we're outside a render pass.
    fn is_outside_render_pass(&self) -> bool;
}

/// Extension trait for both `CommandsList` and `CommandsListOutput` that indicates that we're
/// possibly inside a render pass.
///
/// In other words, if this trait is *not* implemented then we're guaranteed *not* to be inside
/// a render pass. If it is implemented, then we maybe are but that's not sure.
// TODO: make all return values optional, since we're possibly not in a render pass
pub unsafe trait CommandsListPossibleInsideRenderPass {
    type RenderPass: RenderPass;

    /// Returns the number of the subpass we're in. The value is 0-indexed, so immediately after
    /// calling `begin_render_pass` the value will be `0`.
    ///
    /// The value should always be strictly inferior to the number of subpasses in the render pass.
    fn current_subpass_num(&self) -> u32;

    /// If true, only secondary command buffers can be added inside the subpass. If false, only
    /// inline draw commands can be added.
    fn secondary_subpass(&self) -> bool;

    /// Returns the description of the render pass we're in.
    // TODO: return a trait object instead?
    fn render_pass(&self) -> &Self::RenderPass;

    //fn current_subpass(&self) -> Subpass<&Self::RenderPass>;
}

pub unsafe trait CommandsListOutput<S = StatesManager> {
    /// Returns the inner object.
    // TODO: crappy API
    fn inner(&self) -> vk::CommandBuffer;

    /// Returns the device this object belongs to.
    fn device(&self) -> &Arc<Device>;

    unsafe fn on_submit(&self, states: &S, queue: &Arc<Queue>,
                        fence: &mut FnMut() -> Arc<Fence>) -> SubmitInfo;
}

pub struct CommandBuffer<C = Box<CommandsListOutput>> {
    states: StatesManager,
    commands: C,
}

unsafe impl<C> Submit for CommandBuffer<C> where C: CommandsListOutput {
    #[inline]
    fn inner(&self) -> vk::CommandBuffer {
        self.commands.inner()
    }

    #[inline]
    fn device(&self) -> &Arc<Device> {
        self.commands.device()
    }

    #[inline]
    unsafe fn on_submit<F>(&self, queue: &Arc<Queue>, mut fence: F) -> SubmitInfo
        where F: FnMut() -> Arc<Fence>
    {
        self.commands.on_submit(&self.states, queue, &mut fence)
    }
}
<commit_msg>Impl CommandsList for Box<CommandsList><commit_after>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::iter;
use std::sync::Arc;

use buffer::TrackedBuffer;
use command_buffer::DynamicState;
use command_buffer::pool::CommandPool;
use command_buffer::states_manager::StatesManager;
use command_buffer::submit::Submit;
use command_buffer::submit::SubmitInfo;
use command_buffer::sys::PipelineBarrierBuilder;
use command_buffer::sys::UnsafeCommandBufferBuilder;
use descriptor::PipelineLayout;
use descriptor::descriptor_set::collection::TrackedDescriptorSetsCollection;
use device::Device;
use device::Queue;
use framebuffer::traits::TrackedFramebuffer;
use framebuffer::RenderPass;
use framebuffer::RenderPassClearValues;
use instance::QueueFamily;
use pipeline::ComputePipeline;
use pipeline::GraphicsPipeline;
use pipeline::vertex::Source;
use sync::Fence;
use vk;

pub use self::empty::PrimaryCb;
pub use self::empty::PrimaryCbBuilder;

pub mod dispatch;
pub mod draw;
pub mod empty;
pub mod execute;
pub mod fill_buffer;
pub mod render_pass;
pub mod update_buffer;

/// A list of commands that can be turned into a command buffer.
pub unsafe trait CommandsList {
    /// Adds a command that writes the content of a buffer.
    ///
    /// After this command is executed, the content of `buffer` will become `data`.
    #[inline]
    fn update_buffer<'a, B, D: ?Sized>(self, buffer: B, data: &'a D)
                                       -> update_buffer::UpdateCommand<'a, Self, B, D>
        where Self: Sized + CommandsListPossibleOutsideRenderPass, B: TrackedBuffer, D: Copy + 'static
    {
        update_buffer::UpdateCommand::new(self, buffer, data)
    }

    /// Adds a command that writes the content of a buffer.
    #[inline]
    fn fill_buffer<B>(self, buffer: B, data: u32) -> fill_buffer::FillCommand<Self, B>
        where Self: Sized + CommandsListPossibleOutsideRenderPass, B: TrackedBuffer
    {
        fill_buffer::FillCommand::new(self, buffer, data)
    }

    /// Adds a command that executes a secondary command buffer.
    ///
    /// When you create a command buffer, you have the possibility to create either a primary
    /// command buffer or a secondary command buffer. Secondary command buffers can't be executed
    /// directly, but can be executed from a primary command buffer.
    ///
    /// A secondary command buffer can't execute another secondary command buffer. The only way
    /// you can use `execute` is to make a primary command buffer call a secondary command buffer.
    #[inline]
    fn execute<Cb>(self, command_buffer: Cb) -> execute::ExecuteCommand<Cb, Self>
        where Self: Sized, Cb: CommandsListOutput       /* FIXME: */
    {
        execute::ExecuteCommand::new(self, command_buffer)
    }

    /// Adds a command that executes a compute shader.
    ///
    /// The `dimensions` are the number of working groups to start. The GPU will execute the
    /// compute shader `dimensions[0] * dimensions[1] * dimensions[2]` times.
    ///
    /// The `pipeline` is the compute pipeline that will be executed, and the sets and push
    /// constants will be accessible to all the invocations.
    #[inline]
    fn dispatch<'a, Pl, S, Pc>(self, pipeline: Arc<ComputePipeline<Pl>>, sets: S,
                               dimensions: [u32; 3], push_constants: &'a Pc)
                               -> dispatch::DispatchCommand<'a, Self, Pl, S, Pc>
        where Self: Sized + CommandsList + CommandsListPossibleOutsideRenderPass, Pl: PipelineLayout,
              S: TrackedDescriptorSetsCollection, Pc: 'a
    {
        dispatch::DispatchCommand::new(self, pipeline, sets, dimensions, push_constants)
    }

    /// Adds a command that starts a render pass.
    ///
    /// If `secondary` is true, then you will only be able to add secondary command buffers while
    /// you're inside the first subpass on the render pass. If `secondary` is false, you will only
    /// be able to add inline draw commands and not secondary command buffers.
    ///
    /// You must call this before you can add draw commands.
    #[inline]
    fn begin_render_pass<F, C>(self, framebuffer: F, secondary: bool, clear_values: C)
                               -> render_pass::BeginRenderPassCommand<Self, F::RenderPass, F>
        where Self: Sized + CommandsListPossibleOutsideRenderPass,
              F: TrackedFramebuffer, F::RenderPass: RenderPass + RenderPassClearValues<C>
    {
        render_pass::BeginRenderPassCommand::new(self, framebuffer, secondary, clear_values)
    }

    /// Adds a command that jumps to the next subpass of the current render pass.
    #[inline]
    fn next_subpass(self, secondary: bool) -> render_pass::NextSubpassCommand<Self>
        where Self: Sized + CommandsListPossibleInsideRenderPass
    {
        render_pass::NextSubpassCommand::new(self, secondary)
    }

    /// Adds a command that ends the current render pass.
    ///
    /// This must be called after you went through all the subpasses and before you can build
    /// the command buffer or add further commands.
    #[inline]
    fn end_render_pass(self) -> render_pass::EndRenderPassCommand<Self>
        where Self: Sized + CommandsListPossibleInsideRenderPass
    {
        render_pass::EndRenderPassCommand::new(self)
    }

    /// Adds a command that draws.
    ///
    /// Can only be used from inside a render pass.
    #[inline]
    fn draw<'a, Pv, Pl, Prp, S, Pc, V>(self, pipeline: Arc<GraphicsPipeline<Pv, Pl, Prp>>,
                                       dynamic: &DynamicState, vertices: V, sets: S,
                                       push_constants: &'a Pc)
                                       -> draw::DrawCommand<'a, Self, Pv, Pl, Prp, S, Pc>
        where Self: Sized + CommandsList + CommandsListPossibleInsideRenderPass, Pl: PipelineLayout,
              S: TrackedDescriptorSetsCollection, Pc: 'a, Pv: Source<V>
    {
        draw::DrawCommand::regular(self, pipeline, dynamic, vertices, sets, push_constants)
    }

    /// Turns the commands list into a command buffer that can be submitted.
    #[inline]
    fn build(self) -> CommandBuffer<Self::Output> where Self: Sized + CommandsListConcrete {
        CommandsListConcrete::build(self)
    }

    /// Builds a boxed command buffer object.
    #[inline]
    fn abstract_build(self) -> CommandBuffer where Self: Sized + CommandsListAbstract {
        CommandsListAbstract::abstract_build(self)
    }

    /// Returns true if the command buffer can be built. This function should always return true,
    /// except when we're building a primary command buffer that is inside a render pass.
    fn buildable_state(&self) -> bool;

    /// Returns the number of commands in the commands list.
    ///
    /// Note that multiple actual commands may count for just 1.
    fn num_commands(&self) -> usize;

    /// Checks whether the command can be executed on the given queue family.
    // TODO: error type?
    fn check_queue_validity(&self, queue: QueueFamily) -> Result<(), ()>;

    /// Extracts the object that contains the states of all the resources of the commands list.
    ///
    /// Panics if the states were already extracted.
    fn extract_states(&mut self) -> StatesManager;

    /// Returns true if the given compute pipeline is currently binded in the commands list.
    // TODO: better API?
    fn is_compute_pipeline_bound(&self, pipeline: vk::Pipeline) -> bool;

    /// Returns true if the given graphics pipeline is currently binded in the commands list.
    // TODO: better API?
    fn is_graphics_pipeline_bound(&self, pipeline: vk::Pipeline) -> bool;
}

unsafe impl CommandsList for Box<CommandsList> {
    #[inline]
    fn buildable_state(&self) -> bool {
        (**self).buildable_state()
    }

    #[inline]
    fn num_commands(&self) -> usize {
        (**self).num_commands()
    }

    #[inline]
    fn check_queue_validity(&self, queue: QueueFamily) -> Result<(), ()> {
        (**self).check_queue_validity(queue)
    }

    #[inline]
    fn extract_states(&mut self) -> StatesManager {
        (**self).extract_states()
    }

    #[inline]
    fn is_compute_pipeline_bound(&self, pipeline: vk::Pipeline) -> bool {
        (**self).is_compute_pipeline_bound(pipeline)
    }

    #[inline]
    fn is_graphics_pipeline_bound(&self, pipeline: vk::Pipeline) -> bool {
        (**self).is_graphics_pipeline_bound(pipeline)
    }
}

pub unsafe trait CommandsListConcrete: CommandsList {
    type Pool: CommandPool;
    /// The type of the command buffer that will be generated.
    type Output: CommandsListOutput;

    /// Turns the commands list into a command buffer.
    ///
    /// This function accepts additional arguments that will customize the output:
    ///
    /// - `additional_elements` is a closure that must be called on the command buffer builder
    ///   after it has finished building and before `final_barrier` are added.
    /// - `barriers` is a list of pipeline barriers accompanied by a command number. The
    ///   pipeline barrier must happen after the given command number. Usually you want all the
    ///   the command numbers to be inferior to `num_commands`.
    /// - `final_barrier` is a pipeline barrier that must be added at the end of the
    ///   command buffer builder.
    ///
    /// This function doesn't check that `buildable_state` returns true.
    unsafe fn raw_build<I, F>(self, in_s: &mut StatesManager, out: &mut StatesManager,
                              additional_elements: F, barriers: I,
                              final_barrier: PipelineBarrierBuilder) -> Self::Output
        where F: FnOnce(&mut UnsafeCommandBufferBuilder<Self::Pool>),
              I: Iterator<Item = (usize, PipelineBarrierBuilder)>;

    /// Turns the commands list into a command buffer that can be submitted.
    // This function isn't inline because `raw_build` implementations usually are inline.
    fn build(mut self) -> CommandBuffer<Self::Output> where Self: Sized {
        assert!(self.buildable_state(), "Tried to build a command buffer still inside a \
                                         render pass");

        let mut states_in = self.extract_states();
        let mut states_out = StatesManager::new();

        let output = unsafe {
            self.raw_build(&mut states_in, &mut states_out, |_| {},
                           iter::empty(), PipelineBarrierBuilder::new())
        };

        CommandBuffer {
            states: states_out,
            commands: output,
        }
    }
}

pub unsafe trait CommandsListAbstract: CommandsList {
    /// Turns the commands list into a command buffer that can be submitted.
    fn abstract_build(self) -> CommandBuffer where Self: Sized;
}

unsafe impl<C> CommandsListAbstract for C
    where C: CommandsListConcrete, <C as CommandsListConcrete>::Output: 'static
{
    #[inline]
    fn abstract_build(self) -> CommandBuffer {
        let tmp = CommandsListConcrete::build(self);

        CommandBuffer {
            states: tmp.states,
            commands: Box::new(tmp.commands) as Box<_>,
        }
    }
}

/// Extension trait for both `CommandsList` and `CommandsListOutput` that indicates that we're
/// possibly outside a render pass.
///
/// In other words, if this trait is *not* implemented then we're guaranteed *not* to be outside
/// of a render pass. If it is implemented, then we maybe are but that's not sure.
pub unsafe trait CommandsListPossibleOutsideRenderPass {
    /// Returns `true` if we're outside a render pass.
    fn is_outside_render_pass(&self) -> bool;
}

/// Extension trait for both `CommandsList` and `CommandsListOutput` that indicates that we're
/// possibly inside a render pass.
///
/// In other words, if this trait is *not* implemented then we're guaranteed *not* to be inside
/// a render pass. If it is implemented, then we maybe are but that's not sure.
// TODO: make all return values optional, since we're possibly not in a render pass
pub unsafe trait CommandsListPossibleInsideRenderPass {
    type RenderPass: RenderPass;

    /// Returns the number of the subpass we're in. The value is 0-indexed, so immediately after
    /// calling `begin_render_pass` the value will be `0`.
    ///
    /// The value should always be strictly inferior to the number of subpasses in the render pass.
    fn current_subpass_num(&self) -> u32;

    /// If true, only secondary command buffers can be added inside the subpass. If false, only
    /// inline draw commands can be added.
    fn secondary_subpass(&self) -> bool;

    /// Returns the description of the render pass we're in.
    // TODO: return a trait object instead?
    fn render_pass(&self) -> &Self::RenderPass;

    //fn current_subpass(&self) -> Subpass<&Self::RenderPass>;
}

pub unsafe trait CommandsListOutput<S = StatesManager> {
    /// Returns the inner object.
    // TODO: crappy API
    fn inner(&self) -> vk::CommandBuffer;

    /// Returns the device this object belongs to.
    fn device(&self) -> &Arc<Device>;

    unsafe fn on_submit(&self, states: &S, queue: &Arc<Queue>,
                        fence: &mut FnMut() -> Arc<Fence>) -> SubmitInfo;
}

pub struct CommandBuffer<C = Box<CommandsListOutput>> {
    states: StatesManager,
    commands: C,
}

unsafe impl<C> Submit for CommandBuffer<C> where C: CommandsListOutput {
    #[inline]
    fn inner(&self) -> vk::CommandBuffer {
        self.commands.inner()
    }

    #[inline]
    fn device(&self) -> &Arc<Device> {
        self.commands.device()
    }

    #[inline]
    unsafe fn on_submit<F>(&self, queue: &Arc<Queue>, mut fence: F) -> SubmitInfo
        where F: FnMut() -> Arc<Fence>
    {
        self.commands.on_submit(&self.states, queue, &mut fence)
    }
}
<|endoftext|>


<commit_before>//! Accept connections from clients over TCP.
//!
//! This is a (still incomplete) rewrite of the original mio-based code.

use byteorder::{BigEndian, ByteOrder, ReadBytesExt, WriteBytesExt};
use bytes::BufMut;
use cache::Cache;
use client_query::*;
use dns::{self, NormalizedQuestion};
use futures::future::{self, Future};
use futures::Sink;
use futures::stream::Stream;
use futures::sync::mpsc::{channel, Sender};
use std::cell::RefCell;
use std::io::{self, Read, Write};
use std::net::{self, SocketAddr};
use std::rc::Rc;
use std::time;
use std::sync::{Arc, mpsc};
use std::thread;
use super::EdgeDNSContext;
use super::{DNS_QUERY_MAX_SIZE, DNS_QUERY_MIN_SIZE, MAX_TCP_IDLE_MS};
use tcp_arbitrator::TcpArbitrator;
use tokio_core::net::{TcpListener, TcpStream};
use tokio_core::reactor::{Core, Handle};
use tokio_io::io::{read_exact, write_all, ReadHalf, WriteHalf};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_timer::{wheel, Timer};
use varz::Varz;

struct TcpAcceptor {
    timer: Timer,
    handle: Handle,
    net_tcp_listener: net::TcpListener,
    resolver_tx: Sender<ClientQuery>,
    cache: Cache,
    varz: Arc<Varz>,
    tcp_arbitrator: TcpArbitrator,
}

pub struct TcpAcceptorCore {
    timer: Timer,
    handle: Handle,
    net_tcp_listener: net::TcpListener,
    resolver_tx: Sender<ClientQuery>,
    cache: Cache,
    varz: Arc<Varz>,
    service_ready_tx: Option<mpsc::SyncSender<u8>>,
    tcp_arbitrator: TcpArbitrator,
}

struct TcpClientQuery {
    timer: Timer,
    wh: WriteHalf<TcpStream>,
    handle: Handle,
    resolver_tx: Sender<ClientQuery>,
    cache: Cache,
    varz: Arc<Varz>,
}

impl TcpClientQuery {
    pub fn new(tcp_acceptor: &TcpAcceptor, wh: WriteHalf<TcpStream>) -> Self {
        TcpClientQuery {
            timer: tcp_acceptor.timer.clone(),
            wh: wh,
            handle: tcp_acceptor.handle.clone(),
            resolver_tx: tcp_acceptor.resolver_tx.clone(),
            cache: tcp_acceptor.cache.clone(),
            varz: tcp_acceptor.varz.clone(),
        }
    }

    fn fut_process_query(mut self,
                         normalized_question: NormalizedQuestion)
                         -> Box<Future<Item = (), Error = io::Error>> {
        let (tcpclient_tx, tcpclient_rx) = channel(1);
        let cache_entry = self.cache.get2(&normalized_question);
        let client_query = ClientQuery::tcp(tcpclient_tx, normalized_question);
        let wh_cell = RefCell::new(self.wh);
        let fut = tcpclient_rx
            .into_future()
            .map_err(|_| {})
            .map(|(resolver_response, _)| resolver_response)
            .and_then(move |resolver_response| match resolver_response {
                          None => {
                warn!("No resolver response - TX part of the channel closed");
                future::err(())
            }
                          Some(resolver_response) => future::ok(resolver_response),
                      })
            .and_then(|resolver_response| {
                          let wh = wh_cell.into_inner();
                          write_all(wh, resolver_response.packet)
                              .map(|_| {})
                              .map_err(|_| {})
                      });
        if let Some(mut cache_entry) = cache_entry {
            if !cache_entry.is_expired() {
                self.varz.client_queries_cached.inc();
                self.handle.spawn(fut.map_err(|_| {}));
                return client_query.response_send(&mut cache_entry.packet, None);
            }
            debug!("expired");
            self.varz.client_queries_expired.inc();
        }
        let fut_send = self.resolver_tx.send(client_query).map_err(|_| {});
        let futs = fut.join(fut_send);
        Box::new(futs.map(|_| {}).map_err(|_| io::Error::last_os_error()))
    }
}

impl TcpAcceptor {
    fn new(tcp_acceptor_core: &TcpAcceptorCore) -> Self {
        TcpAcceptor {
            timer: tcp_acceptor_core.timer.clone(),
            handle: tcp_acceptor_core.handle.clone(),
            net_tcp_listener: tcp_acceptor_core
                .net_tcp_listener
                .try_clone()
                .expect("Couldn't clone a TCP socket"),
            resolver_tx: tcp_acceptor_core.resolver_tx.clone(),
            cache: tcp_acceptor_core.cache.clone(),
            varz: tcp_acceptor_core.varz.clone(),
            tcp_arbitrator: tcp_acceptor_core.tcp_arbitrator.clone(),
        }
    }

    fn fut_process_client(&mut self,
                          client: TcpStream,
                          client_addr: SocketAddr)
                          -> Box<Future<Item = (), Error = io::Error>> {
        let (session_rx, session_idx) = match self.tcp_arbitrator.new_session(&client_addr) {
            Ok(r) => r,
            Err(_) => return Box::new(future::err(io::Error::last_os_error())),
        };
        debug!("Incoming connection using TCP, session index {}",
               session_idx);
        let varz = self.varz.clone();
        varz.client_queries_tcp.inc();
        let (rh, wh) = client.split();
        let fut_expected_len = read_exact(rh, vec![0u8; 2]).and_then(move |(rh, len_buf)| {
            let expected_len = BigEndian::read_u16(&len_buf) as usize;
            if expected_len < DNS_QUERY_MIN_SIZE || expected_len > DNS_QUERY_MAX_SIZE {
                info!("Suspicious query length: {}", expected_len);
                varz.client_queries_errors.inc();
                return future::err(io::Error::new(io::ErrorKind::InvalidInput,
                                                  "Suspicious query length"));
            }
            debug!("Expected length: {}", expected_len);
            future::ok((rh, expected_len))
        });
        let fut_packet_read =
            fut_expected_len.and_then(|(rh, expected_len)| read_exact(rh, vec![0u8; expected_len]));
        let varz = self.varz.clone();
        let tcp_client_query = TcpClientQuery::new(self, wh);
        let fut_packet = fut_packet_read.and_then(move |(rh, packet)| {
            let normalized_question = match dns::normalize(&packet, true) {
                Ok(normalized_question) => normalized_question,
                Err(e) => {
                    debug!("Error while parsing the question: {}", e);
                    varz.client_queries_errors.inc();
                    return Box::new(future::err(io::Error::new(io::ErrorKind::InvalidInput,
                                                               "Suspicious query"))) as
                           Box<Future<Item = _, Error = _>>;
                }
            };
            tcp_client_query.fut_process_query(normalized_question)
        });
        let fut_timeout = self.timer
            .timeout(fut_packet, time::Duration::from_millis(MAX_TCP_IDLE_MS));
        let mut tcp_arbitrator = self.tcp_arbitrator.clone();
        let fut_with_timeout =
            fut_timeout.then(move |_| {
                                 debug!("Closing TCP connection with session index {}",
                                        session_idx);
                                 tcp_arbitrator.delete_session(session_idx);
                                 future::ok(())
                             });
        let fut_session_rx = session_rx.map(|_| {});
        let fut = fut_session_rx
            .select(fut_with_timeout)
            .map(|_| {})
            .map_err(|_| io::Error::last_os_error());
        Box::new(fut) as Box<Future<Item = _, Error = _>>
    }

    fn fut_process_stream<'a>(mut self,
                              handle: &Handle)
                              -> impl Future<Item = (), Error = io::Error> + 'a {
        let tcp_listener = TcpListener::from_listener(self.net_tcp_listener
                                                          .try_clone()
                                                          .expect("Unable to clone a TCP socket"),
                                                      &self.net_tcp_listener.local_addr().unwrap(),
                                                      handle)
                .expect("Unable to create a tokio TCP listener");
        let handle = handle.clone();
        tcp_listener
            .incoming()
            .for_each(move |(client, client_addr)| {
                          let fut_client = self.fut_process_client(client, client_addr);
                          handle.spawn(fut_client.map_err(|_| {}));
                          Ok(())
                      })
            .map_err(|_| io::Error::last_os_error())
    }
}

impl TcpAcceptorCore {
    fn run(mut self, mut event_loop: Core, tcp_acceptor: TcpAcceptor) -> io::Result<()> {
        let service_ready_tx = self.service_ready_tx.take().unwrap();
        let handle = event_loop.handle();
        let stream = tcp_acceptor.fut_process_stream(&handle);
        handle.spawn(stream.map_err(|_| {}).map(|_| {}));
        service_ready_tx
            .send(0)
            .map_err(|_| io::Error::last_os_error())?;
        loop {
            event_loop.turn(None)
        }
    }

    pub fn spawn(edgedns_context: &EdgeDNSContext,
                 resolver_tx: Sender<ClientQuery>,
                 service_ready_tx: mpsc::SyncSender<u8>)
                 -> io::Result<(thread::JoinHandle<()>)> {
        let net_tcp_listener = edgedns_context.tcp_listener.try_clone()?;
        let cache = edgedns_context.cache.clone();
        let varz = edgedns_context.varz.clone();
        let tcp_arbitrator = edgedns_context.tcp_arbitrator.clone();
        let timer = wheel()
            .tick_duration(time::Duration::from_millis(MAX_TCP_IDLE_MS / 2))
            .max_timeout(time::Duration::from_millis(MAX_TCP_IDLE_MS))
            .build();
        let tcp_acceptor_th = thread::Builder::new()
            .name("tcp_acceptor".to_string())
            .spawn(move || {
                let event_loop = Core::new().unwrap();
                let tcp_acceptor_core = TcpAcceptorCore {
                    timer: timer,
                    handle: event_loop.handle(),
                    net_tcp_listener: net_tcp_listener,
                    cache: cache,
                    resolver_tx: resolver_tx,
                    service_ready_tx: Some(service_ready_tx),
                    varz: varz,
                    tcp_arbitrator: tcp_arbitrator,
                };
                let tcp_acceptor = TcpAcceptor::new(&tcp_acceptor_core);
                tcp_acceptor_core
                    .run(event_loop, tcp_acceptor)
                    .expect("Unable to spawn a tcp listener");
            })
            .unwrap();
        info!("TCP listener is ready");
        Ok(tcp_acceptor_th)
    }
}
<commit_msg>Format<commit_after>//! Accept connections from clients over TCP.
//!
//! This is a (still incomplete) rewrite of the original mio-based code.

use byteorder::{BigEndian, ByteOrder, ReadBytesExt, WriteBytesExt};
use bytes::BufMut;
use cache::Cache;
use client_query::*;
use dns::{self, NormalizedQuestion};
use futures::future::{self, Future};
use futures::Sink;
use futures::stream::Stream;
use futures::sync::mpsc::{channel, Sender};
use std::cell::RefCell;
use std::io::{self, Read, Write};
use std::net::{self, SocketAddr};
use std::rc::Rc;
use std::time;
use std::sync::{Arc, mpsc};
use std::thread;
use super::EdgeDNSContext;
use super::{DNS_QUERY_MAX_SIZE, DNS_QUERY_MIN_SIZE, MAX_TCP_IDLE_MS};
use tcp_arbitrator::TcpArbitrator;
use tokio_core::net::{TcpListener, TcpStream};
use tokio_core::reactor::{Core, Handle};
use tokio_io::io::{read_exact, write_all, ReadHalf, WriteHalf};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_timer::{wheel, Timer};
use varz::Varz;

struct TcpAcceptor {
    timer: Timer,
    handle: Handle,
    net_tcp_listener: net::TcpListener,
    resolver_tx: Sender<ClientQuery>,
    cache: Cache,
    varz: Arc<Varz>,
    tcp_arbitrator: TcpArbitrator,
}

pub struct TcpAcceptorCore {
    timer: Timer,
    handle: Handle,
    net_tcp_listener: net::TcpListener,
    resolver_tx: Sender<ClientQuery>,
    cache: Cache,
    varz: Arc<Varz>,
    service_ready_tx: Option<mpsc::SyncSender<u8>>,
    tcp_arbitrator: TcpArbitrator,
}

struct TcpClientQuery {
    timer: Timer,
    wh: WriteHalf<TcpStream>,
    handle: Handle,
    resolver_tx: Sender<ClientQuery>,
    cache: Cache,
    varz: Arc<Varz>,
}

impl TcpClientQuery {
    pub fn new(tcp_acceptor: &TcpAcceptor, wh: WriteHalf<TcpStream>) -> Self {
        TcpClientQuery {
            timer: tcp_acceptor.timer.clone(),
            wh: wh,
            handle: tcp_acceptor.handle.clone(),
            resolver_tx: tcp_acceptor.resolver_tx.clone(),
            cache: tcp_acceptor.cache.clone(),
            varz: tcp_acceptor.varz.clone(),
        }
    }

    fn fut_process_query(mut self,
                         normalized_question: NormalizedQuestion)
                         -> Box<Future<Item = (), Error = io::Error>> {
        let (tcpclient_tx, tcpclient_rx) = channel(1);
        let cache_entry = self.cache.get2(&normalized_question);
        let client_query = ClientQuery::tcp(tcpclient_tx, normalized_question);
        let wh_cell = RefCell::new(self.wh);
        let fut = tcpclient_rx
            .into_future()
            .map_err(|_| {})
            .map(|(resolver_response, _)| resolver_response)
            .and_then(move |resolver_response| match resolver_response {
                          None => {
                              warn!("No resolver response - TX part of the channel closed");
                              future::err(())
                          }
                          Some(resolver_response) => future::ok(resolver_response),
                      })
            .and_then(|resolver_response| {
                          let wh = wh_cell.into_inner();
                          write_all(wh, resolver_response.packet)
                              .map(|_| {})
                              .map_err(|_| {})
                      });
        if let Some(mut cache_entry) = cache_entry {
            if !cache_entry.is_expired() {
                self.varz.client_queries_cached.inc();
                self.handle.spawn(fut.map_err(|_| {}));
                return client_query.response_send(&mut cache_entry.packet, None);
            }
            debug!("expired");
            self.varz.client_queries_expired.inc();
        }
        let fut_send = self.resolver_tx.send(client_query).map_err(|_| {});
        let futs = fut.join(fut_send);
        Box::new(futs.map(|_| {}).map_err(|_| io::Error::last_os_error()))
    }
}

impl TcpAcceptor {
    fn new(tcp_acceptor_core: &TcpAcceptorCore) -> Self {
        TcpAcceptor {
            timer: tcp_acceptor_core.timer.clone(),
            handle: tcp_acceptor_core.handle.clone(),
            net_tcp_listener: tcp_acceptor_core
                .net_tcp_listener
                .try_clone()
                .expect("Couldn't clone a TCP socket"),
            resolver_tx: tcp_acceptor_core.resolver_tx.clone(),
            cache: tcp_acceptor_core.cache.clone(),
            varz: tcp_acceptor_core.varz.clone(),
            tcp_arbitrator: tcp_acceptor_core.tcp_arbitrator.clone(),
        }
    }

    fn fut_process_client(&mut self,
                          client: TcpStream,
                          client_addr: SocketAddr)
                          -> Box<Future<Item = (), Error = io::Error>> {
        let (session_rx, session_idx) = match self.tcp_arbitrator.new_session(&client_addr) {
            Ok(r) => r,
            Err(_) => return Box::new(future::err(io::Error::last_os_error())),
        };
        debug!("Incoming connection using TCP, session index {}",
               session_idx);
        let varz = self.varz.clone();
        varz.client_queries_tcp.inc();
        let (rh, wh) = client.split();
        let fut_expected_len = read_exact(rh, vec![0u8; 2]).and_then(move |(rh, len_buf)| {
            let expected_len = BigEndian::read_u16(&len_buf) as usize;
            if expected_len < DNS_QUERY_MIN_SIZE || expected_len > DNS_QUERY_MAX_SIZE {
                info!("Suspicious query length: {}", expected_len);
                varz.client_queries_errors.inc();
                return future::err(io::Error::new(io::ErrorKind::InvalidInput,
                                                  "Suspicious query length"));
            }
            debug!("Expected length: {}", expected_len);
            future::ok((rh, expected_len))
        });
        let fut_packet_read =
            fut_expected_len.and_then(|(rh, expected_len)| read_exact(rh, vec![0u8; expected_len]));
        let varz = self.varz.clone();
        let tcp_client_query = TcpClientQuery::new(self, wh);
        let fut_packet = fut_packet_read.and_then(move |(rh, packet)| {
            let normalized_question = match dns::normalize(&packet, true) {
                Ok(normalized_question) => normalized_question,
                Err(e) => {
                    debug!("Error while parsing the question: {}", e);
                    varz.client_queries_errors.inc();
                    return Box::new(future::err(io::Error::new(io::ErrorKind::InvalidInput,
                                                               "Suspicious query"))) as
                           Box<Future<Item = _, Error = _>>;
                }
            };
            tcp_client_query.fut_process_query(normalized_question)
        });
        let fut_timeout = self.timer
            .timeout(fut_packet, time::Duration::from_millis(MAX_TCP_IDLE_MS));
        let mut tcp_arbitrator = self.tcp_arbitrator.clone();
        let fut_with_timeout = fut_timeout
            .then(move |_| {
                      debug!("Closing TCP connection with session index {}", session_idx);
                      tcp_arbitrator.delete_session(session_idx);
                      future::ok(())
                  });
        let fut_session_rx = session_rx.map(|_| {});
        let fut = fut_session_rx
            .select(fut_with_timeout)
            .map(|_| {})
            .map_err(|_| io::Error::last_os_error());
        Box::new(fut) as Box<Future<Item = _, Error = _>>
    }

    fn fut_process_stream<'a>(mut self,
                              handle: &Handle)
                              -> impl Future<Item = (), Error = io::Error> + 'a {
        let tcp_listener = TcpListener::from_listener(self.net_tcp_listener
                                                          .try_clone()
                                                          .expect("Unable to clone a TCP socket"),
                                                      &self.net_tcp_listener.local_addr().unwrap(),
                                                      handle)
                .expect("Unable to create a tokio TCP listener");
        let handle = handle.clone();
        tcp_listener
            .incoming()
            .for_each(move |(client, client_addr)| {
                          let fut_client = self.fut_process_client(client, client_addr);
                          handle.spawn(fut_client.map_err(|_| {}));
                          Ok(())
                      })
            .map_err(|_| io::Error::last_os_error())
    }
}

impl TcpAcceptorCore {
    fn run(mut self, mut event_loop: Core, tcp_acceptor: TcpAcceptor) -> io::Result<()> {
        let service_ready_tx = self.service_ready_tx.take().unwrap();
        let handle = event_loop.handle();
        let stream = tcp_acceptor.fut_process_stream(&handle);
        handle.spawn(stream.map_err(|_| {}).map(|_| {}));
        service_ready_tx
            .send(0)
            .map_err(|_| io::Error::last_os_error())?;
        loop {
            event_loop.turn(None)
        }
    }

    pub fn spawn(edgedns_context: &EdgeDNSContext,
                 resolver_tx: Sender<ClientQuery>,
                 service_ready_tx: mpsc::SyncSender<u8>)
                 -> io::Result<(thread::JoinHandle<()>)> {
        let net_tcp_listener = edgedns_context.tcp_listener.try_clone()?;
        let cache = edgedns_context.cache.clone();
        let varz = edgedns_context.varz.clone();
        let tcp_arbitrator = edgedns_context.tcp_arbitrator.clone();
        let timer = wheel()
            .tick_duration(time::Duration::from_millis(MAX_TCP_IDLE_MS / 2))
            .max_timeout(time::Duration::from_millis(MAX_TCP_IDLE_MS))
            .build();
        let tcp_acceptor_th = thread::Builder::new()
            .name("tcp_acceptor".to_string())
            .spawn(move || {
                let event_loop = Core::new().unwrap();
                let tcp_acceptor_core = TcpAcceptorCore {
                    timer: timer,
                    handle: event_loop.handle(),
                    net_tcp_listener: net_tcp_listener,
                    cache: cache,
                    resolver_tx: resolver_tx,
                    service_ready_tx: Some(service_ready_tx),
                    varz: varz,
                    tcp_arbitrator: tcp_arbitrator,
                };
                let tcp_acceptor = TcpAcceptor::new(&tcp_acceptor_core);
                tcp_acceptor_core
                    .run(event_loop, tcp_acceptor)
                    .expect("Unable to spawn a tcp listener");
            })
            .unwrap();
        info!("TCP listener is ready");
        Ok(tcp_acceptor_th)
    }
}
<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2015 Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) 2020 nicoo            <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use std::cmp::min;
use std::mem::swap;

pub fn gcd(mut u: u64, mut v: u64) -> u64 {
    // Stein's binary GCD algorithm
    // Base cases: gcd(n, 0) = gcd(0, n) = n
    if u == 0 {
        return v;
    } else if v == 0 {
        return u;
    }

    // gcd(2ⁱ u, 2ʲ v) = 2ᵏ gcd(u, v) with u, v odd and k = min(i, j)
    // 2ᵏ is the greatest power of two that divides both u and v
    let k = {
        let i = u.trailing_zeros();
        let j = v.trailing_zeros();
        u >>= i;
        v >>= j;
        min(i, j)
    };

    loop {
        // Loop invariant: u and v are odd
        debug_assert!(u % 2 == 1, "u = {} is even", u);
        debug_assert!(v % 2 == 1, "v = {} is even", v);

        // gcd(u, v) = gcd(|u - v|, min(u, v))
        if u > v {
            swap(&mut u, &mut v);
        }
        v -= u;

        if v == 0 {
            // Reached the base case; gcd is 2ᵏ u
            return u << k;
        }

        // gcd(u, 2ʲ v) = gcd(u, v) as u is odd
        v >>= v.trailing_zeros();
    }
}

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

    quickcheck! {
        fn euclidean(a: u64, b: u64) -> bool {
            // Test against the Euclidean algorithm
            let g = {
                let (mut a, mut b) = (a, b);
                while b > 0 {
                    a %= b;
                    swap(&mut a, &mut b);
                }
                a
            };
            gcd(a, b) == g
        }

        fn one(a: u64) -> bool {
            gcd(1, a) == 1
        }

        fn divisor(a: u64, b: u64) -> bool {
            // Test that gcd(a, b) divides a and b
            let g = gcd(a, b);
            (a == 0 && b == 0) || (a % g == 0 && b % g == 0)
        }

        fn commutative(a: u64, b: u64) -> bool {
            gcd(a, b) == gcd(b, a)
        }

        fn associative(a: u64, b: u64, c: u64) -> bool {
            gcd(a, gcd(b, c)) == gcd(gcd(a, b), c)
        }

        fn scalar_mult(a: u64, b: u64, k: u64) -> bool {
            gcd(k * a, k * b) == k * gcd(a, b)
        }

        fn multiplicative(a: u64, b: u64, c: u64) -> bool {
            // gcd(ab, c) = gcd(a, c) gcd(b, c) when a and b coprime
            gcd(a, b) != 1 || gcd(a * b, c) == gcd(a, c) * gcd(b, c)
        }

        fn linearity(a: u64, b: u64, k: u64) -> bool {
            gcd(a + k * b, b) == gcd(a, b)
        }
    }
}
<commit_msg>factor::numeric::gcd: Add explicit test for the 0 case<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2015 Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) 2020 nicoo            <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use std::cmp::min;
use std::mem::swap;

pub fn gcd(mut u: u64, mut v: u64) -> u64 {
    // Stein's binary GCD algorithm
    // Base cases: gcd(n, 0) = gcd(0, n) = n
    if u == 0 {
        return v;
    } else if v == 0 {
        return u;
    }

    // gcd(2ⁱ u, 2ʲ v) = 2ᵏ gcd(u, v) with u, v odd and k = min(i, j)
    // 2ᵏ is the greatest power of two that divides both u and v
    let k = {
        let i = u.trailing_zeros();
        let j = v.trailing_zeros();
        u >>= i;
        v >>= j;
        min(i, j)
    };

    loop {
        // Loop invariant: u and v are odd
        debug_assert!(u % 2 == 1, "u = {} is even", u);
        debug_assert!(v % 2 == 1, "v = {} is even", v);

        // gcd(u, v) = gcd(|u - v|, min(u, v))
        if u > v {
            swap(&mut u, &mut v);
        }
        v -= u;

        if v == 0 {
            // Reached the base case; gcd is 2ᵏ u
            return u << k;
        }

        // gcd(u, 2ʲ v) = gcd(u, v) as u is odd
        v >>= v.trailing_zeros();
    }
}

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

    quickcheck! {
        fn euclidean(a: u64, b: u64) -> bool {
            // Test against the Euclidean algorithm
            let g = {
                let (mut a, mut b) = (a, b);
                while b > 0 {
                    a %= b;
                    swap(&mut a, &mut b);
                }
                a
            };
            gcd(a, b) == g
        }

        fn one(a: u64) -> bool {
            gcd(1, a) == 1
        }

        fn zero(a: u64) -> bool {
            gcd(0, a) == a
        }

        fn divisor(a: u64, b: u64) -> bool {
            // Test that gcd(a, b) divides a and b
            let g = gcd(a, b);
            (a == 0 && b == 0) || (a % g == 0 && b % g == 0)
        }

        fn commutative(a: u64, b: u64) -> bool {
            gcd(a, b) == gcd(b, a)
        }

        fn associative(a: u64, b: u64, c: u64) -> bool {
            gcd(a, gcd(b, c)) == gcd(gcd(a, b), c)
        }

        fn scalar_mult(a: u64, b: u64, k: u64) -> bool {
            gcd(k * a, k * b) == k * gcd(a, b)
        }

        fn multiplicative(a: u64, b: u64, c: u64) -> bool {
            // gcd(ab, c) = gcd(a, c) gcd(b, c) when a and b coprime
            gcd(a, b) != 1 || gcd(a * b, c) == gcd(a, c) * gcd(b, c)
        }

        fn linearity(a: u64, b: u64, k: u64) -> bool {
            gcd(a + k * b, b) == gcd(a, b)
        }
    }
}
<|endoftext|>


<commit_before>extern crate rlisp;

use self::rlisp::Rlisp;

#[test]
fn test_rlisp() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(def {a} 12)"), "()");

    assert_eq!(rlisp.execute("(+ a a)"), "24");

    assert_eq!(rlisp.execute("(- a 4)"), "8");

    assert_eq!(rlisp.execute("(eval (head {(+ 1 2) (+ 10 20)}))"), "3");

    assert_eq!(rlisp.execute("((lambda {a} {+ a 1}) 2)"), "3");

    assert_eq!(rlisp.execute("((lambda {a} {+ a 1}) a)"), "13");

    assert_eq!(rlisp.execute("(def {addgen} (lambda {a} {lambda {} {+ a 1}}))"), "()");
    assert_eq!(rlisp.execute("(def {adder} (addgen 1))"), "()");
    assert_eq!(rlisp.execute("(adder)"), "2");

    assert_eq!(rlisp.execute("(def {defgen} (lambda {a} {lambda {} {def {o} a}}))"), "()");
    assert_eq!(rlisp.execute("(def {defer} (defgen 10))"), "()");
    assert_eq!(rlisp.execute("(defer)"), "()");
    assert_eq!(rlisp.execute("o"), "10");

    assert_eq!(rlisp.execute("((lambda {x y} {+ x y}) 10 20)"), "30");

    assert_eq!(rlisp.execute("(def {add-mul} (lambda {x y} {+ x (* x y)}))"), "()");
    assert_eq!(rlisp.execute("(add-mul 10 20)"), "210");

    assert_eq!(rlisp.execute("(def {add-mul-10} (add-mul 10))"), "()");
    assert_eq!(rlisp.execute("(add-mul-10 50)"), "510");

    assert_eq!(rlisp.execute("(def {f x y} {+ x (/ (+ x x) y)})"), "()");
    assert_eq!(rlisp.execute("(f 10 20)"), "11");

    assert_eq!(rlisp.execute("(def {hof f} {f 12 12})"), "()");
    assert_eq!(rlisp.execute("(hof +)"), "24");
    assert_eq!(rlisp.execute("(hof (lambda {a b} {+ a (/ b 2)}))"), "18");

    assert_eq!(rlisp.execute("(def {elipse ...} {head (list ...)})"), "()");
    assert_eq!(rlisp.execute("(elipse 1 2 3)"), "{1 2 3}");
    assert_eq!(rlisp.execute("(elipse 1 2 3 4 5)"), "{1 2 3 4 5}");
 }

#[test]
fn test_elipsis_empty_arg_list() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(def {a ...} {head (list ...)})"), "()");

    assert_eq!(rlisp.execute("(a)"), "{}");
}

#[test]
fn test_curried_builtin() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(def {a} (* 5))"), "()");

    assert_eq!(rlisp.execute("(a 10)"), "50");
}

#[test]
fn test_no_input() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute(""), "()");
}

#[test]
fn test_utf8_string() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("\"ߚ\""), "\"ߚ\"");
}

#[test]
fn test_unpack() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(unpack + {1 2 3 4 5 6})"), "21");

    assert_eq!(rlisp.execute("(def {add-unpacked} (unpack +))"), "()");

    assert_eq!(rlisp.execute("(add-unpacked {1 2 3 4 5 6})"), "21");
}

#[test]
fn test_pack() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(pack head 5 6 7)"), "5");
}

#[test]
fn test_no_extra_args() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("((* 1))"), "Error: func: (* Integer|Float) got no arguments");

    assert_eq!(rlisp.execute("(def {mul a b} {* a b})"), "()");

    assert_eq!(rlisp.execute("((mul 5))"), "Error: (lambda {b} {* a b}) got no arguments");
}

#[test]
fn test_comparisons() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(> 1 2)"), "#f");
    assert_eq!(rlisp.execute("(> 2 1)"), "#t");
    assert_eq!(rlisp.execute("(> 2 2)"), "#f");

    assert_eq!(rlisp.execute("(< 1 2)"), "#t");
    assert_eq!(rlisp.execute("(< 2 1)"), "#f");
    assert_eq!(rlisp.execute("(< 2 2)"), "#f");

    assert_eq!(rlisp.execute("(>= 1 2)"), "#f");
    assert_eq!(rlisp.execute("(>= 2 1)"), "#t");
    assert_eq!(rlisp.execute("(>= 2 2)"), "#t");

    assert_eq!(rlisp.execute("(<= 1 2)"), "#t");
    assert_eq!(rlisp.execute("(<= 2 1)"), "#f");
    assert_eq!(rlisp.execute("(<= 2 2)"), "#t");

    assert_eq!(rlisp.execute("(and #f #f)"), "#f");
    assert_eq!(rlisp.execute("(and #t #f)"), "#f");
    assert_eq!(rlisp.execute("(and #t #t)"), "#t");
    assert_eq!(rlisp.execute("(and #t #t #f)"), "#f");
    assert_eq!(rlisp.execute("(and #t #t #t)"), "#t");

    assert_eq!(rlisp.execute("(or #f #f)"), "#f");
    assert_eq!(rlisp.execute("(or #t #f)"), "#t");
    assert_eq!(rlisp.execute("(or #t #t)"), "#t");
    assert_eq!(rlisp.execute("(or #t #t #f)"), "#t");
    assert_eq!(rlisp.execute("(or #f #f #f)"), "#f");
}

#[test]
fn test_if() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(if #t 1 2)"), "1");
    assert_eq!(rlisp.execute("(if #f 1 2)"), "2");

    assert_eq!(rlisp.execute("(if (== 1 1) {+ 1 1} {+ 2 2})"), "2");
    assert_eq!(rlisp.execute("(if (== 1 2) {+ 1 1} {+ 2 2})"), "4");
}<commit_msg>Added some more if tests<commit_after>extern crate rlisp;

use self::rlisp::Rlisp;

#[test]
fn test_rlisp() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(def {a} 12)"), "()");

    assert_eq!(rlisp.execute("(+ a a)"), "24");

    assert_eq!(rlisp.execute("(- a 4)"), "8");

    assert_eq!(rlisp.execute("(eval (head {(+ 1 2) (+ 10 20)}))"), "3");

    assert_eq!(rlisp.execute("((lambda {a} {+ a 1}) 2)"), "3");

    assert_eq!(rlisp.execute("((lambda {a} {+ a 1}) a)"), "13");

    assert_eq!(rlisp.execute("(def {addgen} (lambda {a} {lambda {} {+ a 1}}))"), "()");
    assert_eq!(rlisp.execute("(def {adder} (addgen 1))"), "()");
    assert_eq!(rlisp.execute("(adder)"), "2");

    assert_eq!(rlisp.execute("(def {defgen} (lambda {a} {lambda {} {def {o} a}}))"), "()");
    assert_eq!(rlisp.execute("(def {defer} (defgen 10))"), "()");
    assert_eq!(rlisp.execute("(defer)"), "()");
    assert_eq!(rlisp.execute("o"), "10");

    assert_eq!(rlisp.execute("((lambda {x y} {+ x y}) 10 20)"), "30");

    assert_eq!(rlisp.execute("(def {add-mul} (lambda {x y} {+ x (* x y)}))"), "()");
    assert_eq!(rlisp.execute("(add-mul 10 20)"), "210");

    assert_eq!(rlisp.execute("(def {add-mul-10} (add-mul 10))"), "()");
    assert_eq!(rlisp.execute("(add-mul-10 50)"), "510");

    assert_eq!(rlisp.execute("(def {f x y} {+ x (/ (+ x x) y)})"), "()");
    assert_eq!(rlisp.execute("(f 10 20)"), "11");

    assert_eq!(rlisp.execute("(def {hof f} {f 12 12})"), "()");
    assert_eq!(rlisp.execute("(hof +)"), "24");
    assert_eq!(rlisp.execute("(hof (lambda {a b} {+ a (/ b 2)}))"), "18");

    assert_eq!(rlisp.execute("(def {elipse ...} {head (list ...)})"), "()");
    assert_eq!(rlisp.execute("(elipse 1 2 3)"), "{1 2 3}");
    assert_eq!(rlisp.execute("(elipse 1 2 3 4 5)"), "{1 2 3 4 5}");
 }

#[test]
fn test_elipsis_empty_arg_list() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(def {a ...} {head (list ...)})"), "()");

    assert_eq!(rlisp.execute("(a)"), "{}");
}

#[test]
fn test_curried_builtin() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(def {a} (* 5))"), "()");

    assert_eq!(rlisp.execute("(a 10)"), "50");
}

#[test]
fn test_no_input() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute(""), "()");
}

#[test]
fn test_utf8_string() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("\"ߚ\""), "\"ߚ\"");
}

#[test]
fn test_unpack() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(unpack + {1 2 3 4 5 6})"), "21");

    assert_eq!(rlisp.execute("(def {add-unpacked} (unpack +))"), "()");

    assert_eq!(rlisp.execute("(add-unpacked {1 2 3 4 5 6})"), "21");
}

#[test]
fn test_pack() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(pack head 5 6 7)"), "5");
}

#[test]
fn test_no_extra_args() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("((* 1))"), "Error: func: (* Integer|Float) got no arguments");

    assert_eq!(rlisp.execute("(def {mul a b} {* a b})"), "()");

    assert_eq!(rlisp.execute("((mul 5))"), "Error: (lambda {b} {* a b}) got no arguments");
}

#[test]
fn test_comparisons() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(> 1 2)"), "#f");
    assert_eq!(rlisp.execute("(> 2 1)"), "#t");
    assert_eq!(rlisp.execute("(> 2 2)"), "#f");

    assert_eq!(rlisp.execute("(< 1 2)"), "#t");
    assert_eq!(rlisp.execute("(< 2 1)"), "#f");
    assert_eq!(rlisp.execute("(< 2 2)"), "#f");

    assert_eq!(rlisp.execute("(>= 1 2)"), "#f");
    assert_eq!(rlisp.execute("(>= 2 1)"), "#t");
    assert_eq!(rlisp.execute("(>= 2 2)"), "#t");

    assert_eq!(rlisp.execute("(<= 1 2)"), "#t");
    assert_eq!(rlisp.execute("(<= 2 1)"), "#f");
    assert_eq!(rlisp.execute("(<= 2 2)"), "#t");

    assert_eq!(rlisp.execute("(and #f #f)"), "#f");
    assert_eq!(rlisp.execute("(and #t #f)"), "#f");
    assert_eq!(rlisp.execute("(and #t #t)"), "#t");
    assert_eq!(rlisp.execute("(and #t #t #f)"), "#f");
    assert_eq!(rlisp.execute("(and #t #t #t)"), "#t");

    assert_eq!(rlisp.execute("(or #f #f)"), "#f");
    assert_eq!(rlisp.execute("(or #t #f)"), "#t");
    assert_eq!(rlisp.execute("(or #t #t)"), "#t");
    assert_eq!(rlisp.execute("(or #t #t #f)"), "#t");
    assert_eq!(rlisp.execute("(or #f #f #f)"), "#f");

    assert_eq!(rlisp.execute("(== {1 2 3 {4 5}} {1 2 3 {4 5}})"), "#t");
    assert_eq!(rlisp.execute("(== {1 2 3 {4 5}} {1 2 3 {4 6}})"), "#f");

    assert_eq!(rlisp.execute("(!= {1 2 3 {4 5}} {1 2 3 {4 5}})"), "#f");
    assert_eq!(rlisp.execute("(!= {1 2 3 {4 5}} {1 2 3 {4 6}})"), "#t");
}

#[test]
fn test_if() {
    let mut rlisp = Rlisp::new();

    assert_eq!(rlisp.execute("(if #t 1 2)"), "1");
    assert_eq!(rlisp.execute("(if #f 1 2)"), "2");

    assert_eq!(rlisp.execute("(if (== 1 1) {+ 1 1} {+ 2 2})"), "2");
    assert_eq!(rlisp.execute("(if (== 1 2) {+ 1 1} {+ 2 2})"), "4");

    assert_eq!(rlisp.execute("(def {x} 100)"), "()");
    assert_eq!(rlisp.execute("(def {y} 200)"), "()");

    assert_eq!(rlisp.execute("(if (== x y) {+ x y} {- x y})"), "-100");
}<|endoftext|>


<commit_before>// Copyright 2016-2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Cauchy distribution.

use Rng;
use distributions::Distribution;
use std::f64::consts::PI;

/// The Cauchy distribution `Cauchy(median, scale)`.
///
/// This distribution has a density function:
/// `f(x) = 1 / (pi * scale * (1 + ((x - median) / scale)^2))`
///
/// # Example
///
/// ```
/// use rand::distributions::{Cauchy, Distribution};
///
/// let cau = Cauchy::new(2.0, 5.0);
/// let v = cau.sample(&mut rand::thread_rng());
/// println!("{} is from a Cauchy(2, 5) distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Cauchy {
    median: f64,
    scale: f64
}

impl Cauchy {
    /// Construct a new `Cauchy` with the given shape parameters
    /// `median` the peak location and `scale` the scale factor.
    /// Panics if `scale <= 0`.
    pub fn new(median: f64, scale: f64) -> Cauchy {
        assert!(scale > 0.0, "Cauchy::new called with scale factor <= 0");
        Cauchy {
            median,
            scale
        }
    }
}

impl Distribution<f64> for Cauchy {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
        let mut x = rng.gen::<f64>();
        // guard against the extremely unlikely event we get 0 or 1 from the generator
        while x <= 0.0 || x >= 1.0 {
            x = rng.gen::<f64>();
        }
        // get standard cauchy random number
        let comp_dev = (2.0 * PI * x).tan();
        // shift and scale according to parameters
        let result = self.median + self.scale * comp_dev;
        result
    }
}

#[cfg(test)]
mod test {
    use distributions::Distribution;
    use super::Cauchy;

    fn median(mut numbers: &mut [f64]) -> f64 {
        sort(&mut numbers);
        let mid = numbers.len() / 2;
        numbers[mid]
    }

    fn sort(numbers: &mut [f64]) {
        numbers.sort_by(|a, b| a.partial_cmp(b).unwrap());
    }

    #[test]
    fn test_cauchy_median() {
        let cauchy = Cauchy::new(10.0, 5.0);
        let mut rng = ::test::rng(123);
        let mut numbers: [f64; 1000] = [0.0; 1000];
        for i in 0..1000 {
            numbers[i] = cauchy.sample(&mut rng);
        }
        let median = median(&mut numbers);
        println!("Cauchy median: {}", median);
        assert!((median - 10.0).abs() < 0.5); // not 100% certain, but probable enough
    }

    #[test]
    fn test_cauchy_mean() {
        let cauchy = Cauchy::new(10.0, 5.0);
        let mut rng = ::test::rng(123);
        let mut sum = 0.0;
        for _ in 0..1000 {
            sum += cauchy.sample(&mut rng);
        }
        let mean = sum / 1000.0;
        println!("Cauchy mean: {}", mean);
        // for a Cauchy distribution the mean should not converge
        assert!((mean - 10.0).abs() > 0.5); // not 100% certain, but probable enough
    }

    #[test]
    #[should_panic]
    fn test_cauchy_invalid_scale_zero() {
        Cauchy::new(0.0, 0.0);
    }

    #[test]
    #[should_panic]
    fn test_cauchy_invalid_scale_neg() {
        Cauchy::new(0.0, -10.0);
    }
}
<commit_msg>Fixed mistake in the domain getting passed into the tangent function<commit_after>// Copyright 2016-2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Cauchy distribution.

use Rng;
use distributions::Distribution;
use std::f64::consts::PI;

/// The Cauchy distribution `Cauchy(median, scale)`.
///
/// This distribution has a density function:
/// `f(x) = 1 / (pi * scale * (1 + ((x - median) / scale)^2))`
///
/// # Example
///
/// ```
/// use rand::distributions::{Cauchy, Distribution};
///
/// let cau = Cauchy::new(2.0, 5.0);
/// let v = cau.sample(&mut rand::thread_rng());
/// println!("{} is from a Cauchy(2, 5) distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Cauchy {
    median: f64,
    scale: f64
}

impl Cauchy {
    /// Construct a new `Cauchy` with the given shape parameters
    /// `median` the peak location and `scale` the scale factor.
    /// Panics if `scale <= 0`.
    pub fn new(median: f64, scale: f64) -> Cauchy {
        assert!(scale > 0.0, "Cauchy::new called with scale factor <= 0");
        Cauchy {
            median,
            scale
        }
    }
}

impl Distribution<f64> for Cauchy {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
        // sample from [0, 1)
        let mut x: f64 = rng.gen::<f64>();
        // guard against the extremely unlikely case we get the invalid 0.5
        while x == 0.5 {
            x = rng.gen::<f64>();
        }
        // get standard cauchy random number
        let comp_dev = (PI * x).tan();
        // shift and scale according to parameters
        let result = self.median + self.scale * comp_dev;
        result
    }
}

#[cfg(test)]
mod test {
    use distributions::Distribution;
    use super::Cauchy;

    fn median(mut numbers: &mut [f64]) -> f64 {
        sort(&mut numbers);
        let mid = numbers.len() / 2;
        numbers[mid]
    }

    fn sort(numbers: &mut [f64]) {
        numbers.sort_by(|a, b| a.partial_cmp(b).unwrap());
    }

    #[test]
    fn test_cauchy_median() {
        let cauchy = Cauchy::new(10.0, 5.0);
        let mut rng = ::test::rng(123);
        let mut numbers: [f64; 1000] = [0.0; 1000];
        for i in 0..1000 {
            numbers[i] = cauchy.sample(&mut rng);
        }
        let median = median(&mut numbers);
        println!("Cauchy median: {}", median);
        assert!((median - 10.0).abs() < 0.5); // not 100% certain, but probable enough
    }

    #[test]
    fn test_cauchy_mean() {
        let cauchy = Cauchy::new(10.0, 5.0);
        let mut rng = ::test::rng(123);
        let mut sum = 0.0;
        for _ in 0..1000 {
            sum += cauchy.sample(&mut rng);
        }
        let mean = sum / 1000.0;
        println!("Cauchy mean: {}", mean);
        // for a Cauchy distribution the mean should not converge
        assert!((mean - 10.0).abs() > 0.5); // not 100% certain, but probable enough
    }

    #[test]
    #[should_panic]
    fn test_cauchy_invalid_scale_zero() {
        Cauchy::new(0.0, 0.0);
    }

    #[test]
    #[should_panic]
    fn test_cauchy_invalid_scale_neg() {
        Cauchy::new(0.0, -10.0);
    }
}
<|endoftext|>


<commit_before>///Script to calculate checksums
extern crate crypto;

use crypto::digest::Digest;
use crypto::md5::Md5;
use crypto::sha1::Sha1;
use crypto::sha2::{Sha256, Sha512};

use std::env::args as cmd_args;
use std::fs::{read_dir, metadata};
use std::io::{Read, BufReader};

#[inline(always)]
fn usage() {
println!("Usage: hash_sum [options] <input>...

Options:
    --md5       Enables md5 calculation.
    --sha1      Enables sha1 calculation.
    --sha256    Enables sha256 calculation.
    --sha512    Enables sha512 calculation.
");
}

//Buffer size 10kbs
const BUFFER_SIZE: usize = 10000;

macro_rules! is_file {
    ($path:expr) => { std::fs::metadata($path).unwrap().is_file() }
}

macro_rules! input {
    ($msg:expr) => {{
        use std::io::{Read, Write};
        print!($msg);
        std::io::stdout().flush().unwrap();
        std::io::stdin().bytes().next();
    }}
}

struct Checksum(String, Box<Digest>);

impl Checksum {
    #[inline(always)]
    pub fn input(&mut self, slice_content: &[u8]) { self.1.input(slice_content); }
    #[inline(always)]
    pub fn result(&mut self) -> String { format!("{:8}{}", self.0, self.1.result_str()) }
    #[inline(always)]
    pub fn reset(&mut self) { self.1.reset(); }
}

impl PartialEq for Checksum {
    fn eq(&self, right: &Checksum) -> bool {
        self.0 == right.0
    }

    fn ne(&self, right: &Checksum) -> bool {
        self.0 != right.0
    }
}

struct HashSum(Vec<String>, Vec<Checksum>);

impl HashSum {
    #[inline(always)]
    fn default_algos(&mut self) {
        self.1.push(Checksum("MD5:".to_string(), Box::new(Md5::new())));
        self.1.push(Checksum("SHA1:".to_string(), Box::new(Sha1::new())));
        self.1.push(Checksum("SHA256:".to_string(), Box::new(Sha256::new())));
        self.1.push(Checksum("SHA512:".to_string(), Box::new(Sha512::new())));
    }

    #[inline(always)]
    fn reset_algos(&mut self) {
        for algo in self.1.iter_mut() {
            algo.reset();
        }
    }

    pub fn run_from_args() {
        let mut hash_sum = HashSum(Vec::new(), Vec::new());

        for arg in cmd_args().skip(1) {
            if arg.starts_with("-") {
                match arg.as_ref() {
                    "--md5" => hash_sum.1.push(Checksum("MD5:".to_string(), Box::new(Md5::new()) as Box<Digest>)),
                    "--sha1" => hash_sum.1.push(Checksum("SHA1:".to_string(), Box::new(Sha1::new()) as Box<Digest>)),
                    "--sha256" => hash_sum.1.push(Checksum("SHA256:".to_string(), Box::new(Sha256::new()) as Box<Digest>)),
                    "--sha512" => hash_sum.1.push(Checksum("SHA512:".to_string(), Box::new(Sha512::new()) as Box<Digest>)),
                    arg @ _ => println!(">>>Invalid flag {}", arg),
                }
            }
            else { hash_sum.0.push(arg) }
        }

        if hash_sum.1.len() == 0 { hash_sum.default_algos(); }

        hash_sum.run()
    }

    fn run(&mut self) {
        for arg in self.0.clone() {
            if let Ok(arg_meta_data) = metadata(&arg) {
                if arg_meta_data.is_file() {
                    self.calc(&arg);
                }
                else if let Ok(dir_content) = read_dir(&arg) {
                    for entry_path in dir_content.map(|entry| entry.unwrap().path()) {
                        if is_file!(&entry_path) { self.calc(&entry_path.to_str().unwrap()); }
                    }
                }
                else {
                    println!(">>>{}: cannot access content of this dir", &arg);
                    continue;
                }
            }
            else {
                println!(">>>{}: cannot access file/directory", &arg);
                continue;
            }
        }
    }

    fn calc(&mut self, path: &str) {
        self.reset_algos();
        if let Ok(file) = std::fs::File::open(&path) {
            let mut file = BufReader::new(file);
            println!(">>>{}:", &path);
            loop {
                let mut file_content: [u8; BUFFER_SIZE] = [0; BUFFER_SIZE];
                let read_bytes = file.read(&mut file_content).unwrap();
                if read_bytes == 0 { break; }

                //In case if read is unable to fill buffer completely.
                let slice_content = &file_content[..read_bytes];
                for algo in self.1.iter_mut() {
                    algo.input(slice_content);
                }
            }

            for algo in self.1.iter_mut() {
                println!("{}", algo.result());
            }
            println!("=======================================================\n");
        }
        else {
            println!(">>>{}: failed to open", &path);
        }
    }
}

fn main() {
    if cmd_args().len() < 2 { return usage(); }

    HashSum::run_from_args();

    //this is done mostly for convenient drag and drop.
    if cfg!(windows) {
        input!("Press Enter to exit...");
    }
}
<commit_msg>Rust: hash_sum new flags and functionality<commit_after>///Script to calculate checksums
extern crate crypto;

use crypto::digest::Digest;
use crypto::md5::Md5;
use crypto::sha1::Sha1;
use crypto::sha2::{Sha256, Sha512};

use std::env::args as cmd_args;
use std::fs::{read_dir, metadata};
use std::io::{Read, Write, BufReader};

#[inline(always)]
fn usage() {
println!("Usage: hash_sum [-c | -o | -p] [options] <input>...

Options:
    --md5       Enables md5 calculation.
    --sha1      Enables sha1 calculation.
    --sha256    Enables sha256 calculation.
    --sha512    Enables sha512 calculation.

Flags:
    Flags are mutually exclusive.
    -p/--print  Prints checksums to stdout. Default.
    -c/--check  Verifies checksum from files.
    -o/--output Write calculations into files with corresponding extension.
");
}

//Buffer size 10kbs
const BUFFER_SIZE: usize = 10000;

macro_rules! is_file {
    ($path:expr) => { std::fs::metadata($path).unwrap().is_file() }
}

macro_rules! input {
    ($msg:expr) => {{
        use std::io::{Read, Write};
        print!($msg);
        std::io::stdout().flush().unwrap();
        std::io::stdin().bytes().next();
    }}
}

struct Checksum(String, Box<Digest>);

impl Checksum {
    #[inline(always)]
    pub fn input(&mut self, slice_content: &[u8]) { self.1.input(slice_content); }
    #[inline(always)]
    pub fn result(&mut self) -> String { format!("{:8}{}", self.0, self.1.result_str()) }
    #[inline(always)]
    pub fn get_checksum(&mut self) -> String { self.1.result_str() }
    #[inline(always)]
    pub fn reset(&mut self) { self.1.reset(); }
    #[inline(always)]
    pub fn get_file_ext(&self) -> String { self.0[0..self.0.len()-1].chars().map(|elem| elem.to_lowercase().next().unwrap()).collect() }
    #[inline(always)]
    pub fn get_type_string(&self) -> String { format!("{:8}", self.0) }
}

impl PartialEq for Checksum {
    fn eq(&self, right: &Checksum) -> bool {
        self.0 == right.0
    }

    fn ne(&self, right: &Checksum) -> bool {
        self.0 != right.0
    }
}

enum FlagType {
    Print,
    Output,
    Check,
}

struct HashSum(Vec<String>, Vec<Checksum>, FlagType);

impl HashSum {
    #[inline(always)]
    fn default_algos(&mut self) {
        self.1.push(Checksum("MD5:".to_string(), Box::new(Md5::new())));
        self.1.push(Checksum("SHA1:".to_string(), Box::new(Sha1::new())));
        self.1.push(Checksum("SHA256:".to_string(), Box::new(Sha256::new())));
        self.1.push(Checksum("SHA512:".to_string(), Box::new(Sha512::new())));
    }

    #[inline(always)]
    fn reset_algos(&mut self) {
        for algo in self.1.iter_mut() {
            algo.reset();
        }
    }

    pub fn run_from_args() {
        let mut hash_sum = HashSum(Vec::new(), Vec::new(), FlagType::Print);

        for arg in cmd_args().skip(1) {
            if arg.starts_with("-") {
                match arg.as_ref() {
                    "--md5" => hash_sum.1.push(Checksum("MD5:".to_string(), Box::new(Md5::new()) as Box<Digest>)),
                    "--sha1" => hash_sum.1.push(Checksum("SHA1:".to_string(), Box::new(Sha1::new()) as Box<Digest>)),
                    "--sha256" => hash_sum.1.push(Checksum("SHA256:".to_string(), Box::new(Sha256::new()) as Box<Digest>)),
                    "--sha512" => hash_sum.1.push(Checksum("SHA512:".to_string(), Box::new(Sha512::new()) as Box<Digest>)),
                    "-o" | "--output" => hash_sum.2 = FlagType::Output,
                    "-c" | "--check" => hash_sum.2 = FlagType::Check,
                    "-p" | "--print" => hash_sum.2 = FlagType::Print,
                    arg @ _ => println!(">>>Invalid flag {}", arg),
                }
            }
            else { hash_sum.0.push(arg) }
        }

        if hash_sum.1.len() == 0 { hash_sum.default_algos(); }

        if hash_sum.0.len() == 0 {
            println!(">>>No input is given!");
            return;
        }

        hash_sum.run()
    }

    fn run(&mut self) {
        for arg in self.0.clone() {
            if let Ok(arg_meta_data) = metadata(&arg) {
                if arg_meta_data.is_file() {
                    self.calc(&arg);
                }
                else if let Ok(dir_content) = read_dir(&arg) {
                    for entry_path in dir_content.map(|entry| entry.unwrap().path()) {
                        if is_file!(&entry_path) { self.calc(&entry_path.to_str().unwrap()); }
                    }
                }
                else {
                    println!(">>>{}: cannot access content of this dir", &arg);
                    continue;
                }
            }
            else {
                println!(">>>{}: cannot access file/directory", &arg);
                continue;
            }
        }
    }

    fn calc(&mut self, path: &str) {
        self.reset_algos();
        if let Ok(file) = std::fs::File::open(&path) {
            let mut file = BufReader::new(file);
            println!(">>>{}:", &path);
            loop {
                let mut file_content: [u8; BUFFER_SIZE] = [0; BUFFER_SIZE];
                let read_bytes = file.read(&mut file_content).unwrap();
                if read_bytes == 0 { break; }

                //In case if read is unable to fill buffer completely.
                let slice_content = &file_content[..read_bytes];
                for algo in self.1.iter_mut() {
                    algo.input(slice_content);
                }
            }

            for algo in self.1.iter_mut() {
                match self.2 {
                    FlagType::Output => {
                        let file_name = format!("{}.{}", &path, algo.get_file_ext());
                        if let Ok(mut file) = std::fs::File::create(&file_name) {
                            file.write_fmt(format_args!("{}\n", algo.get_checksum())).unwrap();
                            println!("{}{}", algo.get_type_string(), &file_name);
                        }
                        else {
                            println!("{}Unable to create file with checksum!", algo.get_type_string());
                        }
                    },
                    FlagType::Check => {
                        let file_name = format!("{}.{}", &path, algo.get_file_ext());
                        if let Ok(mut file) = std::fs::File::open(&file_name) {
                            let mut expected_checksum = String::new();
                            if file.read_to_string(&mut expected_checksum).is_ok() {
                                if expected_checksum.trim() == algo.get_checksum() {
                                    println!("{}OK", algo.get_type_string());
                                }
                                else {
                                    println!("{}NOT_OK", algo.get_type_string());
                                }
                            }
                            else {
                                println!("{}Failed to get checksum from file!", algo.get_type_string());
                            }
                        }
                        else {
                            println!("{}No checksum file!", algo.get_type_string());
                        }
                    },
                    FlagType::Print => {
                        println!("{}", algo.result());
                    },

                }
            }
            println!("=======================================================\n");
        }
        else {
            println!(">>>{}: failed to open", &path);
        }
    }
}

fn main() {
    if cmd_args().len() < 2 { return usage(); }

    HashSum::run_from_args();

    //this is done mostly for convenient drag and drop.
    if cfg!(windows) {
        input!("Press Enter to exit...");
    }
}
<|endoftext|>