<commit_before>extern crate num_cpus;
extern crate itertools;
extern crate rand;
extern crate crossbeam;

use self::rand::{thread_rng, Rng};
use self::itertools::Itertools;
use self::crossbeam::{scope, ScopedJoinHandle};

use std::ops::Range;
use std::fmt::{Debug, Formatter, Result as FmtResult};
use std::sync::{Mutex, RwLock, MutexGuard};
use std::sync::mpsc::{Sender, Receiver, channel};
use std::thread::spawn;

use task::{TaskGenerator, Task};
use candidate::{WorkingCandidate, Candidate};
use solution::Solution;
use scaling::{ScalingFunction, proportionate};
use result::{Result as AbcResult, Error as AbcError};

/// Manages the parameters of the ABC algorithm.
pub struct HiveBuilder<S: Solution> {
    workers: usize,
    observers: usize,
    retries: usize,
    builder: Mutex<S::Builder>,
    threads: usize,
    scale: Box<ScalingFunction>,
}

impl<S: Solution> HiveBuilder<S> {
    /// Creates a new hive.
    ///
    /// * `builder` - Factory-like state that can be used while generating solutions.
    /// * `workers` - Number of working solution candidates to maintain at a time.
    pub fn new(builder: S::Builder, workers: usize) -> HiveBuilder<S> {
        if workers == 0 {
            panic!("HiveBuilder must have at least one worker.");
        }

        HiveBuilder {
            workers: workers,
            observers: workers,
            retries: workers,

            builder: Mutex::new(builder),
            threads: num_cpus::get(),
            scale: proportionate(),
        }
    }

    /// Sets the number of "bees" that will pick a candidate to work on at random.
    ///
    /// This defaults to the number of workers.
    pub fn set_observers(mut self, observers: usize) -> HiveBuilder<S> {
        self.observers = observers;
        self
    }

    /// Sets the number of times a candidate can go unimproved before being reinitialized.
    ///
    /// This defaults to the number of workers.
    pub fn set_retries(mut self, retries: usize) -> HiveBuilder<S> {
        self.retries = retries;
        self
    }

    /// Sets the number of worker threads to use while running.
    pub fn set_threads(mut self, threads: usize) -> HiveBuilder<S> {
        self.threads = threads;
        self
    }

    /// Sets the scaling function for observers to use.
    pub fn set_scaling(mut self, scale: Box<ScalingFunction>) -> HiveBuilder<S> {
        self.scale = scale;
        self
    }

    /// Activates the `HiveBuilder` to create a runnable object.
    pub fn build(self) -> AbcResult<Hive<S>> {
        Hive::new(self)
    }
}

/// Runs the ABC algorithm, maintaining any necessary state.
pub struct Hive<S: Solution> {
    hive: HiveBuilder<S>,

    working: Vec<RwLock<WorkingCandidate<S>>>,
    best: Mutex<Candidate<S>>,

    tasks: Mutex<Option<TaskGenerator>>,
    streaming: Option<Mutex<Sender<Candidate<S>>>>,
}

impl<S: Solution> Hive<S> {
    fn new(hive: HiveBuilder<S>) -> AbcResult<Hive<S>> {
        let tokens: Mutex<Range<usize>> = Mutex::new(0..hive.workers);
        let candidates = Mutex::new(Vec::with_capacity(hive.workers));
        let mut handles = Vec::with_capacity(hive.threads);

        try!(crossbeam::scope(|scope| {
            for _ in 0..hive.threads {
                handles.push(scope.spawn(|| {
                    while let Some(_) = tokens.lock().unwrap().next() {
                        let mut builder = match hive.builder.lock() {
                            Ok(b) => b,
                            Err(err) => return Err(AbcError::from(err)),
                        };
                        let solution = S::make(&mut builder);
                        drop(builder);
                        let candidate = Candidate::new(solution);
                        try!(candidates.lock()).push(candidate);
                    }
                    Ok(())
                }));
            }

            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
        }));

        let mut candidates = try!(candidates.into_inner());

        let best = {
            let best_candidate = candidates.iter()
                                           .fold1(|best, next| {
                                               if next.fitness > best.fitness {
                                                   next
                                               } else {
                                                   best
                                               }
                                           })
                                           .unwrap();
            Mutex::new(best_candidate.clone())
        };

        let working = candidates.drain(..)
                                .map(|c| RwLock::new(WorkingCandidate::new(c, hive.retries)))
                                .collect::<Vec<_>>();

        Ok(Hive {
            working: working,
            best: best,
            hive: hive,
            tasks: Mutex::new(None),
            streaming: None,
        })
    }

    fn current_working(&self) -> AbcResult<Vec<Candidate<S>>> {
        let mut current_working = Vec::with_capacity(self.working.len());
        for candidate_mutex in &self.working {
            let read_guard = try!(candidate_mutex.read());
            current_working.push(read_guard.candidate.clone())
        }
        Ok(current_working)
    }

    /// Returns a guard for the current best solution found by the hive.
    ///
    /// If the hive is running, you should drop the guard returned by this
    /// function as soon as convenient, since the logic of the hive can block
    /// on the availability of the associated mutex. If you plan on performing
    /// expensive computations, you should `drop` the guard as soon as
    /// possible, or acquire and clone it within a small block.
    pub fn get(&self) -> AbcResult<MutexGuard<Candidate<S>>> {
        self.best.lock().map_err(AbcError::from)
    }

    fn consider_improvement(&self, candidate: &Candidate<S>) -> AbcResult<()> {
        let mut best_guard = try!(self.best.lock());
        if candidate.fitness > best_guard.fitness {
            *best_guard = candidate.clone();
            if let Some(mutex) = self.streaming.as_ref() {
                // We're streaming, so we need to post the improved candidate.
                let sender_guard = try!(mutex.lock());
                // If this errors, the receiver was dropped, so we're done.
                if let Err(_) = sender_guard.send(candidate.clone()) {
                    try!(self.stop());
                }
            }
        }
        Ok(())
    }

    fn work_on(&self, current_working: &[Candidate<S>], n: usize) -> AbcResult<()> {
        let variant = Candidate::new(S::explore(current_working, n));
        let mut write_guard = try!(self.working[n].write());
        if variant.fitness > write_guard.candidate.fitness {
            *write_guard = WorkingCandidate::new(variant, self.hive.retries);
            try!(self.consider_improvement(&write_guard.candidate));
        } else {
            write_guard.deplete();
            // Scouting has been folded into the working process
            if write_guard.expired() {
                let mut builder = try!(self.hive.builder.lock());
                let candidate = Candidate::new(S::make(&mut builder));
                drop(builder);
                *write_guard = WorkingCandidate::new(candidate, self.hive.retries);
                try!(self.consider_improvement(&write_guard.candidate));
            }
        }
        Ok(())
    }

    fn choose(&self, current_working: &[Candidate<S>]) -> usize {
        let fitnesses = (self.hive.scale)(current_working.iter()
                                                         .map(|candidate| candidate.fitness)
                                                         .collect::<Vec<f64>>());

        let running_totals = fitnesses.iter()
                                      .scan(0f64, |total, fitness| {
                                          *total += *fitness;
                                          Some(*total)
                                      })
                                      .collect::<Vec<f64>>();

        // Multiplying the choice point is equivalent to, and more efficient than, normalizing
        // all of the scaled fitnesses and having a choice point in [0,1)
        let total_fitness = running_totals.last().unwrap();
        let choice_point = thread_rng().next_f64() * total_fitness;

        for (i, total) in running_totals.iter().enumerate() {
            if *total > choice_point {
                return i;
            }
        }
        unreachable!();
    }

    fn execute(&self, task: &Task) -> AbcResult<()> {
        let current_working = try!(self.current_working());
        let index = match *task {
            Task::Worker(n) => n,
            Task::Observer(_) => self.choose(&current_working),
        };
        self.work_on(&current_working, index)
    }

    fn run(&self, tasks: TaskGenerator) -> AbcResult<()> {
        let mut guard = try!(self.tasks.lock());
        *guard = Some(tasks);
        drop(guard);

        let mut handles: Vec<ScopedJoinHandle<AbcResult<()>>> = Vec::new();

        scope(|scope| {
            for _ in 0..self.hive.threads {
                handles.push(scope.spawn(|| {
                    loop {
                        let mut guard = try!(self.tasks.lock());
                        let task = guard.as_mut().and_then(|gen| gen.next());
                        drop(guard);

                        match task {
                            Some(t) => try!(self.execute(&t)),
                            None => return Ok(()),
                        };
                    }
                }));
            }

            // Returns `Ok(())` only if all threads join cleanly, and the task
            // cycle is successfully cleared away.
            //
            // We avoid `try!` because we want all of the following logic to
            // execute unconditionally.
            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
                   .and(self.tasks
                            .lock()
                            .map(|mut tasks_guard| *tasks_guard = None)
                            .map_err(AbcError::from))
        })
    }

    /// Runs for a fixed number of rounds, then return the best solution found.
    ///
    /// If one of the worker threads panics while working, this will return
    /// `Err(abc::Error)`. Otherwise, it will return `Ok` with a `Candidate`.
    pub fn run_for_rounds(&self, rounds: usize) -> AbcResult<Candidate<S>> {
        let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers).max_rounds(rounds);
        try!(self.run(tasks));
        self.get().map(|guard| guard.clone())
    }

    /// Runs indefinitely in the background, providing a stream of results.
    ///
    /// This method consumes the hive, which will run until the `HiveBuilder` object
    /// is dropped. It returns an `mpsc::Receiver`, which receives a
    /// `Candidate` each time the hive improves on its best solution.
    pub fn stream(mut self) -> Receiver<Candidate<S>> {
        let (sender, receiver) = channel();
        spawn(move || {
            let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers);
            self.streaming = Some(Mutex::new(sender));
            self.run(tasks)
        });
        receiver
    }

    /// Stops a running hive.
    ///
    /// If a worker thread has panicked, this returns `Err(abc::Error)`.
    pub fn stop(&self) -> AbcResult<()> {
        let mut tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_mut().map_or((), |t| t.stop()))
    }

    /// Returns the current round of a running hive.
    ///
    /// If a worker thread has panicked and poisoned the task generator lock,
    /// `get_round` will return `Err(abc::Error)`.
    ///
    /// If the hive has not been run, `get_round` will return `Ok(None)`.
    pub fn get_round(&self) -> AbcResult<Option<usize>> {
        let tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_ref().map(|tasks| tasks.round))
    }
}

impl<S: Solution + Debug> Debug for Hive<S> {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        for mutex in (&self.working).iter() {
            let working = mutex.read().unwrap();
            try!(write!(f, "..{:?}..\n", working.candidate));
        }
        let best_candidate = self.get().unwrap();
        write!(f, ">>{:?}<<", *best_candidate)
    }
}

impl<S: Solution> Drop for Hive<S> {
    fn drop(&mut self) {
        self.stop().unwrap()
    }
}
<commit_msg>Ignore panics on stream drop<commit_after>extern crate num_cpus;
extern crate itertools;
extern crate rand;
extern crate crossbeam;

use self::rand::{thread_rng, Rng};
use self::itertools::Itertools;
use self::crossbeam::{scope, ScopedJoinHandle};

use std::ops::Range;
use std::fmt::{Debug, Formatter, Result as FmtResult};
use std::sync::{Mutex, RwLock, MutexGuard};
use std::sync::mpsc::{Sender, Receiver, channel};
use std::thread::spawn;

use task::{TaskGenerator, Task};
use candidate::{WorkingCandidate, Candidate};
use solution::Solution;
use scaling::{ScalingFunction, proportionate};
use result::{Result as AbcResult, Error as AbcError};

/// Manages the parameters of the ABC algorithm.
pub struct HiveBuilder<S: Solution> {
    workers: usize,
    observers: usize,
    retries: usize,
    builder: Mutex<S::Builder>,
    threads: usize,
    scale: Box<ScalingFunction>,
}

impl<S: Solution> HiveBuilder<S> {
    /// Creates a new hive.
    ///
    /// * `builder` - Factory-like state that can be used while generating solutions.
    /// * `workers` - Number of working solution candidates to maintain at a time.
    pub fn new(builder: S::Builder, workers: usize) -> HiveBuilder<S> {
        if workers == 0 {
            panic!("HiveBuilder must have at least one worker.");
        }

        HiveBuilder {
            workers: workers,
            observers: workers,
            retries: workers,

            builder: Mutex::new(builder),
            threads: num_cpus::get(),
            scale: proportionate(),
        }
    }

    /// Sets the number of "bees" that will pick a candidate to work on at random.
    ///
    /// This defaults to the number of workers.
    pub fn set_observers(mut self, observers: usize) -> HiveBuilder<S> {
        self.observers = observers;
        self
    }

    /// Sets the number of times a candidate can go unimproved before being reinitialized.
    ///
    /// This defaults to the number of workers.
    pub fn set_retries(mut self, retries: usize) -> HiveBuilder<S> {
        self.retries = retries;
        self
    }

    /// Sets the number of worker threads to use while running.
    pub fn set_threads(mut self, threads: usize) -> HiveBuilder<S> {
        self.threads = threads;
        self
    }

    /// Sets the scaling function for observers to use.
    pub fn set_scaling(mut self, scale: Box<ScalingFunction>) -> HiveBuilder<S> {
        self.scale = scale;
        self
    }

    /// Activates the `HiveBuilder` to create a runnable object.
    pub fn build(self) -> AbcResult<Hive<S>> {
        Hive::new(self)
    }
}

/// Runs the ABC algorithm, maintaining any necessary state.
pub struct Hive<S: Solution> {
    hive: HiveBuilder<S>,

    working: Vec<RwLock<WorkingCandidate<S>>>,
    best: Mutex<Candidate<S>>,

    tasks: Mutex<Option<TaskGenerator>>,
    streaming: Option<Mutex<Sender<Candidate<S>>>>,
}

impl<S: Solution> Hive<S> {
    fn new(hive: HiveBuilder<S>) -> AbcResult<Hive<S>> {
        let tokens: Mutex<Range<usize>> = Mutex::new(0..hive.workers);
        let candidates = Mutex::new(Vec::with_capacity(hive.workers));
        let mut handles = Vec::with_capacity(hive.threads);

        try!(crossbeam::scope(|scope| {
            for _ in 0..hive.threads {
                handles.push(scope.spawn(|| {
                    while let Some(_) = tokens.lock().unwrap().next() {
                        let mut builder = match hive.builder.lock() {
                            Ok(b) => b,
                            Err(err) => return Err(AbcError::from(err)),
                        };
                        let solution = S::make(&mut builder);
                        drop(builder);
                        let candidate = Candidate::new(solution);
                        try!(candidates.lock()).push(candidate);
                    }
                    Ok(())
                }));
            }

            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
        }));

        let mut candidates = try!(candidates.into_inner());

        let best = {
            let best_candidate = candidates.iter()
                                           .fold1(|best, next| {
                                               if next.fitness > best.fitness {
                                                   next
                                               } else {
                                                   best
                                               }
                                           })
                                           .unwrap();
            Mutex::new(best_candidate.clone())
        };

        let working = candidates.drain(..)
                                .map(|c| RwLock::new(WorkingCandidate::new(c, hive.retries)))
                                .collect::<Vec<_>>();

        Ok(Hive {
            working: working,
            best: best,
            hive: hive,
            tasks: Mutex::new(None),
            streaming: None,
        })
    }

    fn current_working(&self) -> AbcResult<Vec<Candidate<S>>> {
        let mut current_working = Vec::with_capacity(self.working.len());
        for candidate_mutex in &self.working {
            let read_guard = try!(candidate_mutex.read());
            current_working.push(read_guard.candidate.clone())
        }
        Ok(current_working)
    }

    /// Returns a guard for the current best solution found by the hive.
    ///
    /// If the hive is running, you should drop the guard returned by this
    /// function as soon as convenient, since the logic of the hive can block
    /// on the availability of the associated mutex. If you plan on performing
    /// expensive computations, you should `drop` the guard as soon as
    /// possible, or acquire and clone it within a small block.
    pub fn get(&self) -> AbcResult<MutexGuard<Candidate<S>>> {
        self.best.lock().map_err(AbcError::from)
    }

    fn consider_improvement(&self, candidate: &Candidate<S>) -> AbcResult<()> {
        let mut best_guard = try!(self.best.lock());
        if candidate.fitness > best_guard.fitness {
            *best_guard = candidate.clone();
            if let Some(mutex) = self.streaming.as_ref() {
                // We're streaming, so we need to post the improved candidate.
                let sender_guard = try!(mutex.lock());
                // If this errors, the receiver was dropped, so we're done.
                if let Err(_) = sender_guard.send(candidate.clone()) {
                    try!(self.stop());
                }
            }
        }
        Ok(())
    }

    fn work_on(&self, current_working: &[Candidate<S>], n: usize) -> AbcResult<()> {
        let variant = Candidate::new(S::explore(current_working, n));
        let mut write_guard = try!(self.working[n].write());
        if variant.fitness > write_guard.candidate.fitness {
            *write_guard = WorkingCandidate::new(variant, self.hive.retries);
            try!(self.consider_improvement(&write_guard.candidate));
        } else {
            write_guard.deplete();
            // Scouting has been folded into the working process
            if write_guard.expired() {
                let mut builder = try!(self.hive.builder.lock());
                let candidate = Candidate::new(S::make(&mut builder));
                drop(builder);
                *write_guard = WorkingCandidate::new(candidate, self.hive.retries);
                try!(self.consider_improvement(&write_guard.candidate));
            }
        }
        Ok(())
    }

    fn choose(&self, current_working: &[Candidate<S>]) -> usize {
        let fitnesses = (self.hive.scale)(current_working.iter()
                                                         .map(|candidate| candidate.fitness)
                                                         .collect::<Vec<f64>>());

        let running_totals = fitnesses.iter()
                                      .scan(0f64, |total, fitness| {
                                          *total += *fitness;
                                          Some(*total)
                                      })
                                      .collect::<Vec<f64>>();

        // Multiplying the choice point is equivalent to, and more efficient than, normalizing
        // all of the scaled fitnesses and having a choice point in [0,1)
        let total_fitness = running_totals.last().unwrap();
        let choice_point = thread_rng().next_f64() * total_fitness;

        for (i, total) in running_totals.iter().enumerate() {
            if *total > choice_point {
                return i;
            }
        }
        unreachable!();
    }

    fn execute(&self, task: &Task) -> AbcResult<()> {
        let current_working = try!(self.current_working());
        let index = match *task {
            Task::Worker(n) => n,
            Task::Observer(_) => self.choose(&current_working),
        };
        self.work_on(&current_working, index)
    }

    fn run(&self, tasks: TaskGenerator) -> AbcResult<()> {
        let mut guard = try!(self.tasks.lock());
        *guard = Some(tasks);
        drop(guard);

        let mut handles: Vec<ScopedJoinHandle<AbcResult<()>>> = Vec::new();

        scope(|scope| {
            for _ in 0..self.hive.threads {
                handles.push(scope.spawn(|| {
                    loop {
                        let mut guard = try!(self.tasks.lock());
                        let task = guard.as_mut().and_then(|gen| gen.next());
                        drop(guard);

                        match task {
                            Some(t) => try!(self.execute(&t)),
                            None => return Ok(()),
                        };
                    }
                }));
            }

            // Returns `Ok(())` only if all threads join cleanly, and the task
            // cycle is successfully cleared away.
            //
            // We avoid `try!` because we want all of the following logic to
            // execute unconditionally.
            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
                   .and(self.tasks
                            .lock()
                            .map(|mut tasks_guard| *tasks_guard = None)
                            .map_err(AbcError::from))
        })
    }

    /// Runs for a fixed number of rounds, then return the best solution found.
    ///
    /// If one of the worker threads panics while working, this will return
    /// `Err(abc::Error)`. Otherwise, it will return `Ok` with a `Candidate`.
    pub fn run_for_rounds(&self, rounds: usize) -> AbcResult<Candidate<S>> {
        let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers).max_rounds(rounds);
        try!(self.run(tasks));
        self.get().map(|guard| guard.clone())
    }

    /// Runs indefinitely in the background, providing a stream of results.
    ///
    /// This method consumes the hive, which will run until the `HiveBuilder` object
    /// is dropped. It returns an `mpsc::Receiver`, which receives a
    /// `Candidate` each time the hive improves on its best solution.
    pub fn stream(mut self) -> Receiver<Candidate<S>> {
        let (sender, receiver) = channel();
        spawn(move || {
            let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers);
            self.streaming = Some(Mutex::new(sender));
            self.run(tasks)
        });
        receiver
    }

    /// Stops a running hive.
    ///
    /// If a worker thread has panicked, this returns `Err(abc::Error)`.
    pub fn stop(&self) -> AbcResult<()> {
        let mut tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_mut().map_or((), |t| t.stop()))
    }

    /// Returns the current round of a running hive.
    ///
    /// If a worker thread has panicked and poisoned the task generator lock,
    /// `get_round` will return `Err(abc::Error)`.
    ///
    /// If the hive has not been run, `get_round` will return `Ok(None)`.
    pub fn get_round(&self) -> AbcResult<Option<usize>> {
        let tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_ref().map(|tasks| tasks.round))
    }
}

impl<S: Solution + Debug> Debug for Hive<S> {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        for mutex in (&self.working).iter() {
            let working = mutex.read().unwrap();
            try!(write!(f, "..{:?}..\n", working.candidate));
        }
        let best_candidate = self.get().unwrap();
        write!(f, ">>{:?}<<", *best_candidate)
    }
}

impl<S: Solution> Drop for Hive<S> {
    fn drop(&mut self) {
        self.stop().unwrap_or(())
    }
}
<|endoftext|>


<commit_before>use std::str::Chars;
use std::iter::Peekable;
use std::iter::Enumerate;

use ::syntax::lexer::{Token, Lexer, Span};
use ::syntax::next_is;

pub struct RustSyntax;

impl Lexer for RustSyntax {
    fn handle_char(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize, idx: usize) -> Option<Span> {
        match ch {
            '#' => {
                let st = idx;
                let mut end = idx;
                if next_is(&mut iter, '!') || next_is(&mut iter, '[') {
                    let mut s = String::from("#");
                    while let Some(&(e, c)) = iter.peek() {
                        if c == '\n' { break }
                        end = e;
                        s.push(iter.next().unwrap().1)
                    }
                    return Some(Span {
                        y_pos: y_pos,
                        start: st,
                        end: end,
                        token: Token::Attribute(s),
                    });
                }
                return Some(Span{y_pos: y_pos, start: idx, end: idx + 1, token: Token::Hash})
            }

            '/' => {
                let st = idx;
                let mut end = idx;
                if next_is(&mut iter, '/') {
                    let mut s = String::from("/");
                    s.push(iter.next().unwrap().1);

                    let mut doc_comment = false;
                    if next_is(&mut iter, '/') || next_is(&mut iter, '!') {
                        doc_comment = true;
                    }

                    while let Some(&(e, c)) = iter.peek() {
                        if c == '\n' { break }
                        end = e;
                        s.push(iter.next().unwrap().1)
                    }

                    return Some(Span {
                        y_pos: y_pos,
                        start: st,
                        end: end,
                        token: if doc_comment{ Token::DocComment(s) } else { Token::SingleLineComment(s) },
                    });
                }
                return Some(Span {
                    y_pos: y_pos,
                    start: st,
                    end: end,
                    token: Token::ForwardSlash,
                });
            }

            '"' => {
                let st = idx;
                let mut end = idx;
                let mut s = String::from("\"");
                while let Some(&(e, c)) = iter.peek() {
                    end = e;
                    s.push(iter.next().unwrap().1);
                    if c == '\\' {
                        if let Some(&(e, c_)) = iter.peek() {
                            // this is handling escaped single quotes...
                            if c_ == '"' {
                                s.push(iter.next().unwrap().1);
                                continue;
                            }
                        }
                    }
                    if c == '"' {
                        break;
                    }
                }
                return Some(Span {
                    y_pos: y_pos,
                    start: st,
                    end: end,
                    token: Token::String(s),
                });
            }

            '\'' => {
                let st = idx;
                let mut end = idx;
                let mut s = String::from("\'");
                while let Some(&(e, c)) = iter.peek() {
                    end = e;
                    s.push(iter.next().unwrap().1);
                    if c == '\\' {
                        if let Some(&(e, c_)) = iter.peek() {
                            // this is handling escaped single quotes...
                            if c_ == '\'' {
                                s.push(iter.next().unwrap().1);
                                continue;
                            }
                        }
                    }
                    if c == '\'' {
                        break;
                    }
                }
                return Some(Span {
                    y_pos: y_pos,
                    start: st,
                    end: end,
                    token: Token::Special(s),
                });
            }

            _ => None,
        }

    }

    fn handle_ident(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize) -> Span {
        let mut ident = String::new();
        ident.push(ch);
        let mut start = 0;

        if let Some(&(x, c)) = iter.peek() {
            start = x - 1;
        }

        while self.is_ident(iter.peek()) {
            ident.push(iter.next().unwrap().1)
        }

        let end = ident.len() - 1;
        let token;

        if next_is(&mut iter, '(') {
            token = Token::FunctionCallDef(ident);
        } else {
            token = Token::Ident(ident);
        }

        Span {
            y_pos: y_pos,
            start: start,
            end: start + end,
            token: token,
        }
    }
}

<commit_msg>Highlight macro calls in rust<commit_after>use std::str::Chars;
use std::iter::Peekable;
use std::iter::Enumerate;

use ::syntax::lexer::{Token, Lexer, Span};
use ::syntax::next_is;

pub struct RustSyntax;

impl Lexer for RustSyntax {
    fn handle_char(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize, idx: usize) -> Option<Span> {
        match ch {
            '#' => {
                let st = idx;
                let mut end = idx;
                if next_is(&mut iter, '!') || next_is(&mut iter, '[') {
                    let mut s = String::from("#");
                    while let Some(&(e, c)) = iter.peek() {
                        if c == '\n' { break }
                        end = e;
                        s.push(iter.next().unwrap().1)
                    }
                    return Some(Span {
                        y_pos: y_pos,
                        start: st,
                        end: end,
                        token: Token::Attribute(s),
                    });
                }
                return Some(Span{y_pos: y_pos, start: idx, end: idx + 1, token: Token::Hash})
            }

            '/' => {
                let st = idx;
                let mut end = idx;
                if next_is(&mut iter, '/') {
                    let mut s = String::from("/");
                    s.push(iter.next().unwrap().1);

                    let mut doc_comment = false;
                    if next_is(&mut iter, '/') || next_is(&mut iter, '!') {
                        doc_comment = true;
                    }

                    while let Some(&(e, c)) = iter.peek() {
                        if c == '\n' { break }
                        end = e;
                        s.push(iter.next().unwrap().1)
                    }

                    return Some(Span {
                        y_pos: y_pos,
                        start: st,
                        end: end,
                        token: if doc_comment{ Token::DocComment(s) } else { Token::SingleLineComment(s) },
                    });
                }
                return Some(Span {
                    y_pos: y_pos,
                    start: st,
                    end: end,
                    token: Token::ForwardSlash,
                });
            }

            '"' => {
                let st = idx;
                let mut end = idx;
                let mut s = String::from("\"");
                while let Some(&(e, c)) = iter.peek() {
                    end = e;
                    s.push(iter.next().unwrap().1);
                    if c == '\\' {
                        if let Some(&(e, c_)) = iter.peek() {
                            // this is handling escaped single quotes...
                            if c_ == '"' {
                                s.push(iter.next().unwrap().1);
                                continue;
                            }
                        }
                    }
                    if c == '"' {
                        break;
                    }
                }
                return Some(Span {
                    y_pos: y_pos,
                    start: st,
                    end: end,
                    token: Token::String(s),
                });
            }

            '\'' => {
                let st = idx;
                let mut end = idx;
                let mut s = String::from("\'");
                while let Some(&(e, c)) = iter.peek() {
                    end = e;
                    s.push(iter.next().unwrap().1);
                    if c == '\\' {
                        if let Some(&(e, c_)) = iter.peek() {
                            // this is handling escaped single quotes...
                            if c_ == '\'' {
                                s.push(iter.next().unwrap().1);
                                continue;
                            }
                        }
                    }
                    if c == '\'' {
                        break;
                    }
                }
                return Some(Span {
                    y_pos: y_pos,
                    start: st,
                    end: end,
                    token: Token::Special(s),
                });
            }

            _ => None,
        }

    }

    fn handle_ident(&self, ch: char, mut iter: &mut Peekable<Enumerate<Chars>>, y_pos: usize) -> Span {
        let mut ident = String::new();
        ident.push(ch);
        let mut start = 0;

        if let Some(&(x, c)) = iter.peek() {
            start = x - 1;
        }

        while self.is_ident(iter.peek()) {
            ident.push(iter.next().unwrap().1)
        }

        let end = ident.len() - 1;
        let token;

        if next_is(&mut iter, '(') {
            // function calls or definitions
            token = Token::FunctionCallDef(ident);
        } else if next_is(&mut iter, '!') {
            // macro calls
            ident.push(iter.next().unwrap().1);
            token = Token::Special(ident);
        } else {
            // regular idents
            token = Token::Ident(ident);
        }

        Span {
            y_pos: y_pos,
            start: start,
            end: start + end,
            token: token,
        }
    }
}

<|endoftext|>


<commit_before>//! Low level and internal http and https implementation.

use std::sync::Arc;
use std::time::Duration;
use std::io::{Read, Write, stderr};
use std::net::TcpStream;

use rustls::{Session, ClientConfig, ClientSession, ProtocolVersion};
use webpki::DNSNameRef;
use webpki_roots;
use log::LogLevel::Info;

use super::request::Request;
use super::results::{CabotResult, CabotError};
use super::dns::Resolver;
use super::constants;

const BUFFER_PAGE_SIZE: usize = 1024;
const RESPONSE_BUFFER_SIZE: usize = 1024;

fn log_request(request: &[u8], verbose: bool) {
    if !log_enabled!(Info) && !verbose {
        return;
    }
    let request: Vec<&[u8]> = constants::SPLIT_HEADERS_RE.splitn(request, 2).collect();
    let headers = String::from_utf8_lossy(request.get(0).unwrap());
    let headers: Vec<&str> = constants::SPLIT_HEADER_RE.split(&headers).collect();
    let bodylen = if request.len() == 2 {
        let body = request.get(1).unwrap();
        body.len()
    } else {
        0
    };
    if log_enabled!(Info) {
        for header in headers {
            info!("> {}", header);
        }
        if bodylen > 0 {
            info!("> [{} bytes]", bodylen);
        }
        info!(">");
    } else if verbose {
        for header in headers {
            writeln!(&mut stderr(), "> {}", header).unwrap();
        }
        if bodylen > 0 {
            writeln!(&mut stderr(), "> [{} bytes]", bodylen).unwrap();
        }
        writeln!(&mut stderr(), ">").unwrap();
    }
}

fn read_buf<T>(client: &mut T, buf: &mut [u8]) -> Vec<u8>
    where T: Read + Sized
{
    let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
    loop {
        match client.read(&mut buf[..]) {
            Ok(count) => {
                if count > 0 {
                    response.extend_from_slice(&buf[0..count]);
                } else {
                    break;
                }
            }
            Err(_) => break, // connection is closed by client
        }
    }
    response
}

fn from_http(request: &Request,
             client: &mut TcpStream,
             out: &mut Write,
             verbose: bool)
             -> CabotResult<()> {

    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose);

    debug!("Sending request...");
    client.write(&raw_request).unwrap();
    let mut buf = [0; BUFFER_PAGE_SIZE];
    let response = read_buf(client, &mut buf);
    out.write_all(response.as_slice()).unwrap();
    Ok(())
}

fn from_https(request: &Request,
              mut client: &mut TcpStream,
              out: &mut Write,
              verbose: bool)
              -> CabotResult<()> {

    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
    let mut buf = [0; BUFFER_PAGE_SIZE];

    let mut config = ClientConfig::new();
    config.root_store.add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
    let rc_config = Arc::new(config);
    let host = DNSNameRef::try_from_ascii_str(request.host());
    if host.is_err() {
        return Err(CabotError::HostnameParseError(request.host().to_owned()))
    }
    let mut tlsclient = ClientSession::new(&rc_config, host.unwrap());
    let mut is_handshaking = true;
    loop {
        while tlsclient.wants_write() {
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }
        if is_handshaking && !tlsclient.is_handshaking() {
            info!("Handshake complete");
            is_handshaking = false;
            let protocol = tlsclient.get_protocol_version();
            match protocol {
                Some(ProtocolVersion::SSLv2) => {
                    info!("Protocol SSL v2 negociated");
                }
                Some(ProtocolVersion::SSLv3) => {
                    info!("Protocol SSL v3 negociated");
                }
                Some(ProtocolVersion::TLSv1_0) => {
                    info!("Protocol TLS v1.0 negociated");
                }
                Some(ProtocolVersion::TLSv1_1) => {
                    info!("Protocol TLS v1.1 negociated");
                }
                Some(ProtocolVersion::TLSv1_2) => {
                    info!("Protocol TLS v1.2 negociated");
                }
                Some(ProtocolVersion::TLSv1_3) => {
                    info!("Protocol TLS v1.3 negociated");
                }
                Some(ProtocolVersion::Unknown(num)) => {
                    info!("Unknown TLS Protocol negociated: {}", num);
                }
                None => {
                    info!("No TLS Protocol negociated");
                }
            }
            log_request(&raw_request, verbose);
            tlsclient.write_all(&raw_request).unwrap();
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }

        if tlsclient.wants_read() {
            let count = tlsclient.read_tls(&mut client);
            if let Err(err) = count {
                error!("{:?}", err);
                return Err(CabotError::IOError(format!("{}", err)));
            }

            let count = count.unwrap();
            debug!("Read {} TLS bytes", count);
            if count == 0 {
                break;
            }

            if let Err(err) = tlsclient.process_new_packets() {
                return Err(CabotError::CertificateError(format!("{}", err)));
            }

            let mut part: Vec<u8> = read_buf(&mut tlsclient, &mut buf);
            response.append(&mut part);
        } else {
            break;
        }
    }
    out.write_all(response.as_slice()).unwrap();
    Ok(())
}

pub fn http_query(request: &Request, mut out: &mut Write, verbose: bool) -> CabotResult<()> {
    debug!("HTTP Query {} {}",
           request.http_method(),
           request.request_uri());

    let resolver = Resolver::new(verbose);
    let authority = request.authority();
    let addr = resolver.get_addr(authority)?;

    info!("Connecting to {}", addr);
    let mut client = TcpStream::connect(addr).unwrap();
    client.set_read_timeout(Some(Duration::new(5, 0))).unwrap();

    match request.scheme() {
        "http" => from_http(request, &mut client, &mut out, verbose)?,
        "https" => from_https(request, &mut client, &mut out, verbose)?,
        _ => {
            return Err(CabotError::SchemeError(format!("Unrecognized scheme {}", request.scheme())))
        }
    };

    out.flush().unwrap();

    Ok(())

}
<commit_msg>Don't panic if the server socket refuse the connection<commit_after>//! Low level and internal http and https implementation.

use std::sync::Arc;
use std::time::Duration;
use std::io::{Read, Write, stderr};
use std::net::TcpStream;

use rustls::{Session, ClientConfig, ClientSession, ProtocolVersion};
use webpki::DNSNameRef;
use webpki_roots;
use log::LogLevel::Info;

use super::request::Request;
use super::results::{CabotResult, CabotError};
use super::dns::Resolver;
use super::constants;

const BUFFER_PAGE_SIZE: usize = 1024;
const RESPONSE_BUFFER_SIZE: usize = 1024;

fn log_request(request: &[u8], verbose: bool) {
    if !log_enabled!(Info) && !verbose {
        return;
    }
    let request: Vec<&[u8]> = constants::SPLIT_HEADERS_RE.splitn(request, 2).collect();
    let headers = String::from_utf8_lossy(request.get(0).unwrap());
    let headers: Vec<&str> = constants::SPLIT_HEADER_RE.split(&headers).collect();
    let bodylen = if request.len() == 2 {
        let body = request.get(1).unwrap();
        body.len()
    } else {
        0
    };
    if log_enabled!(Info) {
        for header in headers {
            info!("> {}", header);
        }
        if bodylen > 0 {
            info!("> [{} bytes]", bodylen);
        }
        info!(">");
    } else if verbose {
        for header in headers {
            writeln!(&mut stderr(), "> {}", header).unwrap();
        }
        if bodylen > 0 {
            writeln!(&mut stderr(), "> [{} bytes]", bodylen).unwrap();
        }
        writeln!(&mut stderr(), ">").unwrap();
    }
}

fn read_buf<T>(client: &mut T, buf: &mut [u8]) -> Vec<u8>
    where T: Read + Sized
{
    let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
    loop {
        match client.read(&mut buf[..]) {
            Ok(count) => {
                if count > 0 {
                    response.extend_from_slice(&buf[0..count]);
                } else {
                    break;
                }
            }
            Err(_) => break, // connection is closed by client
        }
    }
    response
}

fn from_http(request: &Request,
             client: &mut TcpStream,
             out: &mut Write,
             verbose: bool)
             -> CabotResult<()> {

    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose);

    debug!("Sending request...");
    client.write(&raw_request).unwrap();
    let mut buf = [0; BUFFER_PAGE_SIZE];
    let response = read_buf(client, &mut buf);
    out.write_all(response.as_slice()).unwrap();
    Ok(())
}

fn from_https(request: &Request,
              mut client: &mut TcpStream,
              out: &mut Write,
              verbose: bool)
              -> CabotResult<()> {

    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
    let mut buf = [0; BUFFER_PAGE_SIZE];

    let mut config = ClientConfig::new();
    config.root_store.add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
    let rc_config = Arc::new(config);
    let host = DNSNameRef::try_from_ascii_str(request.host());
    if host.is_err() {
        return Err(CabotError::HostnameParseError(request.host().to_owned()))
    }
    let mut tlsclient = ClientSession::new(&rc_config, host.unwrap());
    let mut is_handshaking = true;
    loop {
        while tlsclient.wants_write() {
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }
        if is_handshaking && !tlsclient.is_handshaking() {
            info!("Handshake complete");
            is_handshaking = false;
            let protocol = tlsclient.get_protocol_version();
            match protocol {
                Some(ProtocolVersion::SSLv2) => {
                    info!("Protocol SSL v2 negociated");
                }
                Some(ProtocolVersion::SSLv3) => {
                    info!("Protocol SSL v3 negociated");
                }
                Some(ProtocolVersion::TLSv1_0) => {
                    info!("Protocol TLS v1.0 negociated");
                }
                Some(ProtocolVersion::TLSv1_1) => {
                    info!("Protocol TLS v1.1 negociated");
                }
                Some(ProtocolVersion::TLSv1_2) => {
                    info!("Protocol TLS v1.2 negociated");
                }
                Some(ProtocolVersion::TLSv1_3) => {
                    info!("Protocol TLS v1.3 negociated");
                }
                Some(ProtocolVersion::Unknown(num)) => {
                    info!("Unknown TLS Protocol negociated: {}", num);
                }
                None => {
                    info!("No TLS Protocol negociated");
                }
            }
            log_request(&raw_request, verbose);
            tlsclient.write_all(&raw_request).unwrap();
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }

        if tlsclient.wants_read() {
            let count = tlsclient.read_tls(&mut client);
            if let Err(err) = count {
                error!("{:?}", err);
                return Err(CabotError::IOError(format!("{}", err)));
            }

            let count = count.unwrap();
            debug!("Read {} TLS bytes", count);
            if count == 0 {
                break;
            }

            if let Err(err) = tlsclient.process_new_packets() {
                return Err(CabotError::CertificateError(format!("{}", err)));
            }

            let mut part: Vec<u8> = read_buf(&mut tlsclient, &mut buf);
            response.append(&mut part);
        } else {
            break;
        }
    }
    out.write_all(response.as_slice()).unwrap();
    Ok(())
}

pub fn http_query(request: &Request, mut out: &mut Write, verbose: bool) -> CabotResult<()> {
    debug!("HTTP Query {} {}",
           request.http_method(),
           request.request_uri());

    let resolver = Resolver::new(verbose);
    let authority = request.authority();
    let addr = resolver.get_addr(authority)?;

    info!("Connecting to {}", addr);
    let client = TcpStream::connect(addr);
    if client.is_err() {
        return Err(CabotError::IOError(format!("{}", client.unwrap_err())));
    }
    let mut client = client.unwrap();
    client.set_read_timeout(Some(Duration::new(5, 0))).unwrap();

    match request.scheme() {
        "http" => from_http(request, &mut client, &mut out, verbose)?,
        "https" => from_https(request, &mut client, &mut out, verbose)?,
        _ => {
            return Err(CabotError::SchemeError(format!("Unrecognized scheme {}", request.scheme())))
        }
    };

    out.flush().unwrap();

    Ok(())

}
<|endoftext|>


<commit_before>// TODO: use Impls https://github.com/iron/iron/blob/master/src/iron.rs

extern crate collections;

use collections::dlist::DList;
use collections::Deque;

use std::cmp;
use std::mem;

use std::rc::{mod, Rc};

static CHB_MIN_SIZE:uint = 32u;

// TODO: move to utils
fn blit<T:Clone>(src: &[T], dst: &mut [T], dst_ofs: uint) {
    let len = src.len();
    let sd = dst.slice_mut(dst_ofs, dst_ofs + len);
    if len > sd.len() {
        fail!("blit: source larger than destination");
    }

    let _ = sd.clone_from_slice(src);
}

/// Chained buffer of bytes.
/// Consists of linked list of nodes.
pub struct Chain {
    head: DList<Box<Node>>,
    length: uint
}

impl Chain {
    pub fn new() -> Chain {
        return Chain{
            head: DList::new(),
            length: 0
        }
    }

    pub fn from_foreign(src: Chain) -> Chain {
        let mut ch = Chain::new();
        ch.concat(src);
        ch
    }

    pub fn len(&self) -> uint {
        self.length
    }

    // XXX: maybe DEDUP append/prepend?
    pub fn append_bytes(&mut self, data: &[u8]) {
        let size = data.len();
        // XXX: Damn, https://github.com/rust-lang/rust/issues/6393
        let should_create = match self.head.back() {
            Some(nd) => {
                (nd.room() < size) || !rc::is_unique(&nd.dh)
            }
            None => {
                true
            }
        };
        // We either not the only owner of DH or don't have enough room
        if should_create {
            let nsize = if size < CHB_MIN_SIZE { size << 1 } else { size };
            let node = Node::new(DataHolder::new(nsize));
            self.add_node_tail(node);
        }
        // infailable: added node above
        let node = self.head.back_mut().unwrap();
        // XXX: Damn, https://github.com/rust-lang/rust/issues/6268
        // XXX: we need additional var and scope only to fight borrow checker
        {
            let node_end = node.end;
            // we should be sole owner of data holder inside node here
            let dh = rc::get_mut(&mut node.dh).unwrap();
            dh.copy_data_from(data, node_end);
        }
        node.end += size;
        self.length += size;
    }

    pub fn prepend_bytes(&mut self, data: &[u8]) {
        let size = data.len();
        // XXX: Damn, https://github.com/rust-lang/rust/issues/6393
        let should_create = match self.head.front() {
            Some(nd) => {
                (size > nd.start || !rc::is_unique(&nd.dh))
            }
            None => {
                true
            }
        };
        if should_create {
            let nsize = if size < CHB_MIN_SIZE { size << 1 } else { size };
            let mut node = Node::new(DataHolder::new(nsize)); // Box<Node>
            let r = node.room();
            node.start = r;
            node.end = r;
            self.add_node_head(node);
        }
        // See comments in `append_bytes`
        let node = self.head.front_mut().unwrap();
        {
            let node_start = node.start;
            let dh = rc::get_mut(&mut node.dh).unwrap();
            dh.copy_data_from(data, node_start - size);
        }
        node.start -= size;
        self.length += size;
    }

    pub fn pullup(&mut self, size: uint) -> Option<&[u8]> {
        if size == 0 || size > self.len() {
            return None
        }
        // could not fail, because self.size() > 0 => has node
        if self.head.front().unwrap().size() >= size {
            let node = self.head.front().unwrap();
            return Some(node.data(size));
        }

        let mut newn = Node::new(DataHolder::new(size));
        // XXX: we need this scope to be able to move newn inside our list
        {
            // we just created new data holder, so we have unique ownership
            let mut msize = size;
            while msize > 0 {
                {
                    let node = self.head.front_mut().unwrap();
                    let csize = cmp::min(node.size(), size);
                    // XXX: we need this scope only to beat borrow checker
                    {
                        let node_end = newn.end;
                        let dh = rc::get_mut(&mut newn.dh).unwrap();
                        dh.copy_data_from(node.data(csize), node_end);
                    }
                    newn.end += csize;

                    if node.size() > size {
                        node.start += size;
                        self.length -= size;
                        break
                    }
                }
                // infailable
                let n = self.head.pop_front().unwrap();
                msize -= n.size();
                // XXX: free node?
            }
        }
        self.add_node_head(newn);
        // Now first node.size >= size, so we recurse
        return self.pullup(size)
    }

    pub fn concat(&mut self, src: Chain) {
        self.length += src.length;
        self.head.append(src.head);
        // No need to cleanup `src`, because it has moved and cannot be used
    }

    // XXX: chb_drop; `drop` is the sole method of built-in Drop trait,
    // so use another name
    pub fn reset(&mut self) {
        self.head = DList::new();
        self.length = 0;
    }

    // XXX: deprecated & experimental

    // TODO: maybe we do not need this method?
    // TODO: `concat` better express move semantics
    // XXX: to delete...
    pub fn concat_ptr(&mut self, src: &mut Chain) {
        self.length += src.length;
        let sh = mem::replace(&mut src.head, DList::new());
        self.head.append(sh);
        src.length = 0;
    }

    // XXX: private
    fn add_node_tail(&mut self, node: Box<Node>) {
        self.length += node.size();
        self.head.push(node);
    }

    fn add_node_head(&mut self, node: Box<Node>) {
        self.length += node.size();
        self.head.push_front(node);
    }
}

/// Node of chain buffer.
/// Owned by Chain.
struct Node {
    dh: Rc<DataHolder>, // можно заменить на RC
    start: uint,
    end: uint
}

impl Node {
    fn new(dh: Rc<DataHolder>) -> Box<Node> {
        let n = box Node {
            dh: dh.clone(),
            start: 0,
            end: 0
        };
        // TODO: ref dh ? auto, when using RC
        return n;
    }

    fn size(&self) -> uint {
        self.end - self.start
    }

    fn room(&self) -> uint {
        self.dh.size - self.end
    }

    fn data(&self, size:uint) -> &[u8] {
        self.dh.data.slice(self.start, self.start + size)
    }

}

/// Refcounted data holder
/// TODO: can be shared among different chains
/// TODO: implement other storages: shmem, mmap
struct DataHolder{
    size: uint,
    data: Vec<u8>
}

impl DataHolder {
    fn new(size: uint) -> Rc<DataHolder> {
        let dh = Rc::new(DataHolder {
            size: size,
            data: Vec::from_elem(size, 0)
        });
        return dh;
    }

    fn copy_data_from(&mut self, src: &[u8], dst_offs: uint) {
        blit(src,
             self.data.as_mut_slice(), dst_offs);
    }
}

#[cfg(test)]
mod test {
    use super::Chain;
    use std::rand::{task_rng, Rng};

    #[test]
    fn test_append_bytes_changes_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld";
        let ls = s.len();
        chain.append_bytes(s.as_bytes());
        assert_eq!(chain.len(), ls);
    }

    #[test]
    fn test_prepend_bytes_changes_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.prepend_bytes(s);
        assert_eq!(chain.len(), ls);
    }

    #[test]
    fn test_from_foreign_moves_all_data() {
        let mut orig = Chain::new();
        orig.append_bytes("HelloWorld".as_bytes());
        let sz = orig.len();
        let new = Chain::from_foreign(orig);
        assert_eq!(new.len(), sz);
    }

    #[test]
    fn test_pullup_return_none_on_empty_chain() {
        let mut chain = Chain::new();
        assert!(chain.pullup(1).is_none());
    }

    #[test]
    fn test_pullup_return_what_has_been_appended() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.append_bytes(s);
        let res = chain.pullup(ls);
        assert!(res.is_some());
        assert_eq!(res.unwrap(), s);
    }

    #[test]
    fn test_pullup_does_not_change_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.append_bytes(s);
        let olds = chain.len();
        chain.pullup(ls / 2);
        assert_eq!(chain.len(), olds);
    }

    #[test]
    fn test_pullup_works_on_large_sequences() {
        let mut chain = Chain::new();
        let total = 2048u;
        let mut t = total;
        let one_seq = 128u;
        let mut buf = Vec::new();
        while t > 0 {
            let s:String = task_rng().gen_ascii_chars().take(one_seq).collect();
            let b = s.as_bytes();
            buf = buf.append(b);
            chain.append_bytes(b);
            t -= one_seq;
        }
        let ret = chain.pullup(total);
        assert!(ret.is_some());
        assert_eq!(ret.unwrap(), buf.as_slice());
    }

    #[test]
    fn test_concat_increase_dst_length() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        chain1.append_bytes("HelloWorld".as_bytes());
        let l1 = chain1.len();
        chain2.append_bytes("HelloWorld".as_bytes());
        let l2 = chain2.len();
        chain1.concat(chain2);
        assert_eq!(chain1.len(), l1+l2);
    }

    #[test]
    fn test_concat_appends_content() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let b = "HelloWorld".as_bytes();
        let bl = b.len();
        chain2.append_bytes(b);
        chain1.concat(chain2);
        let res = chain1.pullup(bl);

        assert!(res.is_some());
        assert_eq!(res.unwrap(), b);
    }

    #[test]
    fn test_reset_empties_chain() {
        let mut chain = Chain::new();
        chain.append_bytes("HelloWorld".as_bytes());
        chain.reset();
        assert!(chain.pullup(1).is_none());
        assert_eq!(chain.len(), 0);
    }

    // XXX: do not need to test it for `concat`, because it is done for us
    // XXX: by type-system.
    #[test]
    fn test_concat_ptr_destroy_src() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        chain2.append_bytes("HelloWorld".as_bytes());
        chain1.concat_ptr(&mut chain2);
        assert_eq!(chain2.len(), 0);
    }

}
<commit_msg>[change] do not box Nodes, DList already box its own nodes<commit_after>// TODO: use Impls https://github.com/iron/iron/blob/master/src/iron.rs

extern crate collections;

use collections::dlist::DList;
use collections::Deque;

use std::cmp;
use std::mem;

use std::rc::{mod, Rc};

static CHB_MIN_SIZE:uint = 32u;

// TODO: move to utils
fn blit<T:Clone>(src: &[T], dst: &mut [T], dst_ofs: uint) {
    let len = src.len();
    let sd = dst.slice_mut(dst_ofs, dst_ofs + len);
    if len > sd.len() {
        fail!("blit: source larger than destination");
    }

    let _ = sd.clone_from_slice(src);
}

/// Chained buffer of bytes.
/// Consists of linked list of nodes.
pub struct Chain {
    head: DList<Node>,
    length: uint
}

impl Chain {
    pub fn new() -> Chain {
        return Chain{
            head: DList::new(),
            length: 0
        }
    }

    pub fn from_foreign(src: Chain) -> Chain {
        let mut ch = Chain::new();
        ch.concat(src);
        ch
    }

    pub fn len(&self) -> uint {
        self.length
    }

    // XXX: maybe DEDUP append/prepend?
    pub fn append_bytes(&mut self, data: &[u8]) {
        let size = data.len();
        // XXX: Damn, https://github.com/rust-lang/rust/issues/6393
        let should_create = match self.head.back() {
            Some(nd) => {
                (nd.room() < size) || !rc::is_unique(&nd.dh)
            }
            None => {
                true
            }
        };
        // We either not the only owner of DH or don't have enough room
        if should_create {
            let nsize = if size < CHB_MIN_SIZE { size << 1 } else { size };
            let node = Node::new(DataHolder::new(nsize));
            self.add_node_tail(node);
        }
        // infailable: added node above
        let node = self.head.back_mut().unwrap();
        // XXX: Damn, https://github.com/rust-lang/rust/issues/6268
        // XXX: we need additional var and scope only to fight borrow checker
        {
            let node_end = node.end;
            // we should be sole owner of data holder inside node here
            let dh = rc::get_mut(&mut node.dh).unwrap();
            dh.copy_data_from(data, node_end);
        }
        node.end += size;
        self.length += size;
    }

    pub fn prepend_bytes(&mut self, data: &[u8]) {
        let size = data.len();
        // XXX: Damn, https://github.com/rust-lang/rust/issues/6393
        let should_create = match self.head.front() {
            Some(nd) => {
                (size > nd.start || !rc::is_unique(&nd.dh))
            }
            None => {
                true
            }
        };
        if should_create {
            let nsize = if size < CHB_MIN_SIZE { size << 1 } else { size };
            let mut node = Node::new(DataHolder::new(nsize)); // Box<Node>
            let r = node.room();
            node.start = r;
            node.end = r;
            self.add_node_head(node);
        }
        // See comments in `append_bytes`
        let node = self.head.front_mut().unwrap();
        {
            let node_start = node.start;
            let dh = rc::get_mut(&mut node.dh).unwrap();
            dh.copy_data_from(data, node_start - size);
        }
        node.start -= size;
        self.length += size;
    }

    pub fn pullup(&mut self, size: uint) -> Option<&[u8]> {
        if size == 0 || size > self.len() {
            return None
        }
        // could not fail, because self.size() > 0 => has node
        if self.head.front().unwrap().size() >= size {
            let node = self.head.front().unwrap();
            return Some(node.data(size));
        }

        let mut newn = Node::new(DataHolder::new(size));
        // XXX: we need this scope to be able to move newn inside our list
        {
            // we just created new data holder, so we have unique ownership
            let mut msize = size;
            while msize > 0 {
                {
                    let node = self.head.front_mut().unwrap();
                    let csize = cmp::min(node.size(), size);
                    // XXX: we need this scope only to beat borrow checker
                    {
                        let node_end = newn.end;
                        let dh = rc::get_mut(&mut newn.dh).unwrap();
                        dh.copy_data_from(node.data(csize), node_end);
                    }
                    newn.end += csize;

                    if node.size() > size {
                        node.start += size;
                        self.length -= size;
                        break
                    }
                }
                // infailable
                let n = self.head.pop_front().unwrap();
                msize -= n.size();
                // XXX: free node?
            }
        }
        self.add_node_head(newn);
        // Now first node.size >= size, so we recurse
        return self.pullup(size)
    }

    pub fn concat(&mut self, src: Chain) {
        self.length += src.length;
        self.head.append(src.head);
        // No need to cleanup `src`, because it has moved and cannot be used
    }

    // XXX: chb_drop; `drop` is the sole method of built-in Drop trait,
    // so use another name
    pub fn reset(&mut self) {
        self.head = DList::new();
        self.length = 0;
    }

    // XXX: deprecated & experimental

    // TODO: maybe we do not need this method?
    // TODO: `concat` better express move semantics
    // XXX: to delete...
    pub fn concat_ptr(&mut self, src: &mut Chain) {
        self.length += src.length;
        let sh = mem::replace(&mut src.head, DList::new());
        self.head.append(sh);
        src.length = 0;
    }

    // XXX: private
    fn add_node_tail(&mut self, node: Node) {
        self.length += node.size();
        self.head.push(node);
    }

    fn add_node_head(&mut self, node: Node) {
        self.length += node.size();
        self.head.push_front(node);
    }
}

/// Node of chain buffer.
/// Owned by Chain.
struct Node {
    dh: Rc<DataHolder>, // можно заменить на RC
    start: uint,
    end: uint
}

impl Node {
    fn new(dh: Rc<DataHolder>) -> Node {
        let n = Node {
            dh: dh.clone(),
            start: 0,
            end: 0
        };
        // TODO: ref dh ? auto, when using RC
        return n;
    }

    fn size(&self) -> uint {
        self.end - self.start
    }

    fn room(&self) -> uint {
        self.dh.size - self.end
    }

    fn data(&self, size:uint) -> &[u8] {
        self.dh.data.slice(self.start, self.start + size)
    }

}

/// Refcounted data holder
/// TODO: can be shared among different chains
/// TODO: implement other storages: shmem, mmap
struct DataHolder{
    size: uint,
    data: Vec<u8>
}

impl DataHolder {
    fn new(size: uint) -> Rc<DataHolder> {
        let dh = Rc::new(DataHolder {
            size: size,
            data: Vec::from_elem(size, 0)
        });
        return dh;
    }

    fn copy_data_from(&mut self, src: &[u8], dst_offs: uint) {
        blit(src,
             self.data.as_mut_slice(), dst_offs);
    }
}

#[cfg(test)]
mod test {
    use super::Chain;
    use std::rand::{task_rng, Rng};

    #[test]
    fn test_append_bytes_changes_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld";
        let ls = s.len();
        chain.append_bytes(s.as_bytes());
        assert_eq!(chain.len(), ls);
    }

    #[test]
    fn test_prepend_bytes_changes_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.prepend_bytes(s);
        assert_eq!(chain.len(), ls);
    }

    #[test]
    fn test_from_foreign_moves_all_data() {
        let mut orig = Chain::new();
        orig.append_bytes("HelloWorld".as_bytes());
        let sz = orig.len();
        let new = Chain::from_foreign(orig);
        assert_eq!(new.len(), sz);
    }

    #[test]
    fn test_pullup_return_none_on_empty_chain() {
        let mut chain = Chain::new();
        assert!(chain.pullup(1).is_none());
    }

    #[test]
    fn test_pullup_return_what_has_been_appended() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.append_bytes(s);
        let res = chain.pullup(ls);
        assert!(res.is_some());
        assert_eq!(res.unwrap(), s);
    }

    #[test]
    fn test_pullup_does_not_change_length() {
        let mut chain = Chain::new();
        let s = "HelloWorld".as_bytes();
        let ls = s.len();
        chain.append_bytes(s);
        let olds = chain.len();
        chain.pullup(ls / 2);
        assert_eq!(chain.len(), olds);
    }

    #[test]
    fn test_pullup_works_on_large_sequences() {
        let mut chain = Chain::new();
        let total = 2048u;
        let mut t = total;
        let one_seq = 128u;
        let mut buf = Vec::new();
        while t > 0 {
            let s:String = task_rng().gen_ascii_chars().take(one_seq).collect();
            let b = s.as_bytes();
            buf = buf.append(b);
            chain.append_bytes(b);
            t -= one_seq;
        }
        let ret = chain.pullup(total);
        assert!(ret.is_some());
        assert_eq!(ret.unwrap(), buf.as_slice());
    }

    #[test]
    fn test_concat_increase_dst_length() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        chain1.append_bytes("HelloWorld".as_bytes());
        let l1 = chain1.len();
        chain2.append_bytes("HelloWorld".as_bytes());
        let l2 = chain2.len();
        chain1.concat(chain2);
        assert_eq!(chain1.len(), l1+l2);
    }

    #[test]
    fn test_concat_appends_content() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        let b = "HelloWorld".as_bytes();
        let bl = b.len();
        chain2.append_bytes(b);
        chain1.concat(chain2);
        let res = chain1.pullup(bl);

        assert!(res.is_some());
        assert_eq!(res.unwrap(), b);
    }

    #[test]
    fn test_reset_empties_chain() {
        let mut chain = Chain::new();
        chain.append_bytes("HelloWorld".as_bytes());
        chain.reset();
        assert!(chain.pullup(1).is_none());
        assert_eq!(chain.len(), 0);
    }

    // XXX: do not need to test it for `concat`, because it is done for us
    // XXX: by type-system.
    #[test]
    fn test_concat_ptr_destroy_src() {
        let mut chain1 = Chain::new();
        let mut chain2 = Chain::new();
        chain2.append_bytes("HelloWorld".as_bytes());
        chain1.concat_ptr(&mut chain2);
        assert_eq!(chain2.len(), 0);
    }

}
<|endoftext|>