<commit_before>use luminance::{Dim2, Flat, RGBA32F, Texture};
use serde_json::from_reader;
use std::collections::HashMap;
use std::fs::File;
use std::path::Path;

use resource::{Load, LoadError, Result, ResCache};

/// Time.
pub type Time = f32;

/// A clip is an object that implements a visual sequence.
pub struct Clip<'a, 'b> where 'a: 'b {
  act: Box<Fn(Time) -> &'a Texture<Flat, Dim2, RGBA32F> + 'b>
}

impl<'a, 'b> Clip<'a, 'b> {
  pub fn new<F>(act: F) -> Self where F: 'b + Fn(Time) -> &'a Texture<Flat, Dim2, RGBA32F> {
    Clip {
      act: Box::new(act)
    }
  }
}

/// A cut is an object that slices a `Clip` at an *input time* and *output time*. It is instantiated
/// in a `Track` at a given *instance time*.
#[derive(Clone)]
pub struct Cut<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  pub in_time: Time,
  pub out_time: Time,
  pub inst_time: Time,
  pub clip: &'c Clip<'a, 'b>
}

impl<'a, 'b, 'c> Cut<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  pub fn new(in_time: Time, out_time: Time, inst_time: Time, clip: &'c Clip<'a, 'b>) -> Self {
    assert!(in_time <= out_time);

    Cut {
      in_time: in_time,
      out_time: out_time,
      inst_time: inst_time,
      clip: clip
    }
  }

  /// Duration of the cut.
  pub fn dur(&self) -> Time {
    self.out_time - self.in_time
  }
}

/// A track gathers `Cut`s and its purpose is to be used inside a `Timeline`.
#[derive(Clone)]
pub struct Track<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  cuts: Vec<Cut<'a, 'b, 'c>>
}

impl<'a, 'b, 'c> Track<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  pub fn new() -> Self {
    Track {
      cuts: Vec::new()
    }
  }

  pub fn add_cut(&mut self, cut: Cut<'a, 'b, 'c>) {
    self.cuts.push(cut);
  }
}

impl<'a, 'b, 'c, 'd> From<&'d [Cut<'a, 'b, 'c>]> for Track<'a, 'b, 'c> {
  fn from(cuts: &'d [Cut<'a, 'b, 'c>]) -> Self {
    Track {
      cuts: cuts.to_vec()
    }
  }
}

/// A timeline gathers tracks used to build up the visual aspect of the demo.
#[derive(Clone)]
pub struct Timeline<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  tracks: Vec<Track<'a, 'b, 'c>>,
}

impl<'a, 'b, 'c> Timeline<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  pub fn new() -> Self {
    Timeline {
      tracks: Vec::new(),
    }
  }

  /// Turn a TimelineManifest into a Timeline by providing a mapping between clips’ names and real
  /// clips.
  pub fn from_manifest(manifest: &TimelineManifest, mapping: &HashMap<String, &'c Clip<'a, 'b>>) -> Self {
    let mut timeline = Self::new();

    for track_manifest in &manifest.tracks {
      let mut track = Track::new();

      for cut_manifest in &track_manifest.cuts {
        let in_time = cut_manifest.in_time;
        let out_time = cut_manifest.out_time;
        let inst_time = cut_manifest.inst_time;
        let clip = mapping.get(&cut_manifest.clip).cloned().unwrap();

        track.add_cut(Cut::new(in_time, out_time, inst_time, clip));
      }

      timeline.add_track(track);
    }

    timeline
  }

  pub fn add_track(&mut self, track: Track<'a, 'b, 'c>) {
    self.tracks.push(track);
  }

  // TODO: currently, we play the first track we find; add transition support
  pub fn play(&self, t: Time) -> Option<&'a Texture<Flat, Dim2, RGBA32F>> {
    for track in &self.tracks {
      for cut in &track.cuts {
        if cut.inst_time <= t && t <= cut.inst_time + cut.dur() {
          return Some((cut.clip.act)(t));
        }
      }
    }

    None
  }
}

impl<'a, 'b, 'c, 'd> From<&'d [Track<'a, 'b, 'c>]> for Timeline<'a, 'b, 'c> {
  fn from(tracks: &'d [Track<'a, 'b, 'c>]) -> Self {
    Timeline {
      tracks: tracks.to_vec(),
    }
  }
}

#[derive(Debug, Deserialize, PartialEq, Serialize)]
pub struct TimelineManifest {
  pub tracks: Vec<TrackManifest>
}

impl Load for TimelineManifest {
  type Args = ();

  const TY_STR: &'static str = "edit/timelines";

  fn load<P>(path: P, cache: &mut ResCache, args: Self::Args) -> Result<Self> where P: AsRef<Path> {
    let path = path.as_ref();

    info!("loading timeline: {:?}", path);

    let file = File::open(path).map_err(|e| LoadError::FileNotFound(path.to_path_buf(), format!("{:?}", e)))?;
    from_reader(file).map_err(|e| LoadError::ParseFailed(format!("{:?}", e)))
  }
}

#[derive(Debug, Deserialize, PartialEq, Serialize)]
pub struct TrackManifest {
  pub cuts: Vec<CutManifest>
}

impl Load for TrackManifest {
  type Args = ();

  const TY_STR: &'static str = "edit/tracks";

  fn load<P>(path: P, cache: &mut ResCache, args: Self::Args) -> Result<Self> where P: AsRef<Path> {
    let path = path.as_ref();

    info!("loading track: {:?}", path);

    let file = File::open(path).map_err(|e| LoadError::FileNotFound(path.to_path_buf(), format!("{:?}", e)))?;
    from_reader(file).map_err(|e| LoadError::ParseFailed(format!("{:?}", e)))
  }
}

#[derive(Debug, Deserialize, PartialEq, Serialize)]
pub struct CutManifest {
  pub in_time: Time,
  pub out_time: Time,
  pub inst_time: Time,
  pub clip: String
}

impl Load for CutManifest {
  type Args = ();

  const TY_STR: &'static str = "edit/cuts";

  fn load<P>(path: P, cache: &mut ResCache, args: Self::Args) -> Result<Self> where P: AsRef<Path> {
    let path = path.as_ref();

    info!("loading cut: {:?}", path);

    let file = File::open(path).map_err(|e| LoadError::FileNotFound(path.to_path_buf(), format!("{:?}", e)))?;
    from_reader(file).map_err(|e| LoadError::ParseFailed(format!("{:?}", e)))
  }
}
<commit_msg>Removed unnecessary Load impls.<commit_after>use luminance::{Dim2, Flat, RGBA32F, Texture};
use serde_json::from_reader;
use std::collections::HashMap;
use std::fs::File;
use std::path::Path;

use resource::{Load, LoadError, Result, ResCache};

/// Time.
pub type Time = f32;

/// A clip is an object that implements a visual sequence.
pub struct Clip<'a, 'b> where 'a: 'b {
  act: Box<Fn(Time) -> &'a Texture<Flat, Dim2, RGBA32F> + 'b>
}

impl<'a, 'b> Clip<'a, 'b> {
  pub fn new<F>(act: F) -> Self where F: 'b + Fn(Time) -> &'a Texture<Flat, Dim2, RGBA32F> {
    Clip {
      act: Box::new(act)
    }
  }
}

/// A cut is an object that slices a `Clip` at an *input time* and *output time*. It is instantiated
/// in a `Track` at a given *instance time*.
#[derive(Clone)]
pub struct Cut<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  pub in_time: Time,
  pub out_time: Time,
  pub inst_time: Time,
  pub clip: &'c Clip<'a, 'b>
}

impl<'a, 'b, 'c> Cut<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  pub fn new(in_time: Time, out_time: Time, inst_time: Time, clip: &'c Clip<'a, 'b>) -> Self {
    assert!(in_time <= out_time);

    Cut {
      in_time: in_time,
      out_time: out_time,
      inst_time: inst_time,
      clip: clip
    }
  }

  /// Duration of the cut.
  pub fn dur(&self) -> Time {
    self.out_time - self.in_time
  }
}

/// A track gathers `Cut`s and its purpose is to be used inside a `Timeline`.
#[derive(Clone)]
pub struct Track<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  cuts: Vec<Cut<'a, 'b, 'c>>
}

impl<'a, 'b, 'c> Track<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  pub fn new() -> Self {
    Track {
      cuts: Vec::new()
    }
  }

  pub fn add_cut(&mut self, cut: Cut<'a, 'b, 'c>) {
    self.cuts.push(cut);
  }
}

impl<'a, 'b, 'c, 'd> From<&'d [Cut<'a, 'b, 'c>]> for Track<'a, 'b, 'c> {
  fn from(cuts: &'d [Cut<'a, 'b, 'c>]) -> Self {
    Track {
      cuts: cuts.to_vec()
    }
  }
}

/// A timeline gathers tracks used to build up the visual aspect of the demo.
#[derive(Clone)]
pub struct Timeline<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  tracks: Vec<Track<'a, 'b, 'c>>,
}

impl<'a, 'b, 'c> Timeline<'a, 'b, 'c> where 'a: 'b, 'b: 'c {
  pub fn new() -> Self {
    Timeline {
      tracks: Vec::new(),
    }
  }

  /// Turn a TimelineManifest into a Timeline by providing a mapping between clips’ names and real
  /// clips.
  pub fn from_manifest(manifest: &TimelineManifest, mapping: &HashMap<String, &'c Clip<'a, 'b>>) -> Self {
    let mut timeline = Self::new();

    for track_manifest in &manifest.tracks {
      let mut track = Track::new();

      for cut_manifest in &track_manifest.cuts {
        let in_time = cut_manifest.in_time;
        let out_time = cut_manifest.out_time;
        let inst_time = cut_manifest.inst_time;
        let clip = mapping.get(&cut_manifest.clip).cloned().unwrap();

        track.add_cut(Cut::new(in_time, out_time, inst_time, clip));
      }

      timeline.add_track(track);
    }

    timeline
  }

  pub fn add_track(&mut self, track: Track<'a, 'b, 'c>) {
    self.tracks.push(track);
  }

  // TODO: currently, we play the first track we find; add transition support
  pub fn play(&self, t: Time) -> Option<&'a Texture<Flat, Dim2, RGBA32F>> {
    for track in &self.tracks {
      for cut in &track.cuts {
        if cut.inst_time <= t && t <= cut.inst_time + cut.dur() {
          return Some((cut.clip.act)(t));
        }
      }
    }

    None
  }
}

impl<'a, 'b, 'c, 'd> From<&'d [Track<'a, 'b, 'c>]> for Timeline<'a, 'b, 'c> {
  fn from(tracks: &'d [Track<'a, 'b, 'c>]) -> Self {
    Timeline {
      tracks: tracks.to_vec(),
    }
  }
}

#[derive(Debug, Deserialize, PartialEq, Serialize)]
pub struct TimelineManifest {
  pub tracks: Vec<TrackManifest>
}

impl Load for TimelineManifest {
  type Args = ();

  const TY_STR: &'static str = "edit";

  fn load<P>(path: P, cache: &mut ResCache, args: Self::Args) -> Result<Self> where P: AsRef<Path> {
    let path = path.as_ref();

    info!("loading timeline: {:?}", path);

    let file = File::open(path).map_err(|e| LoadError::FileNotFound(path.to_path_buf(), format!("{:?}", e)))?;
    from_reader(file).map_err(|e| LoadError::ParseFailed(format!("{:?}", e)))
  }
}

#[derive(Debug, Deserialize, PartialEq, Serialize)]
pub struct TrackManifest {
  pub cuts: Vec<CutManifest>
}

#[derive(Debug, Deserialize, PartialEq, Serialize)]
pub struct CutManifest {
  pub in_time: Time,
  pub out_time: Time,
  pub inst_time: Time,
  pub clip: String
}
<|endoftext|>


<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 std::collections::BTreeSet;

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)
    }

    fn new_candidate(&self) -> AbcResult<Candidate<S>> {
        let mut builder = try!(self.builder.lock());
        let solution = S::make(&mut builder);
        drop(builder);
        Ok(Candidate::new(solution))
    }
}

/// 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>>,
    scouting: RwLock<BTreeSet<usize>>,

    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::<ScopedJoinHandle<AbcResult<()>>>::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 candidate = try!(hive.new_candidate());
                        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 {
            hive: hive,
            working: working,
            best: best,
            scouting: RwLock::new(BTreeSet::new()),
            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 scouting_guard = try!(self.scouting.write());
                scouting_guard.insert(n);
                drop(scouting_guard);

                let candidate = try!(self.hive.new_candidate());
                *write_guard = WorkingCandidate::new(candidate, self.hive.retries);
                try!(self.consider_improvement(&write_guard.candidate));
                drop(write_guard);

                let mut scouting_guard = try!(self.scouting.write());
                scouting_guard.remove(&n);
            }
        }
        Ok(())
    }

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

        // Avoid observing candidates that are being scouted.
        let scouting_guard = try!(self.scouting.read());
        let running_totals = fitnesses.iter()
                                      .enumerate()
                                      .filter(|&(ref i, _)| !scouting_guard.contains(i))
                                      .scan(0f64, |total, (i, fitness)| {
                                          *total += *fitness;
                                          Some((i, *total))
                                      })
                                      .collect::<Vec<(usize, f64)>>();
        drop(scouting_guard);

        // 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)
        match running_totals.last() {
            Some(&(_, total_fitness)) => {
                let choice_point = thread_rng().next_f64() * total_fitness;
                for &(i, total) in running_totals.iter() {
                    if total > choice_point {
                        return Ok(i);
                    }
                }
                unreachable!();
            }

            // If we are currently scouting all of the solutions, pick one at random.
            None => Ok(thread_rng().gen_range::<usize>(0, fitnesses.len()))
        }
    }

    fn execute(&self, task: &Task) -> AbcResult<()> {
        let current_working = try!(self.current_working());
        let index = match *task {
            Task::Worker(n) => n,
            Task::Observer(_) => try!(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)`.
    ///
    /// If the hive is running, this will return `Ok(Some(n))`. `n` will start
    /// at 0, and increment each time every task in the round has been claimed
    /// (though not necessarily completed) by a worker thread.
    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(())
    }
}
<commit_msg>Workers skip while their candidate is scouted<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 std::collections::BTreeSet;

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)
    }

    fn new_candidate(&self) -> AbcResult<Candidate<S>> {
        let mut builder = try!(self.builder.lock());
        let solution = S::make(&mut builder);
        drop(builder);
        Ok(Candidate::new(solution))
    }
}

/// 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>>,
    scouting: RwLock<BTreeSet<usize>>,

    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::<ScopedJoinHandle<AbcResult<()>>>::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 candidate = try!(hive.new_candidate());
                        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 {
            hive: hive,
            working: working,
            best: best,
            scouting: RwLock::new(BTreeSet::new()),
            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 scouting_guard = try!(self.scouting.write());
                scouting_guard.insert(n);
                drop(scouting_guard);

                let candidate = try!(self.hive.new_candidate());
                *write_guard = WorkingCandidate::new(candidate, self.hive.retries);
                try!(self.consider_improvement(&write_guard.candidate));
                drop(write_guard);

                let mut scouting_guard = try!(self.scouting.write());
                scouting_guard.remove(&n);
            }
        }
        Ok(())
    }

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

        // Avoid observing candidates that are being scouted.
        let scouting_guard = try!(self.scouting.read());
        let running_totals = fitnesses.iter()
                                      .enumerate()
                                      .filter(|&(ref i, _)| !scouting_guard.contains(i))
                                      .scan(0f64, |total, (i, fitness)| {
                                          *total += *fitness;
                                          Some((i, *total))
                                      })
                                      .collect::<Vec<(usize, f64)>>();
        drop(scouting_guard);

        // 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)
        match running_totals.last() {
            Some(&(_, total_fitness)) => {
                let choice_point = thread_rng().next_f64() * total_fitness;
                for &(i, total) in running_totals.iter() {
                    if total > choice_point {
                        return Ok(i);
                    }
                }
                unreachable!();
            }

            // If we are currently scouting all of the solutions, pick one at random.
            None => Ok(thread_rng().gen_range::<usize>(0, fitnesses.len()))
        }
    }

    fn execute(&self, task: &Task) -> AbcResult<()> {
        let current_working = try!(self.current_working());
        let index = match *task {
            Task::Worker(n) => {
                // If the worker's candidate is in the middle of being replaced, just skip it.
                let scouting_guard = try!(self.scouting.read());
                if scouting_guard.contains(&n) { return Ok(()) }
                n
            }
            Task::Observer(_) => try!(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)`.
    ///
    /// If the hive is running, this will return `Ok(Some(n))`. `n` will start
    /// at 0, and increment each time every task in the round has been claimed
    /// (though not necessarily completed) by a worker thread.
    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>// This file is part of Grust, GObject introspection bindings for Rust
//
// Copyright (C) 2013, 2014  Mikhail Zabaluev <mikhail.zabaluev@gmail.com>
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

use ffi;
use types::{gchar,gpointer};
use util::is_false;

use libc;
use std::error::Error;
use std::fmt;
use std::mem;
use std::ptr;
use std::slice;
use std::str;

const NUL: u8 = 0;

pub struct GStr {
    ptr: *const gchar,
}

pub unsafe fn parse_as_bytes(raw: & *const gchar) -> &[u8] {
    assert!(raw.is_not_null());
    let r = mem::copy_lifetime(raw, &(*raw as *const u8));
    slice::from_raw_buf(r, libc::strlen(*raw) as uint)
}

#[inline]
pub unsafe fn parse_as_utf8(raw: & *const gchar)
                           -> Result<&str, str::Utf8Error>
{
    str::from_utf8(parse_as_bytes(raw))
}

#[inline]
pub unsafe fn parse_as_utf8_unchecked(raw: & *const gchar) -> &str {
    str::from_utf8_unchecked(parse_as_bytes(raw))
}

impl GStr {

    pub unsafe fn from_raw_buf(ptr: *mut gchar) -> GStr {
        assert!(ptr.is_not_null());
        GStr { ptr: ptr as *const gchar }
    }

    pub fn parse_as_bytes<'a>(&'a self) -> &'a [u8] {
        unsafe {
            let r = mem::copy_lifetime(self, &(self.ptr as *const u8));
            slice::from_raw_buf(r, libc::strlen(self.ptr) as uint)
        }
    }

    #[inline]
    pub fn parse_as_utf8<'a>(&'a self) -> Result<&'a str, str::Utf8Error> {
        str::from_utf8(self.parse_as_bytes())
    }

    #[inline]
    pub unsafe fn parse_as_utf8_unchecked<'a>(&'a self) -> &'a str {
        str::from_utf8_unchecked(self.parse_as_bytes())
    }
}

impl Drop for GStr {
    fn drop(&mut self) {
        unsafe { ffi::g_free(self.ptr as gpointer) }
    }
}

impl Clone for GStr {
    fn clone(&self) -> GStr {
        unsafe {
            GStr::from_raw_buf(ffi::g_strdup(self.ptr))
        }
    }
}

impl PartialEq for GStr {
    fn eq(&self, other: &GStr) -> bool {
        unsafe { libc::strcmp(self.ptr, other.ptr) == 0 }
    }
}

impl Eq for GStr { }

#[deriving(Copy)]
pub enum StrDataError {
    ContainsNul(uint),
    InvalidUtf8(uint)
}

impl Error for StrDataError {

    fn description(&self) -> &str {
        match *self {
            StrDataError::ContainsNul(_)
                => "invalid data for C string: contains a zero byte",
            StrDataError::InvalidUtf8(_)
                => "invalid UTF-8 sequence",
        }
    }

    fn detail(&self) -> Option<String> {
        match *self {
            StrDataError::ContainsNul(pos)
                => Some(format!("NUL at position {}", pos)),
            StrDataError::InvalidUtf8(pos)
                => Some(format!("Invalid sequence at {}", pos)),
        }
    }
}

impl fmt::Show for StrDataError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            StrDataError::ContainsNul(pos)
                => write!(f, "invalid data for C string: NUL at position {}", pos),
            StrDataError::InvalidUtf8(pos)
                => write!(f, "invalid UTF-8 sequence at position {}", pos),
        }
    }
}

enum GStrData {
    Static(&'static [u8]),
    Owned(Vec<u8>),
    GLib(GStr)
}

impl GStrData {

    fn as_ptr(&self) -> *const gchar {
        match *self {
            GStrData::Static(s) => s.as_ptr() as *const gchar,
            GStrData::Owned(ref v) => v.as_ptr() as *const gchar,
            GStrData::GLib(ref g) => g.ptr
        }
    }

    fn from_static_str(s: &'static str) -> GStrData {
        if !s.ends_with("\0") {
            panic!("static string is not null-terminated: \"{}\"", s);
        }
        GStrData::Static(s.as_bytes())
    }

    fn from_static_bytes(bytes: &'static [u8]) -> GStrData {
        assert!(!bytes.is_empty());
        if bytes[bytes.len() - 1] != 0 {
            panic!("static byte string is not null-terminated: \"{}\"", bytes);
        }
        GStrData::Static(bytes)
    }
}

fn vec_into_g_str_data(mut v: Vec<u8>) -> GStrData {
    v.push(NUL);
    GStrData::Owned(v)
}

macro_rules! g_str_data_from_bytes(
    ($inp:expr) => {
        {
            let bytes = $inp;
            if let Some(pos) = bytes.position_elem(&NUL) {
                return Err(StrDataError::ContainsNul(pos));
            }
            vec_into_g_str_data(bytes.to_vec())
        }
    }
);

pub struct GStrArg {
    data: GStrData
}

impl GStrArg {

    #[inline]
    pub fn from_str(s: &str) -> Result<GStrArg, StrDataError> {
        GStrArg::from_bytes(s.as_bytes())
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<GStrArg, StrDataError> {
        Ok(GStrArg { data: g_str_data_from_bytes!(bytes) })
    }

    pub fn from_static_str(s: &'static str) -> GStrArg {
        GStrArg { data: GStrData::from_static_str(s) }
    }

    pub fn from_static_bytes(bytes: &'static [u8]) -> GStrArg {
        GStrArg { data: GStrData::from_static_bytes(bytes) }
    }

    pub fn as_ptr(&self) -> *const gchar {
        self.data.as_ptr()
    }
}

pub struct Utf8Arg {
    data: GStrData
}

impl Utf8Arg {

    pub fn from_str(s: &str) -> Result<Utf8Arg, StrDataError> {
        Ok(Utf8Arg { data: g_str_data_from_bytes!(s.as_bytes()) })
    }

    pub fn from_static_str(s: &'static str) -> Utf8Arg {
        Utf8Arg { data: GStrData::from_static_str(s) }
    }

    #[inline]
    pub fn as_ptr(&self) -> *const gchar {
        self.data.as_ptr()
    }
}

pub trait IntoUtf8 {

    fn into_utf8(self) -> Result<Utf8Arg, StrDataError>;

    unsafe fn into_utf8_unchecked(self) -> Utf8Arg;
}

fn vec_into_utf8(v: Vec<u8>) -> Utf8Arg {
    Utf8Arg { data: vec_into_g_str_data(v) }
}

impl<'a> IntoUtf8 for &'a str {

    #[inline]
    fn into_utf8(self) -> Result<Utf8Arg, StrDataError> {
        Utf8Arg::from_str(self)
    }

    unsafe fn into_utf8_unchecked(self) -> Utf8Arg {
        vec_into_utf8(self.as_bytes().to_vec())
    }
}

impl IntoUtf8 for String {

    fn into_utf8(self) -> Result<Utf8Arg, StrDataError> {
        Ok(Utf8Arg { data: g_str_data_from_bytes!(self.into_bytes()) })
    }

    unsafe fn into_utf8_unchecked(self) -> Utf8Arg {
        vec_into_utf8(self.into_bytes())
    }
}

impl IntoUtf8 for GStr {

    fn into_utf8(self) -> Result<Utf8Arg, StrDataError> {
        let mut end: *const gchar = ptr::null_mut();
        let valid = unsafe { ffi::g_utf8_validate(self.ptr, -1, &mut end) };
        if is_false(valid) {
            let pos = end as uint - self.ptr as uint;
            return Err(StrDataError::InvalidUtf8(pos));
        }
        Ok(Utf8Arg { data: GStrData::GLib(self) })
    }

    unsafe fn into_utf8_unchecked(self) -> Utf8Arg {
        Utf8Arg { data: GStrData::GLib(self) }
    }
}
<commit_msg>gstr: worked {GStr,Utf8}Arg::from_static_* panics into assert!() macro<commit_after>// This file is part of Grust, GObject introspection bindings for Rust
//
// Copyright (C) 2013, 2014  Mikhail Zabaluev <mikhail.zabaluev@gmail.com>
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

use ffi;
use types::{gchar,gpointer};
use util::is_false;

use libc;
use std::error::Error;
use std::fmt;
use std::mem;
use std::ptr;
use std::slice;
use std::str;

const NUL: u8 = 0;

pub struct GStr {
    ptr: *const gchar,
}

pub unsafe fn parse_as_bytes(raw: & *const gchar) -> &[u8] {
    assert!(raw.is_not_null());
    let r = mem::copy_lifetime(raw, &(*raw as *const u8));
    slice::from_raw_buf(r, libc::strlen(*raw) as uint)
}

#[inline]
pub unsafe fn parse_as_utf8(raw: & *const gchar)
                           -> Result<&str, str::Utf8Error>
{
    str::from_utf8(parse_as_bytes(raw))
}

#[inline]
pub unsafe fn parse_as_utf8_unchecked(raw: & *const gchar) -> &str {
    str::from_utf8_unchecked(parse_as_bytes(raw))
}

impl GStr {

    pub unsafe fn from_raw_buf(ptr: *mut gchar) -> GStr {
        assert!(ptr.is_not_null());
        GStr { ptr: ptr as *const gchar }
    }

    pub fn parse_as_bytes<'a>(&'a self) -> &'a [u8] {
        unsafe {
            let r = mem::copy_lifetime(self, &(self.ptr as *const u8));
            slice::from_raw_buf(r, libc::strlen(self.ptr) as uint)
        }
    }

    #[inline]
    pub fn parse_as_utf8<'a>(&'a self) -> Result<&'a str, str::Utf8Error> {
        str::from_utf8(self.parse_as_bytes())
    }

    #[inline]
    pub unsafe fn parse_as_utf8_unchecked<'a>(&'a self) -> &'a str {
        str::from_utf8_unchecked(self.parse_as_bytes())
    }
}

impl Drop for GStr {
    fn drop(&mut self) {
        unsafe { ffi::g_free(self.ptr as gpointer) }
    }
}

impl Clone for GStr {
    fn clone(&self) -> GStr {
        unsafe {
            GStr::from_raw_buf(ffi::g_strdup(self.ptr))
        }
    }
}

impl PartialEq for GStr {
    fn eq(&self, other: &GStr) -> bool {
        unsafe { libc::strcmp(self.ptr, other.ptr) == 0 }
    }
}

impl Eq for GStr { }

#[deriving(Copy)]
pub enum StrDataError {
    ContainsNul(uint),
    InvalidUtf8(uint)
}

impl Error for StrDataError {

    fn description(&self) -> &str {
        match *self {
            StrDataError::ContainsNul(_)
                => "invalid data for C string: contains a zero byte",
            StrDataError::InvalidUtf8(_)
                => "invalid UTF-8 sequence",
        }
    }

    fn detail(&self) -> Option<String> {
        match *self {
            StrDataError::ContainsNul(pos)
                => Some(format!("NUL at position {}", pos)),
            StrDataError::InvalidUtf8(pos)
                => Some(format!("Invalid sequence at {}", pos)),
        }
    }
}

impl fmt::Show for StrDataError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            StrDataError::ContainsNul(pos)
                => write!(f, "invalid data for C string: NUL at position {}", pos),
            StrDataError::InvalidUtf8(pos)
                => write!(f, "invalid UTF-8 sequence at position {}", pos),
        }
    }
}

enum GStrData {
    Static(&'static [u8]),
    Owned(Vec<u8>),
    GLib(GStr)
}

impl GStrData {

    fn as_ptr(&self) -> *const gchar {
        match *self {
            GStrData::Static(s) => s.as_ptr() as *const gchar,
            GStrData::Owned(ref v) => v.as_ptr() as *const gchar,
            GStrData::GLib(ref g) => g.ptr
        }
    }

    fn from_static_str(s: &'static str) -> GStrData {
        assert!(s.ends_with("\0"),
                "static string is not null-terminated: \"{}\"", s);
        GStrData::Static(s.as_bytes())
    }

    fn from_static_bytes(bytes: &'static [u8]) -> GStrData {
        assert!(!bytes.is_empty());
        assert!(bytes[bytes.len() - 1] == 0,
                "static byte string is not null-terminated: {}", bytes);
        GStrData::Static(bytes)
    }
}

fn vec_into_g_str_data(mut v: Vec<u8>) -> GStrData {
    v.push(NUL);
    GStrData::Owned(v)
}

macro_rules! g_str_data_from_bytes(
    ($inp:expr) => {
        {
            let bytes = $inp;
            if let Some(pos) = bytes.position_elem(&NUL) {
                return Err(StrDataError::ContainsNul(pos));
            }
            vec_into_g_str_data(bytes.to_vec())
        }
    }
);

pub struct GStrArg {
    data: GStrData
}

impl GStrArg {

    #[inline]
    pub fn from_str(s: &str) -> Result<GStrArg, StrDataError> {
        GStrArg::from_bytes(s.as_bytes())
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<GStrArg, StrDataError> {
        Ok(GStrArg { data: g_str_data_from_bytes!(bytes) })
    }

    pub fn from_static_str(s: &'static str) -> GStrArg {
        GStrArg { data: GStrData::from_static_str(s) }
    }

    pub fn from_static_bytes(bytes: &'static [u8]) -> GStrArg {
        GStrArg { data: GStrData::from_static_bytes(bytes) }
    }

    pub fn as_ptr(&self) -> *const gchar {
        self.data.as_ptr()
    }
}

pub struct Utf8Arg {
    data: GStrData
}

impl Utf8Arg {

    pub fn from_str(s: &str) -> Result<Utf8Arg, StrDataError> {
        Ok(Utf8Arg { data: g_str_data_from_bytes!(s.as_bytes()) })
    }

    pub fn from_static_str(s: &'static str) -> Utf8Arg {
        Utf8Arg { data: GStrData::from_static_str(s) }
    }

    #[inline]
    pub fn as_ptr(&self) -> *const gchar {
        self.data.as_ptr()
    }
}

pub trait IntoUtf8 {

    fn into_utf8(self) -> Result<Utf8Arg, StrDataError>;

    unsafe fn into_utf8_unchecked(self) -> Utf8Arg;
}

fn vec_into_utf8(v: Vec<u8>) -> Utf8Arg {
    Utf8Arg { data: vec_into_g_str_data(v) }
}

impl<'a> IntoUtf8 for &'a str {

    #[inline]
    fn into_utf8(self) -> Result<Utf8Arg, StrDataError> {
        Utf8Arg::from_str(self)
    }

    unsafe fn into_utf8_unchecked(self) -> Utf8Arg {
        vec_into_utf8(self.as_bytes().to_vec())
    }
}

impl IntoUtf8 for String {

    fn into_utf8(self) -> Result<Utf8Arg, StrDataError> {
        Ok(Utf8Arg { data: g_str_data_from_bytes!(self.into_bytes()) })
    }

    unsafe fn into_utf8_unchecked(self) -> Utf8Arg {
        vec_into_utf8(self.into_bytes())
    }
}

impl IntoUtf8 for GStr {

    fn into_utf8(self) -> Result<Utf8Arg, StrDataError> {
        let mut end: *const gchar = ptr::null_mut();
        let valid = unsafe { ffi::g_utf8_validate(self.ptr, -1, &mut end) };
        if is_false(valid) {
            let pos = end as uint - self.ptr as uint;
            return Err(StrDataError::InvalidUtf8(pos));
        }
        Ok(Utf8Arg { data: GStrData::GLib(self) })
    }

    unsafe fn into_utf8_unchecked(self) -> Utf8Arg {
        Utf8Arg { data: GStrData::GLib(self) }
    }
}
<|endoftext|>


<commit_before>use parser::{parse};

#[derive(Clone, Debug, PartialEq)]
pub enum UnOp {
  Not,
  Neg,
}

#[derive(Clone, Debug, PartialEq)]
pub enum BinOp {
  Plus,
  Minus,
  Times,
  Div,
  Eq,
  Ne,
  Leq,
  Geq,
  Lt,
  Gt,
  And,
  Or,
  Mod,
  Seq,
}

#[derive(Clone, Debug, PartialEq)]
pub enum Expr {
  Int(isize),
  Bool(bool),
  Var(String),
  BinOp(BinOp, Box<Expr>, Box<Expr>),
  UnOp(UnOp, Box<Expr>),
  Ternary(Box<Expr>, Box<Expr>, Box<Expr>),
  Let(Box<Expr>, Box<Expr>, Box<Expr>),
  Func(Option<Box<Expr>>, Box<Expr>, Vec<Expr>),
  FnCall(Box<Expr>, Vec<Expr>),
}

fn is_int(e: &Expr) -> bool {
  match *e {
    Expr::Int(_) => true,
    _ => false,
  }
}

fn is_bool(e: &Expr) -> bool {
  match *e {
    Expr::Bool(_) => true,
    _ => false,
  }
}

fn is_func(e: &Expr) -> bool {
  match *e {
    Expr::Func(_, _, _) => true,
    _ => false,
  }
}

fn is_value(e: &Expr) -> bool {
  match *e {
    Expr::Int(_) | Expr::Bool(_) | Expr::Func(_, _, _) => true,
    _ => false,
  }
}

fn to_int(e: &Expr) -> isize {
  match *e {
    Expr::Int(n) => n,
    _ => {
      debug!("cant turn into int: {:?}", e);
      panic!()
    }
  }
}

fn to_bool(e: &Expr) -> bool {
  match *e {
    Expr::Bool(b) => b,
    _ => {
      debug!("cant turn into bool: {:?}", e);
      panic!()
    }
  }
}

fn subst(e: Expr, x: Expr, v: Expr) -> Expr {
  let sub = |e1: Expr| subst(e1.clone(), x.clone(), v.clone());

  match (e.clone(), x.clone()) {
    (Expr::Var(ref s1), Expr::Var(ref s2)) if s1 == s2 => v.clone(),
    (Expr::FnCall(v1, xs), _) => {
      let xs2 = xs.iter().map(|xn| sub(xn.clone())).collect();

      if *v1 == x {
        Expr::FnCall(Box::new(v.clone()), xs2)
      } else {
        Expr::FnCall(v1, xs2)
      }
    },
    (Expr::Var(_), _) => e,
    (Expr::Int(_), _) => e,
    (Expr::Bool(_), _) => e,
    (Expr::BinOp(op, e1, e2), _) => {
      Expr::BinOp(
        op,
        Box::new(sub(*e1)),
        Box::new(sub(*e2))
      )
    },
    (Expr::UnOp(op, e1), _) => Expr::UnOp(op, Box::new(sub(*e1))),
    (Expr::Ternary(e1, e2, e3), _) => {
      Expr::Ternary(
        Box::new(sub(*e1)),
        Box::new(sub(*e2)),
        Box::new(sub(*e3))
      )
    },
    (Expr::Let(e1, e2, e3), _) => {
      let e3s = if *e1 == x {
        *e3
      } else {
        sub(*e3)
      };

      Expr::Let(
        Box::new(*e1),
        Box::new(sub(*e2)),
        Box::new(e3s)
      )
    },
    (Expr::Func(name, e1, xs), _) => {
      let xs2 = xs.iter().map(|xn| sub(xn.clone())).collect();
      Expr::Func(name, Box::new(sub(*e1)), xs2)
    },
  }
}

pub fn step(e: Expr) -> Expr {
  match e {
    // TODO 'use Expr::*;
    //
    /**
     * Values are ineligible for step
     */
    Expr::Int(_) | Expr::Bool(_) | Expr::Func(_, _, _) | Expr::Var(_) => {
      debug!("stepping on a value {:?}", e);
      panic!("stepping on a value");
    }
    /**
     * Base cases
     */
    Expr::UnOp(UnOp::Not, ref e1) if is_bool(e1) => {
      Expr::Bool(!to_bool(&*e1))
    },
    Expr::UnOp(UnOp::Neg, ref e1) if is_int(e1) => {
      Expr::Int(-1 * to_int(&*e1))
    },
    Expr::BinOp(BinOp::And, ref e1, ref e2) if is_bool(e1) && is_bool(e2) => {
      Expr::Bool(to_bool(&*e1) && to_bool(&*e2))
    },
    Expr::BinOp(BinOp::Or, ref e1, ref e2) if is_bool(e1) && is_bool(e2) => {
      Expr::Bool(to_bool(&*e1) || to_bool(&*e2))
    },
    Expr::BinOp(BinOp::Eq, ref e1, ref e2) if is_value(e1) && is_value(e2) => {
      Expr::Bool(*e1 == *e2)
    },
    Expr::BinOp(BinOp::Ne, ref e1, ref e2) if is_value(e1) && is_value(e2) => {
      Expr::Bool(*e1 != *e2)
    },
    Expr::BinOp(BinOp::Mod, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      let n1 = to_int(&*e1);
      let n2 = to_int(&*e2);

      // rust % gives the remainder, not modulus
      let result = ((n1 % n2) + n2) % n2;

      Expr::Int(result)
    },
    Expr::BinOp(BinOp::Lt, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Bool(to_int(&*e1) < to_int(&*e2))
    },
    Expr::BinOp(BinOp::Gt, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Bool(to_int(&*e1) > to_int(&*e2))
    },
    Expr::BinOp(BinOp::Leq, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Bool(to_int(&*e1) <= to_int(&*e2))
    },
    Expr::BinOp(BinOp::Geq, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Bool(to_int(&*e1) >= to_int(&*e2))
    },
    Expr::BinOp(BinOp::Plus, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Int(to_int(&*e1) + to_int(&*e2))
    },
    Expr::BinOp(BinOp::Minus, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Int(to_int(&*e1) - to_int(&*e2))
    },
    Expr::BinOp(BinOp::Times, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Int(to_int(&*e1) * to_int(&*e2))
    },
    Expr::BinOp(BinOp::Div, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Int(to_int(&*e1) / to_int(&*e2))
    },
    Expr::BinOp(BinOp::Seq, ref e1, ref e2) if is_value(e1) => {
      *e2.clone()
    },
    Expr::Ternary(ref e1, ref e2, ref e3) if is_value(e1) => {
      match to_bool(&*e1) {
        true => *e2.clone(),
        false => *e3.clone(),
      }
    },
    Expr::Let(ref x, ref e1, ref e2) if is_value(e1) => {
      subst(*e2.clone(), *x.clone(), *e1.clone())
    },
    Expr::FnCall(ref v1, ref es) if is_func(v1) => {
      match **v1 {
        Expr::Func(ref name, ref e1, ref xs) => {
          // sub the params
          let exp = xs.iter().zip(es.iter())
            .fold(*e1.clone(), |exp, (xn, en)| subst(exp, xn.clone(), en.clone()));

          // sub the fn body for named functions
          match *name {
            None => exp,
            Some(ref s) => subst(exp, *s.clone(), *v1.clone())
          }
        },
        _ => {
          debug!("expected a Func, got {:?}", v1);
          panic!()
        },
      }
    },
    /**
     * Search Cases
     */
    Expr::BinOp(ref op, ref v1, ref e2) if is_value(v1) => {
      Expr::BinOp(op.clone(), Box::new(*v1.clone()), Box::new(step(*e2.clone())))
    },
    Expr::BinOp(op, e1, e2) => {
      Expr::BinOp(op, Box::new(step(*e1)), e2)
    },
    Expr::UnOp(op, e1) => {
      Expr::UnOp(op, Box::new(step(*e1)))
    },
    Expr::Ternary(e1, e2, e3) => {
      Expr::Ternary(Box::new(step(*e1)), e2, e3)
    },
    Expr::Let(ref e1, ref e2, ref e3) if is_value(e1) => {
      Expr::Let(Box::new(step(*e1.clone())), Box::new(*e2.clone()), Box::new(*e3.clone()))
    },
    Expr::Let(e1, e2, e3) => {
      Expr::Let(Box::new(step(*e1)), e2, e3)
    },
    Expr::FnCall(e1, mut xs) => {
      let mut found_first = true;

      for x in xs.iter_mut() {
        if is_value(x) && found_first == true {
          found_first = false;
          *x = step(x.clone());
        }
      }

      Expr::FnCall(e1, xs)
    }
  }
}

pub fn boxx(input: &str) -> Expr {
  let mut exp = parse(input);

  loop {
    if is_value(&exp) {
      return exp
    } else {
      exp = step(exp);
    }
  }
}
<commit_msg>tests passing<commit_after>use parser::{parse};

#[derive(Clone, Debug, PartialEq)]
pub enum UnOp {
  Not,
  Neg,
}

#[derive(Clone, Debug, PartialEq)]
pub enum BinOp {
  Plus,
  Minus,
  Times,
  Div,
  Eq,
  Ne,
  Leq,
  Geq,
  Lt,
  Gt,
  And,
  Or,
  Mod,
  Seq,
}

#[derive(Clone, Debug, PartialEq)]
pub enum Expr {
  Int(isize),
  Bool(bool),
  Var(String),
  BinOp(BinOp, Box<Expr>, Box<Expr>),
  UnOp(UnOp, Box<Expr>),
  Ternary(Box<Expr>, Box<Expr>, Box<Expr>),
  Let(Box<Expr>, Box<Expr>, Box<Expr>),
  Func(Option<Box<Expr>>, Box<Expr>, Vec<Expr>),
  FnCall(Box<Expr>, Vec<Expr>),
}

fn is_int(e: &Expr) -> bool {
  match *e {
    Expr::Int(_) => true,
    _ => false,
  }
}

fn is_bool(e: &Expr) -> bool {
  match *e {
    Expr::Bool(_) => true,
    _ => false,
  }
}

fn is_func(e: &Expr) -> bool {
  match *e {
    Expr::Func(_, _, _) => true,
    _ => false,
  }
}

fn is_value(e: &Expr) -> bool {
  match *e {
    Expr::Int(_) | Expr::Bool(_) | Expr::Var(_) | Expr::Func(_, _, _) => true,
    _ => false,
  }
}

fn to_int(e: &Expr) -> isize {
  match *e {
    Expr::Int(n) => n,
    _ => {
      debug!("cant turn into int: {:?}", e);
      panic!()
    }
  }
}

fn to_bool(e: &Expr) -> bool {
  match *e {
    Expr::Bool(b) => b,
    _ => {
      debug!("cant turn into bool: {:?}", e);
      panic!()
    }
  }
}

fn subst(e: Expr, x: Expr, v: Expr) -> Expr {
  let sub = |e1: Expr| subst(e1.clone(), x.clone(), v.clone());

  match (e.clone(), x.clone()) {
    (Expr::Var(ref s1), Expr::Var(ref s2)) if s1 == s2 => v.clone(),
    (Expr::FnCall(v1, xs), _) => {
      let xs2 = xs.iter().map(|xn| sub(xn.clone())).collect();

      if *v1 == x {
        Expr::FnCall(Box::new(v.clone()), xs2)
      } else {
        Expr::FnCall(v1, xs2)
      }
    },
    (Expr::Var(_), _) => e,
    (Expr::Int(_), _) => e,
    (Expr::Bool(_), _) => e,
    (Expr::BinOp(op, e1, e2), _) => {
      Expr::BinOp(
        op,
        Box::new(sub(*e1)),
        Box::new(sub(*e2))
      )
    },
    (Expr::UnOp(op, e1), _) => Expr::UnOp(op, Box::new(sub(*e1))),
    (Expr::Ternary(e1, e2, e3), _) => {
      Expr::Ternary(
        Box::new(sub(*e1)),
        Box::new(sub(*e2)),
        Box::new(sub(*e3))
      )
    },
    (Expr::Let(e1, e2, e3), _) => {
      let e3s = if *e1 == x {
        *e3
      } else {
        sub(*e3)
      };

      Expr::Let(
        Box::new(*e1),
        Box::new(sub(*e2)),
        Box::new(e3s)
      )
    },
    (Expr::Func(name, e1, xs), _) => {
      let xs2 = xs.iter().map(|xn| sub(xn.clone())).collect();
      Expr::Func(name, Box::new(sub(*e1)), xs2)
    },
  }
}

pub fn step(e: Expr) -> Expr {
  debug!("step(e) : {:?}", e);
  match e {
    // TODO 'use Expr::*;
    //
    /**
     * Values are ineligible for step
     */
    Expr::Int(_) | Expr::Bool(_) | Expr::Func(_, _, _) | Expr::Var(_) => {
      debug!("stepping on a value {:?}", e);
      panic!("stepping on a value");
    }
    /**
     * Base cases
     */
    Expr::UnOp(UnOp::Not, ref e1) if is_bool(e1) => {
      Expr::Bool(!to_bool(&*e1))
    },
    Expr::UnOp(UnOp::Neg, ref e1) if is_int(e1) => {
      Expr::Int(-1 * to_int(&*e1))
    },
    Expr::BinOp(BinOp::And, ref e1, ref e2) if is_bool(e1) && is_bool(e2) => {
      Expr::Bool(to_bool(&*e1) && to_bool(&*e2))
    },
    Expr::BinOp(BinOp::Or, ref e1, ref e2) if is_bool(e1) && is_bool(e2) => {
      Expr::Bool(to_bool(&*e1) || to_bool(&*e2))
    },
    Expr::BinOp(BinOp::Eq, ref e1, ref e2) if is_value(e1) && is_value(e2) => {
      Expr::Bool(*e1 == *e2)
    },
    Expr::BinOp(BinOp::Ne, ref e1, ref e2) if is_value(e1) && is_value(e2) => {
      Expr::Bool(*e1 != *e2)
    },
    Expr::BinOp(BinOp::Mod, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      let n1 = to_int(&*e1);
      let n2 = to_int(&*e2);

      // rust % gives the remainder, not modulus
      let result = ((n1 % n2) + n2) % n2;

      Expr::Int(result)
    },
    Expr::BinOp(BinOp::Lt, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Bool(to_int(&*e1) < to_int(&*e2))
    },
    Expr::BinOp(BinOp::Gt, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Bool(to_int(&*e1) > to_int(&*e2))
    },
    Expr::BinOp(BinOp::Leq, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Bool(to_int(&*e1) <= to_int(&*e2))
    },
    Expr::BinOp(BinOp::Geq, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Bool(to_int(&*e1) >= to_int(&*e2))
    },
    Expr::BinOp(BinOp::Plus, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Int(to_int(&*e1) + to_int(&*e2))
    },
    Expr::BinOp(BinOp::Minus, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Int(to_int(&*e1) - to_int(&*e2))
    },
    Expr::BinOp(BinOp::Times, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Int(to_int(&*e1) * to_int(&*e2))
    },
    Expr::BinOp(BinOp::Div, ref e1, ref e2) if is_int(e1) && is_int(e2) => {
      Expr::Int(to_int(&*e1) / to_int(&*e2))
    },
    Expr::BinOp(BinOp::Seq, ref e1, ref e2) if is_value(e1) => {
      *e2.clone()
    },
    Expr::Ternary(ref e1, ref e2, ref e3) if is_value(e1) => {
      match to_bool(&*e1) {
        true => *e2.clone(),
        false => *e3.clone(),
      }
    },
    Expr::Let(ref x, ref e1, ref e2) if is_value(e1) => {
      subst(*e2.clone(), *x.clone(), *e1.clone())
    },
    Expr::FnCall(ref v1, ref es) if is_func(v1) => {
      match **v1 {
        Expr::Func(ref name, ref e1, ref xs) => {
          // sub the params
          let exp = xs.iter().zip(es.iter())
            .fold(*e1.clone(), |exp, (xn, en)| subst(exp, xn.clone(), en.clone()));

          // sub the fn body for named functions
          match *name {
            None => exp,
            Some(ref s) => subst(exp, *s.clone(), *v1.clone())
          }
        },
        _ => {
          debug!("expected a Func, got {:?}", v1);
          panic!()
        },
      }
    },
    /**
     * Search Cases
     */
    Expr::BinOp(ref op, ref v1, ref e2) if is_value(v1) => {
      Expr::BinOp(op.clone(), Box::new(*v1.clone()), Box::new(step(*e2.clone())))
    },
    Expr::BinOp(op, e1, e2) => {
      Expr::BinOp(op, Box::new(step(*e1)), e2)
    },
    Expr::UnOp(op, e1) => {
      Expr::UnOp(op, Box::new(step(*e1)))
    },
    Expr::Ternary(e1, e2, e3) => {
      Expr::Ternary(Box::new(step(*e1)), e2, e3)
    },
    Expr::Let(ref e1, ref e2, ref e3) if is_value(e1) => {
      Expr::Let(Box::new(*e1.clone()), Box::new(step(*e2.clone())), Box::new(*e3.clone()))
    },
    Expr::Let(e1, e2, e3) => {
      Expr::Let(Box::new(step(*e1)), e2, e3)
    },
    Expr::FnCall(e1, mut xs) => {
      let mut found_first = true;

      for x in xs.iter_mut() {
        if is_value(x) && found_first == true {
          found_first = false;
          *x = step(x.clone());
        }
      }

      Expr::FnCall(e1, xs)
    }
  }
}

pub fn boxx(input: &str) -> Expr {
  let mut exp = parse(input);

  loop {
    if is_value(&exp) {
      return exp
    } else {
      exp = step(exp);
    }
  }
}
<|endoftext|>


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

use std::collections::HashMap;
use std::time::Duration;
use std::vec::Vec;

use async_std::io::{self, stderr, Read, Result as IoResult, Write};
use async_std::net::{SocketAddr, TcpStream};
use async_std::prelude::*;
use log::Level::Info;

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

#[derive(Debug, PartialEq)]
enum TransferEncoding {
    ContentLength(usize),
    Chunked,
    Unkown,
    None,
}

#[derive(Debug, PartialEq)]
enum TransferEncodingStatus {
    ReadingHeader,
    ChunkHeader,
    ReadingBody(usize),
}

fn drain_buffer<T>(buffer: &mut Vec<T>, size: usize) -> Vec<T> {
    if buffer.len() >= size {
        buffer.drain(size..).collect()
    } else {
        error!("invalid size in drained buffer");
        buffer.clear();
        buffer.drain(..).collect()
    }
}

impl From<&[u8]> for TransferEncoding {
    fn from(hdr: &[u8]) -> Self {
        let hdr = String::from_utf8_lossy(hdr);
        let hdrup = hdr.to_ascii_uppercase();
        match hdrup.as_str() {
            "CHUNKED" => TransferEncoding::Chunked,
            _ => TransferEncoding::Unkown,
        }
    }
}

struct HttpDecoder<'a> {
    reader: &'a mut (dyn Read + Unpin), // tls require Write
    writer: &'a mut (dyn Write + Unpin),
    buffer: Vec<u8>,
    transfer_encoding: TransferEncoding,
    transfer_encoding_status: TransferEncodingStatus,
    read_timeout: Duration,
}

impl<'a> HttpDecoder<'a> {
    fn new(
        writer: &'a mut (dyn Write + Unpin),
        reader: &'a mut (dyn Read + Unpin),
        read_timeout: u64,
    ) -> Self {
        HttpDecoder {
            writer,
            reader,
            buffer: Vec::with_capacity(constants::BUFFER_PAGE_SIZE),
            transfer_encoding: TransferEncoding::None,
            transfer_encoding_status: TransferEncodingStatus::ReadingHeader,
            read_timeout: Duration::from_millis(read_timeout),
        }
    }

    async fn chunk_read(&mut self) -> IoResult<usize> {
        let ret = io::timeout(self.read_timeout, async {
            let mut buf = [0; constants::BUFFER_PAGE_SIZE];
            let ret = self.reader.read(&mut buf[..]).await;
            if let Ok(count) = ret {
                if count > 0 {
                    self.buffer.extend_from_slice(&buf[..count]);
                }
            }
            ret
        });
        ret.await.map_err(|err| match err.kind() {
            io::ErrorKind::TimedOut => io::Error::new(err.kind(), "Read Timeout".to_owned()),
            _ => err,
        })
    }
    async fn read_write_headers(&mut self) -> IoResult<()> {
        info!("Reading headers");
        loop {
            let _count = self.chunk_read().await?;
            if let Some(_) = self.process_headers().await {
                break;
            }
        }
        Ok(())
    }

    async fn read_write_no_transfer_encoding(&mut self) -> IoResult<()> {
        loop {
            self.writer.write(self.buffer.as_slice()).await.unwrap();
            self.buffer.clear();
            let cnt = self.chunk_read().await?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    async fn read_content_length(&mut self, size: usize) -> IoResult<()> {
        let mut read_count = self.buffer.len();
        loop {
            self.writer.write(self.buffer.as_slice()).await.unwrap();
            self.buffer.clear();

            if read_count >= size {
                break;
            }
            read_count = read_count + self.chunk_read().await?;
            warn!("< {}", read_count);
        }
        Ok(())
    }

    async fn read_write_chunk(&mut self) -> IoResult<()> {
        loop {
            let done = self.process_chunk().await?;
            if done {
                break;
            }
            let cnt = self.chunk_read().await?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    async fn read_write_body(&mut self) -> IoResult<()> {
        info!("Reading body");
        match self.transfer_encoding {
            TransferEncoding::ContentLength(size) => {
                self.read_content_length(size).await?;
            }
            TransferEncoding::Chunked => {
                self.read_write_chunk().await?;
            }
            _ => {
                warn!("Neither Content-Length, not chunk is response header");
                self.read_write_no_transfer_encoding().await?;
            }
        }

        if self.buffer.len() > 0 {
            let b = String::from_utf8_lossy(self.buffer.as_slice());
            error!("Buffer not clear: {}", b);
        }

        Ok(())
    }

    async fn read_write(&mut self) -> IoResult<()> {
        self.read_write_headers().await?;
        self.read_write_body().await?;
        self.writer.flush().await
    }

    async fn process_headers(&mut self) -> Option<usize> {
        let ret = {
            let resp_header: Vec<&[u8]> = constants::SPLIT_HEADERS_RE
                .splitn(self.buffer.as_slice(), 2)
                .collect();
            if resp_header.len() == 2 {
                // We have headers
                let headers = resp_header.get(0).unwrap();
                if let Some(header) = constants::TRANSFER_ENCODING.captures(headers) {
                    if let Some(tenc) = header.get(1) {
                        self.transfer_encoding = TransferEncoding::from(tenc.as_bytes());
                    }
                } else if let Some(header) = constants::CONTENT_LENGTH.captures(headers) {
                    if let Some(clength) = header.get(1) {
                        let clength = String::from_utf8_lossy(clength.as_bytes()).into_owned();
                        let clength = usize::from_str_radix(clength.as_str(), 10).unwrap();
                        self.transfer_encoding = TransferEncoding::ContentLength(clength);
                    }
                }
                let resp = headers.len();
                self.writer.write(headers).await.unwrap();
                Some(resp + 4) // + CRLF CRLF
            } else {
                None
            }
        };
        if let Some(to_drain) = ret {
            self.buffer = drain_buffer(&mut self.buffer, to_drain);
            info!("End of Headers readched");
            debug!("Transfer encoding: {:?}", self.transfer_encoding);
            debug!("{:?}", String::from_utf8_lossy(self.buffer.as_slice()));
        }
        ret
    }
    async fn process_chunk(&mut self) -> IoResult<bool> {
        debug!(
            "transfer_encoding_status: {:?}",
            self.transfer_encoding_status
        );
        let mut can_process_buffer = self.buffer.len() > 0;
        let mut body_chunk_size = 0;
        let mut header_len: usize;
        while can_process_buffer {
            header_len = 0;
            if self.transfer_encoding_status == TransferEncodingStatus::ReadingHeader {
                debug!("Reading header in Transfer-Encoding chunked");
                // we read the chunk size to drain
                let header: Vec<&[u8]> = constants::SPLIT_HEADER_BRE
                    .splitn(self.buffer.as_slice(), 2)
                    .collect();
                if header.len() == 2 {
                    if let Some(header) = constants::GET_CHUNK_SIZE.captures(header[0]) {
                        if let Some(size) = header.get(1) {
                            let size = size.as_bytes();
                            header_len = size.len() + 2;
                            let size = String::from_utf8_lossy(size).into_owned();
                            body_chunk_size = usize::from_str_radix(size.as_str(), 16).unwrap();
                            debug!("Chunk Size to read: {}", body_chunk_size);
                            self.transfer_encoding_status = TransferEncodingStatus::ChunkHeader;
                        } else {
                            error!("Chunk Header is invalid");
                        }
                    } else {
                        error!("Chunk Header has improper size");
                        // else return Error
                        // break;
                    }
                } else {
                    debug!(
                        "Chunked not complete: {}",
                        String::from_utf8_lossy(self.buffer.as_slice())
                    );
                }
            }

            if self.transfer_encoding_status == TransferEncodingStatus::ChunkHeader {
                if body_chunk_size == 0 {
                    debug!(
                        "0 chunked size received: {}",
                        String::from_utf8_lossy(self.buffer.as_slice())
                    );
                    self.buffer.clear(); // should we check that it is '0\r\n' ?
                    return Ok(true);
                }

                debug!(
                    "Before header cleanup: {}",
                    String::from_utf8_lossy(self.buffer.as_slice())
                );
                self.buffer = drain_buffer(&mut self.buffer, header_len);
                debug!(
                    "After header cleanup: {}",
                    String::from_utf8_lossy(self.buffer.as_slice())
                );
                self.transfer_encoding_status =
                    TransferEncodingStatus::ReadingBody(body_chunk_size);
            }

            if let TransferEncodingStatus::ReadingBody(buf_size) = self.transfer_encoding_status {
                if self.buffer.len() > (buf_size + 2) {
                    let mut buffer: Vec<u8> = self.buffer.drain(buf_size..).collect();
                    self.writer.write(self.buffer.as_slice()).await?;
                    self.buffer = drain_buffer(&mut buffer, 2); // CRLF
                    self.transfer_encoding_status = TransferEncodingStatus::ReadingHeader;
                    body_chunk_size = 0;
                    if self.buffer.len() < 4 {
                        can_process_buffer = false;
                    }
                } else {
                    can_process_buffer = false;
                }
            } else {
                can_process_buffer = false;
            }
        }
        return Ok(false);
    }
}

async 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[0]);
    let headers: Vec<&str> = constants::SPLIT_HEADER_RE.split(&headers).collect();
    let bodylen = if request.len() == 2 {
        let body = &request[1];
        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).await.unwrap();
        }
        if bodylen > 0 {
            writeln!(&mut stderr(), "> [{} bytes]", bodylen)
                .await
                .unwrap();
        }
        writeln!(&mut stderr(), ">").await.unwrap();
    }
}

async fn from_http(
    request: &Request,
    client: &mut TcpStream,
    out: &mut (dyn Write + Unpin),
    verbose: bool,
    read_timeout: u64,
    request_timeout: u64,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose).await;

    debug!("Sending request...");
    client.write_all(&raw_request).await?;

    debug!("Reading response headers...");

    let mut http_decoder = HttpDecoder::new(out, client, read_timeout);

    if request_timeout > 0 {
        io::timeout(Duration::from_millis(request_timeout), async {
            http_decoder.read_write().await
        })
        .await
        .map_err(|err| match err.kind() {
            io::ErrorKind::TimedOut => {
                if err.to_string() == "Read Timeout" {
                    err
                } else {
                    io::Error::new(err.kind(), "Request Timeout".to_owned())
                }
            }
            _ => err,
        })?;
    } else {
        http_decoder.read_write().await?;
    }
    Ok(())
}

async fn from_https(
    request: &Request,
    client: &mut TcpStream,
    out: &mut (dyn Write + Unpin),
    verbose: bool,
    read_timeout: u64,
    request_timeout: u64,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose).await;

    let mut tls_client = TLSStream::new(client, request.host())?;
    tls_client.starttls().await?;

    debug!("Sending request...");
    tls_client.write_all(&raw_request).await?;
    debug!("Request sent");

    debug!("Decoding response...");
    let mut http_decoder = HttpDecoder::new(out, &mut tls_client, read_timeout);

    if request_timeout > 0 {
        io::timeout(Duration::from_millis(request_timeout), async {
            http_decoder.read_write().await
        })
        .await
        .map_err(|err| match err.kind() {
            io::ErrorKind::TimedOut => {
                if err.to_string() == "Read Timeout" {
                    err
                } else {
                    io::Error::new(err.kind(), "Request Timeout".to_owned())
                }
            }
            _ => err,
        })?;
    } else {
        http_decoder.read_write().await?;
    }

    Ok(())
}

pub async fn http_query(
    request: &Request,
    mut out: &mut (dyn Write + Unpin),
    authorities: &HashMap<String, SocketAddr>,
    verbose: bool,
    ipv4: bool,
    ipv6: bool,
    dns_timeout: u64,
    connect_timeout: u64,
    read_timeout: u64,
    request_timeout: u64,
) -> CabotResult<()> {
    debug!(
        "HTTP Query {} {}",
        request.http_method(),
        request.request_uri()
    );

    let authority = request.authority();

    let addr = match authorities.get(authority) {
        Some(val) => {
            info!("Fetch authority {} using autorities map", authority);
            *val
        }
        None => {
            info!("Fetch authority {} using resolver", authority);
            let resolver = Resolver::new(verbose);
            resolver
                .get_addr(authority, ipv4, ipv6, dns_timeout)
                .await?
        }
    };

    info!("Connecting to {}", addr);
    let mut client = io::timeout(Duration::from_millis(connect_timeout), async {
        TcpStream::connect(addr).await
    })
    .await
    .map_err(|err| match err.kind() {
        io::ErrorKind::TimedOut => io::Error::new(err.kind(), "Connection Timeout".to_owned()),
        _ => err,
    })?;

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

    out.flush().await.unwrap();

    Ok(())
}
<commit_msg>Fix when the chunk is bigger than buffer size<commit_after>//! Low level and internal http and https implementation.

use std::collections::HashMap;
use std::time::Duration;
use std::vec::Vec;

use async_std::io::{self, stderr, Read, Result as IoResult, Write};
use async_std::net::{SocketAddr, TcpStream};
use async_std::prelude::*;
use log::Level::Info;

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

#[derive(Debug, PartialEq)]
enum TransferEncoding {
    ContentLength(usize),
    Chunked,
    Unkown,
    None,
}

#[derive(Debug, PartialEq)]
enum TransferEncodingStatus {
    ReadingHeader,
    ChunkHeader,
    ReadingBody(usize),
}

fn drain_buffer<T>(buffer: &mut Vec<T>, size: usize) -> Vec<T> {
    if buffer.len() >= size {
        buffer.drain(size..).collect()
    } else {
        error!("invalid size in drained buffer");
        buffer.clear();
        buffer.drain(..).collect()
    }
}

impl From<&[u8]> for TransferEncoding {
    fn from(hdr: &[u8]) -> Self {
        let hdr = String::from_utf8_lossy(hdr);
        let hdrup = hdr.to_ascii_uppercase();
        match hdrup.as_str() {
            "CHUNKED" => TransferEncoding::Chunked,
            _ => TransferEncoding::Unkown,
        }
    }
}

struct HttpDecoder<'a> {
    reader: &'a mut (dyn Read + Unpin), // tls require Write
    writer: &'a mut (dyn Write + Unpin),
    buffer: Vec<u8>,
    transfer_encoding: TransferEncoding,
    transfer_encoding_status: TransferEncodingStatus,
    read_timeout: Duration,
}

impl<'a> HttpDecoder<'a> {
    fn new(
        writer: &'a mut (dyn Write + Unpin),
        reader: &'a mut (dyn Read + Unpin),
        read_timeout: u64,
    ) -> Self {
        HttpDecoder {
            writer,
            reader,
            buffer: Vec::with_capacity(constants::BUFFER_PAGE_SIZE),
            transfer_encoding: TransferEncoding::None,
            transfer_encoding_status: TransferEncodingStatus::ReadingHeader,
            read_timeout: Duration::from_millis(read_timeout),
        }
    }

    async fn chunk_read(&mut self) -> IoResult<usize> {
        let ret = io::timeout(self.read_timeout, async {
            let mut buf = [0; constants::BUFFER_PAGE_SIZE];
            let ret = self.reader.read(&mut buf[..]).await;
            if let Ok(count) = ret {
                if count > 0 {
                    self.buffer.extend_from_slice(&buf[..count]);
                }
            }
            ret
        });
        ret.await.map_err(|err| match err.kind() {
            io::ErrorKind::TimedOut => io::Error::new(err.kind(), "Read Timeout".to_owned()),
            _ => err,
        })
    }
    async fn read_write_headers(&mut self) -> IoResult<()> {
        info!("Reading headers");
        loop {
            let _count = self.chunk_read().await?;
            if let Some(_) = self.process_headers().await {
                break;
            }
        }
        Ok(())
    }

    async fn read_write_no_transfer_encoding(&mut self) -> IoResult<()> {
        loop {
            self.writer.write(self.buffer.as_slice()).await.unwrap();
            self.buffer.clear();
            let cnt = self.chunk_read().await?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    async fn read_content_length(&mut self, size: usize) -> IoResult<()> {
        let mut read_count = self.buffer.len();
        loop {
            self.writer.write(self.buffer.as_slice()).await.unwrap();
            self.buffer.clear();

            if read_count >= size {
                break;
            }
            read_count = read_count + self.chunk_read().await?;
            warn!("< {}", read_count);
        }
        Ok(())
    }

    async fn read_write_chunk(&mut self) -> IoResult<()> {
        loop {
            let done = self.process_chunk().await?;
            if done {
                break;
            }
            let cnt = self.chunk_read().await?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    async fn read_write_body(&mut self) -> IoResult<()> {
        info!("Reading body");
        match self.transfer_encoding {
            TransferEncoding::ContentLength(size) => {
                self.read_content_length(size).await?;
            }
            TransferEncoding::Chunked => {
                self.read_write_chunk().await?;
            }
            _ => {
                warn!("Neither Content-Length, not chunk is response header");
                self.read_write_no_transfer_encoding().await?;
            }
        }

        if self.buffer.len() > 0 {
            let b = String::from_utf8_lossy(self.buffer.as_slice());
            error!("Buffer not clear: {}", b);
        }

        Ok(())
    }

    async fn read_write(&mut self) -> IoResult<()> {
        self.read_write_headers().await?;
        self.read_write_body().await?;
        self.writer.flush().await
    }

    async fn process_headers(&mut self) -> Option<usize> {
        let ret = {
            let resp_header: Vec<&[u8]> = constants::SPLIT_HEADERS_RE
                .splitn(self.buffer.as_slice(), 2)
                .collect();
            if resp_header.len() == 2 {
                // We have headers
                let headers = resp_header.get(0).unwrap();
                if let Some(header) = constants::TRANSFER_ENCODING.captures(headers) {
                    if let Some(tenc) = header.get(1) {
                        self.transfer_encoding = TransferEncoding::from(tenc.as_bytes());
                    }
                } else if let Some(header) = constants::CONTENT_LENGTH.captures(headers) {
                    if let Some(clength) = header.get(1) {
                        let clength = String::from_utf8_lossy(clength.as_bytes()).into_owned();
                        let clength = usize::from_str_radix(clength.as_str(), 10).unwrap();
                        self.transfer_encoding = TransferEncoding::ContentLength(clength);
                    }
                }
                let resp = headers.len();
                self.writer.write(headers).await.unwrap();
                Some(resp + 4) // + CRLF CRLF
            } else {
                None
            }
        };
        if let Some(to_drain) = ret {
            self.buffer = drain_buffer(&mut self.buffer, to_drain);
            info!("End of Headers readched");
            debug!("Transfer encoding: {:?}", self.transfer_encoding);
            debug!("{:?}", String::from_utf8_lossy(self.buffer.as_slice()));
        }
        ret
    }
    async fn process_chunk(&mut self) -> IoResult<bool> {
        debug!(
            "transfer_encoding_status: {:?}",
            self.transfer_encoding_status
        );
        if self.buffer.len() == 0 {
            return Ok(false);
        }
        let mut body_chunk_size = 0;
        let mut header_len: usize;
        loop {
            header_len = 0;
            if self.transfer_encoding_status == TransferEncodingStatus::ReadingHeader {
                debug!("Reading header in Transfer-Encoding chunked");
                // we read the chunk size to drain
                let header: Vec<&[u8]> = constants::SPLIT_HEADER_BRE
                    .splitn(self.buffer.as_slice(), 2)
                    .collect();
                if header.len() == 2 {
                    if let Some(header) = constants::GET_CHUNK_SIZE.captures(header[0]) {
                        if let Some(size) = header.get(1) {
                            let size = size.as_bytes();
                            header_len = size.len() + 2;
                            let size = String::from_utf8_lossy(size).into_owned();
                            body_chunk_size = usize::from_str_radix(size.as_str(), 16).unwrap();
                            debug!("Chunk Size to read: {}", body_chunk_size);
                            self.transfer_encoding_status = TransferEncodingStatus::ChunkHeader;
                        } else {
                            error!("Chunk Header is invalid");
                        }
                    } else {
                        error!("Chunk Header has improper size");
                        // else return Error
                        // break;
                    }
                } else {
                    debug!(
                        "Chunked not complete: {}",
                        String::from_utf8_lossy(self.buffer.as_slice())
                    );
                }
            }

            if self.transfer_encoding_status == TransferEncodingStatus::ChunkHeader {
                if body_chunk_size == 0 {
                    debug!(
                        "0 chunked size received: {}",
                        String::from_utf8_lossy(self.buffer.as_slice())
                    );
                    self.buffer.clear(); // should we check that it is '0\r\n' ?
                    return Ok(true);
                }

                debug!(
                    "Before header cleanup: {}",
                    String::from_utf8_lossy(self.buffer.as_slice())
                );
                self.buffer = drain_buffer(&mut self.buffer, header_len);
                debug!(
                    "After header cleanup: {}",
                    String::from_utf8_lossy(self.buffer.as_slice())
                );
                self.transfer_encoding_status =
                    TransferEncodingStatus::ReadingBody(body_chunk_size);
            }

            if let TransferEncodingStatus::ReadingBody(buf_size) = self.transfer_encoding_status {
                if buf_size >= constants::BUFFER_PAGE_SIZE {
                    let buf_size = buf_size - self.buffer.len();
                    self.writer.write(self.buffer.as_slice()).await?;
                    self.buffer.clear();
                    self.transfer_encoding_status = TransferEncodingStatus::ReadingBody(buf_size);
                    break;
                }
                if self.buffer.len() > (buf_size + 2) {
                    let mut buffer: Vec<u8> = self.buffer.drain(buf_size..).collect();
                    self.writer.write(self.buffer.as_slice()).await?;
                    self.buffer = drain_buffer(&mut buffer, 2); // CRLF
                    self.transfer_encoding_status = TransferEncodingStatus::ReadingHeader;
                    body_chunk_size = 0;
                    if self.buffer.len() < 4 {
                        break;
                    }
                } else {
                    break;
                }
            } else {
                break;
            }
        }
        return Ok(false);
    }
}

async 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[0]);
    let headers: Vec<&str> = constants::SPLIT_HEADER_RE.split(&headers).collect();
    let bodylen = if request.len() == 2 {
        let body = &request[1];
        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).await.unwrap();
        }
        if bodylen > 0 {
            writeln!(&mut stderr(), "> [{} bytes]", bodylen)
                .await
                .unwrap();
        }
        writeln!(&mut stderr(), ">").await.unwrap();
    }
}

async fn from_http(
    request: &Request,
    client: &mut TcpStream,
    out: &mut (dyn Write + Unpin),
    verbose: bool,
    read_timeout: u64,
    request_timeout: u64,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose).await;

    debug!("Sending request...");
    client.write_all(&raw_request).await?;

    debug!("Reading response headers...");

    let mut http_decoder = HttpDecoder::new(out, client, read_timeout);

    if request_timeout > 0 {
        io::timeout(Duration::from_millis(request_timeout), async {
            http_decoder.read_write().await
        })
        .await
        .map_err(|err| match err.kind() {
            io::ErrorKind::TimedOut => {
                if err.to_string() == "Read Timeout" {
                    err
                } else {
                    io::Error::new(err.kind(), "Request Timeout".to_owned())
                }
            }
            _ => err,
        })?;
    } else {
        http_decoder.read_write().await?;
    }
    Ok(())
}

async fn from_https(
    request: &Request,
    client: &mut TcpStream,
    out: &mut (dyn Write + Unpin),
    verbose: bool,
    read_timeout: u64,
    request_timeout: u64,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose).await;

    let mut tls_client = TLSStream::new(client, request.host())?;
    tls_client.starttls().await?;

    debug!("Sending request...");
    tls_client.write_all(&raw_request).await?;
    debug!("Request sent");

    debug!("Decoding response...");
    let mut http_decoder = HttpDecoder::new(out, &mut tls_client, read_timeout);

    if request_timeout > 0 {
        io::timeout(Duration::from_millis(request_timeout), async {
            http_decoder.read_write().await
        })
        .await
        .map_err(|err| match err.kind() {
            io::ErrorKind::TimedOut => {
                if err.to_string() == "Read Timeout" {
                    err
                } else {
                    io::Error::new(err.kind(), "Request Timeout".to_owned())
                }
            }
            _ => err,
        })?;
    } else {
        http_decoder.read_write().await?;
    }

    Ok(())
}

pub async fn http_query(
    request: &Request,
    mut out: &mut (dyn Write + Unpin),
    authorities: &HashMap<String, SocketAddr>,
    verbose: bool,
    ipv4: bool,
    ipv6: bool,
    dns_timeout: u64,
    connect_timeout: u64,
    read_timeout: u64,
    request_timeout: u64,
) -> CabotResult<()> {
    debug!(
        "HTTP Query {} {}",
        request.http_method(),
        request.request_uri()
    );

    let authority = request.authority();

    let addr = match authorities.get(authority) {
        Some(val) => {
            info!("Fetch authority {} using autorities map", authority);
            *val
        }
        None => {
            info!("Fetch authority {} using resolver", authority);
            let resolver = Resolver::new(verbose);
            resolver
                .get_addr(authority, ipv4, ipv6, dns_timeout)
                .await?
        }
    };

    info!("Connecting to {}", addr);
    let mut client = io::timeout(Duration::from_millis(connect_timeout), async {
        TcpStream::connect(addr).await
    })
    .await
    .map_err(|err| match err.kind() {
        io::ErrorKind::TimedOut => io::Error::new(err.kind(), "Connection Timeout".to_owned()),
        _ => err,
    })?;

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

    out.flush().await.unwrap();

    Ok(())
}
<|endoftext|>