<commit_before>//! Trait solving using Chalk.
use std::{panic, sync::Arc};

use chalk_ir::cast::Cast;
use hir_def::{expr::ExprId, DefWithBodyId, ImplId, TraitId, TypeAliasId};
use ra_db::{impl_intern_key, salsa, CrateId};
use ra_prof::profile;
use rustc_hash::FxHashSet;

use crate::{db::HirDatabase, DebruijnIndex};

use super::{Canonical, GenericPredicate, HirDisplay, ProjectionTy, TraitRef, Ty, TypeWalk};

use self::chalk::{from_chalk, Interner, ToChalk};

pub(crate) mod chalk;
mod builtin;

/// This controls the maximum size of types Chalk considers. If we set this too
/// high, we can run into slow edge cases; if we set it too low, Chalk won't
/// find some solutions.
const CHALK_SOLVER_MAX_SIZE: usize = 10;
/// This controls how much 'time' we give the Chalk solver before giving up.
const CHALK_SOLVER_FUEL: i32 = 100;

#[derive(Debug, Copy, Clone)]
struct ChalkContext<'a> {
    db: &'a dyn HirDatabase,
    krate: CrateId,
}

fn create_chalk_solver() -> chalk_solve::Solver<Interner> {
    let solver_choice =
        chalk_solve::SolverChoice::SLG { max_size: CHALK_SOLVER_MAX_SIZE, expected_answers: None };
    solver_choice.into_solver()
}

/// Collects impls for the given trait in the whole dependency tree of `krate`.
pub(crate) fn impls_for_trait_query(
    db: &dyn HirDatabase,
    krate: CrateId,
    trait_: TraitId,
) -> Arc<[ImplId]> {
    let mut impls = FxHashSet::default();
    // We call the query recursively here. On the one hand, this means we can
    // reuse results from queries for different crates; on the other hand, this
    // will only ever get called for a few crates near the root of the tree (the
    // ones the user is editing), so this may actually be a waste of memory. I'm
    // doing it like this mainly for simplicity for now.
    for dep in &db.crate_graph()[krate].dependencies {
        impls.extend(db.impls_for_trait(dep.crate_id, trait_).iter());
    }
    let crate_impl_defs = db.impls_in_crate(krate);
    impls.extend(crate_impl_defs.lookup_impl_defs_for_trait(trait_));
    impls.into_iter().collect()
}

/// A set of clauses that we assume to be true. E.g. if we are inside this function:
/// ```rust
/// fn foo<T: Default>(t: T) {}
/// ```
/// we assume that `T: Default`.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TraitEnvironment {
    pub predicates: Vec<GenericPredicate>,
}

impl TraitEnvironment {
    /// Returns trait refs with the given self type which are supposed to hold
    /// in this trait env. E.g. if we are in `foo<T: SomeTrait>()`, this will
    /// find that `T: SomeTrait` if we call it for `T`.
    pub(crate) fn trait_predicates_for_self_ty<'a>(
        &'a self,
        ty: &'a Ty,
    ) -> impl Iterator<Item = &'a TraitRef> + 'a {
        self.predicates.iter().filter_map(move |pred| match pred {
            GenericPredicate::Implemented(tr) if tr.self_ty() == ty => Some(tr),
            _ => None,
        })
    }
}

/// Something (usually a goal), along with an environment.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct InEnvironment<T> {
    pub environment: Arc<TraitEnvironment>,
    pub value: T,
}

impl<T> InEnvironment<T> {
    pub fn new(environment: Arc<TraitEnvironment>, value: T) -> InEnvironment<T> {
        InEnvironment { environment, value }
    }
}

/// Something that needs to be proven (by Chalk) during type checking, e.g. that
/// a certain type implements a certain trait. Proving the Obligation might
/// result in additional information about inference variables.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum Obligation {
    /// Prove that a certain type implements a trait (the type is the `Self` type
    /// parameter to the `TraitRef`).
    Trait(TraitRef),
    Projection(ProjectionPredicate),
}

impl Obligation {
    pub fn from_predicate(predicate: GenericPredicate) -> Option<Obligation> {
        match predicate {
            GenericPredicate::Implemented(trait_ref) => Some(Obligation::Trait(trait_ref)),
            GenericPredicate::Projection(projection_pred) => {
                Some(Obligation::Projection(projection_pred))
            }
            GenericPredicate::Error => None,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct ProjectionPredicate {
    pub projection_ty: ProjectionTy,
    pub ty: Ty,
}

impl TypeWalk for ProjectionPredicate {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        self.projection_ty.walk(f);
        self.ty.walk(f);
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        self.projection_ty.walk_mut_binders(f, binders);
        self.ty.walk_mut_binders(f, binders);
    }
}

/// Solve a trait goal using Chalk.
pub(crate) fn trait_solve_query(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: Canonical<InEnvironment<Obligation>>,
) -> Option<Solution> {
    let _p = profile("trait_solve_query").detail(|| match &goal.value.value {
        Obligation::Trait(it) => db.trait_data(it.trait_).name.to_string(),
        Obligation::Projection(_) => "projection".to_string(),
    });
    eprintln!("trait_solve_query({})", goal.value.value.display(db));

    if let Obligation::Projection(pred) = &goal.value.value {
        if let Ty::Bound(_) = &pred.projection_ty.parameters[0] {
            // Hack: don't ask Chalk to normalize with an unknown self type, it'll say that's impossible
            return Some(Solution::Ambig(Guidance::Unknown));
        }
    }

    let canonical = goal.to_chalk(db).cast(&Interner);

    // We currently don't deal with universes (I think / hope they're not yet
    // relevant for our use cases?)
    let u_canonical = chalk_ir::UCanonical { canonical, universes: 1 };
    let solution = solve(db, krate, &u_canonical);
    solution.map(|solution| solution_from_chalk(db, solution))
}

fn solve(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal<Interner>>>,
) -> Option<chalk_solve::Solution<Interner>> {
    let context = ChalkContext { db, krate };
    log::debug!("solve goal: {:?}", goal);
    let mut solver = create_chalk_solver();

    let fuel = std::cell::Cell::new(CHALK_SOLVER_FUEL);

    let solution = solver.solve_limited(&context, goal, || {
        context.db.check_canceled();
        let remaining = fuel.get();
        fuel.set(remaining - 1);
        if remaining == 0 {
            log::debug!("fuel exhausted");
        }
        remaining > 0
    });

    log::debug!("solve({:?}) => {:?}", goal, solution);
    solution
}

fn solution_from_chalk(
    db: &dyn HirDatabase,
    solution: chalk_solve::Solution<Interner>,
) -> Solution {
    let convert_subst = |subst: chalk_ir::Canonical<chalk_ir::Substitution<Interner>>| {
        let value = subst
            .value
            .iter(&Interner)
            .map(|p| match p.ty(&Interner) {
                Some(ty) => from_chalk(db, ty.clone()),
                None => unimplemented!(),
            })
            .collect();
        let result = Canonical { value, num_vars: subst.binders.len() };
        SolutionVariables(result)
    };
    match solution {
        chalk_solve::Solution::Unique(constr_subst) => {
            let subst = chalk_ir::Canonical {
                value: constr_subst.value.subst,
                binders: constr_subst.binders,
            };
            Solution::Unique(convert_subst(subst))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Definite(subst)) => {
            Solution::Ambig(Guidance::Definite(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Suggested(subst)) => {
            Solution::Ambig(Guidance::Suggested(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Unknown) => {
            Solution::Ambig(Guidance::Unknown)
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SolutionVariables(pub Canonical<Vec<Ty>>);

#[derive(Clone, Debug, PartialEq, Eq)]
/// A (possible) solution for a proposed goal.
pub enum Solution {
    /// The goal indeed holds, and there is a unique value for all existential
    /// variables.
    Unique(SolutionVariables),

    /// The goal may be provable in multiple ways, but regardless we may have some guidance
    /// for type inference. In this case, we don't return any lifetime
    /// constraints, since we have not "committed" to any particular solution
    /// yet.
    Ambig(Guidance),
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// When a goal holds ambiguously (e.g., because there are multiple possible
/// solutions), we issue a set of *guidance* back to type inference.
pub enum Guidance {
    /// The existential variables *must* have the given values if the goal is
    /// ever to hold, but that alone isn't enough to guarantee the goal will
    /// actually hold.
    Definite(SolutionVariables),

    /// There are multiple plausible values for the existentials, but the ones
    /// here are suggested as the preferred choice heuristically. These should
    /// be used for inference fallback only.
    Suggested(SolutionVariables),

    /// There's no useful information to feed back to type inference
    Unknown,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum FnTrait {
    FnOnce,
    FnMut,
    Fn,
}

impl FnTrait {
    fn lang_item_name(self) -> &'static str {
        match self {
            FnTrait::FnOnce => "fn_once",
            FnTrait::FnMut => "fn_mut",
            FnTrait::Fn => "fn",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ClosureFnTraitImplData {
    def: DefWithBodyId,
    expr: ExprId,
    fn_trait: FnTrait,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct UnsizeToSuperTraitObjectData {
    trait_: TraitId,
    super_trait: TraitId,
}

/// An impl. Usually this comes from an impl block, but some built-in types get
/// synthetic impls.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Impl {
    /// A normal impl from an impl block.
    ImplDef(ImplId),
    /// Closure types implement the Fn traits synthetically.
    ClosureFnTraitImpl(ClosureFnTraitImplData),
    /// [T; n]: Unsize<[T]>
    UnsizeArray,
    /// T: Unsize<dyn Trait> where T: Trait
    UnsizeToTraitObject(TraitId),
    /// dyn Trait: Unsize<dyn SuperTrait> if Trait: SuperTrait
    UnsizeToSuperTraitObject(UnsizeToSuperTraitObjectData),
}
/// This exists just for Chalk, because our ImplIds are only unique per module.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct GlobalImplId(salsa::InternId);
impl_intern_key!(GlobalImplId);

/// An associated type value. Usually this comes from a `type` declaration
/// inside an impl block, but for built-in impls we have to synthesize it.
/// (We only need this because Chalk wants a unique ID for each of these.)
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum AssocTyValue {
    /// A normal assoc type value from an impl block.
    TypeAlias(TypeAliasId),
    /// The output type of the Fn trait implementation.
    ClosureFnTraitImplOutput(ClosureFnTraitImplData),
}
/// This exists just for Chalk, because it needs a unique ID for each associated
/// type value in an impl (even synthetic ones).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct AssocTyValueId(salsa::InternId);
impl_intern_key!(AssocTyValueId);
<commit_msg>Use log info in trait_solve_query<commit_after>//! Trait solving using Chalk.
use std::{panic, sync::Arc};

use chalk_ir::cast::Cast;
use hir_def::{expr::ExprId, DefWithBodyId, ImplId, TraitId, TypeAliasId};
use ra_db::{impl_intern_key, salsa, CrateId};
use ra_prof::profile;
use rustc_hash::FxHashSet;

use crate::{db::HirDatabase, DebruijnIndex};

use super::{Canonical, GenericPredicate, HirDisplay, ProjectionTy, TraitRef, Ty, TypeWalk};

use self::chalk::{from_chalk, Interner, ToChalk};

pub(crate) mod chalk;
mod builtin;

/// This controls the maximum size of types Chalk considers. If we set this too
/// high, we can run into slow edge cases; if we set it too low, Chalk won't
/// find some solutions.
const CHALK_SOLVER_MAX_SIZE: usize = 10;
/// This controls how much 'time' we give the Chalk solver before giving up.
const CHALK_SOLVER_FUEL: i32 = 100;

#[derive(Debug, Copy, Clone)]
struct ChalkContext<'a> {
    db: &'a dyn HirDatabase,
    krate: CrateId,
}

fn create_chalk_solver() -> chalk_solve::Solver<Interner> {
    let solver_choice =
        chalk_solve::SolverChoice::SLG { max_size: CHALK_SOLVER_MAX_SIZE, expected_answers: None };
    solver_choice.into_solver()
}

/// Collects impls for the given trait in the whole dependency tree of `krate`.
pub(crate) fn impls_for_trait_query(
    db: &dyn HirDatabase,
    krate: CrateId,
    trait_: TraitId,
) -> Arc<[ImplId]> {
    let mut impls = FxHashSet::default();
    // We call the query recursively here. On the one hand, this means we can
    // reuse results from queries for different crates; on the other hand, this
    // will only ever get called for a few crates near the root of the tree (the
    // ones the user is editing), so this may actually be a waste of memory. I'm
    // doing it like this mainly for simplicity for now.
    for dep in &db.crate_graph()[krate].dependencies {
        impls.extend(db.impls_for_trait(dep.crate_id, trait_).iter());
    }
    let crate_impl_defs = db.impls_in_crate(krate);
    impls.extend(crate_impl_defs.lookup_impl_defs_for_trait(trait_));
    impls.into_iter().collect()
}

/// A set of clauses that we assume to be true. E.g. if we are inside this function:
/// ```rust
/// fn foo<T: Default>(t: T) {}
/// ```
/// we assume that `T: Default`.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TraitEnvironment {
    pub predicates: Vec<GenericPredicate>,
}

impl TraitEnvironment {
    /// Returns trait refs with the given self type which are supposed to hold
    /// in this trait env. E.g. if we are in `foo<T: SomeTrait>()`, this will
    /// find that `T: SomeTrait` if we call it for `T`.
    pub(crate) fn trait_predicates_for_self_ty<'a>(
        &'a self,
        ty: &'a Ty,
    ) -> impl Iterator<Item = &'a TraitRef> + 'a {
        self.predicates.iter().filter_map(move |pred| match pred {
            GenericPredicate::Implemented(tr) if tr.self_ty() == ty => Some(tr),
            _ => None,
        })
    }
}

/// Something (usually a goal), along with an environment.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct InEnvironment<T> {
    pub environment: Arc<TraitEnvironment>,
    pub value: T,
}

impl<T> InEnvironment<T> {
    pub fn new(environment: Arc<TraitEnvironment>, value: T) -> InEnvironment<T> {
        InEnvironment { environment, value }
    }
}

/// Something that needs to be proven (by Chalk) during type checking, e.g. that
/// a certain type implements a certain trait. Proving the Obligation might
/// result in additional information about inference variables.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum Obligation {
    /// Prove that a certain type implements a trait (the type is the `Self` type
    /// parameter to the `TraitRef`).
    Trait(TraitRef),
    Projection(ProjectionPredicate),
}

impl Obligation {
    pub fn from_predicate(predicate: GenericPredicate) -> Option<Obligation> {
        match predicate {
            GenericPredicate::Implemented(trait_ref) => Some(Obligation::Trait(trait_ref)),
            GenericPredicate::Projection(projection_pred) => {
                Some(Obligation::Projection(projection_pred))
            }
            GenericPredicate::Error => None,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct ProjectionPredicate {
    pub projection_ty: ProjectionTy,
    pub ty: Ty,
}

impl TypeWalk for ProjectionPredicate {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        self.projection_ty.walk(f);
        self.ty.walk(f);
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        self.projection_ty.walk_mut_binders(f, binders);
        self.ty.walk_mut_binders(f, binders);
    }
}

/// Solve a trait goal using Chalk.
pub(crate) fn trait_solve_query(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: Canonical<InEnvironment<Obligation>>,
) -> Option<Solution> {
    let _p = profile("trait_solve_query").detail(|| match &goal.value.value {
        Obligation::Trait(it) => db.trait_data(it.trait_).name.to_string(),
        Obligation::Projection(_) => "projection".to_string(),
    });
    log::info!("trait_solve_query({})", goal.value.value.display(db));

    if let Obligation::Projection(pred) = &goal.value.value {
        if let Ty::Bound(_) = &pred.projection_ty.parameters[0] {
            // Hack: don't ask Chalk to normalize with an unknown self type, it'll say that's impossible
            return Some(Solution::Ambig(Guidance::Unknown));
        }
    }

    let canonical = goal.to_chalk(db).cast(&Interner);

    // We currently don't deal with universes (I think / hope they're not yet
    // relevant for our use cases?)
    let u_canonical = chalk_ir::UCanonical { canonical, universes: 1 };
    let solution = solve(db, krate, &u_canonical);
    solution.map(|solution| solution_from_chalk(db, solution))
}

fn solve(
    db: &dyn HirDatabase,
    krate: CrateId,
    goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal<Interner>>>,
) -> Option<chalk_solve::Solution<Interner>> {
    let context = ChalkContext { db, krate };
    log::debug!("solve goal: {:?}", goal);
    let mut solver = create_chalk_solver();

    let fuel = std::cell::Cell::new(CHALK_SOLVER_FUEL);

    let solution = solver.solve_limited(&context, goal, || {
        context.db.check_canceled();
        let remaining = fuel.get();
        fuel.set(remaining - 1);
        if remaining == 0 {
            log::debug!("fuel exhausted");
        }
        remaining > 0
    });

    log::debug!("solve({:?}) => {:?}", goal, solution);
    solution
}

fn solution_from_chalk(
    db: &dyn HirDatabase,
    solution: chalk_solve::Solution<Interner>,
) -> Solution {
    let convert_subst = |subst: chalk_ir::Canonical<chalk_ir::Substitution<Interner>>| {
        let value = subst
            .value
            .iter(&Interner)
            .map(|p| match p.ty(&Interner) {
                Some(ty) => from_chalk(db, ty.clone()),
                None => unimplemented!(),
            })
            .collect();
        let result = Canonical { value, num_vars: subst.binders.len() };
        SolutionVariables(result)
    };
    match solution {
        chalk_solve::Solution::Unique(constr_subst) => {
            let subst = chalk_ir::Canonical {
                value: constr_subst.value.subst,
                binders: constr_subst.binders,
            };
            Solution::Unique(convert_subst(subst))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Definite(subst)) => {
            Solution::Ambig(Guidance::Definite(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Suggested(subst)) => {
            Solution::Ambig(Guidance::Suggested(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Unknown) => {
            Solution::Ambig(Guidance::Unknown)
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SolutionVariables(pub Canonical<Vec<Ty>>);

#[derive(Clone, Debug, PartialEq, Eq)]
/// A (possible) solution for a proposed goal.
pub enum Solution {
    /// The goal indeed holds, and there is a unique value for all existential
    /// variables.
    Unique(SolutionVariables),

    /// The goal may be provable in multiple ways, but regardless we may have some guidance
    /// for type inference. In this case, we don't return any lifetime
    /// constraints, since we have not "committed" to any particular solution
    /// yet.
    Ambig(Guidance),
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// When a goal holds ambiguously (e.g., because there are multiple possible
/// solutions), we issue a set of *guidance* back to type inference.
pub enum Guidance {
    /// The existential variables *must* have the given values if the goal is
    /// ever to hold, but that alone isn't enough to guarantee the goal will
    /// actually hold.
    Definite(SolutionVariables),

    /// There are multiple plausible values for the existentials, but the ones
    /// here are suggested as the preferred choice heuristically. These should
    /// be used for inference fallback only.
    Suggested(SolutionVariables),

    /// There's no useful information to feed back to type inference
    Unknown,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum FnTrait {
    FnOnce,
    FnMut,
    Fn,
}

impl FnTrait {
    fn lang_item_name(self) -> &'static str {
        match self {
            FnTrait::FnOnce => "fn_once",
            FnTrait::FnMut => "fn_mut",
            FnTrait::Fn => "fn",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ClosureFnTraitImplData {
    def: DefWithBodyId,
    expr: ExprId,
    fn_trait: FnTrait,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct UnsizeToSuperTraitObjectData {
    trait_: TraitId,
    super_trait: TraitId,
}

/// An impl. Usually this comes from an impl block, but some built-in types get
/// synthetic impls.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Impl {
    /// A normal impl from an impl block.
    ImplDef(ImplId),
    /// Closure types implement the Fn traits synthetically.
    ClosureFnTraitImpl(ClosureFnTraitImplData),
    /// [T; n]: Unsize<[T]>
    UnsizeArray,
    /// T: Unsize<dyn Trait> where T: Trait
    UnsizeToTraitObject(TraitId),
    /// dyn Trait: Unsize<dyn SuperTrait> if Trait: SuperTrait
    UnsizeToSuperTraitObject(UnsizeToSuperTraitObjectData),
}
/// This exists just for Chalk, because our ImplIds are only unique per module.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct GlobalImplId(salsa::InternId);
impl_intern_key!(GlobalImplId);

/// An associated type value. Usually this comes from a `type` declaration
/// inside an impl block, but for built-in impls we have to synthesize it.
/// (We only need this because Chalk wants a unique ID for each of these.)
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum AssocTyValue {
    /// A normal assoc type value from an impl block.
    TypeAlias(TypeAliasId),
    /// The output type of the Fn trait implementation.
    ClosureFnTraitImplOutput(ClosureFnTraitImplData),
}
/// This exists just for Chalk, because it needs a unique ID for each associated
/// type value in an impl (even synthetic ones).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct AssocTyValueId(salsa::InternId);
impl_intern_key!(AssocTyValueId);
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use capnp::capability::Promise;
use capnp_rpc::{rpc_twoparty_capnp, twoparty, RpcSystem};

use crate::hello_world_capnp::hello_world;

use futures::task::LocalSpawn;
use futures::{AsyncReadExt, FutureExt};
use std::net::ToSocketAddrs;

struct HelloWorldImpl;

impl hello_world::Server for HelloWorldImpl {
    fn say_hello(
        &mut self,
        params: hello_world::SayHelloParams,
        mut results: hello_world::SayHelloResults,
    ) -> Promise<(), ::capnp::Error> {

        let request = params.get().unwrap().get_request().unwrap();
        let name = request.get_name().unwrap();
        let message = format!("Hello, {}!", name);

        results.get().init_reply().set_message(&message);

        Promise::ok(())
    }
}

pub fn main() {
    let args: Vec<String> = ::std::env::args().collect();
    if args.len() != 3 {
        println!("usage: {} server ADDRESS[:PORT]", args[0]);
        return;
    }

    let addr = args[2]
        .to_socket_addrs()
        .unwrap()
        .next()
        .expect("could not parse address");

    let mut exec = futures::executor::LocalPool::new();
    let spawner = exec.spawner();

    exec.run_until(async move {
        let listener = async_std::net::TcpListener::bind(&addr).await.unwrap();
        let hello_world_impl =
            hello_world::ToClient::new(HelloWorldImpl).into_client::<::capnp_rpc::Server>();

        loop {
            let (stream, _) = listener.accept().await.unwrap();
            stream.set_nodelay(true).unwrap();
            let (reader, writer) = stream.split();
            let network = twoparty::VatNetwork::new(
                reader,
                writer,
                rpc_twoparty_capnp::Side::Server,
                Default::default(),
            );

            let rpc_system =
                RpcSystem::new(Box::new(network), Some(hello_world_impl.clone().client));

            spawner
                .spawn_local_obj(Box::pin(rpc_system.map(|_| ())).into())
                .unwrap();
        }
    });
}
<commit_msg>update Hello World example to use capnp_rpc::new_client()<commit_after>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use capnp::capability::Promise;
use capnp_rpc::{rpc_twoparty_capnp, twoparty, RpcSystem};

use crate::hello_world_capnp::hello_world;

use futures::task::LocalSpawn;
use futures::{AsyncReadExt, FutureExt};
use std::net::ToSocketAddrs;

struct HelloWorldImpl;

impl hello_world::Server for HelloWorldImpl {
    fn say_hello(
        &mut self,
        params: hello_world::SayHelloParams,
        mut results: hello_world::SayHelloResults,
    ) -> Promise<(), ::capnp::Error> {

        let request = params.get().unwrap().get_request().unwrap();
        let name = request.get_name().unwrap();
        let message = format!("Hello, {}!", name);

        results.get().init_reply().set_message(&message);

        Promise::ok(())
    }
}

pub fn main() {
    let args: Vec<String> = ::std::env::args().collect();
    if args.len() != 3 {
        println!("usage: {} server ADDRESS[:PORT]", args[0]);
        return;
    }

    let addr = args[2]
        .to_socket_addrs()
        .unwrap()
        .next()
        .expect("could not parse address");

    let mut exec = futures::executor::LocalPool::new();
    let spawner = exec.spawner();

    exec.run_until(async move {
        let listener = async_std::net::TcpListener::bind(&addr).await.unwrap();
        let hello_world_client: hello_world::Client = capnp_rpc::new_client(HelloWorldImpl);

        loop {
            let (stream, _) = listener.accept().await.unwrap();
            stream.set_nodelay(true).unwrap();
            let (reader, writer) = stream.split();
            let network = twoparty::VatNetwork::new(
                reader,
                writer,
                rpc_twoparty_capnp::Side::Server,
                Default::default(),
            );

            let rpc_system =
                RpcSystem::new(Box::new(network), Some(hello_world_client.clone().client));

            spawner
                .spawn_local_obj(Box::pin(rpc_system.map(|_| ())).into())
                .unwrap();
        }
    });
}
<|endoftext|>


<commit_before>#![recursion_limit = "128"]

#[macro_use]
extern crate lazy_static;
extern crate futures;
extern crate tokio;
extern crate trust_dns_resolver;

#[cfg(feature = "tokio-runtime")]
use std::io;
#[cfg(feature = "tokio-runtime")]
use std::net::SocketAddr;

#[cfg(feature = "tokio-runtime")]
use futures::{Future, FutureExt};
#[cfg(feature = "tokio-runtime")]
use trust_dns_resolver::TokioAsyncResolver;
#[cfg(feature = "tokio-runtime")]
use trust_dns_resolver::{IntoName, TryParseIp};

// This is an example of registering a static global resolver into any system.
//
// We may want to create a GlobalResolver as part of the Resolver library
//   in the mean time, this example has the necessary steps to do so.
//
// Thank you to @zonyitoo for the original example.
// TODO: this example can probably be made much simpler with the new
//      `AsyncResolver`.
#[cfg(feature = "tokio-runtime")]
lazy_static! {
    // First we need to setup the global Resolver
    static ref GLOBAL_DNS_RESOLVER: TokioAsyncResolver = {
        use std::sync::{Arc, Mutex, Condvar};
        use std::thread;

        // We'll be using this condvar to get the Resolver from the thread...
        let pair = Arc::new((Mutex::new(None::<TokioAsyncResolver>), Condvar::new()));
        let pair2 = pair.clone();

        // Spawn the runtime to a new thread...
        //
        // This thread will manage the actual resolution runtime
        thread::spawn(move || {
            // A runtime for this new thread
            let mut runtime = tokio::runtime::Runtime::new().expect("failed to launch Runtime");

            // our platform independent future, result, see next blocks
            let resolver = {

                // To make this independent, if targeting macOS, BSD, Linux, or Windows, we can use the system's configuration:
                #[cfg(any(unix, windows))]
                {
                    // use the system resolver configuration
                    TokioAsyncResolver::from_system_conf(runtime.handle().clone())
                }

                // For other operating systems, we can use one of the preconfigured definitions
                #[cfg(not(any(unix, windows)))]
                {
                    // Directly reference the config types
                    use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};

                    // Get a new resolver with the google nameservers as the upstream recursive resolvers
                    TokioAsyncResolver::new(ResolverConfig::google(), ResolverOpts::default(), runtime.handle().clone())
                }
            };

            let &(ref lock, ref cvar) = &*pair2;
            let mut started = lock.lock().unwrap();

            let resolver = runtime.block_on(resolver).expect("failed to create trust-dns-resolver");

            *started = Some(resolver);
            cvar.notify_one();
            drop(started);
        });

        // Wait for the thread to start up.
        let &(ref lock, ref cvar) = &*pair;
        let mut resolver = lock.lock().unwrap();
        while resolver.is_none() {
            resolver = cvar.wait(resolver).unwrap();
        }

        // take the started resolver
        let resolver = std::mem::replace(&mut *resolver, None);

        // set the global resolver
        resolver.expect("resolver should not be none")
    };
}

/// Provide a general purpose resolution function.
///
/// This looks up the `host` (a `&str` or `String` is good), and combines that with the provided port
///   this mimics the lookup functions of `std::net`.
#[cfg(feature = "tokio-runtime")]
pub fn resolve<N: IntoName + TryParseIp + 'static>(
    host: N,
    port: u16,
) -> impl Future<Output = io::Result<Vec<SocketAddr>>> {
    // Now we use the global resolver to perform a lookup_ip.
    GLOBAL_DNS_RESOLVER.lookup_ip(host).map(move |result| {
        // map the result into what we want...
        result
            .map_err(move |err| {
                // we transform the error into a standard IO error for convenience
                io::Error::new(
                    io::ErrorKind::AddrNotAvailable,
                    format!("dns resolution error: {}", err),
                )
            })
            .map(move |lookup_ip| {
                // we take all the IPs returned, and then send back the set of IPs
                lookup_ip
                    .iter()
                    .map(|ip| SocketAddr::new(ip, port))
                    .collect::<Vec<_>>()
            })
    })
}

#[cfg(feature = "tokio-runtime")]
fn main() {
    use std::thread;

    // Let's resolve some names, we should be able to do it across threads
    let names = &["www.google.com", "www.reddit.com", "www.wikipedia.org"];

    // spawn all the threads to do the lookups
    let threads = names
        .iter()
        .map(|name| {
            let join = thread::spawn(move || {
                let mut runtime = tokio::runtime::Runtime::new().expect("failed to launch Runtime");
                runtime.block_on(resolve(*name, 443))
            });

            (name, join)
        })
        .collect::<Vec<_>>();

    // print the resolved IPs
    for (name, join) in threads {
        let result = join
            .join()
            .unwrap_or_else(|_| panic!("error resolving: {}", name));
        println!("{} resolved to {:?}", name, result);
    }
}

#[cfg(not(feature = "tokio-runtime"))]
fn main() {
    println!("tokio-runtime feature must be enabled")
}
<commit_msg>simplify the global resolver example<commit_after>#![recursion_limit = "128"]

#[macro_use]
extern crate lazy_static;
extern crate futures;
extern crate tokio;
extern crate trust_dns_resolver;

#[cfg(feature = "tokio-runtime")]
use std::io;
#[cfg(feature = "tokio-runtime")]
use std::net::SocketAddr;

#[cfg(feature = "tokio-runtime")]
use trust_dns_resolver::TokioAsyncResolver;
#[cfg(feature = "tokio-runtime")]
use trust_dns_resolver::{IntoName, TryParseIp};

// This is an example of registering a static global resolver into any system.
//
// We may want to create a GlobalResolver as part of the Resolver library
//   in the mean time, this example has the necessary steps to do so.
//
// Thank you to @zonyitoo for the original example.
// TODO: this example can probably be made much simpler with the new
//      `AsyncResolver`.
#[cfg(feature = "tokio-runtime")]
lazy_static! {
    // First we need to setup the global Resolver
    static ref GLOBAL_DNS_RESOLVER: TokioAsyncResolver = {
        use std::sync::{Arc, Mutex, Condvar};
        use std::thread;

        // We'll be using this condvar to get the Resolver from the thread...
        let pair = Arc::new((Mutex::new(None::<TokioAsyncResolver>), Condvar::new()));
        let pair2 = pair.clone();

        // Spawn the runtime to a new thread...
        //
        // This thread will manage the actual resolution runtime
        thread::spawn(move || {
            // A runtime for this new thread
            let mut runtime = tokio::runtime::Runtime::new().expect("failed to launch Runtime");

            // our platform independent future, result, see next blocks
            let resolver = {

                // To make this independent, if targeting macOS, BSD, Linux, or Windows, we can use the system's configuration:
                #[cfg(any(unix, windows))]
                {
                    // use the system resolver configuration
                    TokioAsyncResolver::from_system_conf(runtime.handle().clone())
                }

                // For other operating systems, we can use one of the preconfigured definitions
                #[cfg(not(any(unix, windows)))]
                {
                    // Directly reference the config types
                    use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};

                    // Get a new resolver with the google nameservers as the upstream recursive resolvers
                    TokioAsyncResolver::new(ResolverConfig::google(), ResolverOpts::default(), runtime.handle().clone())
                }
            };

            let &(ref lock, ref cvar) = &*pair2;
            let mut started = lock.lock().unwrap();

            let resolver = runtime.block_on(resolver).expect("failed to create trust-dns-resolver");

            *started = Some(resolver);
            cvar.notify_one();
            drop(started);
        });

        // Wait for the thread to start up.
        let &(ref lock, ref cvar) = &*pair;
        let mut resolver = lock.lock().unwrap();
        while resolver.is_none() {
            resolver = cvar.wait(resolver).unwrap();
        }

        // take the started resolver
        let resolver = std::mem::replace(&mut *resolver, None);

        // set the global resolver
        resolver.expect("resolver should not be none")
    };
}

/// Provide a general purpose resolution function.
///
/// This looks up the `host` (a `&str` or `String` is good), and combines that with the provided port
///   this mimics the lookup functions of `std::net`.
#[cfg(feature = "tokio-runtime")]
pub async fn resolve<N: IntoName + TryParseIp + 'static>(
    host: N,
    port: u16,
) -> io::Result<Vec<SocketAddr>> {
    // Now we use the global resolver to perform a lookup_ip.
    let result = GLOBAL_DNS_RESOLVER.lookup_ip(host).await;
    // map the result into what we want...
    result
        .map_err(move |err| {
            // we transform the error into a standard IO error for convenience
            io::Error::new(
                io::ErrorKind::AddrNotAvailable,
                format!("dns resolution error: {}", err),
            )
        })
        .map(move |lookup_ip| {
            // we take all the IPs returned, and then send back the set of IPs
            lookup_ip
                .iter()
                .map(|ip| SocketAddr::new(ip, port))
                .collect::<Vec<_>>()
        })
}

#[cfg(feature = "tokio-runtime")]
fn main() {
    use std::thread;

    // Let's resolve some names, we should be able to do it across threads
    let names = &["www.google.com", "www.reddit.com", "www.wikipedia.org"];

    // spawn all the threads to do the lookups
    let threads = names
        .iter()
        .map(|name| {
            let join = thread::spawn(move || {
                let mut runtime = tokio::runtime::Runtime::new().expect("failed to launch Runtime");
                runtime.block_on(resolve(*name, 443))
            });

            (name, join)
        })
        .collect::<Vec<_>>();

    // print the resolved IPs
    for (name, join) in threads {
        let result = join
            .join()
            .unwrap_or_else(|_| panic!("error resolving: {}", name));
        println!("{} resolved to {:?}", name, result);
    }
}

#[cfg(not(feature = "tokio-runtime"))]
fn main() {
    println!("tokio-runtime feature must be enabled")
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::fmt;

use {Args, Context, MessagePart};

/// A message that has been localized and can be formatted in a
/// locale-aware manner.
///
/// While a `Message` can be created directly from [`MessagePart`]
/// components, it is easiest to create it from [`icu::parse`].
///
/// [`MessagePart`]: trait.MessagePart.html
/// [`icu::parse`]: icu/fn.parse.html
#[derive(Debug)]
pub struct Message {
    parts: Vec<Box<MessagePart>>,
}

impl Message {
    /// Construct a message from constituent parts.
    pub fn new(parts: Vec<Box<MessagePart>>) -> Self {
        Message { parts: parts }
    }

    /// Write a message to a stream.
    ///
    /// This shouldn't be called directly in the usual case.
    /// Use `Context::write` or `Context::format` instead.
    pub fn write_message<'f>(&self,
                             ctx: &Context,
                             stream: &mut fmt::Write,
                             args: Option<&Args<'f>>)
                             -> fmt::Result {
        for part in &self.parts {
            try!(part.apply_format(&ctx, stream, args));
        }
        Ok(())
    }
}
<commit_msg>Fix clippy warning: needless borrow.<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::fmt;

use {Args, Context, MessagePart};

/// A message that has been localized and can be formatted in a
/// locale-aware manner.
///
/// While a `Message` can be created directly from [`MessagePart`]
/// components, it is easiest to create it from [`icu::parse`].
///
/// [`MessagePart`]: trait.MessagePart.html
/// [`icu::parse`]: icu/fn.parse.html
#[derive(Debug)]
pub struct Message {
    parts: Vec<Box<MessagePart>>,
}

impl Message {
    /// Construct a message from constituent parts.
    pub fn new(parts: Vec<Box<MessagePart>>) -> Self {
        Message { parts: parts }
    }

    /// Write a message to a stream.
    ///
    /// This shouldn't be called directly in the usual case.
    /// Use `Context::write` or `Context::format` instead.
    pub fn write_message<'f>(&self,
                             ctx: &Context,
                             stream: &mut fmt::Write,
                             args: Option<&Args<'f>>)
                             -> fmt::Result {
        for part in &self.parts {
            try!(part.apply_format(ctx, stream, args));
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Song analysis module.
//!
//! Use features-extraction functions from other modules
//! e.g. tempo features, spectral features, etc to build an
//! Analysis of a Song
#[cfg(feature = "aubio-lib")]
extern crate aubio_lib;

extern crate crossbeam;
extern crate ndarray;
extern crate ndarray_npy;

use std::f64::consts::PI;

use aubio_rs::vec::CVec;
use aubio_rs::FFT;
use ndarray::ShapeBuilder;
use ndarray::{arr1, s, Array, Array1, Array2};

use crate::chroma::ChromaDesc;
use crate::decode::decode_song;
use crate::misc::LoudnessDesc;
use crate::temporal::BPMDesc;
use crate::timbral::{SpectralDesc, ZeroCrossingRateDesc};
use crate::{Analysis, Song};
use crossbeam::thread;

pub fn decode_and_analyze(path: &str) -> Result<Song, String> {
    // TODO error handling here
    let mut song = decode_song(&path)?;

    song.analysis = analyze(&song);
    Ok(song)
}

fn reflect_pad(array: &[f32], pad: usize) -> Vec<f32> {
    let mut prefix = array[1..=pad].iter().rev().copied().collect::<Vec<f32>>();
    let suffix = array[(array.len() - 2) - pad + 1..array.len() - 1]
        .iter()
        .rev()
        .copied()
        .collect::<Vec<f32>>();
    prefix.extend(array);
    prefix.extend(suffix);
    prefix
}

pub fn stft(signal: &[f32], window_length: usize, hop_length: usize) -> Array2<f64> {
    let mut fft = FFT::new(window_length).unwrap();

    let signal = reflect_pad(&signal, window_length / 2);
    //let mut stft = Array2::zeros((window_length / 2 + 1, 0));
    let mut stft: Vec<f32> = Vec::new();

    // TODO actually have it constant - no need to compute it everytime
    // Periodic, so window_size + 1
    let mut hann_window = Array::zeros(window_length + 1);
    for n in 0..window_length {
        hann_window[[n]] = 0.5 - 0.5 * f64::cos(2. * n as f64 * PI / (window_length as f64));
    }
    // TODO no need for count
    let mut count = 0;
    hann_window = hann_window.slice_move(s![0..window_length]);
    for i in 1..signal.len() {
        if i >= window_length && (i - window_length) % hop_length == 0 {
            let beginning = i - window_length;
            let end = i;
            let mut fftgrain: Vec<f32> = vec![0.0; window_length + 2];
            // TODO directly apply hann window maybe
            let signal = (arr1(&signal[beginning..end]).mapv(f64::from) * &hann_window)
                .mapv(|x| x as f32)
                .to_vec();
            fft.do_(&signal, fftgrain.as_mut_slice()).unwrap();
            let cvec: CVec = fftgrain.as_slice().into();
            let norm: Array1<f32> = arr1(&cvec.norm());
            stft.append(&mut norm.to_vec());
            count += 1;
        }
    }
    let stft = Array::from_shape_vec(
        (window_length / 2 + 1, count).strides((1, window_length / 2 + 1)),
        stft,
    )
    .unwrap();
    stft.mapv(f64::from)
}

pub fn analyze(song: &Song) -> Analysis {
    let mut spectral_desc = SpectralDesc::new(song.sample_rate);
    let mut zcr_desc = ZeroCrossingRateDesc::default();
    let mut tempo_desc = BPMDesc::new(song.sample_rate);
    let mut loudness_desc = LoudnessDesc::default();

    thread::scope(|s| {
        let child = s.spawn(|_| {
        let mut chroma_desc = ChromaDesc::new(song.sample_rate, 12);
            chroma_desc.do_(&song.sample_array);
            chroma_desc.get_values()
        });

        // These descriptors can be streamed
        for i in 1..song.sample_array.len() {
            if (i % SpectralDesc::HOP_SIZE) == 0 {
                let beginning = (i / SpectralDesc::HOP_SIZE - 1) * SpectralDesc::HOP_SIZE;
                let end = i;
                spectral_desc.do_(&song.sample_array[beginning..end]);
                zcr_desc.do_(&song.sample_array[beginning..end]);
            }

            if (i % BPMDesc::HOP_SIZE) == 0 {
                let beginning = (i / BPMDesc::HOP_SIZE - 1) * BPMDesc::HOP_SIZE;
                let end = i;
                tempo_desc.do_(&song.sample_array[beginning..end]);
            }

            // Contiguous windows, so WINDOW_SIZE here
            if (i % LoudnessDesc::WINDOW_SIZE) == 0 {
                let beginning = (i / LoudnessDesc::WINDOW_SIZE - 1) * LoudnessDesc::WINDOW_SIZE;
                let end = i;
                loudness_desc.do_(&song.sample_array[beginning..end]);
            }
        }
        // Non-streaming approach for that one
        let (is_major, fifth) = child.join().unwrap();

        Analysis {
            tempo: tempo_desc.get_value(),
            spectral_centroid: spectral_desc.get_centroid(),
            zero_crossing_rate: zcr_desc.get_value(),
            spectral_rolloff: spectral_desc.get_rolloff(),
            spectral_flatness: spectral_desc.get_flatness(),
            loudness: loudness_desc.get_value(),
            is_major,
            fifth,
        }
    }).unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decode::decode_song;
    use ndarray::Array2;
    use ndarray_npy::ReadNpyExt;
    use std::f32::consts::PI;
    use std::fs::File;

    #[test]
    fn test_analyze() {
        let song = decode_song("data/s16_mono_22_5kHz.flac").unwrap();
        let expected_analysis = Analysis {
            tempo: 0.37860596,
            spectral_centroid: -0.75483,
            zero_crossing_rate: -0.85036564,
            spectral_rolloff: -0.6326486,
            spectral_flatness: -0.77610075,
            loudness: 0.27126348,
            is_major: -1.,
            fifth: (f32::cos(5. * PI / 3.), f32::sin(5. * PI / 3.)),
        };
        assert!(expected_analysis.approx_eq(&analyze(&song)));
    }

    #[test]
    fn test_compute_stft() {
        let file = File::open("data/librosa-stft.npy").unwrap();
        let expected_stft = Array2::<f32>::read_npy(file).unwrap().mapv(|x| x as f64);

        let song = decode_song("data/piano.flac").unwrap();

        let stft = stft(&song.sample_array, 2048, 512);

        assert!(!stft.is_empty() && !expected_stft.is_empty());
        for (expected, actual) in expected_stft.iter().zip(stft.iter()) {
            assert!(0.0001 > (expected - actual).abs());
        }
    }
}
<commit_msg>Optimize stft computation<commit_after>//! Song analysis module.
//!
//! Use features-extraction functions from other modules
//! e.g. tempo features, spectral features, etc to build an
//! Analysis of a Song
#[cfg(feature = "aubio-lib")]
extern crate aubio_lib;

extern crate crossbeam;
extern crate ndarray;
extern crate ndarray_npy;

use std::f32::consts::PI as PI;

use ndarray::{arr1, s, Array, Array2};
use rustfft::num_complex::Complex;
use rustfft::num_traits::Zero;
use rustfft::FFTplanner;

use crate::chroma::ChromaDesc;
use crate::decode::decode_song;
use crate::misc::LoudnessDesc;
use crate::temporal::BPMDesc;
use crate::timbral::{SpectralDesc, ZeroCrossingRateDesc};
use crate::{Analysis, Song};
use crossbeam::thread;

pub fn decode_and_analyze(path: &str) -> Result<Song, String> {
    // TODO error handling here
    let mut song = decode_song(&path)?;

    song.analysis = analyze(&song);
    Ok(song)
}

fn reflect_pad(array: &[f32], pad: usize) -> Vec<f32> {
    let mut prefix = array[1..=pad].iter().rev().copied().collect::<Vec<f32>>();
    let suffix = array[(array.len() - 2) - pad + 1..array.len() - 1]
        .iter()
        .rev()
        .copied()
        .collect::<Vec<f32>>();
    prefix.extend(array);
    prefix.extend(suffix);
    prefix
}

pub fn stft(signal: &[f32], window_length: usize, hop_length: usize) -> Array2<f64> {
    let mut stft = Array2::zeros((
        window_length / 2 + 1,
        (signal.len() as f32 / hop_length as f32).ceil() as usize,
    ));
    let signal = reflect_pad(&signal, window_length / 2);

    // TODO actually have it constant - no need to compute it everytime
    // Periodic, so window_size + 1
    let mut hann_window = Array::zeros(window_length + 1);
    for n in 0..window_length {
        hann_window[[n]] = 0.5 - 0.5 * f32::cos(2. * n as f32 * PI / (window_length as f32));
    }
    hann_window = hann_window.slice_move(s![0..window_length]);
    let mut output_window = Array::from_elem(window_length, Complex::zero());
    let mut planner = FFTplanner::new(false);
    let fft = planner.plan_fft(window_length);

    for (window, mut stft_col) in signal
        .windows(window_length)
        .step_by(hop_length)
        .zip(stft.gencolumns_mut())
    {
        let mut signal = (arr1(&window) * &hann_window).mapv(|x| Complex::new(x as f32, 0.));
        fft.process(
            signal.as_slice_mut().unwrap(),
            output_window.as_slice_mut().unwrap(),
        );
        stft_col.assign(
            &output_window
                .slice(s![..window_length / 2 + 1])
                .mapv(|x| x.norm() as f64),
        );
    }
    stft
}

pub fn analyze(song: &Song) -> Analysis {
    let mut spectral_desc = SpectralDesc::new(song.sample_rate);
    let mut zcr_desc = ZeroCrossingRateDesc::default();
    let mut tempo_desc = BPMDesc::new(song.sample_rate);
    let mut loudness_desc = LoudnessDesc::default();

    thread::scope(|s| {
        let child = s.spawn(|_| {
            let mut chroma_desc = ChromaDesc::new(song.sample_rate, 12);
            chroma_desc.do_(&song.sample_array);
            chroma_desc.get_values()
        });

        // These descriptors can be streamed
        for i in 1..song.sample_array.len() {
            if (i % SpectralDesc::HOP_SIZE) == 0 {
                let beginning = (i / SpectralDesc::HOP_SIZE - 1) * SpectralDesc::HOP_SIZE;
                let end = i;
                spectral_desc.do_(&song.sample_array[beginning..end]);
                zcr_desc.do_(&song.sample_array[beginning..end]);
            }

            if (i % BPMDesc::HOP_SIZE) == 0 {
                let beginning = (i / BPMDesc::HOP_SIZE - 1) * BPMDesc::HOP_SIZE;
                let end = i;
                tempo_desc.do_(&song.sample_array[beginning..end]);
            }

            // Contiguous windows, so WINDOW_SIZE here
            if (i % LoudnessDesc::WINDOW_SIZE) == 0 {
                let beginning = (i / LoudnessDesc::WINDOW_SIZE - 1) * LoudnessDesc::WINDOW_SIZE;
                let end = i;
                loudness_desc.do_(&song.sample_array[beginning..end]);
            }
        }
        // Non-streaming approach for that one
        let (is_major, fifth) = child.join().unwrap();

        Analysis {
            tempo: tempo_desc.get_value(),
            spectral_centroid: spectral_desc.get_centroid(),
            zero_crossing_rate: zcr_desc.get_value(),
            spectral_rolloff: spectral_desc.get_rolloff(),
            spectral_flatness: spectral_desc.get_flatness(),
            loudness: loudness_desc.get_value(),
            is_major,
            fifth,
        }
    })
    .unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decode::decode_song;
    use ndarray::{arr1, Array2};
    use ndarray_npy::ReadNpyExt;
    use std::f32::consts::PI;
    use std::fs::File;

    #[test]
    fn test_analyze() {
        let song = decode_song("data/s16_mono_22_5kHz.flac").unwrap();
        let expected_analysis = Analysis {
            tempo: 0.37860596,
            spectral_centroid: -0.75483,
            zero_crossing_rate: -0.85036564,
            spectral_rolloff: -0.6326486,
            spectral_flatness: -0.77610075,
            loudness: 0.27126348,
            is_major: -1.,
            fifth: (f32::cos(5. * PI / 3.), f32::sin(5. * PI / 3.)),
        };
        assert!(expected_analysis.approx_eq(&analyze(&song)));
    }

    #[test]
    fn test_compute_stft() {
        let file = File::open("data/librosa-stft.npy").unwrap();
        let expected_stft = Array2::<f32>::read_npy(file).unwrap().mapv(|x| x as f64);

        let song = decode_song("data/piano.flac").unwrap();

        let stft = stft(&song.sample_array, 2048, 512);

        assert!(!stft.is_empty() && !expected_stft.is_empty());
        for (expected, actual) in expected_stft.iter().zip(stft.iter()) {
            assert!(0.0001 > (expected - actual).abs());
        }
    }

    #[test]
    fn test_foo() {
        let array = arr1(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
        for w in array.windows(4).into_iter().step_by(2) {
            println!("{:?}", w);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Walter Kuppens.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use debugger::debugger::Debugger;
use io::binutils::INESHeader;
use io::errors::*;
use io::log;
use nes::cpu::CPU;
use nes::ppu::PPU;
use std::fs::File;
use std::io::{self, stdin, Read, Write, BufReader, BufRead};
use std::sync::mpsc::{self, SyncSender, Receiver};
use std::{thread, panic};
use rustyline::error::ReadlineError;
use rustyline::Editor;

use nes::memory::{
    Memory,
    TRAINER_START,
    TRAINER_SIZE,
    PRG_ROM_1_START,
    PRG_ROM_2_START,
    PRG_ROM_SIZE
};

const HISTORY_FILE: &'static str = ".nes-rs-history.txt";

/// The NES struct owns all hardware peripherals and lends them when needed. The
/// runtime cost of this should be removed with optimized builds (untested).
pub struct NES {
    pub header: INESHeader,
    pub runtime_options: NESRuntimeOptions,
    pub cpu: CPU,
    pub ppu: PPU,
    pub memory: Memory,
}

impl NES {
    /// Initializes the NES emulator by dumping the ROM into memory and
    /// initializing the initial hardware state.
    pub fn new(rom: Vec<u8>, header: INESHeader, runtime_options: NESRuntimeOptions) -> Self {
        // An offset is used when copying from the ROM into RAM as the presence
        // of a trainer will shift the locations of other structures.
        let mut cursor: usize = 0x10;

        // Spew out some useful metadata about the rom when verbose is on.
        log::log("init", format!("Using {:?} mapper", header.mapper()), &runtime_options);
        log::log("init", format!("Using {:?} mirroring", header.mirror_type()), &runtime_options);

        // Copy the trainer data to 0x7000 if it exists and adjust the cursor
        // size to accommodate. Trainer data will offset the location of ROM
        // data in the INES ROM file.
        let mut memory = Memory::new();
        if header.has_trainer() {
            log::log("init", "Trainer data found", &runtime_options);
            memory.memdump(TRAINER_START, &rom[0x10..0x210]);
            cursor += TRAINER_SIZE;
        }

        // Copy PRG-ROM into memory so it can be addressed by the chosen memory
        // mapper. PRG-ROM bank 1 begins at 0x8000 and bank 2 begins at 0xC000.
        //
        // In the event that there are 2 PRG-ROM banks, make both banks
        // addressable at their respective locations. However if there's only
        // one bank, make PRG-ROM bank 1 addressable starting from both
        // addresses.
        //
        // NOTE: Should this be moved to mapper code?
        if header.prg_rom_size == 2 {
            log::log("init", "2 PRG-ROM banks detected", &runtime_options);
            let prg_rom_1_addr = cursor;
            let prg_rom_2_addr = cursor + PRG_ROM_SIZE;
            memory.memdump(PRG_ROM_1_START, &rom[prg_rom_1_addr..prg_rom_1_addr + PRG_ROM_SIZE]);
            memory.memdump(PRG_ROM_2_START, &rom[prg_rom_2_addr..prg_rom_2_addr + PRG_ROM_SIZE]);
        } else {
            log::log("init", "1 PRG-ROM bank detected", &runtime_options);
            let prg_rom_1_addr = cursor;
            memory.memdump(PRG_ROM_1_START, &rom[prg_rom_1_addr..prg_rom_1_addr + PRG_ROM_SIZE]);
            memory.memdump(PRG_ROM_2_START, &rom[prg_rom_1_addr..prg_rom_1_addr + PRG_ROM_SIZE]);
        }

        // Set the initial program counter to the address stored at 0xFFFC (this
        // allows ROMs to specify entry point). If a program counter was
        // specified on the command-line, use that one instead.
        let pc = match runtime_options.program_counter {
            Some(pc) => pc,
            None => {
                memory.read_u16(0xFFFC)
            },
        };

        NES {
            header: header,
            cpu: CPU::new(runtime_options.clone(), pc),
            ppu: PPU::new(runtime_options.clone()),
            runtime_options: runtime_options,
            memory: memory,
        }
    }

    /// Starts the execution loop and starts executing PRG-ROM.
    pub fn run(&mut self) -> i32 {
        // Put the CPU into testing mode if a cpu log was passed in the runtime
        // options. This is done before execution so the log and the CPU state
        // are kept in sync.
        match self.runtime_options.cpu_log {
            Some(ref filename) => {
                match File::open(filename) {
                    Ok(f) => {
                        self.cpu.begin_testing(BufReader::new(f))
                    },
                    Err(e) => {
                        let mut stderr = io::stderr();
                        writeln!(stderr, "nes-rs: cannot open {}: {}", filename, e).unwrap();
                        return EXIT_CPU_LOG_NOT_FOUND;
                    },
                }
            },
            None => {},
        }

        // Start cycling the CPU and PPU and add a panic catcher so crash
        // information can be shown if the CPU panics.The PPU ticks three times
        // every CPU cycle, though there may need to be changes made for PAL
        // (currently assumes NTSC PPU clock speed).
        //
        // Depending on the runtime environment, execution can go one of two
        // ways. Either the virtual machine step function is called in an
        // infinite loop, or the debugger handles execution if the debug flag is
        // set.
        //
        // In debug mode, there is another step function that wraps the main
        // step function that lets the debugger control execution flow and
        // access virtual machine state. Another thread is also setup that waits
        // for input on stdin that sends input to the debugger for the debugger
        // subshell.
        let result = panic::catch_unwind(panic::AssertUnwindSafe(|| {
            if self.runtime_options.debugging {
                let (tx, rx): (SyncSender<String>, Receiver<String>) = mpsc::sync_channel(1);
                let (mtx, mrx): (SyncSender<u8>, Receiver<u8>) = mpsc::sync_channel(1);

                // Input is read on another thread, so spin one up.
                self.setup_readline_thread(tx, mrx);

                // Execute until shutdown signal is received from debugger.
                let mut debugger = Debugger::new(mtx, rx);
                while !debugger.step(self) {}
            } else {
                loop {
                    self.step();
                }
            }
        }));

        // Unwinding point with shutdown code. In the event of a panic, we want
        // to display some diagnostic information to the user that can be sent
        // to the developer.
        match result {
            Ok(_) => {
                println!("Shutting down nes-rs, happy emulating!");
                return EXIT_SUCCESS; // Success exit code.
            },
            Err(_) => {
                println!("{}", self.cpu);
                return EXIT_RUNTIME_FAILURE; // Runtime failure exit code.
            }
        }
    }

    /// Executes a CPU instruction and steps the PPU 3 times per CPU cycle. This
    /// works since the PPU and CPU clocks are synchronized 1 to 3.
    pub fn step(&mut self) {
        let mut cycles = self.cpu.step(&mut self.memory);
        self.cpu.sleep(cycles);

        while cycles > 0 {
            for _ in 0..3 { // *Should* unroll.
                self.ppu.step(&mut self.memory);
            }
            cycles -= 1;
        }
    }

    /// Creates a readline loop on another thread and sends commands to the
    /// debugger over a synchronous rust channel. Offers quality of life features
    /// such as history built into the library used.
    fn setup_readline_thread(&self, tx: SyncSender<String>, rx: Receiver<u8>) {
        thread::spawn(move || {
            let mut rl = Editor::<()>::new();
            if let Err(_) = rl.load_history(HISTORY_FILE) {
                // No history saved, do nothing.
            }

            loop {
                let readline = rl.readline("(nes-rs) ");
                match readline {
                    Ok(line) => {
                        rl.add_history_entry(&line);
                        tx.send(line).unwrap();

                        // Block until command is run.
                        match rx.recv() {
                            Ok(_) => {},
                            Err(err) => {
                                println!("Error: {:?}", err);
                                tx.send("exit".to_string()).unwrap();
                                break;
                            },
                        }
                    },
                    Err(ReadlineError::Interrupted) => {
                        tx.send("exit".to_string()).unwrap();
                        break;
                    },
                    Err(ReadlineError::Eof) => {
                        tx.send("exit".to_string()).unwrap();
                        break;
                    },
                    Err(err) => {
                        println!("Error: {:?}", err);
                        tx.send("exit".to_string()).unwrap();
                        break;
                    },
                };
            }

            rl.save_history(HISTORY_FILE).unwrap();
        });
    }
}

/// Flags and other information set through command-line arguments.
#[derive(Clone, Debug)]
pub struct NESRuntimeOptions {
    pub program_counter: Option<u16>,
    pub cpu_log:         Option<String>,
    pub verbose:         bool,
    pub debugging:       bool,
}
<commit_msg>Do not send exit after channel receive error<commit_after>// Copyright 2016 Walter Kuppens.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use debugger::debugger::Debugger;
use io::binutils::INESHeader;
use io::errors::*;
use io::log;
use nes::cpu::CPU;
use nes::ppu::PPU;
use std::fs::File;
use std::io::{self, stdin, Read, Write, BufReader, BufRead};
use std::sync::mpsc::{self, SyncSender, Receiver};
use std::{thread, panic};
use std::time::Duration;
use rustyline::error::ReadlineError;
use rustyline::Editor;

use nes::memory::{
    Memory,
    TRAINER_START,
    TRAINER_SIZE,
    PRG_ROM_1_START,
    PRG_ROM_2_START,
    PRG_ROM_SIZE
};

const HISTORY_FILE: &'static str = ".nes-rs-history.txt";

/// The NES struct owns all hardware peripherals and lends them when needed. The
/// runtime cost of this should be removed with optimized builds (untested).
pub struct NES {
    pub header: INESHeader,
    pub runtime_options: NESRuntimeOptions,
    pub cpu: CPU,
    pub ppu: PPU,
    pub memory: Memory,
}

impl NES {
    /// Initializes the NES emulator by dumping the ROM into memory and
    /// initializing the initial hardware state.
    pub fn new(rom: Vec<u8>, header: INESHeader, runtime_options: NESRuntimeOptions) -> Self {
        // An offset is used when copying from the ROM into RAM as the presence
        // of a trainer will shift the locations of other structures.
        let mut cursor: usize = 0x10;

        // Spew out some useful metadata about the rom when verbose is on.
        log::log("init", format!("Using {:?} mapper", header.mapper()), &runtime_options);
        log::log("init", format!("Using {:?} mirroring", header.mirror_type()), &runtime_options);

        // Copy the trainer data to 0x7000 if it exists and adjust the cursor
        // size to accommodate. Trainer data will offset the location of ROM
        // data in the INES ROM file.
        let mut memory = Memory::new();
        if header.has_trainer() {
            log::log("init", "Trainer data found", &runtime_options);
            memory.memdump(TRAINER_START, &rom[0x10..0x210]);
            cursor += TRAINER_SIZE;
        }

        // Copy PRG-ROM into memory so it can be addressed by the chosen memory
        // mapper. PRG-ROM bank 1 begins at 0x8000 and bank 2 begins at 0xC000.
        //
        // In the event that there are 2 PRG-ROM banks, make both banks
        // addressable at their respective locations. However if there's only
        // one bank, make PRG-ROM bank 1 addressable starting from both
        // addresses.
        //
        // NOTE: Should this be moved to mapper code?
        if header.prg_rom_size == 2 {
            log::log("init", "2 PRG-ROM banks detected", &runtime_options);
            let prg_rom_1_addr = cursor;
            let prg_rom_2_addr = cursor + PRG_ROM_SIZE;
            memory.memdump(PRG_ROM_1_START, &rom[prg_rom_1_addr..prg_rom_1_addr + PRG_ROM_SIZE]);
            memory.memdump(PRG_ROM_2_START, &rom[prg_rom_2_addr..prg_rom_2_addr + PRG_ROM_SIZE]);
        } else {
            log::log("init", "1 PRG-ROM bank detected", &runtime_options);
            let prg_rom_1_addr = cursor;
            memory.memdump(PRG_ROM_1_START, &rom[prg_rom_1_addr..prg_rom_1_addr + PRG_ROM_SIZE]);
            memory.memdump(PRG_ROM_2_START, &rom[prg_rom_1_addr..prg_rom_1_addr + PRG_ROM_SIZE]);
        }

        // Set the initial program counter to the address stored at 0xFFFC (this
        // allows ROMs to specify entry point). If a program counter was
        // specified on the command-line, use that one instead.
        let pc = match runtime_options.program_counter {
            Some(pc) => pc,
            None => {
                memory.read_u16(0xFFFC)
            },
        };

        NES {
            header: header,
            cpu: CPU::new(runtime_options.clone(), pc),
            ppu: PPU::new(runtime_options.clone()),
            runtime_options: runtime_options,
            memory: memory,
        }
    }

    /// Starts the execution loop and starts executing PRG-ROM.
    pub fn run(&mut self) -> i32 {
        // Put the CPU into testing mode if a cpu log was passed in the runtime
        // options. This is done before execution so the log and the CPU state
        // are kept in sync.
        match self.runtime_options.cpu_log {
            Some(ref filename) => {
                match File::open(filename) {
                    Ok(f) => {
                        self.cpu.begin_testing(BufReader::new(f))
                    },
                    Err(e) => {
                        let mut stderr = io::stderr();
                        writeln!(stderr, "nes-rs: cannot open {}: {}", filename, e).unwrap();
                        return EXIT_CPU_LOG_NOT_FOUND;
                    },
                }
            },
            None => {},
        }

        // Start cycling the CPU and PPU and add a panic catcher so crash
        // information can be shown if the CPU panics.The PPU ticks three times
        // every CPU cycle, though there may need to be changes made for PAL
        // (currently assumes NTSC PPU clock speed).
        //
        // Depending on the runtime environment, execution can go one of two
        // ways. Either the virtual machine step function is called in an
        // infinite loop, or the debugger handles execution if the debug flag is
        // set.
        //
        // In debug mode, there is another step function that wraps the main
        // step function that lets the debugger control execution flow and
        // access virtual machine state. Another thread is also setup that waits
        // for input on stdin that sends input to the debugger for the debugger
        // subshell.
        let result = panic::catch_unwind(panic::AssertUnwindSafe(|| {
            if self.runtime_options.debugging {
                let (tx, rx): (SyncSender<String>, Receiver<String>) = mpsc::sync_channel(1);
                let (mtx, mrx): (SyncSender<u8>, Receiver<u8>) = mpsc::sync_channel(1);

                // Input is read on another thread, so spin one up.
                self.setup_readline_thread(tx, mrx);

                // Execute until shutdown signal is received from debugger.
                let mut debugger = Debugger::new(mtx, rx);
                while !debugger.step(self) {}
            } else {
                loop {
                    self.step();
                }
            }
        }));

        // Unwinding point with shutdown code. In the event of a panic, we want
        // to display some diagnostic information to the user that can be sent
        // to the developer.
        match result {
            Ok(_) => {
                println!("Shutting down nes-rs, happy emulating!");
                return EXIT_SUCCESS; // Success exit code.
            },
            Err(_) => {
                thread::sleep(Duration::from_millis(16));
                println!("{}", self.cpu);
                return EXIT_RUNTIME_FAILURE; // Runtime failure exit code.
            }
        }
    }

    /// Executes a CPU instruction and steps the PPU 3 times per CPU cycle. This
    /// works since the PPU and CPU clocks are synchronized 1 to 3.
    pub fn step(&mut self) {
        let mut cycles = self.cpu.step(&mut self.memory);
        self.cpu.sleep(cycles);

        while cycles > 0 {
            for _ in 0..3 { // *Should* unroll.
                self.ppu.step(&mut self.memory);
            }
            cycles -= 1;
        }
    }

    /// Creates a readline loop on another thread and sends commands to the
    /// debugger over a synchronous rust channel. Offers quality of life features
    /// such as history built into the library used.
    fn setup_readline_thread(&self, tx: SyncSender<String>, rx: Receiver<u8>) {
        thread::spawn(move || {
            let mut rl = Editor::<()>::new();
            if let Err(_) = rl.load_history(HISTORY_FILE) {
                // No history saved, do nothing.
            }

            loop {
                let readline = rl.readline("(nes-rs) ");
                match readline {
                    Ok(line) => {
                        rl.add_history_entry(&line);
                        tx.send(line).unwrap();

                        // Block until command is run.
                        match rx.recv() {
                            Ok(_) => {},
                            Err(_) => {
                                break;
                            },
                        }
                    },
                    Err(ReadlineError::Interrupted) => {
                        tx.send("exit".to_string()).unwrap();
                        break;
                    },
                    Err(ReadlineError::Eof) => {
                        tx.send("exit".to_string()).unwrap();
                        break;
                    },
                    Err(err) => {
                        println!("Error: {:?}", err);
                        tx.send("exit".to_string()).unwrap();
                        break;
                    },
                };
            }

            rl.save_history(HISTORY_FILE).unwrap();
        });
    }
}

/// Flags and other information set through command-line arguments.
#[derive(Clone, Debug)]
pub struct NESRuntimeOptions {
    pub program_counter: Option<u16>,
    pub cpu_log:         Option<String>,
    pub verbose:         bool,
    pub debugging:       bool,
}
<|endoftext|>


<commit_before>use std::convert::From;
use std::io::{Read, Write};
use std::sync::mpsc::{Sender, channel};
use std::fmt;
use std::fmt::{Display, Formatter};
use std::error::Error;

use VERSION;

use bincode::serde::{DeserializeError, deserialize, serialize};
use bincode::SizeLimit;
use byteorder::{ByteOrder, LittleEndian, WriteBytesExt};
use mio::{EventLoop as MioEventLoop, EventSet, Handler as MioHandler, PollOpt, Token};
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;
use slab::Index;

pub mod client;
pub mod server;
#[cfg(Test)]
pub mod test;

/// The default port for a server to listen on.
pub const STANDARD_PORT: u16 = 25566;
/// Standard number to ensure network connections are syncronized and the same protocol is being used.
/// Reexported incase it is of use for something not-networking.
pub const NET_MAGIC_NUMBER: u32 = 0xCB011043; //0xcafebade + 0x25565, because programming references.
const MAX_CONNECTIONS: usize = 1024;

pub enum HandlerMessage {
    Send(NetworkPacket, Token),
    AddStream(TcpStream, Sender<Token>),
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum NetworkPacket {
    /// Sent on connection to verify everything is in sync.
    Init {
        version: String,
    },
    /// An error that should crash the game and show an error to the user, but only on a client.
    Error(NetworkError),
    /// Used to unit test networking. When recived, increments a value. TODO: Locate that value.
    #[cfg(test)]
    Test,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum NetworkError {
    VersionMismatch(String, String),
}

impl Display for NetworkError {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        match *self {
            NetworkError::VersionMismatch(ref ver1, ref ver2) => {
                write!(fmt,
                       "VersionMismatch: The versions of the client and server attempting to connect mismatch. ver1: {}, ver2: {}",
                       ver1,
                       ver2)
            }
        }
    }
}

impl Error for NetworkError {
    fn description(&self) -> &str {
        match *self {
            NetworkError::VersionMismatch(_, _) => "VersionMismatch: The versions of the client and server attempting to connect mismatch.",
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            NetworkError::VersionMismatch(_, _) => None,
        }
    }
}

pub type EventLoop = MioEventLoop<Handler>;
pub struct Handler {
    connections: Slab<Connection>,
    listener: Option<TcpListener>,
}

impl Handler {
    /// Creates a new instance of a Handler. Does not listen for connections.
    pub fn new() -> Self {
        Handler {
            connections: Slab::new_starting_at(Token::from_usize(1), MAX_CONNECTIONS),
            listener: None,
        }
    }

    /// Creates a new instance of a Handler. Does listen for connections.
    pub fn new_listener() -> Self {
        unimplemented!();
    }
}

#[derive(Debug)]
struct Connection {
    message_queue: Vec<NetworkPacket>,
    stream: TcpStream,
}

impl Connection {
    fn new(stream: TcpStream) -> Connection {
        Connection {
            stream: stream,
            message_queue: Vec::new(),
        }
    }
}

impl MioHandler for Handler {
    type Timeout = ();
    type Message = HandlerMessage;
    fn ready(&mut self, event_loop: &mut MioEventLoop<Handler>, token: Token, events: EventSet) {
        if events.is_writable() {
            // Get all the messages that should be send to that token, seralize all of them, then send all of them.
            // Afterwards, flush the buffer.
            let to_send = self.connections[token].message_queue.clone();
            self.connections[token].message_queue = Vec::new();
            if !to_send.is_empty() {
                for send in to_send {
                    let bytes = seralize_packet(&send);
                    self.connections[token]
                        .stream
                        .write(&bytes)
                        .expect(&format!("Error writing {:?} to connection {:?}/{:?}",
                                         send,
                                         token,
                                         self.connections[token]));
                }
                self.connections[token]
                    .stream
                    .flush()
                    .expect(&format!("Error flushing connection {:?}/{:?}",
                                     token,
                                     self.connections[token]));;
            }
        }

        if events.is_readable() {
            // Read the 6 byte header of each packet, throw it into get_packet_length, then read that number of bytes.
            // Then throw those bytes into deserialize_packet. Afterwards, throw it to handle_packet.

            let mut header = [0; 6];
            self.connections[token]
                .stream
                .read(&mut header)
                .expect(&format!("An error occured reading from socket {:?}", token));    // TODO: Figure out possible errors and take care of them.
            let length = get_packet_length(header).unwrap();    // TODO: Handle error gracefully.
            let mut packet = Vec::new();
            (&mut self.connections[token].stream).take(length as u64).read_to_end(&mut packet).unwrap();
            // I do this because Read.take takes a self, instead of a reasonable alternitive.
            // However &mut Read is it's self a Reader. So I use that instead.
            let dese = deserialize_packet(&packet).unwrap();
            handle_packet(dese, token, event_loop);
        }

        // This may need to be in an if events.is_readable() block.
        if self.listener.is_some() {
            match self.listener.as_ref().unwrap().accept() {
                Ok(Some((socket, address))) => {
                    let token = self.connections.insert(Connection::new(socket)).unwrap();
                    send(event_loop,
                         NetworkPacket::Init { version: VERSION.to_owned() },
                         token);
                    info!("Accepted connection {}.", address);
                }
                Ok(None) => debug!("Recived Ok(None) in listener.accept()."),
                Err(error) => panic!("listener.accept() errored with {:?}!", error),
            }
        }
    }

    fn notify(&mut self, event_loop: &mut MioEventLoop<Handler>, msg: HandlerMessage) {
        match msg {
            HandlerMessage::AddStream(stream, tx) => {
                let token = self.connections.insert(Connection::new(stream)).unwrap();
                event_loop.register(&self.connections[token].stream,
                                    token,
                                    EventSet::all(),
                                    PollOpt::level())
                          .unwrap();
                tx.send(token).unwrap();
            }
            HandlerMessage::Send(packet, token) => {
                self.connections[token].message_queue.push(packet);
            }
        }
    }
}

fn add_socket(event_loop: &EventLoop, socket: TcpStream) -> Token {
    let (tx, rx) = channel();   // How expensive is this? Should I be creating a whole new channel for just 1 message?
    event_loop.channel().send(HandlerMessage::AddStream(socket, tx)).unwrap();
    rx.recv().unwrap() // I should feel bad.
}

/// Send a NetworkPacket over the network.
///
/// This is a static function instead of an impl function because you can't impl on external structs. In this case, mio::EventLoop.
pub fn send(event_loop: &EventLoop, to_send: NetworkPacket, token: Token) {
    event_loop.channel().send(HandlerMessage::Send(to_send, token)).unwrap();
}

fn seralize_packet(to_ser: &NetworkPacket) -> Vec<u8> {
    let mut result: Vec<u8> = Vec::new();
    result.write_u32::<LittleEndian>(NET_MAGIC_NUMBER).unwrap();   // No possible errors here.
    // The NET_MAGIC_NUMBER is used before every packet, so incase the stream is desynced for whatever reason, the game doesn't just read arbratrary data and crash badly.
    // Instead, it can either recover somehow, by disconnecting and reconnecting, or just erroring gracefully.
    let mut encoded = serialize(to_ser, SizeLimit::Infinite).unwrap();
    // Since the size limit is infinite and i'm not encoding to a stream, there is no error and I can safely unwrap();
    result.write_u16::<LittleEndian>(encoded.len() as u16).unwrap();
    result.append(&mut encoded);
    result
}

#[derive(Debug)]
enum PacketDeseError {
    MagicNumberMismatch,
    InvalidEncoding(DeserializeError),
}

impl From<DeserializeError> for PacketDeseError {
    fn from(err: DeserializeError) -> PacketDeseError {
        PacketDeseError::InvalidEncoding(err)
    }
}

/// Returns the length of a given packet, or an error if the first four bytes do not match NET_MAGIC_NUMBER.
fn get_packet_length(to_ln: [u8; 6]) -> Result<u16, PacketDeseError> {
    let (first_four, next_two) = to_ln.split_at(4);
    let should_be_magic_num = LittleEndian::read_u32(&first_four);
    if should_be_magic_num != NET_MAGIC_NUMBER {
        return Err(PacketDeseError::MagicNumberMismatch);
    }
    let length = LittleEndian::read_u16(&next_two);
    Ok(length)
}

fn deserialize_packet(to_de: &[u8]) -> Result<NetworkPacket, PacketDeseError> {
    Ok(try!(deserialize::<NetworkPacket>(to_de)))
}

fn handle_packet(to_handle: NetworkPacket, sender: Token, event_loop: &EventLoop) {
    match to_handle {
        NetworkPacket::Init{version} => {
            if version != VERSION {
                send(event_loop,
                     NetworkPacket::Error(NetworkError::VersionMismatch(version.to_owned(), VERSION.to_owned())),
                     sender)
            }
        }

        #[cfg(Test)]
        NetworkPacket::Test => {
            unimplemented!()
        }

        NetworkPacket::Error(error) => if ::check_should_crash() { panic!(error) } else { unimplemented!() },
    }
}
<commit_msg>cargo fmt<commit_after>use std::convert::From;
use std::io::{Read, Write};
use std::sync::mpsc::{Sender, channel};
use std::fmt;
use std::fmt::{Display, Formatter};
use std::error::Error;

use VERSION;

use bincode::serde::{DeserializeError, deserialize, serialize};
use bincode::SizeLimit;
use byteorder::{ByteOrder, LittleEndian, WriteBytesExt};
use mio::{EventLoop as MioEventLoop, EventSet, Handler as MioHandler, PollOpt, Token};
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;
use slab::Index;

pub mod client;
pub mod server;
#[cfg(Test)]
pub mod test;

/// The default port for a server to listen on.
pub const STANDARD_PORT: u16 = 25566;
/// Standard number to ensure network connections are syncronized and the same protocol is being used.
/// Reexported incase it is of use for something not-networking.
pub const NET_MAGIC_NUMBER: u32 = 0xCB011043; //0xcafebade + 0x25565, because programming references.
const MAX_CONNECTIONS: usize = 1024;

pub enum HandlerMessage {
    Send(NetworkPacket, Token),
    AddStream(TcpStream, Sender<Token>),
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum NetworkPacket {
    /// Sent on connection to verify everything is in sync.
    Init {
        version: String,
    },
    /// An error that should crash the game and show an error to the user, but only on a client.
    Error(NetworkError),
    /// Used to unit test networking. When recived, increments a value. TODO: Locate that value.
    #[cfg(test)]
    Test,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum NetworkError {
    VersionMismatch(String, String),
}

impl Display for NetworkError {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        match *self {
            NetworkError::VersionMismatch(ref ver1, ref ver2) => {
                write!(fmt,
                       "VersionMismatch: The versions of the client and server attempting to connect mismatch. ver1: {}, ver2: {}",
                       ver1,
                       ver2)
            }
        }
    }
}

impl Error for NetworkError {
    fn description(&self) -> &str {
        match *self {
            NetworkError::VersionMismatch(_, _) => "VersionMismatch: The versions of the client and server attempting to connect mismatch.",
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            NetworkError::VersionMismatch(_, _) => None,
        }
    }
}

pub type EventLoop = MioEventLoop<Handler>;
pub struct Handler {
    connections: Slab<Connection>,
    listener: Option<TcpListener>,
}

impl Handler {
    /// Creates a new instance of a Handler. Does not listen for connections.
    pub fn new() -> Self {
        Handler {
            connections: Slab::new_starting_at(Token::from_usize(1), MAX_CONNECTIONS),
            listener: None,
        }
    }

    /// Creates a new instance of a Handler. Does listen for connections.
    pub fn new_listener() -> Self {
        unimplemented!();
    }
}

#[derive(Debug)]
struct Connection {
    message_queue: Vec<NetworkPacket>,
    stream: TcpStream,
}

impl Connection {
    fn new(stream: TcpStream) -> Connection {
        Connection {
            stream: stream,
            message_queue: Vec::new(),
        }
    }
}

impl MioHandler for Handler {
    type Timeout = ();
    type Message = HandlerMessage;
    fn ready(&mut self, event_loop: &mut MioEventLoop<Handler>, token: Token, events: EventSet) {
        if events.is_writable() {
            // Get all the messages that should be send to that token, seralize all of them, then send all of them.
            // Afterwards, flush the buffer.
            let to_send = self.connections[token].message_queue.clone();
            self.connections[token].message_queue = Vec::new();
            if !to_send.is_empty() {
                for send in to_send {
                    let bytes = seralize_packet(&send);
                    self.connections[token]
                        .stream
                        .write(&bytes)
                        .expect(&format!("Error writing {:?} to connection {:?}/{:?}",
                                         send,
                                         token,
                                         self.connections[token]));
                }
                self.connections[token]
                    .stream
                    .flush()
                    .expect(&format!("Error flushing connection {:?}/{:?}",
                                     token,
                                     self.connections[token]));;
            }
        }

        if events.is_readable() {
            // Read the 6 byte header of each packet, throw it into get_packet_length, then read that number of bytes.
            // Then throw those bytes into deserialize_packet. Afterwards, throw it to handle_packet.

            let mut header = [0; 6];
            self.connections[token]
                .stream
                .read(&mut header)
                .expect(&format!("An error occured reading from socket {:?}", token));    // TODO: Figure out possible errors and take care of them.
            let length = get_packet_length(header).unwrap();    // TODO: Handle error gracefully.
            let mut packet = Vec::new();
            (&mut self.connections[token].stream).take(length as u64).read_to_end(&mut packet).unwrap();
            // I do this because Read.take takes a self, instead of a reasonable alternitive.
            // However &mut Read is it's self a Reader. So I use that instead.
            let dese = deserialize_packet(&packet).unwrap();
            handle_packet(dese, token, event_loop);
        }

        // This may need to be in an if events.is_readable() block.
        if self.listener.is_some() {
            match self.listener.as_ref().unwrap().accept() {
                Ok(Some((socket, address))) => {
                    let token = self.connections.insert(Connection::new(socket)).unwrap();
                    send(event_loop,
                         NetworkPacket::Init { version: VERSION.to_owned() },
                         token);
                    info!("Accepted connection {}.", address);
                }
                Ok(None) => debug!("Recived Ok(None) in listener.accept()."),
                Err(error) => panic!("listener.accept() errored with {:?}!", error),
            }
        }
    }

    fn notify(&mut self, event_loop: &mut MioEventLoop<Handler>, msg: HandlerMessage) {
        match msg {
            HandlerMessage::AddStream(stream, tx) => {
                let token = self.connections.insert(Connection::new(stream)).unwrap();
                event_loop.register(&self.connections[token].stream,
                                    token,
                                    EventSet::all(),
                                    PollOpt::level())
                          .unwrap();
                tx.send(token).unwrap();
            }
            HandlerMessage::Send(packet, token) => {
                self.connections[token].message_queue.push(packet);
            }
        }
    }
}

fn add_socket(event_loop: &EventLoop, socket: TcpStream) -> Token {
    let (tx, rx) = channel();   // How expensive is this? Should I be creating a whole new channel for just 1 message?
    event_loop.channel().send(HandlerMessage::AddStream(socket, tx)).unwrap();
    rx.recv().unwrap() // I should feel bad.
}

/// Send a NetworkPacket over the network.
///
/// This is a static function instead of an impl function because you can't impl on external structs. In this case, mio::EventLoop.
pub fn send(event_loop: &EventLoop, to_send: NetworkPacket, token: Token) {
    event_loop.channel().send(HandlerMessage::Send(to_send, token)).unwrap();
}

fn seralize_packet(to_ser: &NetworkPacket) -> Vec<u8> {
    let mut result: Vec<u8> = Vec::new();
    result.write_u32::<LittleEndian>(NET_MAGIC_NUMBER).unwrap();   // No possible errors here.
    // The NET_MAGIC_NUMBER is used before every packet, so incase the stream is desynced for whatever reason, the game doesn't just read arbratrary data and crash badly.
    // Instead, it can either recover somehow, by disconnecting and reconnecting, or just erroring gracefully.
    let mut encoded = serialize(to_ser, SizeLimit::Infinite).unwrap();
    // Since the size limit is infinite and i'm not encoding to a stream, there is no error and I can safely unwrap();
    result.write_u16::<LittleEndian>(encoded.len() as u16).unwrap();
    result.append(&mut encoded);
    result
}

#[derive(Debug)]
enum PacketDeseError {
    MagicNumberMismatch,
    InvalidEncoding(DeserializeError),
}

impl From<DeserializeError> for PacketDeseError {
    fn from(err: DeserializeError) -> PacketDeseError {
        PacketDeseError::InvalidEncoding(err)
    }
}

/// Returns the length of a given packet, or an error if the first four bytes do not match NET_MAGIC_NUMBER.
fn get_packet_length(to_ln: [u8; 6]) -> Result<u16, PacketDeseError> {
    let (first_four, next_two) = to_ln.split_at(4);
    let should_be_magic_num = LittleEndian::read_u32(&first_four);
    if should_be_magic_num != NET_MAGIC_NUMBER {
        return Err(PacketDeseError::MagicNumberMismatch);
    }
    let length = LittleEndian::read_u16(&next_two);
    Ok(length)
}

fn deserialize_packet(to_de: &[u8]) -> Result<NetworkPacket, PacketDeseError> {
    Ok(try!(deserialize::<NetworkPacket>(to_de)))
}

fn handle_packet(to_handle: NetworkPacket, sender: Token, event_loop: &EventLoop) {
    match to_handle {
        NetworkPacket::Init{version} => {
            if version != VERSION {
                send(event_loop,
                     NetworkPacket::Error(NetworkError::VersionMismatch(version.to_owned(), VERSION.to_owned())),
                     sender)
            }
        }

        #[cfg(Test)]
        NetworkPacket::Test => unimplemented!(),

        NetworkPacket::Error(error) => if ::check_should_crash() { panic!(error) } else { unimplemented!() },
    }
}
<|endoftext|>