<commit_before>extern crate time;
extern crate num_cpus;
extern crate scoped_threadpool;

use time::precise_time_s;
use scoped_threadpool::Pool;
use std::mem;

// 1. Calibration
const ALPHA: f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const BETA: f64 = 0.95;           // Discount factor

// Productivity values
const PRODUCTIVITY: [f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
static TRANSITIONS: [[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                     [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                     [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                     [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                     [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const GRID_CAPITAL: usize = 17820;
const GRID_PRODUCTIVITY: usize = 5;

macro_rules! init_array {
    ($f:expr) => (
        {
            let arr: [Vec<f64>; GRID_PRODUCTIVITY] = [
                $f(0), $f(1), $f(2), $f(3), $f(4),
            ];

            arr
        }
    );
    () => ( init_array![|_| vec![0.0; GRID_CAPITAL]] );
}

fn solve(pool: &mut Pool, print: bool) -> f64 {
    // 2. Steady State
    let capital_steady_state = (ALPHA * BETA).powf(1_f64 / (1_f64 - ALPHA));
    let output_steady_state = capital_steady_state.powf(ALPHA);
    let consumption_steady_state = output_steady_state - capital_steady_state;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 output_steady_state,
                 capital_steady_state,
                 consumption_steady_state);
    }

    let mut grid_capital = [0f64; GRID_CAPITAL];

    for (i, val) in grid_capital.iter_mut().enumerate() {
        *val = 0.5 * capital_steady_state + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut values = init_array![];
    let mut policies = init_array![];
    let mut expected_values = init_array![];

    // 4. We pre-build output for each point in the grid
    let outputs = init_array![|i| {
        (0..GRID_CAPITAL).map(|capital_idx|
            PRODUCTIVITY[i] * grid_capital[capital_idx].powf(ALPHA)
        ).collect()
    }];

    // 5. Main iteration
    const TOLERANCE: f64 = 0.0000001; // compiler warn: variable does not need to be mutable

    let mut iteration = 0;
    let mut finished = false;

    // small array to split where the writes are going, so we don't need a mutex
    let mut diffs = [-100000.0; GRID_PRODUCTIVITY];

    while !finished {
        pool.scoped(|scoped| {
            for (expected_values, transitions) in expected_values.iter_mut()
                                                                 .zip(TRANSITIONS.iter()) {
                // Only capture refs
                let values = &values;

                scoped.execute(move || {
                    for (idx, expected_value) in expected_values.iter_mut().enumerate() {
                        *expected_value = transitions.iter()
                                                     .zip(values.iter())
                                                     .fold(0.0f64, |acc, (transition, value_fns)| {
                            acc + (transition * value_fns[idx])
                        });
                    }
                });
            }
        });

        pool.scoped(|scoped| {
            for ((((policies, value_fns), max_diff), expected_values), outputs) in policies.iter_mut()
                                                                                .zip(values.iter_mut())
                                                                                .zip(diffs.iter_mut())
                                                                                .zip(expected_values.iter())
                                                                                .zip(outputs.iter()) {
                scoped.execute(move || {
                    *max_diff = -100000.0;

                    // We start from previous choice (monotonicity of policy function)
                    let mut grid_capital_next_period = 0;

                    for ((policy, output), value_fn) in policies.iter_mut()
                                                                .zip(outputs.iter())
                                                                .zip(value_fns.iter_mut()) {
                        let mut value_high = -100000.0;
                        let mut capital_choice = grid_capital[0];

                        for capital_next_period in grid_capital_next_period..GRID_CAPITAL {
                            let consumption = output - &grid_capital[capital_next_period];
                            let value_provisional = (1.0 - BETA) * (consumption.ln()) +
                                                    BETA *
                                                    expected_values[capital_next_period];

                            if value_provisional > value_high {
                                value_high = value_provisional;
                                capital_choice = grid_capital[capital_next_period];
                                grid_capital_next_period = capital_next_period;
                            } else {
                                break; // We break when we have achieved the max
                            }
                        }

                        let old = mem::replace(value_fn, value_high);
                        let diff = (old - value_high).abs();
                        if diff > *max_diff {
                            *max_diff = diff
                        }

                        *policy = capital_choice;
                    }
                })
            }
        });

        let mut max_diff = diffs[0];
        for &diff in &diffs[1..] {
            if diff > max_diff {
                max_diff = diff
            }
        }

        finished = max_diff <= TOLERANCE;

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1 || finished) {
            println!("Iteration = {}, Sup Diff = {}", iteration, max_diff);
        }
    }

    policies[2][999]
}

fn main() {
    use std::env::args;

    let sample = args().any(|arg| arg == "--sample");
    let thread_count = num_cpus::get();
    let ref mut pool = Pool::new(thread_count as u32);

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(pool, false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(pool, true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}
<commit_msg>Wrap scoped.execute in unsafe so it runs on stable compiler<commit_after>#![allow(unused_unsafe)]
extern crate time;
extern crate num_cpus;
extern crate scoped_threadpool;

use time::precise_time_s;
use scoped_threadpool::Pool;
use std::mem;

// 1. Calibration
const ALPHA: f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const BETA: f64 = 0.95;           // Discount factor

// Productivity values
const PRODUCTIVITY: [f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
static TRANSITIONS: [[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                     [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                     [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                     [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                     [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const GRID_CAPITAL: usize = 17820;
const GRID_PRODUCTIVITY: usize = 5;

macro_rules! init_array {
    ($f:expr) => (
        {
            let arr: [Vec<f64>; GRID_PRODUCTIVITY] = [
                $f(0), $f(1), $f(2), $f(3), $f(4),
            ];

            arr
        }
    );
    () => ( init_array![|_| vec![0.0; GRID_CAPITAL]] );
}

fn solve(pool: &mut Pool, print: bool) -> f64 {
    // 2. Steady State
    let capital_steady_state = (ALPHA * BETA).powf(1_f64 / (1_f64 - ALPHA));
    let output_steady_state = capital_steady_state.powf(ALPHA);
    let consumption_steady_state = output_steady_state - capital_steady_state;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 output_steady_state,
                 capital_steady_state,
                 consumption_steady_state);
    }

    let mut grid_capital = [0f64; GRID_CAPITAL];

    for (i, val) in grid_capital.iter_mut().enumerate() {
        *val = 0.5 * capital_steady_state + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut values = init_array![];
    let mut policies = init_array![];
    let mut expected_values = init_array![];

    // 4. We pre-build output for each point in the grid
    let outputs = init_array![|i| {
        (0..GRID_CAPITAL).map(|capital_idx|
            PRODUCTIVITY[i] * grid_capital[capital_idx].powf(ALPHA)
        ).collect()
    }];

    // 5. Main iteration
    const TOLERANCE: f64 = 0.0000001; // compiler warn: variable does not need to be mutable

    let mut iteration = 0;
    let mut finished = false;

    // small array to split where the writes are going, so we don't need a mutex
    let mut diffs = [-100000.0; GRID_PRODUCTIVITY];

    while !finished {
        pool.scoped(|scoped| {
            for (expected_values, transitions) in expected_values.iter_mut()
                                                                 .zip(TRANSITIONS.iter()) {
                // Only capture refs
                let values = &values;

                unsafe { // scoped_api requires compiler 1.4 to be compiler-proven safe
                scoped.execute(move || {
                    for (idx, expected_value) in expected_values.iter_mut().enumerate() {
                        *expected_value = transitions.iter()
                                                     .zip(values.iter())
                                                     .fold(0.0f64, |acc, (transition, value_fns)| {
                            acc + (transition * value_fns[idx])
                        });
                    }
                });
                }
            }
        });

        pool.scoped(|scoped| {
            for ((((policies, value_fns), max_diff), expected_values), outputs) in policies.iter_mut()
                                                                                .zip(values.iter_mut())
                                                                                .zip(diffs.iter_mut())
                                                                                .zip(expected_values.iter())
                                                                                .zip(outputs.iter()) {
                unsafe { // scoped_api requires compiler 1.4 to be compiler-proven safe
                scoped.execute(move || {
                    *max_diff = -100000.0;

                    // We start from previous choice (monotonicity of policy function)
                    let mut grid_capital_next_period = 0;

                    for ((policy, output), value_fn) in policies.iter_mut()
                                                                .zip(outputs.iter())
                                                                .zip(value_fns.iter_mut()) {
                        let mut value_high = -100000.0;
                        let mut capital_choice = grid_capital[0];

                        for capital_next_period in grid_capital_next_period..GRID_CAPITAL {
                            let consumption = output - &grid_capital[capital_next_period];
                            let value_provisional = (1.0 - BETA) * (consumption.ln()) +
                                                    BETA *
                                                    expected_values[capital_next_period];

                            if value_provisional > value_high {
                                value_high = value_provisional;
                                capital_choice = grid_capital[capital_next_period];
                                grid_capital_next_period = capital_next_period;
                            } else {
                                break; // We break when we have achieved the max
                            }
                        }

                        let old = mem::replace(value_fn, value_high);
                        let diff = (old - value_high).abs();
                        if diff > *max_diff {
                            *max_diff = diff
                        }

                        *policy = capital_choice;
                    }
                })
                }
            }
        });

        let mut max_diff = diffs[0];
        for &diff in &diffs[1..] {
            if diff > max_diff {
                max_diff = diff
            }
        }

        finished = max_diff <= TOLERANCE;

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1 || finished) {
            println!("Iteration = {}, Sup Diff = {}", iteration, max_diff);
        }
    }

    policies[2][999]
}

fn main() {
    use std::env::args;

    let sample = args().any(|arg| arg == "--sample");
    let thread_count = num_cpus::get();
    let ref mut pool = Pool::new(thread_count as u32);

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(pool, false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(pool, true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::cmp;

use rand;
use rand::Rng;
use itertools::Itertools;
use rgsl::randist::multinomial::multinomial_pdf;

use bio::stats::{Prob, LogProb};

use io::codebook::{Codebook, FeatureID};

use model::readout::{Expressions, Counts, Miscalls, Xi};

pub trait AbstractFeatureModel {
    /// FeatureID of this feature.
    fn feature_id(&self) -> FeatureID;

    /// Neighbors of this feature.
    fn neighbors(&self) -> &[FeatureID];

    /// Minimum hamming distance between any pair of features.
    fn min_dist(&self) -> u8;

    /// Number of events (calls and miscalls) except dropout.
    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32;

    /// Dropout probability.
    fn prob_dropout(&self) -> f64;

    /// Probability of exact miscall to given neighbor index.
    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64;

    /// Probability of mismatching miscall to given neighbor index.
    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64;

    /// Calculate maximum likelihood estimate of miscalls given expression and store in the given
    /// arrays.
    /// Returns the total absolute changes to estimated miscall counts.
    fn mle_miscalls(
        &self,
        expressions: &Expressions,
        miscalls_exact: &mut Miscalls,
        miscalls_mismatch: &mut Miscalls,
        rng: &mut rand::StdRng
    ) {
        let x = expressions[self.feature_id()];

        // We need to ensure that the sum of the estimated miscalls does not exceed x.
        // Hence, we iterate randomly over the neighbors such that we can simply stop once that
        // happens without causing a bias.
        let mut idx = (0..self.neighbors().len()).collect_vec();
        rng.shuffle(&mut idx);

        // by stochastic rounding, we ensure that on average \sum x*p_i = x
        // numeric rounding would not work if many of the probabilities are small enough to yield
        // expectations less than 1
        let mut stochastic_round = |v: f64| v.floor() as u32 + (rng.next_f64() <= v % 1.0) as u32;

        let mut rest = x;
        for i in idx {
            let n = self.neighbors()[i];
            let miscall_exact = self.prob_miscall_exact(i) * x as f64;
            miscalls_exact.set(
                self.feature_id(), n, stochastic_round(miscall_exact), &mut rest
            );

            if self.min_dist() == 4 {
                let miscall_mismatch = self.prob_miscall_mismatch(i) * x as f64;
                let miscall_mismatch = stochastic_round(miscall_mismatch);
                miscalls_mismatch.set(
                    self.feature_id(), n, miscall_mismatch, &mut rest
                );
            }
        }
    }

    /// Calculate maximum likelihood estimate of expression given miscalls.
    /// Returns the absolute amount of change to the expression.
    fn mle_expression(
        &self,
        expressions: &mut Expressions,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> i32 {
        let event_count = self.event_count_non_dropout(miscalls_exact, miscalls_mismatch);

        let x_ = cmp::max(
            (event_count as f64 / (1.0 - self.prob_dropout())).floor() as u32,
            // or at least as many as we have events coming from this feature
            // TODO does this allow x to decrease properly?
            event_count
        );

        let x = expressions.get_mut(self.feature_id()).unwrap();

        let change = *x as i32 - x_ as i32;

        *x = x_;

        change
    }
}

pub struct FeatureModel {
    pub(crate) feature_id: FeatureID,
    event_probs: Vec<f64>,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
    pub(crate) counts: Counts,
    event_counts: RefCell<Vec<u32>>
}

impl FeatureModel {
    pub fn new(
        feature: FeatureID,
        counts: Counts,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        let feature_record = codebook.record(feature);

        let xi_call = xi.prob(feature_record.codeword(), feature_record.codeword());

        let prob_call_exact = xi_call[0];
        let prob_call_mismatch = if codebook.min_dist == 4 {
            xi_call[1]
        } else {
            LogProb::ln_zero()
        };

        let neighbors_4 = codebook.neighbors(feature, 4);
        let neighbors_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2)
        } else {
            vec![]
        };

        let mut neighbors = neighbors_4.clone();
        neighbors.extend(&neighbors_2);

        let xi_miscall_4 = codebook.neighbors(feature, 4).iter().map(|&n| {
            xi.prob(feature_record.codeword(), codebook.record(n).codeword())
        }).collect_vec();
        let xi_miscall_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2).iter().map(|&n| {
                xi.prob(feature_record.codeword(), codebook.record(n).codeword())
            }).collect_vec()
        } else {
            vec![]
        };

        let prob_miscall_exact = xi_miscall_4.iter().map(|xi| {
            xi[0]
        }).chain(xi_miscall_2.iter().map(|xi| xi[0])).collect_vec();

        let prob_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall_4.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let mut total = vec![prob_call_exact, prob_call_mismatch];
        total.extend(&prob_miscall_exact);
        total.extend(&prob_miscall_mismatch);

        let prob_total = LogProb::ln_sum_exp(&total);

        let mut event_probs = vec![
            *Prob::from(prob_call_exact),
            *Prob::from(prob_call_mismatch),
            *Prob::from(prob_total.ln_one_minus_exp())
        ];
        event_probs.extend(prob_miscall_exact.iter().map(|&p| *Prob::from(p)));
        event_probs.extend(prob_miscall_mismatch.iter().map(|&p| *Prob::from(p)));

        debug!("{:?}", event_probs);

        FeatureModel {
            feature_id: feature,
            event_probs: event_probs,
            neighbors: neighbors,
            min_dist: codebook.min_dist,
            counts: counts,
            event_counts: RefCell::new(Vec::new())
        }
    }

    /// Calculate likelihood of given expression.
    pub fn likelihood(
        &self,
        x: u32,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> LogProb {
        let count_exact_miscalled = miscalls_exact.total_to(self.feature_id);
        let count_mismatch_miscalled = miscalls_mismatch.total_to(self.feature_id);

        if self.counts.exact < count_exact_miscalled {
            panic!("bug: likelihood cannot be calculated if exact miscalls exceed counts: counts={}, miscalls={}", self.counts.exact, count_exact_miscalled);
        }
        if self.counts.mismatch < count_mismatch_miscalled {
            panic!("bug: likelihood cannot be calculated if mismatch miscalls exceed counts: counts={}, miscalls={}", self.counts.mismatch, count_mismatch_miscalled);
        }
        let total_count = self.counts.exact - count_exact_miscalled +
                          self.counts.mismatch - count_mismatch_miscalled;
        if total_count > x {
            panic!("bug: event count exceeds expression: x={}, event count={}", x, total_count);
        }

        let calls_exact = self.counts.exact - count_exact_miscalled;
        let calls_mismatch = self.counts.mismatch - count_mismatch_miscalled;

        // setup event counts
        {
            let mut event_counts = self.event_counts.borrow_mut();
            event_counts.clear();
            event_counts.push(
                // exact calls
                calls_exact
            );
            event_counts.push(
                // mismatch calls
                calls_mismatch
            );
            event_counts.push(
                // dropouts (fill in later)
                0
            );
            // miscalls
            event_counts.extend(self.neighbors.iter().map(
                |&n| miscalls_exact.get(self.feature_id, n)
            ));
            if self.min_dist == 4 {
                event_counts.extend(
                    self.neighbors.iter().map(
                        |&n| miscalls_mismatch.get(self.feature_id, n)
                    )
                );
            }

            // set dropouts
            let total_counts = event_counts.iter().sum();
            assert!(x >= total_counts, "total event counts exceed expression: x={} {:?}", x, event_counts);
            let dropouts = x - total_counts;
            event_counts[2] = dropouts;
        }

        LogProb::from(Prob(
            multinomial_pdf(&self.event_probs, &self.event_counts.borrow())
        ))
    }

    /// Minimum expression given fixed miscalls.
    pub fn min_expression(&self, miscalls_exact: &Miscalls, miscalls_mismatch: &Miscalls) -> u32 {
        self.counts.exact.saturating_sub(miscalls_exact.total_to(self.feature_id)) +
        self.counts.mismatch.saturating_sub(miscalls_mismatch.total_to(self.feature_id)) +
        miscalls_exact.total_from(self.feature_id) + miscalls_mismatch.total_from(self.feature_id)
    }
}

impl AbstractFeatureModel for FeatureModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let calls_exact = self.counts.exact.saturating_sub(
            miscalls_exact.total_to(self.feature_id)
        );
        let calls_mismatch = self.counts.mismatch.saturating_sub(
            miscalls_mismatch.total_to(self.feature_id)
        );

        calls_exact + calls_mismatch +
        miscalls_exact.total_from(self.feature_id) +
        miscalls_mismatch.total_from(self.feature_id)
    }

    fn prob_dropout(&self) -> f64 {
        self.event_probs[2]
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + self.neighbors.len() + neighbor_index]
    }
}

pub struct NoiseModel {
    pub(crate) feature_id: FeatureID,
    pub(crate) not_expressed_feature_ids: Vec<FeatureID>,
    pub(crate) not_expressed_counts: Vec<Counts>,
    probs_miscall_exact: Vec<Prob>,
    probs_miscall_mismatch: Vec<Prob>,
    prob_dropout: Prob,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
}

impl NoiseModel {
    pub fn new(
        feature_id: FeatureID,
        not_expressed_feature_ids: Vec<FeatureID>,
        not_expressed_counts: Vec<Counts>,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        // miscalls generated from noise
        let xi_miscall = codebook.records().iter().filter_map(|feat_rec| {
            if feat_rec.expressed() {
                Some(xi.prob(codebook.noise().codeword(), feat_rec.codeword()))
            } else { None }
        }).collect_vec();
        let neighbors = codebook.records().iter().filter_map(|rec| {
            if rec.expressed() {
                Some(codebook.get_id(rec.name()))
            } else { None }
        }).collect_vec();

        let probs_miscall_exact = xi_miscall.iter().map(|xi| xi[0]).collect_vec();
        let probs_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let prob_total_miscall_exact = LogProb::ln_sum_exp(&probs_miscall_exact);
        let prob_total_miscall_mismatch = LogProb::ln_sum_exp(&probs_miscall_mismatch);

        // probs for calls of not expressed features
        let xi_not_expressed = not_expressed_feature_ids.iter().map(|&f| {
            xi.prob(codebook.noise().codeword(), codebook.record(f).codeword())
        }).collect_vec();

        let prob_not_expressed_exact = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[0]).collect_vec()
        );
        let prob_not_expressed_mismatch = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[1]).collect_vec()
        );

        // dropout
        let prob_dropout = LogProb::ln_sum_exp(&[
            prob_not_expressed_exact,
            prob_not_expressed_mismatch,
            prob_total_miscall_exact,
            prob_total_miscall_mismatch
        ]).ln_one_minus_exp();

        NoiseModel {
            feature_id: feature_id,
            not_expressed_feature_ids: not_expressed_feature_ids,
            not_expressed_counts: not_expressed_counts,
            probs_miscall_exact: probs_miscall_exact.iter().map(|&p| Prob::from(p)).collect_vec(),
            probs_miscall_mismatch: probs_miscall_mismatch.iter().map(|&p| Prob::from(p)).collect_vec(),
            prob_dropout: Prob::from(prob_dropout),
            neighbors: neighbors,
            min_dist: codebook.min_dist
        }
    }
}

impl AbstractFeatureModel for NoiseModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let mut calls_exact = 0;
        let mut calls_mismatch = 0;
        for (&id, counts) in self.not_expressed_feature_ids.iter().zip(&self.not_expressed_counts) {
            calls_exact += counts.exact.saturating_sub(miscalls_exact.total_to(id));
            calls_mismatch += counts.mismatch.saturating_sub(miscalls_mismatch.total_to(id));
        }

        let event_count = calls_exact + calls_mismatch +
                          miscalls_exact.total_from(self.feature_id) +
                          miscalls_mismatch.total_from(self.feature_id);

        event_count
    }

    fn prob_dropout(&self) -> f64 {
        *self.prob_dropout
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_exact[neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_mismatch[neighbor_index]
    }
}
<commit_msg>Assertion.<commit_after>use std::cell::RefCell;
use std::cmp;

use rand;
use rand::Rng;
use itertools::Itertools;
use rgsl::randist::multinomial::multinomial_pdf;

use bio::stats::{Prob, LogProb};

use io::codebook::{Codebook, FeatureID};

use model::readout::{Expressions, Counts, Miscalls, Xi};

pub trait AbstractFeatureModel {
    /// FeatureID of this feature.
    fn feature_id(&self) -> FeatureID;

    /// Neighbors of this feature.
    fn neighbors(&self) -> &[FeatureID];

    /// Minimum hamming distance between any pair of features.
    fn min_dist(&self) -> u8;

    /// Number of events (calls and miscalls) except dropout.
    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32;

    /// Dropout probability.
    fn prob_dropout(&self) -> f64;

    /// Probability of exact miscall to given neighbor index.
    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64;

    /// Probability of mismatching miscall to given neighbor index.
    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64;

    /// Calculate maximum likelihood estimate of miscalls given expression and store in the given
    /// arrays.
    /// Returns the total absolute changes to estimated miscall counts.
    fn mle_miscalls(
        &self,
        expressions: &Expressions,
        miscalls_exact: &mut Miscalls,
        miscalls_mismatch: &mut Miscalls,
        rng: &mut rand::StdRng
    ) {
        let x = expressions[self.feature_id()];

        // We need to ensure that the sum of the estimated miscalls does not exceed x.
        // Hence, we iterate randomly over the neighbors such that we can simply stop once that
        // happens without causing a bias.
        let mut idx = (0..self.neighbors().len()).collect_vec();
        rng.shuffle(&mut idx);

        // by stochastic rounding, we ensure that on average \sum x*p_i = x
        // numeric rounding would not work if many of the probabilities are small enough to yield
        // expectations less than 1
        let mut stochastic_round = |v: f64| v.floor() as u32 + (rng.next_f64() <= v % 1.0) as u32;

        let mut rest = x;
        for i in idx {
            let n = self.neighbors()[i];
            let miscall_exact = self.prob_miscall_exact(i) * x as f64;
            miscalls_exact.set(
                self.feature_id(), n, stochastic_round(miscall_exact), &mut rest
            );

            if self.min_dist() == 4 {
                let miscall_mismatch = self.prob_miscall_mismatch(i) * x as f64;
                let miscall_mismatch = stochastic_round(miscall_mismatch);
                miscalls_mismatch.set(
                    self.feature_id(), n, miscall_mismatch, &mut rest
                );
            }
        }
    }

    /// Calculate maximum likelihood estimate of expression given miscalls.
    /// Returns the absolute amount of change to the expression.
    fn mle_expression(
        &self,
        expressions: &mut Expressions,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> i32 {
        let event_count = self.event_count_non_dropout(miscalls_exact, miscalls_mismatch);

        let x_ = cmp::max(
            (event_count as f64 / (1.0 - self.prob_dropout())).floor() as u32,
            // or at least as many as we have events coming from this feature
            // TODO does this allow x to decrease properly?
            event_count
        );

        let x = expressions.get_mut(self.feature_id()).unwrap();

        let change = *x as i32 - x_ as i32;

        *x = x_;

        change
    }
}

pub struct FeatureModel {
    pub(crate) feature_id: FeatureID,
    event_probs: Vec<f64>,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
    pub(crate) counts: Counts,
    event_counts: RefCell<Vec<u32>>
}

impl FeatureModel {
    pub fn new(
        feature: FeatureID,
        counts: Counts,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        let feature_record = codebook.record(feature);

        let xi_call = xi.prob(feature_record.codeword(), feature_record.codeword());

        let prob_call_exact = xi_call[0];
        let prob_call_mismatch = if codebook.min_dist == 4 {
            xi_call[1]
        } else {
            LogProb::ln_zero()
        };

        let neighbors_4 = codebook.neighbors(feature, 4);
        let neighbors_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2)
        } else {
            vec![]
        };

        let mut neighbors = neighbors_4.clone();
        neighbors.extend(&neighbors_2);

        let xi_miscall_4 = codebook.neighbors(feature, 4).iter().map(|&n| {
            xi.prob(feature_record.codeword(), codebook.record(n).codeword())
        }).collect_vec();
        let xi_miscall_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2).iter().map(|&n| {
                xi.prob(feature_record.codeword(), codebook.record(n).codeword())
            }).collect_vec()
        } else {
            vec![]
        };

        let prob_miscall_exact = xi_miscall_4.iter().map(|xi| {
            xi[0]
        }).chain(xi_miscall_2.iter().map(|xi| xi[0])).collect_vec();

        let prob_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall_4.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let mut total = vec![prob_call_exact, prob_call_mismatch];
        total.extend(&prob_miscall_exact);
        total.extend(&prob_miscall_mismatch);

        let prob_total = LogProb::ln_sum_exp(&total);
        assert!(*prob_total <= 0, "prob_total > 0 ({})", prob_total);

        let mut event_probs = vec![
            *Prob::from(prob_call_exact),
            *Prob::from(prob_call_mismatch),
            *Prob::from(prob_total.ln_one_minus_exp())
        ];
        event_probs.extend(prob_miscall_exact.iter().map(|&p| *Prob::from(p)));
        event_probs.extend(prob_miscall_mismatch.iter().map(|&p| *Prob::from(p)));

        debug!("{:?}", event_probs);

        FeatureModel {
            feature_id: feature,
            event_probs: event_probs,
            neighbors: neighbors,
            min_dist: codebook.min_dist,
            counts: counts,
            event_counts: RefCell::new(Vec::new())
        }
    }

    /// Calculate likelihood of given expression.
    pub fn likelihood(
        &self,
        x: u32,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> LogProb {
        let count_exact_miscalled = miscalls_exact.total_to(self.feature_id);
        let count_mismatch_miscalled = miscalls_mismatch.total_to(self.feature_id);

        if self.counts.exact < count_exact_miscalled {
            panic!("bug: likelihood cannot be calculated if exact miscalls exceed counts: counts={}, miscalls={}", self.counts.exact, count_exact_miscalled);
        }
        if self.counts.mismatch < count_mismatch_miscalled {
            panic!("bug: likelihood cannot be calculated if mismatch miscalls exceed counts: counts={}, miscalls={}", self.counts.mismatch, count_mismatch_miscalled);
        }
        let total_count = self.counts.exact - count_exact_miscalled +
                          self.counts.mismatch - count_mismatch_miscalled;
        if total_count > x {
            panic!("bug: event count exceeds expression: x={}, event count={}", x, total_count);
        }

        let calls_exact = self.counts.exact - count_exact_miscalled;
        let calls_mismatch = self.counts.mismatch - count_mismatch_miscalled;

        // setup event counts
        {
            let mut event_counts = self.event_counts.borrow_mut();
            event_counts.clear();
            event_counts.push(
                // exact calls
                calls_exact
            );
            event_counts.push(
                // mismatch calls
                calls_mismatch
            );
            event_counts.push(
                // dropouts (fill in later)
                0
            );
            // miscalls
            event_counts.extend(self.neighbors.iter().map(
                |&n| miscalls_exact.get(self.feature_id, n)
            ));
            if self.min_dist == 4 {
                event_counts.extend(
                    self.neighbors.iter().map(
                        |&n| miscalls_mismatch.get(self.feature_id, n)
                    )
                );
            }

            // set dropouts
            let total_counts = event_counts.iter().sum();
            assert!(x >= total_counts, "total event counts exceed expression: x={} {:?}", x, event_counts);
            let dropouts = x - total_counts;
            event_counts[2] = dropouts;
        }

        LogProb::from(Prob(
            multinomial_pdf(&self.event_probs, &self.event_counts.borrow())
        ))
    }

    /// Minimum expression given fixed miscalls.
    pub fn min_expression(&self, miscalls_exact: &Miscalls, miscalls_mismatch: &Miscalls) -> u32 {
        self.counts.exact.saturating_sub(miscalls_exact.total_to(self.feature_id)) +
        self.counts.mismatch.saturating_sub(miscalls_mismatch.total_to(self.feature_id)) +
        miscalls_exact.total_from(self.feature_id) + miscalls_mismatch.total_from(self.feature_id)
    }
}

impl AbstractFeatureModel for FeatureModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let calls_exact = self.counts.exact.saturating_sub(
            miscalls_exact.total_to(self.feature_id)
        );
        let calls_mismatch = self.counts.mismatch.saturating_sub(
            miscalls_mismatch.total_to(self.feature_id)
        );

        calls_exact + calls_mismatch +
        miscalls_exact.total_from(self.feature_id) +
        miscalls_mismatch.total_from(self.feature_id)
    }

    fn prob_dropout(&self) -> f64 {
        self.event_probs[2]
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + self.neighbors.len() + neighbor_index]
    }
}

pub struct NoiseModel {
    pub(crate) feature_id: FeatureID,
    pub(crate) not_expressed_feature_ids: Vec<FeatureID>,
    pub(crate) not_expressed_counts: Vec<Counts>,
    probs_miscall_exact: Vec<Prob>,
    probs_miscall_mismatch: Vec<Prob>,
    prob_dropout: Prob,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
}

impl NoiseModel {
    pub fn new(
        feature_id: FeatureID,
        not_expressed_feature_ids: Vec<FeatureID>,
        not_expressed_counts: Vec<Counts>,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        // miscalls generated from noise
        let xi_miscall = codebook.records().iter().filter_map(|feat_rec| {
            if feat_rec.expressed() {
                Some(xi.prob(codebook.noise().codeword(), feat_rec.codeword()))
            } else { None }
        }).collect_vec();
        let neighbors = codebook.records().iter().filter_map(|rec| {
            if rec.expressed() {
                Some(codebook.get_id(rec.name()))
            } else { None }
        }).collect_vec();

        let probs_miscall_exact = xi_miscall.iter().map(|xi| xi[0]).collect_vec();
        let probs_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let prob_total_miscall_exact = LogProb::ln_sum_exp(&probs_miscall_exact);
        let prob_total_miscall_mismatch = LogProb::ln_sum_exp(&probs_miscall_mismatch);

        // probs for calls of not expressed features
        let xi_not_expressed = not_expressed_feature_ids.iter().map(|&f| {
            xi.prob(codebook.noise().codeword(), codebook.record(f).codeword())
        }).collect_vec();

        let prob_not_expressed_exact = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[0]).collect_vec()
        );
        let prob_not_expressed_mismatch = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[1]).collect_vec()
        );

        // dropout
        let prob_dropout = LogProb::ln_sum_exp(&[
            prob_not_expressed_exact,
            prob_not_expressed_mismatch,
            prob_total_miscall_exact,
            prob_total_miscall_mismatch
        ]).ln_one_minus_exp();

        NoiseModel {
            feature_id: feature_id,
            not_expressed_feature_ids: not_expressed_feature_ids,
            not_expressed_counts: not_expressed_counts,
            probs_miscall_exact: probs_miscall_exact.iter().map(|&p| Prob::from(p)).collect_vec(),
            probs_miscall_mismatch: probs_miscall_mismatch.iter().map(|&p| Prob::from(p)).collect_vec(),
            prob_dropout: Prob::from(prob_dropout),
            neighbors: neighbors,
            min_dist: codebook.min_dist
        }
    }
}

impl AbstractFeatureModel for NoiseModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let mut calls_exact = 0;
        let mut calls_mismatch = 0;
        for (&id, counts) in self.not_expressed_feature_ids.iter().zip(&self.not_expressed_counts) {
            calls_exact += counts.exact.saturating_sub(miscalls_exact.total_to(id));
            calls_mismatch += counts.mismatch.saturating_sub(miscalls_mismatch.total_to(id));
        }

        let event_count = calls_exact + calls_mismatch +
                          miscalls_exact.total_from(self.feature_id) +
                          miscalls_mismatch.total_from(self.feature_id);

        event_count
    }

    fn prob_dropout(&self) -> f64 {
        *self.prob_dropout
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_exact[neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_mismatch[neighbor_index]
    }
}
<|endoftext|>


<commit_before>use std::vec::Vec;
use collections::ring_buf::RingBuf;

pub struct WrapLineIterator<T> {
    /// Input iterator
    iter: T,
    /// Maximum line length
    line_len: u32,
    /// Characters output since last newline
    current_line_len: u32,
    /// Incoming elements
    buffer: RingBuf<char>,
    /// Peek window
    peek: Option<char>,
}

impl<T: Iterator<char>> WrapLineIterator<T> {
    fn fill_buffer(&mut self) {
        assert!(self.buffer.is_empty());
        let mut seen_text = false;
        let mut text_length = 0;
        match self.peek {
            Some(c) => {
                assert!(c.is_whitespace());
                self.buffer.push_back(c);
                self.peek = None;
            },
            None => ()
        }
        loop {
            let c = match self.iter.next() {
                None => break,
                Some(c) => c
            };

            // Break before non-newline whitespace after
            // non-whitespace text has been seen. (Newlines will be
            // included in the buffer.)
            if seen_text && c.is_whitespace() && c != '\n' {
                self.peek = Some(c);
                break;
            }
            self.buffer.push_back(c);
            if !c.is_whitespace() {
                seen_text = true;
                text_length += 1;
            }

            if text_length >= self.line_len {
                // The word is longer than the allowed line length,
                // need to break it.
                break;
            }
            // Break after hyphen or newline.
            if c == '-' || c == '\n' {
                break;
            }
        }
    }

    fn trim_buffer_left(&mut self) {
        while !self.buffer.is_empty() && self.buffer.front().unwrap().is_whitespace() {
            self.buffer.pop_front();
        }
    }
}

impl<T: Iterator<char>> Iterator<char> for WrapLineIterator<T> {
    fn next(&mut self) -> Option<char> {
        if self.buffer.is_empty() {
            return None;
        }
        if self.current_line_len + self.buffer.len() > self.line_len {
            // Next word won't fit, insert a newline.
            self.current_line_len = 0;
            self.trim_buffer_left();
            return Some('\n');
        }
        let c = self.buffer.pop_front().unwrap();
        self.current_line_len += 1;
        if c == '\n' {
            self.current_line_len = 0;
        }
        if self.buffer.is_empty() {
            self.fill_buffer();
        }
        Some(c)
    }
}

// All char iterators get this trait and the new method.
pub trait WrapUtil {
    fn wrap(self, line_len: u32) -> WrapLineIterator<Self>;
}

impl<T: Iterator<Item=char>> WrapUtil for T {
    fn wrap(self, line_len: u32) -> WrapLineIterator<T> {
        assert!(line_len > 0);
        let mut ret = WrapLineIterator{
            iter: self,
            line_len: line_len,
            current_line_len: 0,
            buffer: RingBuf::new(),
            peek: None,
        };
        ret.fill_buffer();
        ret
    }
}

pub fn wrap_lines(line_len: u32, s: &str) -> String {
    s.chars().wrap(line_len).collect::<Vec<char>>().as_slice().collect()
}

pub struct Map2DIterator<T> {
    /// Input iterator
    iter: T,
    x: i32,
    y: i32,
}

impl<T: Iterator<Item=char>> Iterator for Map2DIterator<T> {
    type Item = (char, i32, i32);

    fn next(&mut self) -> Option<(char, i32, i32)> {
        loop {
            match self.iter.next() {
                None => { return None }
                Some(c) if c == '\n' => { self.y += 1; self.x = 0; }
                Some(c) if (c as u32) < 32 => { }
                Some(c) => { self.x += 1; return Some((c, self.x - 1, self.y)) }
            }
        }
    }
}

pub trait Map2DUtil {
    fn map2d(self) -> Map2DIterator<Self>;
}

impl<T: Iterator<Item=char>> Map2DUtil for T {
    fn map2d(self) -> Map2DIterator<T> {
        Map2DIterator{ iter: self, x: 0, y: 0 }
    }
}
<commit_msg>int fix<commit_after>use std::vec::Vec;
use collections::ring_buf::RingBuf;

pub struct WrapLineIterator<T> {
    /// Input iterator
    iter: T,
    /// Maximum line length
    line_len: usize,
    /// Characters output since last newline
    current_line_len: usize,
    /// Incoming elements
    buffer: RingBuf<char>,
    /// Peek window
    peek: Option<char>,
}

impl<T: Iterator<Item=char>> WrapLineIterator<T> {
    fn fill_buffer(&mut self) {
        assert!(self.buffer.is_empty());
        let mut seen_text = false;
        let mut text_length = 0;
        match self.peek {
            Some(c) => {
                assert!(c.is_whitespace());
                self.buffer.push_back(c);
                self.peek = None;
            },
            None => ()
        }
        loop {
            let c = match self.iter.next() {
                None => break,
                Some(c) => c
            };

            // Break before non-newline whitespace after
            // non-whitespace text has been seen. (Newlines will be
            // included in the buffer.)
            if seen_text && c.is_whitespace() && c != '\n' {
                self.peek = Some(c);
                break;
            }
            self.buffer.push_back(c);
            if !c.is_whitespace() {
                seen_text = true;
                text_length += 1;
            }

            if text_length >= self.line_len {
                // The word is longer than the allowed line length,
                // need to break it.
                break;
            }
            // Break after hyphen or newline.
            if c == '-' || c == '\n' {
                break;
            }
        }
    }

    fn trim_buffer_left(&mut self) {
        while !self.buffer.is_empty() && self.buffer.front().unwrap().is_whitespace() {
            self.buffer.pop_front();
        }
    }
}

impl<T: Iterator<Item=char>> Iterator for WrapLineIterator<T> {
    type Item = char;
    fn next(&mut self) -> Option<char> {
        if self.buffer.is_empty() {
            return None;
        }
        if self.current_line_len + self.buffer.len() > self.line_len {
            // Next word won't fit, insert a newline.
            self.current_line_len = 0;
            self.trim_buffer_left();
            return Some('\n');
        }
        let c = self.buffer.pop_front().unwrap();
        self.current_line_len += 1;
        if c == '\n' {
            self.current_line_len = 0;
        }
        if self.buffer.is_empty() {
            self.fill_buffer();
        }
        Some(c)
    }
}

// All char iterators get this trait and the new method.
pub trait WrapUtil {
    fn wrap(self, line_len: usize) -> WrapLineIterator<Self>;
}

impl<T: Iterator<Item=char>> WrapUtil for T {
    fn wrap(self, line_len: usize) -> WrapLineIterator<T> {
        assert!(line_len > 0);
        let mut ret = WrapLineIterator{
            iter: self,
            line_len: line_len,
            current_line_len: 0,
            buffer: RingBuf::new(),
            peek: None,
        };
        ret.fill_buffer();
        ret
    }
}

pub fn wrap_lines(line_len: usize, s: &str) -> String {
    s.chars().wrap(line_len).collect()
}

pub struct Map2DIterator<T> {
    /// Input iterator
    iter: T,
    x: i32,
    y: i32,
}

impl<T: Iterator<Item=char>> Iterator for Map2DIterator<T> {
    type Item = (char, i32, i32);

    fn next(&mut self) -> Option<(char, i32, i32)> {
        loop {
            match self.iter.next() {
                None => { return None }
                Some(c) if c == '\n' => { self.y += 1; self.x = 0; }
                Some(c) if (c as u32) < 32 => { }
                Some(c) => { self.x += 1; return Some((c, self.x - 1, self.y)) }
            }
        }
    }
}

pub trait Map2DUtil {
    fn map2d(self) -> Map2DIterator<Self>;
}

impl<T: Iterator<Item=char>> Map2DUtil for T {
    fn map2d(self) -> Map2DIterator<T> {
        Map2DIterator{ iter: self, x: 0, y: 0 }
    }
}
<|endoftext|>


<commit_before>use std::env;

use serenity::{
    model::{channel::Message, gateway::Ready},
    prelude::*,
    utils::MessageBuilder,
};

struct Handler;

impl EventHandler for Handler {
    fn message(&self, context: Context, msg: Message) {
        if msg.content == "!ping" {
            let channel = match msg.channel_id.to_channel(&context) {
                Ok(channel) => channel,
                Err(why) => {
                    println!("Error getting channel: {:?}", why);

                    return;
                },
            };

            // The message builder allows for creating a message by
            // mentioning users dynamically, pushing "safe" versions of
            // content (such as bolding normalized content), displaying
            // emojis, and more.
            let response = MessageBuilder::new()
                .push("User ")
                .push_bold_safe(msg.author.name)
                .push(" used the 'ping' command in the ")
                .mention(&channel)
                .push(" channel")
                .build();

            if let Err(why) = msg.channel_id.say(&context.http, &response) {
                println!("Error sending message: {:?}", why);
            }
        }
    }

    fn ready(&self, _: Context, ready: Ready) {
        println!("{} is connected!", ready.user.name);
    }
}

fn main() {
    // Configure the client with your Discord bot token in the environment.
    let token = env::var("DISCORD_TOKEN")
        .expect("Expected a token in the environment");
    let mut client = Client::new(&token, Handler).expect("Err creating client");

    if let Err(why) = client.start() {
        println!("Client error: {:?}", why);
    }
}
<commit_msg>Fix a typo in the message builder example (#775)<commit_after>use std::env;

use serenity::{
    model::{channel::Message, gateway::Ready},
    prelude::*,
    utils::MessageBuilder,
};

struct Handler;

impl EventHandler for Handler {
    fn message(&self, context: Context, msg: Message) {
        if msg.content == "!ping" {
            let channel = match msg.channel_id.to_channel(&context) {
                Ok(channel) => channel,
                Err(why) => {
                    println!("Error getting channel: {:?}", why);

                    return;
                },
            };

            // The message builder allows for creating a message by
            // mentioning users dynamically, pushing "safe" versions of
            // content (such as bolding normalized content), displaying
            // emojis, and more.
            let response = MessageBuilder::new()
                .push("User ")
                .push_bold_safe(&msg.author.name)
                .push(" used the 'ping' command in the ")
                .mention(&channel)
                .push(" channel")
                .build();

            if let Err(why) = msg.channel_id.say(&context.http, &response) {
                println!("Error sending message: {:?}", why);
            }
        }
    }

    fn ready(&self, _: Context, ready: Ready) {
        println!("{} is connected!", ready.user.name);
    }
}

fn main() {
    // Configure the client with your Discord bot token in the environment.
    let token = env::var("DISCORD_TOKEN")
        .expect("Expected a token in the environment");
    let mut client = Client::new(&token, Handler).expect("Err creating client");

    if let Err(why) = client.start() {
        println!("Client error: {:?}", why);
    }
}
<|endoftext|>


<commit_before>use std;
use std::collections::BTreeMap;

use serde_json::{value, ser};
use serde_json::error::Error as SerdeJsonError;
use sodiumoxide::crypto::sign;

use rustc_serialize::base64;
use rustc_serialize::base64::ToBase64;

pub type Object = BTreeMap<String, value::Value>;

quick_error! {
    #[derive(Debug)]
    pub enum SigningJsonError {
        Serialization(err: SerdeJsonError) {
            from()
        }
    }
}

pub type Result<T> = std::result::Result<T, SigningJsonError>;

pub const UNPADDED_BASE64 : base64::Config = base64::Config {
    char_set: base64::CharacterSet::Standard,
    newline: base64::Newline::LF,
    pad: false,
    line_length: None,
};

#[derive(Debug, Clone)]
pub struct SigningKey {
    pub public: sign::PublicKey,
    pub secret: sign::SecretKey,
    pub key_id: String,
}

impl SigningKey {
    pub fn from_seed(seed: &[u8], key_id: String) -> Option<SigningKey> {
        if let Some(seed) = sign::Seed::from_slice(seed) {
            let (public, secret) = sign::keypair_from_seed(&seed);
            Some(SigningKey {
                public: public,
                secret: secret,
                key_id: key_id,
            })
        } else {
            None
        }
    }

    pub fn public_key_b64(&self) -> String {
        self.public.0.to_base64(UNPADDED_BASE64)
    }
}

/// Takes a JSON object, signs it, and adds the signature to the object.
pub fn sign_json(key: &SigningKey, entity_name: String, json: &mut Object) -> Result<()> {
    let b64sig = try!(get_sig_for_json_b64(key, json));

    // json.entry("Signature".to_string()).or_insert();
    let mut sig_obj = get_or_insert(json, "signatures".to_string());
    let mut sig_place = get_or_insert(sig_obj, entity_name);
    sig_place.insert(key.key_id.clone(), value::Value::String(b64sig));
    Ok(())
}

/// Returns the base64 encoded signature for a JSON object
pub fn get_sig_for_json_b64(key: &SigningKey, json: &Object) -> Result<String> {
    let serialized = if json.contains_key("signatures") {
        let mut j = json.clone();
        j.remove("signatures");
        try!(ser::to_string(&j))
    } else {
        try!(ser::to_string(&json))
    };

    let signature = sign::sign_detached(serialized.as_bytes(), &key.secret);
    Ok(signature.0.to_base64(UNPADDED_BASE64))
}

fn get_or_insert<'a>(json: &'a mut Object, s: String) -> &'a mut Object {
    let mut val = json.entry(s).or_insert(value::Value::Object(Object::new()));

    if let &mut value::Value::Object(ref mut obj) = val {
        obj
    } else {
        *val = value::Value::Object(Object::new());
        val.as_object_mut().unwrap()
    }
}

// Tests

#[cfg(test)] use serde_json;
#[cfg(test)] use rustc_serialize::base64::FromBase64;

#[test]
fn test_get_or_insert() {
    let mut map = Object::new();
    get_or_insert(&mut map, "foo".to_string());
    println!("{}", ser::to_string(&map).unwrap());
    map.get_mut("foo").unwrap().as_object_mut().unwrap().insert("test".to_string(), value::Value::U64(1));
    get_or_insert(&mut map, "foo".to_string());
    println!("{}", ser::to_string(&map).unwrap());

    map.insert("foo".to_string(), value::Value::Null);
    println!("{}", ser::to_string(&map).unwrap());
    get_or_insert(&mut map, "foo".to_string());
    println!("{}", ser::to_string(&map).unwrap());
}

#[test]
fn test_sign() {
    let seed_bytes = "YJDBA9Xnr2sVqXD9Vj7XVUnmFZcZrlw8Md7kMW+3XA1".from_base64().unwrap();
    let seed = sign::Seed::from_slice(&seed_bytes[..]).unwrap();
    let (pub_key, priv_key) = sign::keypair_from_seed(&seed);

    let sig_key = SigningKey {
        public: pub_key,
        secret: priv_key,
        key_id: "ed25519:1".to_string(),
    };

    {
        let mut json = serde_json::from_slice::<value::Value>(b"{}").unwrap();
        sign_json(&sig_key, "domain".to_string(), json.as_object_mut().unwrap()).unwrap();

        let signed_json = ser::to_string(&json).unwrap();

        assert_eq!(
            signed_json,
            r#"{"signatures":{"domain":{"ed25519:1":"K8280/U9SSy9IVtjBuVeLr+HpOB4BQFWbg+UZaADMtTdGYI7Geitb76LTrr5QV/7Xg4ahLwYGYZzuHGZKM5ZAQ"}}}"#
        );
    }

    {
        let mut json = serde_json::from_slice::<value::Value>(
            br#"{"one": 1, "two": "Two"}"#
        ).unwrap();
        sign_json(&sig_key, "domain".to_string(), json.as_object_mut().unwrap()).unwrap();

        let signed_json = ser::to_string(&json).unwrap();

        assert_eq!(
            signed_json,
            r#"{"one":1,"signatures":{"domain":{"ed25519:1":"KqmLSbO39/Bzb0QIYE82zqLwsA+PDzYIpIRA2sRQ4sL53+sN6/fpNSoqE7BP7vBZhG6kYdD13EIMJpvhJI+6Bw"}},"two":"Two"}"#
        );
    }

    {
        let mut json = serde_json::from_slice::<value::Value>(
            br#"{"one": 1, "two": "Two", "signatures": {}}"#
        ).unwrap();
        sign_json(&sig_key, "domain".to_string(), json.as_object_mut().unwrap()).unwrap();

        let signed_json = ser::to_string(&json).unwrap();

        assert_eq!(
            signed_json,
            r#"{"one":1,"signatures":{"domain":{"ed25519:1":"KqmLSbO39/Bzb0QIYE82zqLwsA+PDzYIpIRA2sRQ4sL53+sN6/fpNSoqE7BP7vBZhG6kYdD13EIMJpvhJI+6Bw"}},"two":"Two"}"#
        );
    }
}
<commit_msg>Comments<commit_after>use std;
use std::collections::BTreeMap;

use serde_json::{value, ser};
use serde_json::error::Error as SerdeJsonError;
use sodiumoxide::crypto::sign;

use rustc_serialize::base64;
use rustc_serialize::base64::ToBase64;

pub type Object = BTreeMap<String, value::Value>;

quick_error! {
    #[derive(Debug)]
    pub enum SigningJsonError {
        Serialization(err: SerdeJsonError) {
            from()
        }
    }
}

pub type Result<T> = std::result::Result<T, SigningJsonError>;

pub const UNPADDED_BASE64 : base64::Config = base64::Config {
    char_set: base64::CharacterSet::Standard,
    newline: base64::Newline::LF,
    pad: false,
    line_length: None,
};

#[derive(Debug, Clone)]
pub struct SigningKey {
    /// Public part of ED25519 signing key
    pub public: sign::PublicKey,
    /// Secret part of ED25519 signing key
    pub secret: sign::SecretKey,
    /// A unique ID for this signing key.
    pub key_id: String,
}

impl SigningKey {
    /// Create the signing key from a standard ED25519 seed
    pub fn from_seed(seed: &[u8], key_id: String) -> Option<SigningKey> {
        if let Some(seed) = sign::Seed::from_slice(seed) {
            let (public, secret) = sign::keypair_from_seed(&seed);
            Some(SigningKey {
                public: public,
                secret: secret,
                key_id: key_id,
            })
        } else {
            None
        }
    }

    /// Return a unpadded base64 version of the public key.
    pub fn public_key_b64(&self) -> String {
        self.public.0.to_base64(UNPADDED_BASE64)
    }
}

/// Takes a JSON object, signs it, and adds the signature to the object.
pub fn sign_json(key: &SigningKey, entity_name: String, json: &mut Object) -> Result<()> {
    let b64sig = try!(get_sig_for_json_b64(key, json));

    // json.entry("Signature".to_string()).or_insert();
    let mut sig_obj = get_or_insert(json, "signatures".to_string());
    let mut sig_place = get_or_insert(sig_obj, entity_name);
    sig_place.insert(key.key_id.clone(), value::Value::String(b64sig));
    Ok(())
}

/// Returns the base64 encoded signature for a JSON object
pub fn get_sig_for_json_b64(key: &SigningKey, json: &Object) -> Result<String> {
    let serialized = if json.contains_key("signatures") {
        let mut j = json.clone();
        j.remove("signatures");
        try!(ser::to_string(&j))
    } else {
        try!(ser::to_string(&json))
    };

    let signature = sign::sign_detached(serialized.as_bytes(), &key.secret);
    Ok(signature.0.to_base64(UNPADDED_BASE64))
}

fn get_or_insert<'a>(json: &'a mut Object, s: String) -> &'a mut Object {
    let mut val = json.entry(s).or_insert(value::Value::Object(Object::new()));

    if let &mut value::Value::Object(ref mut obj) = val {
        obj
    } else {
        *val = value::Value::Object(Object::new());
        val.as_object_mut().unwrap()
    }
}

// Tests

#[cfg(test)] use serde_json;
#[cfg(test)] use rustc_serialize::base64::FromBase64;

#[test]
fn test_get_or_insert() {
    let mut map = Object::new();
    get_or_insert(&mut map, "foo".to_string());
    println!("{}", ser::to_string(&map).unwrap());
    map.get_mut("foo").unwrap().as_object_mut().unwrap().insert("test".to_string(), value::Value::U64(1));
    get_or_insert(&mut map, "foo".to_string());
    println!("{}", ser::to_string(&map).unwrap());

    map.insert("foo".to_string(), value::Value::Null);
    println!("{}", ser::to_string(&map).unwrap());
    get_or_insert(&mut map, "foo".to_string());
    println!("{}", ser::to_string(&map).unwrap());
}

#[test]
fn test_sign() {
    let seed_bytes = "YJDBA9Xnr2sVqXD9Vj7XVUnmFZcZrlw8Md7kMW+3XA1".from_base64().unwrap();
    let seed = sign::Seed::from_slice(&seed_bytes[..]).unwrap();
    let (pub_key, priv_key) = sign::keypair_from_seed(&seed);

    let sig_key = SigningKey {
        public: pub_key,
        secret: priv_key,
        key_id: "ed25519:1".to_string(),
    };

    {
        let mut json = serde_json::from_slice::<value::Value>(b"{}").unwrap();
        sign_json(&sig_key, "domain".to_string(), json.as_object_mut().unwrap()).unwrap();

        let signed_json = ser::to_string(&json).unwrap();

        assert_eq!(
            signed_json,
            r#"{"signatures":{"domain":{"ed25519:1":"K8280/U9SSy9IVtjBuVeLr+HpOB4BQFWbg+UZaADMtTdGYI7Geitb76LTrr5QV/7Xg4ahLwYGYZzuHGZKM5ZAQ"}}}"#
        );
    }

    {
        let mut json = serde_json::from_slice::<value::Value>(
            br#"{"one": 1, "two": "Two"}"#
        ).unwrap();
        sign_json(&sig_key, "domain".to_string(), json.as_object_mut().unwrap()).unwrap();

        let signed_json = ser::to_string(&json).unwrap();

        assert_eq!(
            signed_json,
            r#"{"one":1,"signatures":{"domain":{"ed25519:1":"KqmLSbO39/Bzb0QIYE82zqLwsA+PDzYIpIRA2sRQ4sL53+sN6/fpNSoqE7BP7vBZhG6kYdD13EIMJpvhJI+6Bw"}},"two":"Two"}"#
        );
    }

    {
        let mut json = serde_json::from_slice::<value::Value>(
            br#"{"one": 1, "two": "Two", "signatures": {}}"#
        ).unwrap();
        sign_json(&sig_key, "domain".to_string(), json.as_object_mut().unwrap()).unwrap();

        let signed_json = ser::to_string(&json).unwrap();

        assert_eq!(
            signed_json,
            r#"{"one":1,"signatures":{"domain":{"ed25519:1":"KqmLSbO39/Bzb0QIYE82zqLwsA+PDzYIpIRA2sRQ4sL53+sN6/fpNSoqE7BP7vBZhG6kYdD13EIMJpvhJI+6Bw"}},"two":"Two"}"#
        );
    }
}
<|endoftext|>


<commit_before>//! In certain situations, rust automatically inserts derefs as necessary: for
//! example, field accesses `foo.bar` still work when `foo` is actually a
//! reference to a type with the field `bar`. This is an approximation of the
//! logic in rustc (which lives in librustc_typeck/check/autoderef.rs).

use std::iter::successors;

use base_db::CrateId;
use chalk_ir::{cast::Cast, fold::Fold, interner::HasInterner, VariableKind};
use hir_def::lang_item::LangItemTarget;
use hir_expand::name::name;
use log::{info, warn};

use crate::{
    db::HirDatabase, static_lifetime, AliasEq, AliasTy, BoundVar, Canonical, CanonicalVarKinds,
    DebruijnIndex, InEnvironment, Interner, ProjectionTyExt, Solution, Substitution, Ty, TyBuilder,
    TyKind,
};

const AUTODEREF_RECURSION_LIMIT: usize = 10;

pub fn autoderef<'a>(
    db: &'a dyn HirDatabase,
    krate: Option<CrateId>,
    ty: InEnvironment<Canonical<Ty>>,
) -> impl Iterator<Item = Canonical<Ty>> + 'a {
    let _p = profile::span("autoderef");
    let InEnvironment { goal: ty, environment } = ty;
    successors(Some(ty), move |ty| {
        let _p = profile::span("autoderef.step");
        deref(db, krate?, InEnvironment { goal: ty, environment: environment.clone() })
    })
    .take(AUTODEREF_RECURSION_LIMIT)
}

pub(crate) fn deref(
    db: &dyn HirDatabase,
    krate: CrateId,
    ty: InEnvironment<&Canonical<Ty>>,
) -> Option<Canonical<Ty>> {
    if let Some(derefed) = builtin_deref(&ty.goal.value) {
        Some(Canonical { value: derefed, binders: ty.goal.binders.clone() })
    } else {
        deref_by_trait(db, krate, ty)
    }
}

fn builtin_deref(ty: &Ty) -> Option<Ty> {
    match ty.kind(&Interner) {
        TyKind::Ref(.., ty) => Some(ty.clone()),
        TyKind::Raw(.., ty) => Some(ty.clone()),
        _ => None,
    }
}

fn deref_by_trait(
    db: &dyn HirDatabase,
    krate: CrateId,
    ty: InEnvironment<&Canonical<Ty>>,
) -> Option<Canonical<Ty>> {
    let _p = profile::span("deref_by_trait");
    let deref_trait = match db.lang_item(krate, "deref".into())? {
        LangItemTarget::TraitId(it) => it,
        _ => return None,
    };
    let target = db.trait_data(deref_trait).associated_type_by_name(&name![Target])?;

    let projection = {
        let b = TyBuilder::assoc_type_projection(db, target);
        if b.remaining() != 1 {
            // the Target type + Deref trait should only have one generic parameter,
            // namely Deref's Self type
            return None;
        }
        b.push(ty.goal.value.clone()).build()
    };

    // FIXME make the Canonical / bound var handling nicer

    // Check that the type implements Deref at all
    let trait_ref = projection.trait_ref(db);
    let implements_goal = Canonical {
        binders: ty.goal.binders.clone(),
        value: InEnvironment {
            goal: trait_ref.cast(&Interner),
            environment: ty.environment.clone(),
        },
    };
    if {
        let _p = profile::span("db.trait_solve");
        db.trait_solve(krate, implements_goal).is_none()
    } {
        return None;
    }

    // Now do the assoc type projection
    let alias_eq = AliasEq {
        alias: AliasTy::Projection(projection),
        ty: TyKind::BoundVar(BoundVar::new(
            DebruijnIndex::INNERMOST,
            ty.goal.binders.len(&Interner),
        ))
        .intern(&Interner),
    };

    let in_env = InEnvironment { goal: alias_eq.cast(&Interner), environment: ty.environment };

    let canonical = Canonical {
        value: in_env,
        binders: CanonicalVarKinds::from_iter(
            &Interner,
            ty.goal.binders.iter(&Interner).cloned().chain(Some(chalk_ir::WithKind::new(
                VariableKind::Ty(chalk_ir::TyVariableKind::General),
                chalk_ir::UniverseIndex::ROOT,
            ))),
        ),
    };

    let solution = db.trait_solve(krate, canonical)?;

    match &solution {
        Solution::Unique(vars) => {
            // FIXME: vars may contain solutions for any inference variables
            // that happened to be inside ty. To correctly handle these, we
            // would have to pass the solution up to the inference context, but
            // that requires a larger refactoring (especially if the deref
            // happens during method resolution). So for the moment, we just
            // check that we're not in the situation we're we would actually
            // need to handle the values of the additional variables, i.e.
            // they're just being 'passed through'. In the 'standard' case where
            // we have `impl<T> Deref for Foo<T> { Target = T }`, that should be
            // the case.

            // FIXME: if the trait solver decides to truncate the type, these
            // assumptions will be broken. We would need to properly introduce
            // new variables in that case

            for i in 1..vars.binders.len(&Interner) {
                if vars.value.subst.at(&Interner, i - 1).assert_ty_ref(&Interner).kind(&Interner)
                    != &TyKind::BoundVar(BoundVar::new(DebruijnIndex::INNERMOST, i - 1))
                {
                    warn!("complex solution for derefing {:?}: {:?}, ignoring", ty.goal, solution);
                    return None;
                }
            }
            // FIXME: we remove lifetime variables here since they can confuse
            // the method resolution code later
            Some(fixup_lifetime_variables(Canonical {
                value: vars
                    .value
                    .subst
                    .at(&Interner, vars.value.subst.len(&Interner) - 1)
                    .assert_ty_ref(&Interner)
                    .clone(),
                binders: vars.binders.clone(),
            }))
        }
        Solution::Ambig(_) => {
            info!("Ambiguous solution for derefing {:?}: {:?}", ty.goal, solution);
            None
        }
    }
}

fn fixup_lifetime_variables<T: Fold<Interner, Result = T> + HasInterner<Interner = Interner>>(
    c: Canonical<T>,
) -> Canonical<T> {
    // Removes lifetime variables from the Canonical, replacing them by static lifetimes.
    let mut i = 0;
    let subst = Substitution::from_iter(
        &Interner,
        c.binders.iter(&Interner).map(|vk| match vk.kind {
            VariableKind::Ty(_) => {
                let index = i;
                i += 1;
                BoundVar::new(DebruijnIndex::INNERMOST, index).to_ty(&Interner).cast(&Interner)
            }
            VariableKind::Lifetime => static_lifetime().cast(&Interner),
            VariableKind::Const(_) => unimplemented!(),
        }),
    );
    let binders = CanonicalVarKinds::from_iter(
        &Interner,
        c.binders.iter(&Interner).filter(|vk| match vk.kind {
            VariableKind::Ty(_) => true,
            VariableKind::Lifetime => false,
            VariableKind::Const(_) => true,
        }),
    );
    let value = subst.apply(c.value, &Interner);
    Canonical { binders, value }
}
<commit_msg>Better places for spans<commit_after>//! In certain situations, rust automatically inserts derefs as necessary: for
//! example, field accesses `foo.bar` still work when `foo` is actually a
//! reference to a type with the field `bar`. This is an approximation of the
//! logic in rustc (which lives in librustc_typeck/check/autoderef.rs).

use std::iter::successors;

use base_db::CrateId;
use chalk_ir::{cast::Cast, fold::Fold, interner::HasInterner, VariableKind};
use hir_def::lang_item::LangItemTarget;
use hir_expand::name::name;
use log::{info, warn};

use crate::{
    db::HirDatabase, static_lifetime, AliasEq, AliasTy, BoundVar, Canonical, CanonicalVarKinds,
    DebruijnIndex, InEnvironment, Interner, ProjectionTyExt, Solution, Substitution, Ty, TyBuilder,
    TyKind,
};

const AUTODEREF_RECURSION_LIMIT: usize = 10;

pub fn autoderef<'a>(
    db: &'a dyn HirDatabase,
    krate: Option<CrateId>,
    ty: InEnvironment<Canonical<Ty>>,
) -> impl Iterator<Item = Canonical<Ty>> + 'a {
    let InEnvironment { goal: ty, environment } = ty;
    successors(Some(ty), move |ty| {
        deref(db, krate?, InEnvironment { goal: ty, environment: environment.clone() })
    })
    .take(AUTODEREF_RECURSION_LIMIT)
}

pub(crate) fn deref(
    db: &dyn HirDatabase,
    krate: CrateId,
    ty: InEnvironment<&Canonical<Ty>>,
) -> Option<Canonical<Ty>> {
    let _p = profile::span("deref");
    if let Some(derefed) = builtin_deref(&ty.goal.value) {
        Some(Canonical { value: derefed, binders: ty.goal.binders.clone() })
    } else {
        deref_by_trait(db, krate, ty)
    }
}

fn builtin_deref(ty: &Ty) -> Option<Ty> {
    match ty.kind(&Interner) {
        TyKind::Ref(.., ty) => Some(ty.clone()),
        TyKind::Raw(.., ty) => Some(ty.clone()),
        _ => None,
    }
}

fn deref_by_trait(
    db: &dyn HirDatabase,
    krate: CrateId,
    ty: InEnvironment<&Canonical<Ty>>,
) -> Option<Canonical<Ty>> {
    let _p = profile::span("deref_by_trait");
    let deref_trait = match db.lang_item(krate, "deref".into())? {
        LangItemTarget::TraitId(it) => it,
        _ => return None,
    };
    let target = db.trait_data(deref_trait).associated_type_by_name(&name![Target])?;

    let projection = {
        let b = TyBuilder::assoc_type_projection(db, target);
        if b.remaining() != 1 {
            // the Target type + Deref trait should only have one generic parameter,
            // namely Deref's Self type
            return None;
        }
        b.push(ty.goal.value.clone()).build()
    };

    // FIXME make the Canonical / bound var handling nicer

    // Check that the type implements Deref at all
    let trait_ref = projection.trait_ref(db);
    let implements_goal = Canonical {
        binders: ty.goal.binders.clone(),
        value: InEnvironment {
            goal: trait_ref.cast(&Interner),
            environment: ty.environment.clone(),
        },
    };
    if {
        let _p = profile::span("db.trait_solve");
        db.trait_solve(krate, implements_goal).is_none()
    } {
        return None;
    }

    // Now do the assoc type projection
    let alias_eq = AliasEq {
        alias: AliasTy::Projection(projection),
        ty: TyKind::BoundVar(BoundVar::new(
            DebruijnIndex::INNERMOST,
            ty.goal.binders.len(&Interner),
        ))
        .intern(&Interner),
    };

    let in_env = InEnvironment { goal: alias_eq.cast(&Interner), environment: ty.environment };

    let canonical = Canonical {
        value: in_env,
        binders: CanonicalVarKinds::from_iter(
            &Interner,
            ty.goal.binders.iter(&Interner).cloned().chain(Some(chalk_ir::WithKind::new(
                VariableKind::Ty(chalk_ir::TyVariableKind::General),
                chalk_ir::UniverseIndex::ROOT,
            ))),
        ),
    };

    let solution = db.trait_solve(krate, canonical)?;

    match &solution {
        Solution::Unique(vars) => {
            // FIXME: vars may contain solutions for any inference variables
            // that happened to be inside ty. To correctly handle these, we
            // would have to pass the solution up to the inference context, but
            // that requires a larger refactoring (especially if the deref
            // happens during method resolution). So for the moment, we just
            // check that we're not in the situation we're we would actually
            // need to handle the values of the additional variables, i.e.
            // they're just being 'passed through'. In the 'standard' case where
            // we have `impl<T> Deref for Foo<T> { Target = T }`, that should be
            // the case.

            // FIXME: if the trait solver decides to truncate the type, these
            // assumptions will be broken. We would need to properly introduce
            // new variables in that case

            for i in 1..vars.binders.len(&Interner) {
                if vars.value.subst.at(&Interner, i - 1).assert_ty_ref(&Interner).kind(&Interner)
                    != &TyKind::BoundVar(BoundVar::new(DebruijnIndex::INNERMOST, i - 1))
                {
                    warn!("complex solution for derefing {:?}: {:?}, ignoring", ty.goal, solution);
                    return None;
                }
            }
            // FIXME: we remove lifetime variables here since they can confuse
            // the method resolution code later
            Some(fixup_lifetime_variables(Canonical {
                value: vars
                    .value
                    .subst
                    .at(&Interner, vars.value.subst.len(&Interner) - 1)
                    .assert_ty_ref(&Interner)
                    .clone(),
                binders: vars.binders.clone(),
            }))
        }
        Solution::Ambig(_) => {
            info!("Ambiguous solution for derefing {:?}: {:?}", ty.goal, solution);
            None
        }
    }
}

fn fixup_lifetime_variables<T: Fold<Interner, Result = T> + HasInterner<Interner = Interner>>(
    c: Canonical<T>,
) -> Canonical<T> {
    // Removes lifetime variables from the Canonical, replacing them by static lifetimes.
    let mut i = 0;
    let subst = Substitution::from_iter(
        &Interner,
        c.binders.iter(&Interner).map(|vk| match vk.kind {
            VariableKind::Ty(_) => {
                let index = i;
                i += 1;
                BoundVar::new(DebruijnIndex::INNERMOST, index).to_ty(&Interner).cast(&Interner)
            }
            VariableKind::Lifetime => static_lifetime().cast(&Interner),
            VariableKind::Const(_) => unimplemented!(),
        }),
    );
    let binders = CanonicalVarKinds::from_iter(
        &Interner,
        c.binders.iter(&Interner).filter(|vk| match vk.kind {
            VariableKind::Ty(_) => true,
            VariableKind::Lifetime => false,
            VariableKind::Const(_) => true,
        }),
    );
    let value = subst.apply(c.value, &Interner);
    Canonical { binders, value }
}
<|endoftext|>


<commit_before>use std::iter;
use std::cmp::max;
use num::{Float, Zero};
use num::traits::{Num, Signed};
use image::{GenericImage, ImageBuffer, Rgba, Pixel};
use calx_layout::Rect;
use img;

/// Constructor object for atlases.
pub struct AtlasBuilder {
    images: Vec<ImageBuffer<Rgba<u8>, Vec<u8>>>,
    draw_offsets: Vec<[i32; 2]>,
}

impl AtlasBuilder {
    pub fn new() -> AtlasBuilder {
        AtlasBuilder {
            images: vec![],
            draw_offsets: vec![],
        }
    }

    /// Add an image to the image atlas with the given draw offset.
    pub fn push<P, I, V>(&mut self, offset: V, image: &I) -> usize
        where P: Pixel<Subpixel = u8> + 'static,
              I: GenericImage<Pixel = P>,
              V: Into<[i32; 2]>
    {
        let Rect { top: pos, size: dim } = img::crop_alpha(image);
        let image = ImageBuffer::from_fn(dim[0] as u32, dim[1] as u32, |x, y| {
            image.get_pixel(pos[0] as u32 + x, pos[1] as u32 + y)
                 .to_rgba()
        });
        self.images.push(image);
        self.draw_offsets.push(pos + offset);
        self.images.len() - 1
    }
}

/// A collection of images packed into a single large image for more efficient
/// caching in graphics hardware.
pub struct Atlas {
    /// The atlas image containing all the subimages.
    pub image: ImageBuffer<Rgba<u8>, Vec<u8>>,
    /// Metadata for the subimages.
    pub items: Vec<AtlasItem>,
}

/// One image stored in a texture atlas
#[derive(Copy, Clone, Debug)]
pub struct AtlasItem {
    /// Vertices for the image rectangle when drawn at origin
    pub pos: Rect<f32>,
    /// Texture coordinates for the image rectangle on the atlas texture
    pub tex: Rect<f32>,
}

impl Atlas {
    pub fn new(builder: &AtlasBuilder) -> Atlas {
        let dims = builder.images
                                        .iter()
                                        .map(|img| {
                                            let (w, h) = img.dimensions();
                                            V2(w as i32, h as i32)
                                        })
                                        .collect::<[i32; 2]>();

        // Add 1 pixel edges to images to prevent texturing artifacts from
        // adjacent pixels in separate subimages.
        let expanded_dims = dims.iter()
                                .map(|&v| v + V2(1, 1))
                                .collect();

        // Guesstimate the size for the atlas container.
        let total_area = dims.iter().map(|dim| dim.0 * dim.1).fold(0, |a, b| a + b);
        let mut d = ((total_area as f64).sqrt() as u32).next_power_of_two();
        let offsets;

        loop {
            assert!(d < 1000000000); // Sanity check
            match pack_rectangles(V2(d as i32, d as i32), &expanded_dims) {
                Some(ret) => {
                    offsets = ret;
                    break;
                }
                None => {
                    d = d * 2;
                }
            }
        }

        // Blit subimages to atlas image.
        let mut image: ImageBuffer<Rgba<u8>, Vec<u8>> = ImageBuffer::new(d, d);
        for (i, &offset) in offsets.iter().enumerate() {
            img::blit(&builder.images[i], &mut image, offset);
        }

        let image_dim = V2(d, d);

        assert!(offsets.len() == builder.draw_offsets.len());

        let items: Vec<AtlasItem> = (0..(offsets.len()))
                                        .map(|i| {
                                            AtlasItem {
                                                pos: Rect(builder.draw_offsets[i].map(|x| x as f32),
                                                          dims[i].map(|x| x as f32)),
                                                tex: Rect(scale_vec(offsets[i], image_dim),
                                                          scale_vec(dims[i], image_dim)),
                                            }
                                        })
                                        .collect();

        return Atlas {
            image: image,
            items: items,
        };

        fn scale_vec(pixel_vec: V2<i32>, image_dim: V2<u32>) -> V2<f32> {
            V2(pixel_vec.0 as f32 / image_dim.0 as f32,
               pixel_vec.1 as f32 / image_dim.1 as f32)
        }
    }
}

/// Try to pack several small rectangles into one large rectangle. Return
/// offsets for the subrectangles within the container if a packing was found.
fn pack_rectangles<T>(container_dim: V2<T>, dims: &Vec<V2<T>>) -> Option<Vec<V2<T>>>
    where T: Num + PartialOrd + Ord + Copy
{
    let init: T = Zero::zero();
    let total_area = dims.iter().map(|dim| dim.0 * dim.1).fold(init, |a, b| a + b);

    // Too much rectangle area to fit in container no matter how you pack it.
    // Fail early.
    if total_area > container_dim.0 * container_dim.1 {
        return None;
    }

    // Take enumeration to keep the original indices around.
    let mut largest_first: Vec<(usize, &V2<T>)> = dims.iter().enumerate().collect();
    largest_first.sort_by(|&(_i, a), &(_j, b)| (b.0 * b.1).cmp(&(a.0 * a.1)));

    let mut slots = vec![Rect(V2(Zero::zero(), Zero::zero()), container_dim)];

    let mut ret: Vec<V2<T>> = iter::repeat(V2(Zero::zero(), Zero::zero()))
                                  .take(dims.len())
                                  .collect();

    for i in 0..(largest_first.len()) {
        let (idx, &dim) = largest_first[i];
        match place(dim, &mut slots) {
            Some(pos) => {
                ret[idx] = pos;
            }
            None => {
                return None;
            }
        }
    }

    return Some(ret);

    /// Find the smallest slot in the slot vector that will fit the given
    /// item.
    fn place<T>(dim: [T; 2], slots: &mut Vec<Rect<T>>) -> Option<[T; 2]>
        where T: Num + PartialOrd + Ord + Signed + Copy
    {
        for i in 0..(slots.len()) {
            let Rect { top: slot_pos, size: slot_dim } = slots[i];
            if fits(dim, slot_dim) {
                // Remove the original slot, it gets the item. Add the two new
                // rectangles that form around the item.
                let (new_1, new_2) = remaining_rects(dim, Rect { top: slot_pos, size: slot_dim });
                slots.swap_remove(i);
                slots.push(new_1);
                slots.push(new_2);
                // Sort by area from smallest to largest.
                slots.sort_by(|&a, &b| a.area().cmp(&b.area()));
                return Some(slot_pos);
            }
        }
        None
    }

    /// Return the two remaining parts of container rect when the dim-sized
    /// item is placed in the top left corner.
    fn remaining_rects<T>(dim: [T; 2], Rect { top: rect_pos, size: rect_dim }: Rect<T>) -> (Rect<T>, Rect<T>)
        where T: Num + PartialOrd + Ord + Copy + Signed
    {
        assert!(fits(dim, rect_dim));

        // Choose between making a vertical or a horizontal split
        // based on which leaves a bigger open rectangle.
        let vert_vol = max(rect_dim[0] * (rect_dim[1] - dim[1]),
                           (rect_dim[0] - dim[0]) * dim[1]);
        let horiz_vol = max(dim[0] * (rect_dim[1] - dim[1]),
                            (rect_dim[0] - dim[0]) * rect_dim[1]);

        if vert_vol > horiz_vol {
            //     |AA
            // ----+--
            // BBBBBBB
            // BBBBBBB
            (Rect { top: [rect_pos[0] + dim[0], rect_pos[1]],
                    size: [rect_dim[0] - dim[0], dim[1]] },
             Rect { top: [rect_pos[0], rect_pos[1] + dim[1]],
                    size: [rect_dim[0], rect_dim[1] - dim[1]] })
        } else {
            //     |BB
            // ----+BB
            // AAAA|BB
            // AAAA|BB
            (Rect { top: [rect_pos[0], rect_pos[1] + dim[1]],
                    size: [dim[0], rect_dim[1] - dim[1]] },
             Rect { top: [rect_pos[0] + dim[0], rect_pos[1]],
                    size: [rect_dim[0] - dim[0], rect_dim[1]] })
        }
    }

    fn fits<T: Ord>(dim: [T; 2], container_dim: [T; 2]) -> bool {
        dim[0] <= container_dim[0] && dim[1] <= container_dim[1]
    }
}
<commit_msg>Cache refactoring<commit_after>use std::iter;
use std::cmp::max;
use num::{Float, Zero};
use num::traits::{Num, Signed};
use image::{GenericImage, ImageBuffer, Rgba, Pixel};
use calx_layout::Rect;
use img;

/// Constructor object for atlases.
pub struct AtlasBuilder {
    images: Vec<ImageBuffer<Rgba<u8>, Vec<u8>>>,
    draw_offsets: Vec<[i32; 2]>,
}

impl AtlasBuilder {
    pub fn new() -> AtlasBuilder {
        AtlasBuilder {
            images: vec![],
            draw_offsets: vec![],
        }
    }

    /// Add an image to the image atlas with the given draw offset.
    pub fn push<P, I, V>(&mut self, offset: V, image: &I) -> usize
        where P: Pixel<Subpixel = u8> + 'static,
              I: GenericImage<Pixel = P>,
              V: Into<[i32; 2]>
    {
        let Rect { top: pos, size: dim } = img::crop_alpha(image);
        let image = ImageBuffer::from_fn(dim[0] as u32,
                                         dim[1] as u32,
                                         |x, y| {
                                             image.get_pixel(pos[0] as u32 + x,
                                                             pos[1] as u32 + y)
                                                  .to_rgba()
                                         });
        self.images.push(image);
        self.draw_offsets.push(pos + offset);
        self.images.len() - 1
    }
}

/// A collection of images packed into a single large image for more efficient
/// caching in graphics hardware.
pub struct Atlas {
    /// The atlas image containing all the subimages.
    pub image: ImageBuffer<Rgba<u8>, Vec<u8>>,
    /// Metadata for the subimages.
    pub items: Vec<AtlasItem>,
}

/// One image stored in a texture atlas
#[derive(Copy, Clone, Debug)]
pub struct AtlasItem {
    /// Vertices for the image rectangle when drawn at origin
    pub pos: Rect<f32>,
    /// Texture coordinates for the image rectangle on the atlas texture
    pub tex: Rect<f32>,
}

impl Atlas {
    pub fn new(builder: &AtlasBuilder) -> Atlas {
        let dims = builder.images
                          .iter()
                          .map(|img| {
                              let (w, h) = img.dimensions();
                              [w as i32, h as i32]
                          })
                          .collect::<[i32; 2]>();

        // Add 1 pixel edges to images to prevent texturing artifacts from
        // adjacent pixels in separate subimages.
        let expanded_dims = dims.iter()
                                .map(|&v| [v[0] + 1, v[1] + 1])
                                .collect::<[i32; 2]>();

        // Guesstimate the size for the atlas container.
        let total_area = dims.iter()
                             .map(|dim| dim[0] * dim[1])
                             .fold(0, |a, b| a + b);
        let mut d = ((total_area as f64).sqrt() as u32).next_power_of_two();
        let offsets;

        loop {
            assert!(d < 1000000000); // Sanity check
            match pack_rectangles(V2(d as i32, d as i32), &expanded_dims) {
                Some(ret) => {
                    offsets = ret;
                    break;
                }
                None => {
                    d = d * 2;
                }
            }
        }

        // Blit subimages to atlas image.
        let mut image: ImageBuffer<Rgba<u8>, Vec<u8>> = ImageBuffer::new(d, d);
        for (i, &offset) in offsets.iter().enumerate() {
            img::blit(&builder.images[i], &mut image, offset);
        }

        let image_dim = [d, d];

        assert!(offsets.len() == builder.draw_offsets.len());

        let items = (0..(offsets.len()))
                        .map(|i| {
                            let top = [builder.draw_offsets[i][0] as f32,
                                       builder.draw_offsets[i][1] as f32];
                            let size = [dims[i][0] as f32, dims[i][1] as f32];
                            AtlasItem {
                                pos: Rect {
                                    top: top,
                                    size: size,
                                },
                                tex: Rect {
                                    top: scale_vec(offsets[i], image_dim),
                                    size: scale_vec(dims[i], image_dim),
                                },
                            }
                        })
                        .collect::<Vec<AtlasItem>>();

        return Atlas {
            image: image,
            items: items,
        };

        fn scale_vec(pixel_vec: [i32; 2], image_dim: [u32; 2]) -> [f32; 2] {
            [pixel_vec[0] as f32 / image_dim[0] as f32,
             pixel_vec[1] as f32 / image_dim[1] as f32]
        }
    }
}

/// Try to pack several small rectangles into one large rectangle. Return
/// offsets for the subrectangles within the container if a packing was found.
fn pack_rectangles<T>(container_dim: [T; 2],
                      dims: &Vec<[T; 2]>)
                      -> Option<Vec<[T; 2]>>
    where T: Num + PartialOrd + Ord + Copy + Signed
{
    let init: T = Zero::zero();
    let total_area = dims.iter()
                         .map(|dim| dim[0] * dim[1])
                         .fold(init, |a, b| a + b);

    // Too much rectangle area to fit in container no matter how you pack it.
    // Fail early.
    if total_area > container_dim[0] * container_dim[1] {
        return None;
    }

    // Take enumeration to keep the original indices around.
    let mut largest_first: Vec<(usize, &[T; 2])> = dims.iter()
                                                       .enumerate()
                                                       .collect();
    largest_first.sort_by(|&(_i, a), &(_j, b)| {
        (b[0] * b[1]).cmp(&(a[0] * a[1]))
    });

    let mut slots = vec![Rect {
                             top: [Zero::zero(), Zero::zero()],
                             size: container_dim,
                         }];

    let mut ret: Vec<[T; 2]> = iter::repeat([Zero::zero(), Zero::zero()])
                                   .take(dims.len())
                                   .collect();

    for i in 0..(largest_first.len()) {
        let (idx, &dim) = largest_first[i];
        match place(dim, &mut slots) {
            Some(pos) => {
                ret[idx] = pos;
            }
            None => {
                return None;
            }
        }
    }

    return Some(ret);

    /// Find the smallest slot in the slot vector that will fit the given
    /// item.
    fn place<T>(dim: [T; 2], slots: &mut Vec<Rect<T>>) -> Option<[T; 2]>
        where T: Num + PartialOrd + Ord + Signed + Copy
    {
        for i in 0..(slots.len()) {
            let Rect { top: slot_pos, size: slot_dim } = slots[i];
            if fits(dim, slot_dim) {
                // Remove the original slot, it gets the item. Add the two new
                // rectangles that form around the item.
                let (new_1, new_2) = remaining_rects(dim,
                                                     Rect {
                                                         top: slot_pos,
                                                         size: slot_dim,
                                                     });
                slots.swap_remove(i);
                slots.push(new_1);
                slots.push(new_2);
                // Sort by area from smallest to largest.
                slots.sort_by(|&a, &b| a.area().cmp(&b.area()));
                return Some(slot_pos);
            }
        }
        None
    }

    /// Return the two remaining parts of container rect when the dim-sized
    /// item is placed in the top left corner.
    fn remaining_rects<T>(dim: [T; 2],
                          Rect { top: rect_pos, size: rect_dim }: Rect<T>)
                          -> (Rect<T>, Rect<T>)
        where T: Num + PartialOrd + Ord + Copy + Signed
    {
        assert!(fits(dim, rect_dim));

        // Choose between making a vertical or a horizontal split
        // based on which leaves a bigger open rectangle.
        let vert_vol = max(rect_dim[0] * (rect_dim[1] - dim[1]),
                           (rect_dim[0] - dim[0]) * dim[1]);
        let horiz_vol = max(dim[0] * (rect_dim[1] - dim[1]),
                            (rect_dim[0] - dim[0]) * rect_dim[1]);

        if vert_vol > horiz_vol {
            //     |AA
            // ----+--
            // BBBBBBB
            // BBBBBBB
            (Rect {
                top: [rect_pos[0] + dim[0], rect_pos[1]],
                size: [rect_dim[0] - dim[0], dim[1]],
            },
             Rect {
                top: [rect_pos[0], rect_pos[1] + dim[1]],
                size: [rect_dim[0], rect_dim[1] - dim[1]],
            })
        } else {
            //     |BB
            // ----+BB
            // AAAA|BB
            // AAAA|BB
            (Rect {
                top: [rect_pos[0], rect_pos[1] + dim[1]],
                size: [dim[0], rect_dim[1] - dim[1]],
            },
             Rect {
                top: [rect_pos[0] + dim[0], rect_pos[1]],
                size: [rect_dim[0] - dim[0], rect_dim[1]],
            })
        }
    }

    fn fits<T: Ord>(dim: [T; 2], container_dim: [T; 2]) -> bool {
        dim[0] <= container_dim[0] && dim[1] <= container_dim[1]
    }
}
<|endoftext|>