<commit_before>use card_strategy::CardStrategy;

use card::Card;
use card::Suit;
use card::Rank;
use game_status::GameStatus;
use player::PlayerName;

use std::collections::BTreeMap;
use std::collections::BTreeSet;

#[derive(Debug)]
pub struct DefensiveCardStrategy;

/*
    Play 2 of clubs if in hand.
    If going to win a deal, do so with highest ranking card.
*/
impl DefensiveCardStrategy {
    fn score_card<'a>(card: &'a Card, game_status: &'a GameStatus) -> (i32, i32, i32, i32, i32, i32, i32) {
        let potential_points = Self::potential_points(card, game_status);
        let definite_points = if Self::will_win_deal(card, game_status) {
            potential_points
        } else {
            0
        };
        let later_potential_points = 0 - Self::later_potential_points(card, game_status);
        let card_penalty_to_me = Self::card_penalty_to_me(card, game_status);
        let card_penalty = 0 - Self::card_penalty(card, game_status);
        let trouble = Self::trouble_score(card, game_status);
        let card_rank = 0 - (u32::from(card.rank) as i32);
        (definite_points, potential_points, later_potential_points, card_penalty_to_me, card_penalty, trouble, card_rank)
    }

    fn card_penalty(card: &Card, game_status: &GameStatus) -> i32 {
        game_status.round_parameters.card_points.get(card)
            .map(|penalty| *penalty)
            .unwrap_or_default()
    }

    fn card_penalty_to_me(card: &Card, game_status: &GameStatus) -> i32 {
        if Self::can_win_deal(card, game_status) {
            Self::card_penalty(card, game_status)
        } else {
            0
        }
    }

    fn trouble_score(card: &Card, game_status: &GameStatus) -> i32 {
        if Self::can_win_deal(card, game_status) {
            let trouble = u32::from(card.rank) as i32;
            if Self::card_penalty_to_me(card, game_status) < 0 {
                0 - trouble
            } else {
                if Self::will_win_deal(card, game_status) {
                    u32::from(Rank::Ace) as i32
                } else {
                    trouble
                }
            }
        } else {
            0
        }
    }

    fn can_win_deal(card: &Card, game_status: &GameStatus) -> bool {
        game_status.my_in_progress_deal.as_ref().and_then(|deal|
            deal.suit.and_then(|suit|
                deal.deal_cards.iter()
                    .map(|deal_card| deal_card.card)
                    .filter(|card| card.suit == suit)
                    .max()
                    .map(|winning_card| card.suit == suit && card.rank > winning_card.rank))).unwrap_or(true)
    }

    fn will_win_deal(card: &Card, game_status: &GameStatus) -> bool {
        Self::can_win_deal(card, game_status) && (Self::plays_left(game_status).len() == 0 ||
            Self::deal_suit(game_status)
                .and_then(|suit| Self::remaining_cards(game_status).iter()
                    .filter(|card| card.suit == *suit)
                    .max()
                    .map(|winning_card| card.rank > winning_card.rank)).unwrap_or(true))
    }

    fn deal_suit(game_status: &GameStatus) -> Option<&Suit> {
        game_status.my_in_progress_deal.as_ref().and_then(|deal| deal.suit.as_ref())
    }

    fn plays_left(game_status: &GameStatus) -> BTreeSet<&PlayerName> {
        let mut players: BTreeSet<&PlayerName> = game_status.my_game_players.iter()
            .filter(|player| !player.has_turn)
            .map(|player| &player.team_name)
            .collect();

        if let &Some(ref deal) = &game_status.my_in_progress_deal {
            for deal_card in &deal.deal_cards {
                players.remove(&deal_card.player_name);
            }
        }

        players
    }

    fn remaining_cards(game_status: &GameStatus) -> BTreeSet<Card> {
        let mut cards = Card::all();

        for deal in &game_status.my_game_deals {
            for deal_card in &deal.deal_cards {
                cards.remove(&deal_card.card);
            }
        }

        if let &Some(ref deal) = &game_status.my_in_progress_deal {
            for deal_card in &deal.deal_cards {
                cards.remove(&deal_card.card);
            }
        }

        for card in &game_status.my_current_hand {
            cards.remove(&card);
        }

        cards
    }

    fn dealt_cards(game_status: &GameStatus) -> BTreeSet<Card> {
        game_status.my_in_progress_deal.as_ref()
            .map(|deal| deal.deal_cards.iter().map(|deal_card| deal_card.card).collect())
            .unwrap_or_default()
    }

    fn potential_points(card: &Card, game_status: &GameStatus) -> i32 {
        if Self::can_win_deal(card, game_status) {
            let mut cards = Self::remaining_cards(game_status);
            cards.extend(Self::dealt_cards(game_status));
            let cards = cards;

            let suit_points: i32 = cards.iter()
                .filter(|other| other.suit == card.suit)
                .filter(|other| other.rank < card.rank)
                .map(|other| Self::card_penalty(other, game_status))
                .sum();

            let other_points: i32 = cards.iter()
                .filter(|other| other.suit != card.suit)
                .map(|other| Self::card_penalty(other, game_status))
                .sum();

            suit_points + ((Self::chance_of_later_win(card, game_status) * (other_points as f32)) as i32)
        } else {
            0
        }
    }

    fn later_potential_points(card: &Card, game_status: &GameStatus) -> i32 {
        let cards = Self::remaining_cards(game_status);

        let suit_points: i32 = cards.iter()
            .filter(|other| other.suit == card.suit)
            .filter(|other| other.rank < card.rank)
            .map(|other| Self::card_penalty(other, game_status))
            .sum();

        let other_points: i32 = cards.iter()
            .filter(|other| other.suit != card.suit)
            .map(|other| Self::card_penalty(other, game_status))
            .sum();

        suit_points + ((Self::chance_of_later_win(card, game_status) * (other_points as f32)) as i32)
    }

    fn chance_of_later_win(card: &Card, game_status: &GameStatus) -> f32 {
        let cards = Self::remaining_cards(game_status);
        let suit_cards = cards.iter().filter(|other| other.suit == card.suit).collect::<Vec<_>>();
        let will_win_count = suit_cards.iter().filter(|other| other.rank < card.rank).collect::<Vec<_>>().len();
        (will_win_count as f32) / (suit_cards.len() as f32)
    }

}

impl CardStrategy for DefensiveCardStrategy {

    fn pass_cards<'a>(&mut self, game_status: &'a GameStatus) -> Vec<&'a Card> {
        // Need to order cards by potential to win lowest penalty hands
        // make sure high risk cards are passed, including those that will win high risk cards (Queen of Spades)
        let mut initial_hand = game_status.my_initial_hand.iter().collect::<Vec<&Card>>();
        initial_hand.reverse();
        initial_hand.into_iter().take(3).collect()
    }

    fn play_card<'a>(&mut self, game_status: &'a GameStatus, player_name: &PlayerName) -> &'a Card {
        let current_suit = game_status.my_in_progress_deal.as_ref().and_then(|deal| deal.suit);
        let mut valid_cards = game_status.my_current_hand.iter()
            .filter(|card| Some(card.suit) == current_suit)
            .collect::<Vec<&Card>>();

        if valid_cards.is_empty() {
            valid_cards.extend(&game_status.my_current_hand);
        }

        let evaluation = valid_cards.into_iter()
            .map(|card| ((Self::score_card(card, game_status), card), card))
            .collect::<BTreeMap<_,&Card>>();

        // println!("Evaluation: {:?}", evaluation);

        evaluation.values()
            .next()
            .expect("No valid cards to play!")
    }

}

#[cfg(test)]
mod tests {
    use super::*;
    use player::PlayerName;
    use card::Card;
    use card::Rank::*;
    use card::Suit::*;
    use game_status::GameStatus;
    use game_status::dto::GameStatusDto;
    use card_strategy::CardStrategy;

    use try_from::TryFrom;
    use error::Error;

    extern crate serde;
    extern crate serde_json;

    use std::fs::File;
    use std::io::Read;

    fn open_scenario(name: &str) -> GameStatus {
        let file_name = format!("samples/scenarios/{}.json", name);
        let mut game_status_file = File::open(file_name).unwrap();
        let mut game_status_string = String::new();
        game_status_file.read_to_string(&mut game_status_string).unwrap();
        let game_status_dto: GameStatusDto = serde_json::from_str(&game_status_string).map_err(Error::from).unwrap();
        GameStatus::try_from(game_status_dto).unwrap()
    }

    fn should_play(name: &str, expected_card: Card) {
        let player_name = PlayerName::new("Derek Williams");
        let game_status = open_scenario(name);
        let card = DefensiveCardStrategy.play_card(&game_status, &player_name).clone();
        assert_eq!(expected_card, card);
    }

    #[test]
    fn normal_1() {
        should_play("normal 1", Jack.of(Diamond));
    }

    #[test]
    fn normal_2() {
        should_play("normal 2", King.of(Club));
    }

    #[test]
    fn normal_3() {
        should_play("normal 3", Six.of(Spade));
    }

    #[test]
    fn should_play_heart_1() {
        should_play("should play heart 1", Seven.of(Heart));
    }

    #[test]
    fn should_play_heart_2() {
        should_play("should play heart 2", Four.of(Heart));
    }

    #[test]
    fn should_play_high_rank_1() {
        should_play("should play high rank 1", King.of(Diamond));
    }

    #[test]
    fn should_play_high_rank_2() {
        should_play("should play high rank 2", Ace.of(Spade));
    }

    // #[test]
    // fn should_play_high_rank_3() { // Should get rid of high risk high rank
    //     should_play("should play high rank 3", Ace.of(Spade));
    // }

}<commit_msg>smarter card passing<commit_after>use card_strategy::CardStrategy;

use card::Card;
use card::Suit;
use card::Rank;
use game_status::GameStatus;
use player::PlayerName;

use std::collections::BTreeMap;
use std::collections::BTreeSet;

#[derive(Debug)]
pub struct DefensiveCardStrategy;

/*
    Play 2 of clubs if in hand.
    If going to win a deal, do so with highest ranking card.
*/
impl DefensiveCardStrategy {
    fn score_card<'a>(card: &'a Card, game_status: &'a GameStatus) -> (i32, i32, i32, i32, i32, i32, i32) {
        let potential_points = Self::potential_points(card, game_status);
        let definite_points = if Self::will_win_deal(card, game_status) {
            potential_points
        } else {
            0
        };
        let later_potential_points = 0 - Self::later_potential_points(card, game_status);
        let card_penalty_to_me = Self::card_penalty_to_me(card, game_status);
        let card_penalty = 0 - Self::card_penalty(card, game_status);
        let trouble = Self::trouble_score(card, game_status);
        let card_rank = 0 - (u32::from(card.rank) as i32);
        (definite_points, potential_points, later_potential_points, card_penalty_to_me, card_penalty, trouble, card_rank)
    }

    fn card_penalty(card: &Card, game_status: &GameStatus) -> i32 {
        game_status.round_parameters.card_points.get(card)
            .map(|penalty| *penalty)
            .unwrap_or_default()
    }

    fn card_penalty_to_me(card: &Card, game_status: &GameStatus) -> i32 {
        if Self::can_win_deal(card, game_status) {
            Self::card_penalty(card, game_status)
        } else {
            0
        }
    }

    fn trouble_score(card: &Card, game_status: &GameStatus) -> i32 {
        if Self::can_win_deal(card, game_status) {
            let trouble = u32::from(card.rank) as i32;
            if Self::card_penalty_to_me(card, game_status) < 0 {
                0 - trouble
            } else {
                if Self::will_win_deal(card, game_status) {
                    u32::from(Rank::Ace) as i32
                } else {
                    trouble
                }
            }
        } else {
            0
        }
    }

    fn can_win_deal(card: &Card, game_status: &GameStatus) -> bool {
        game_status.my_in_progress_deal.as_ref().and_then(|deal|
            deal.suit.and_then(|suit|
                deal.deal_cards.iter()
                    .map(|deal_card| deal_card.card)
                    .filter(|card| card.suit == suit)
                    .max()
                    .map(|winning_card| card.suit == suit && card.rank > winning_card.rank))).unwrap_or(true)
    }

    fn will_win_deal(card: &Card, game_status: &GameStatus) -> bool {
        Self::can_win_deal(card, game_status) && (Self::plays_left(game_status).len() == 0 ||
            Self::deal_suit(game_status)
                .and_then(|suit| Self::remaining_cards(game_status).iter()
                    .filter(|card| card.suit == *suit)
                    .max()
                    .map(|winning_card| card.rank > winning_card.rank)).unwrap_or(true))
    }

    fn deal_suit(game_status: &GameStatus) -> Option<&Suit> {
        game_status.my_in_progress_deal.as_ref().and_then(|deal| deal.suit.as_ref())
    }

    fn plays_left(game_status: &GameStatus) -> BTreeSet<&PlayerName> {
        let mut players: BTreeSet<&PlayerName> = game_status.my_game_players.iter()
            .filter(|player| !player.has_turn)
            .map(|player| &player.team_name)
            .collect();

        if let &Some(ref deal) = &game_status.my_in_progress_deal {
            for deal_card in &deal.deal_cards {
                players.remove(&deal_card.player_name);
            }
        }

        players
    }

    fn remaining_cards(game_status: &GameStatus) -> BTreeSet<Card> {
        let mut cards = Card::all();

        for deal in &game_status.my_game_deals {
            for deal_card in &deal.deal_cards {
                cards.remove(&deal_card.card);
            }
        }

        if let &Some(ref deal) = &game_status.my_in_progress_deal {
            for deal_card in &deal.deal_cards {
                cards.remove(&deal_card.card);
            }
        }

        for card in &game_status.my_current_hand {
            cards.remove(&card);
        }

        cards
    }

    fn dealt_cards(game_status: &GameStatus) -> BTreeSet<Card> {
        game_status.my_in_progress_deal.as_ref()
            .map(|deal| deal.deal_cards.iter().map(|deal_card| deal_card.card).collect())
            .unwrap_or_default()
    }

    fn potential_points(card: &Card, game_status: &GameStatus) -> i32 {
        if Self::can_win_deal(card, game_status) {
            let mut cards = Self::remaining_cards(game_status);
            cards.extend(Self::dealt_cards(game_status));
            let cards = cards;

            let suit_points: i32 = cards.iter()
                .filter(|other| other.suit == card.suit)
                .filter(|other| other.rank < card.rank)
                .map(|other| Self::card_penalty(other, game_status))
                .sum();

            let other_points: i32 = cards.iter()
                .filter(|other| other.suit != card.suit)
                .map(|other| Self::card_penalty(other, game_status))
                .sum();

            suit_points + ((Self::chance_of_later_win(card, game_status) * (other_points as f32)) as i32)
        } else {
            0
        }
    }

    fn later_potential_points(card: &Card, game_status: &GameStatus) -> i32 {
        let cards = Self::remaining_cards(game_status);

        let suit_points: i32 = cards.iter()
            .filter(|other| other.suit == card.suit)
            .filter(|other| other.rank < card.rank)
            .map(|other| Self::card_penalty(other, game_status))
            .sum();

        let other_points: i32 = cards.iter()
            .filter(|other| other.suit != card.suit)
            .map(|other| Self::card_penalty(other, game_status))
            .sum();

        suit_points + ((Self::chance_of_later_win(card, game_status) * (other_points as f32)) as i32)
    }

    fn chance_of_later_win(card: &Card, game_status: &GameStatus) -> f32 {
        let cards = Self::remaining_cards(game_status);
        let suit_cards = cards.iter().filter(|other| other.suit == card.suit).collect::<Vec<_>>();
        let will_win_count = suit_cards.iter().filter(|other| other.rank < card.rank).collect::<Vec<_>>().len();
        (will_win_count as f32) / (suit_cards.len() as f32)
    }

}

impl CardStrategy for DefensiveCardStrategy {

    fn pass_cards<'a>(&mut self, game_status: &'a GameStatus) -> Vec<&'a Card> {
        game_status.my_initial_hand.iter()
            .map(|card| ((0 - Self::later_potential_points(card, game_status), card), card))
            .collect::<BTreeMap<_, &Card>>()
            .into_iter()
            .map(|kv| kv.1)
            .take(3)
            .collect()
    }

    fn play_card<'a>(&mut self, game_status: &'a GameStatus, player_name: &PlayerName) -> &'a Card {
        let current_suit = game_status.my_in_progress_deal.as_ref().and_then(|deal| deal.suit);
        let mut valid_cards = game_status.my_current_hand.iter()
            .filter(|card| Some(card.suit) == current_suit)
            .collect::<Vec<&Card>>();

        if valid_cards.is_empty() {
            valid_cards.extend(&game_status.my_current_hand);
        }

        let evaluation = valid_cards.into_iter()
            .map(|card| ((Self::score_card(card, game_status), card), card))
            .collect::<BTreeMap<_,&Card>>();

        // println!("Evaluation: {:?}", evaluation);

        evaluation.values()
            .next()
            .expect("No valid cards to play!")
    }

}

#[cfg(test)]
mod tests {
    use super::*;
    use player::PlayerName;
    use card::Card;
    use card::Rank::*;
    use card::Suit::*;
    use game_status::GameStatus;
    use game_status::dto::GameStatusDto;
    use card_strategy::CardStrategy;

    use try_from::TryFrom;
    use error::Error;

    extern crate serde;
    extern crate serde_json;

    use std::fs::File;
    use std::io::Read;

    fn open_scenario(name: &str) -> GameStatus {
        let file_name = format!("samples/scenarios/{}.json", name);
        let mut game_status_file = File::open(file_name).unwrap();
        let mut game_status_string = String::new();
        game_status_file.read_to_string(&mut game_status_string).unwrap();
        let game_status_dto: GameStatusDto = serde_json::from_str(&game_status_string).map_err(Error::from).unwrap();
        GameStatus::try_from(game_status_dto).unwrap()
    }

    fn should_play(name: &str, expected_card: Card) {
        let player_name = PlayerName::new("Derek Williams");
        let game_status = open_scenario(name);
        let card = DefensiveCardStrategy.play_card(&game_status, &player_name).clone();
        assert_eq!(expected_card, card);
    }

    #[test]
    fn normal_1() {
        should_play("normal 1", Jack.of(Diamond));
    }

    #[test]
    fn normal_2() {
        should_play("normal 2", King.of(Club));
    }

    #[test]
    fn normal_3() {
        should_play("normal 3", Six.of(Spade));
    }

    #[test]
    fn should_play_heart_1() {
        should_play("should play heart 1", Seven.of(Heart));
    }

    #[test]
    fn should_play_heart_2() {
        should_play("should play heart 2", Four.of(Heart));
    }

    #[test]
    fn should_play_high_rank_1() {
        should_play("should play high rank 1", King.of(Diamond));
    }

    #[test]
    fn should_play_high_rank_2() {
        should_play("should play high rank 2", Ace.of(Spade));
    }

    // #[test]
    // fn should_play_high_rank_3() { // Should get rid of high risk high rank
    //     should_play("should play high rank 3", Ace.of(Spade));
    // }

}<|endoftext|>


<commit_before>extern crate diesel;
extern crate dotenv;
extern crate kuchiki;
extern crate server;
use dotenv::dotenv;
use diesel::pg::PgConnection;
use diesel::dsl::{not, sql};
use diesel::sql_types::Bool;
use diesel::prelude::*;
use kuchiki::traits::*;
use server::models::{get_advertiser, get_targets, Ad};
use server::start_logging;
use server::schema::ads::dsl::*;
use std::env;

fn main() {
    dotenv().ok();
    start_logging();

    let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
    let conn = PgConnection::establish(&database_url).unwrap();
    let dbads: Vec<Ad> = ads.order(created_at.desc())
        .filter(political_probability.gt(0.70))
        .filter(lang.eq("en-US"))
        .filter(not(suppressed.eq(true)))
        .filter(sql::<Bool>("targeting is not null"))
        // .filter(sql::<Bool>("targeting ilike '%part of an audience%'"))
        .load::<Ad>(&conn)
        .unwrap();
    let mut total = 0;
    let mut still_parses_correctly = 0;
    let mut newly_parses_correctly = 0;
    let mut parse_failed = 0;
    let mut parse_got_worse = 0;
    let mut still_doesnt_parse = 0;
    for ad in dbads {
        total += 1;
        let document = kuchiki::parse_html().one(ad.html.clone());
        // println!("{:?}", ad.targeting);
        // println!("{:?}", get_targets(&ad.targeting));
        // println!("");

        // find any cases where the revised targeting parser removes an element
        ;
        if let Some(ref old_targets) = ad.targets {
            if let Some(new_targets) = get_targets(&ad.targeting) {
                if old_targets.to_string().len() > new_targets.to_string().len() {
                    println!("lost target -->");
                    println!("{:?}", ad.targeting);
                    println!("{:?}", get_targets(&ad.targeting));
                    println!("{:?}", old_targets);
                    println!("");
                    parse_got_worse += 1;
                } else {
                    still_parses_correctly += 1;
                }
            } else {
                if old_targets.as_array().unwrap().len() == 0 {
                    still_doesnt_parse += 1;
                } else {
                    println!("failed to parse: ");
                    println!("{:?}", ad.targeting);
                    println!("{:?}", old_targets);
                    println!("");
                    parse_failed += 1;
                }
            }
        } else {
            if let Some(new_targets) = get_targets(&ad.targeting) {
                if new_targets.as_array().unwrap().len() > 0 {
                    newly_parses_correctly += 1;
                } else {
                    still_doesnt_parse += 1;
                }
            } else {
                still_doesnt_parse += 1;
                println!("{:?}", ad.id);
            }
        }

        diesel::update(ads.find(ad.id))
            .set((
                targets.eq(get_targets(&ad.targeting)),
                advertiser.eq(get_advertiser(&ad.targeting, &document)),
            ))
            .execute(&conn)
            .unwrap();
    }
    println!(
        "{:?}/{:?} parse successfully",
        still_parses_correctly + newly_parses_correctly + parse_got_worse,
        total
    );
    println!(
        "{:?}/{:?} now parse correctly",
        newly_parses_correctly, total
    );
    println!(
        "{:?}/{:?} now parse worse than before",
        parse_got_worse, total
    );
    println!("{:?}/{:?} newly fail", parse_failed, total);
    println!("{:?}/{:?} still fail", still_doesnt_parse, total);
}
<commit_msg>comment cleanup in targeting test script<commit_after>extern crate diesel;
extern crate dotenv;
extern crate kuchiki;
extern crate server;
use dotenv::dotenv;
use diesel::pg::PgConnection;
use diesel::dsl::{not, sql};
use diesel::sql_types::Bool;
use diesel::prelude::*;
use kuchiki::traits::*;
use server::models::{get_advertiser, get_targets, Ad};
use server::start_logging;
use server::schema::ads::dsl::*;
use std::env;

fn main() {
    dotenv().ok();
    start_logging();

    let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
    let conn = PgConnection::establish(&database_url).unwrap();
    let dbads: Vec<Ad> = ads.order(created_at.desc())
        .filter(political_probability.gt(0.70))
        .filter(lang.eq("en-US"))
        .filter(not(suppressed.eq(true)))
        .filter(sql::<Bool>("targeting is not null"))
        // .filter(sql::<Bool>("targeting ilike '%part of an audience%'"))
        .load::<Ad>(&conn)
        .unwrap();
    let mut total = 0;
    let mut still_parses_correctly = 0;
    let mut newly_parses_correctly = 0;
    let mut parse_failed = 0;
    let mut parse_got_worse = 0;
    let mut still_doesnt_parse = 0;
    for ad in dbads {
        total += 1;
        let document = kuchiki::parse_html().one(ad.html.clone());
        // find any cases where the revised targeting parser removes an element
        ;
        if let Some(ref old_targets) = ad.targets {
            if let Some(new_targets) = get_targets(&ad.targeting) {
                if old_targets.to_string().len() > new_targets.to_string().len() {
                    println!("lost target -->");
                    println!("{:?}", ad.targeting);
                    println!("{:?}", get_targets(&ad.targeting));
                    println!("{:?}", old_targets);
                    println!("");
                    parse_got_worse += 1;
                } else {
                    still_parses_correctly += 1;
                }
            } else {
                if old_targets.as_array().unwrap().len() == 0 {
                    still_doesnt_parse += 1;
                } else {
                    parse_failed += 1;
                }
            }
        } else {
            if let Some(new_targets) = get_targets(&ad.targeting) {
                if new_targets.as_array().unwrap().len() > 0 {
                    newly_parses_correctly += 1;
                    println!("{:?}", new_targets.as_array().unwrap())
                // how to parse in Ruby.
                // output = open("targeting_output.txt", 'r').read.gsub("\n24865/24877 parse successfully\n1520/24877 now parse correctly\n0/24877 now parse worse than before\n0/24877 newly fail\n12/24877 still fail\n", "")
                // lines = output.gsub("Object(", "").gsub("String(", "").gsub('")', '"').gsub(')}', '').gsub("})", "}").gsub('\\\\"', '"').split("\n")
                // parsed_lines = lines.map{|line| JSON.parse(line)}
                } else {
                    still_doesnt_parse += 1;
                }
            } else {
                still_doesnt_parse += 1;
            }
        }

        diesel::update(ads.find(ad.id))
            .set((
                targets.eq(get_targets(&ad.targeting)),
                advertiser.eq(get_advertiser(&ad.targeting, &document)),
            ))
            .execute(&conn)
            .unwrap();
    }
    println!(
        "{:?}/{:?} parse successfully",
        still_parses_correctly + newly_parses_correctly + parse_got_worse,
        total
    );
    println!(
        "{:?}/{:?} now parse correctly",
        newly_parses_correctly, total
    );
    println!(
        "{:?}/{:?} now parse worse than before",
        parse_got_worse, total
    );
    println!("{:?}/{:?} newly fail", parse_failed, total);
    println!("{:?}/{:?} still fail", still_doesnt_parse, total);
}
<|endoftext|>


<commit_before>use crate::helpers::*;
use crate::response_builder::ResponseBuilder;
use crate::{ApiError, ApiResult, BoxFut, NetworkService, UrlQuery};
use beacon_chain::{BeaconChain, BeaconChainTypes};
use hyper::{Body, Request};
use serde::Serialize;
use ssz_derive::Encode;
use std::sync::Arc;
use store::Store;
use types::{BeaconBlock, BeaconState, Epoch, EthSpec, Hash256, Slot, Validator};

#[derive(Serialize, Encode)]
pub struct HeadResponse {
    pub slot: Slot,
    pub block_root: Hash256,
    pub state_root: Hash256,
    pub finalized_slot: Slot,
    pub finalized_block_root: Hash256,
    pub justified_slot: Slot,
    pub justified_block_root: Hash256,
    pub previous_justified_slot: Slot,
    pub previous_justified_block_root: Hash256,
}

/// HTTP handler to return a `BeaconBlock` at a given `root` or `slot`.
pub fn get_head<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = req
        .extensions()
        .get::<Arc<BeaconChain<T>>>()
        .ok_or_else(|| ApiError::ServerError("Beacon chain extension missing".to_string()))?;

    let chain_head = beacon_chain.head();

    let head = HeadResponse {
        slot: chain_head.beacon_state.slot,
        block_root: chain_head.beacon_block_root,
        state_root: chain_head.beacon_state_root,
        finalized_slot: chain_head
            .beacon_state
            .finalized_checkpoint
            .epoch
            .start_slot(T::EthSpec::slots_per_epoch()),
        finalized_block_root: chain_head.beacon_state.finalized_checkpoint.root,
        justified_slot: chain_head
            .beacon_state
            .current_justified_checkpoint
            .epoch
            .start_slot(T::EthSpec::slots_per_epoch()),
        justified_block_root: chain_head.beacon_state.current_justified_checkpoint.root,
        previous_justified_slot: chain_head
            .beacon_state
            .previous_justified_checkpoint
            .epoch
            .start_slot(T::EthSpec::slots_per_epoch()),
        previous_justified_block_root: chain_head.beacon_state.previous_justified_checkpoint.root,
    };

    ResponseBuilder::new(&req).body(&head)
}

#[derive(Serialize, Encode)]
#[serde(bound = "T: EthSpec")]
pub struct BlockResponse<T: EthSpec> {
    pub root: Hash256,
    pub beacon_block: BeaconBlock<T>,
}

/// HTTP handler to return a `BeaconBlock` at a given `root` or `slot`.
pub fn get_block<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = req
        .extensions()
        .get::<Arc<BeaconChain<T>>>()
        .ok_or_else(|| ApiError::ServerError("Beacon chain extension missing".to_string()))?;

    let query_params = ["root", "slot"];
    let (key, value) = UrlQuery::from_request(&req)?.first_of(&query_params)?;

    let block_root = match (key.as_ref(), value) {
        ("slot", value) => {
            let target = parse_slot(&value)?;

            block_root_at_slot(&beacon_chain, target).ok_or_else(|| {
                ApiError::NotFound(format!("Unable to find BeaconBlock for slot {:?}", target))
            })?
        }
        ("root", value) => parse_root(&value)?,
        _ => return Err(ApiError::ServerError("Unexpected query parameter".into())),
    };

    let block = beacon_chain
        .store
        .get::<BeaconBlock<T::EthSpec>>(&block_root)?
        .ok_or_else(|| {
            ApiError::NotFound(format!(
                "Unable to find BeaconBlock for root {:?}",
                block_root
            ))
        })?;

    let response = BlockResponse {
        root: block_root,
        beacon_block: block,
    };

    ResponseBuilder::new(&req).body(&response)
}

/// HTTP handler to return a `BeaconBlock` root at a given `slot`.
pub fn get_block_root<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;

    let slot_string = UrlQuery::from_request(&req)?.only_one("slot")?;
    let target = parse_slot(&slot_string)?;

    let root = block_root_at_slot(&beacon_chain, target).ok_or_else(|| {
        ApiError::NotFound(format!("Unable to find BeaconBlock for slot {:?}", target))
    })?;

    ResponseBuilder::new(&req).body(&root)
}

/// HTTP handler to return the `Fork` of the current head.
pub fn get_fork<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;
    ResponseBuilder::new(&req).body(&beacon_chain.head().beacon_state)
}

/// HTTP handler to return the set of validators for an `Epoch`
///
/// The `Epoch` parameter can be any epoch number. If it is not specified,
/// the current epoch is assumed.
pub fn get_validators<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;

    let epoch = match UrlQuery::from_request(&req) {
        // We have some parameters, so make sure it's the epoch one and parse it
        Ok(query) => query
            .only_one("epoch")?
            .parse::<u64>()
            .map(Epoch::from)
            .map_err(|e| {
                ApiError::InvalidQueryParams(format!(
                    "Invalid epoch parameter, must be a u64. {:?}",
                    e
                ))
            })?,
        // In this case, our url query did not contain any parameters, so we take the default
        Err(_) => beacon_chain.epoch().map_err(|e| {
            ApiError::ServerError(format!("Unable to determine current epoch: {:?}", e))
        })?,
    };

    let all_validators = &beacon_chain.head().beacon_state.validators;
    let active_vals: Vec<Validator> = all_validators
        .iter()
        .filter(|v| v.is_active_at(epoch))
        .cloned()
        .collect();

    ResponseBuilder::new(&req).body(&active_vals)
}

#[derive(Serialize, Encode)]
#[serde(bound = "T: EthSpec")]
pub struct StateResponse<T: EthSpec> {
    pub root: Hash256,
    pub beacon_state: BeaconState<T>,
}

/// HTTP handler to return a `BeaconState` at a given `root` or `slot`.
///
/// Will not return a state if the request slot is in the future. Will return states higher than
/// the current head by skipping slots.
pub fn get_state<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;
    let head_state = beacon_chain.head().beacon_state;

    let (key, value) = match UrlQuery::from_request(&req) {
        Ok(query) => {
            // We have *some* parameters, just check them.
            let query_params = ["root", "slot"];
            query.first_of(&query_params)?
        }
        Err(ApiError::InvalidQueryParams(_)) => {
            // No parameters provided at all, use current slot.
            (String::from("slot"), head_state.slot.to_string())
        }
        Err(e) => {
            return Err(e);
        }
    };

    let (root, state): (Hash256, BeaconState<T::EthSpec>) = match (key.as_ref(), value) {
        ("slot", value) => state_at_slot(&beacon_chain, parse_slot(&value)?)?,
        ("root", value) => {
            let root = &parse_root(&value)?;

            let state = beacon_chain
                .store
                .get(root)?
                .ok_or_else(|| ApiError::NotFound(format!("No state for root: {:?}", root)))?;

            (*root, state)
        }
        _ => return Err(ApiError::ServerError("Unexpected query parameter".into())),
    };

    let response = StateResponse {
        root,
        beacon_state: state,
    };

    ResponseBuilder::new(&req).body(&response)
}

/// HTTP handler to return a `BeaconState` root at a given `slot`.
///
/// Will not return a state if the request slot is in the future. Will return states higher than
/// the current head by skipping slots.
pub fn get_state_root<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;

    let slot_string = UrlQuery::from_request(&req)?.only_one("slot")?;
    let slot = parse_slot(&slot_string)?;

    let root = state_root_at_slot(&beacon_chain, slot)?;

    ResponseBuilder::new(&req).body(&root)
}

/// HTTP handler to return the highest finalized slot.
pub fn get_current_finalized_checkpoint<T: BeaconChainTypes + 'static>(
    req: Request<Body>,
) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;
    let head_state = beacon_chain.head().beacon_state;

    let checkpoint = head_state.finalized_checkpoint.clone();

    ResponseBuilder::new(&req).body(&checkpoint)
}

/// HTTP handler to return a `BeaconState` at the genesis block.
pub fn get_genesis_state<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;

    let (_root, state) = state_at_slot(&beacon_chain, Slot::new(0))?;

    ResponseBuilder::new(&req).body(&state)
}
<commit_msg>Fixed small bug with get_fork function<commit_after>use crate::helpers::*;
use crate::response_builder::ResponseBuilder;
use crate::{ApiError, ApiResult, BoxFut, NetworkService, UrlQuery};
use beacon_chain::{BeaconChain, BeaconChainTypes};
use hyper::{Body, Request};
use serde::Serialize;
use ssz_derive::Encode;
use std::sync::Arc;
use store::Store;
use types::{BeaconBlock, BeaconState, Epoch, EthSpec, Hash256, Slot, Validator};

#[derive(Serialize, Encode)]
pub struct HeadResponse {
    pub slot: Slot,
    pub block_root: Hash256,
    pub state_root: Hash256,
    pub finalized_slot: Slot,
    pub finalized_block_root: Hash256,
    pub justified_slot: Slot,
    pub justified_block_root: Hash256,
    pub previous_justified_slot: Slot,
    pub previous_justified_block_root: Hash256,
}

/// HTTP handler to return a `BeaconBlock` at a given `root` or `slot`.
pub fn get_head<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = req
        .extensions()
        .get::<Arc<BeaconChain<T>>>()
        .ok_or_else(|| ApiError::ServerError("Beacon chain extension missing".to_string()))?;

    let chain_head = beacon_chain.head();

    let head = HeadResponse {
        slot: chain_head.beacon_state.slot,
        block_root: chain_head.beacon_block_root,
        state_root: chain_head.beacon_state_root,
        finalized_slot: chain_head
            .beacon_state
            .finalized_checkpoint
            .epoch
            .start_slot(T::EthSpec::slots_per_epoch()),
        finalized_block_root: chain_head.beacon_state.finalized_checkpoint.root,
        justified_slot: chain_head
            .beacon_state
            .current_justified_checkpoint
            .epoch
            .start_slot(T::EthSpec::slots_per_epoch()),
        justified_block_root: chain_head.beacon_state.current_justified_checkpoint.root,
        previous_justified_slot: chain_head
            .beacon_state
            .previous_justified_checkpoint
            .epoch
            .start_slot(T::EthSpec::slots_per_epoch()),
        previous_justified_block_root: chain_head.beacon_state.previous_justified_checkpoint.root,
    };

    ResponseBuilder::new(&req).body(&head)
}

#[derive(Serialize, Encode)]
#[serde(bound = "T: EthSpec")]
pub struct BlockResponse<T: EthSpec> {
    pub root: Hash256,
    pub beacon_block: BeaconBlock<T>,
}

/// HTTP handler to return a `BeaconBlock` at a given `root` or `slot`.
pub fn get_block<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = req
        .extensions()
        .get::<Arc<BeaconChain<T>>>()
        .ok_or_else(|| ApiError::ServerError("Beacon chain extension missing".to_string()))?;

    let query_params = ["root", "slot"];
    let (key, value) = UrlQuery::from_request(&req)?.first_of(&query_params)?;

    let block_root = match (key.as_ref(), value) {
        ("slot", value) => {
            let target = parse_slot(&value)?;

            block_root_at_slot(&beacon_chain, target).ok_or_else(|| {
                ApiError::NotFound(format!("Unable to find BeaconBlock for slot {:?}", target))
            })?
        }
        ("root", value) => parse_root(&value)?,
        _ => return Err(ApiError::ServerError("Unexpected query parameter".into())),
    };

    let block = beacon_chain
        .store
        .get::<BeaconBlock<T::EthSpec>>(&block_root)?
        .ok_or_else(|| {
            ApiError::NotFound(format!(
                "Unable to find BeaconBlock for root {:?}",
                block_root
            ))
        })?;

    let response = BlockResponse {
        root: block_root,
        beacon_block: block,
    };

    ResponseBuilder::new(&req).body(&response)
}

/// HTTP handler to return a `BeaconBlock` root at a given `slot`.
pub fn get_block_root<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;

    let slot_string = UrlQuery::from_request(&req)?.only_one("slot")?;
    let target = parse_slot(&slot_string)?;

    let root = block_root_at_slot(&beacon_chain, target).ok_or_else(|| {
        ApiError::NotFound(format!("Unable to find BeaconBlock for slot {:?}", target))
    })?;

    ResponseBuilder::new(&req).body(&root)
}

/// HTTP handler to return the `Fork` of the current head.
pub fn get_fork<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;
    ResponseBuilder::new(&req).body(&beacon_chain.head().beacon_state.fork)
}

/// HTTP handler to return the set of validators for an `Epoch`
///
/// The `Epoch` parameter can be any epoch number. If it is not specified,
/// the current epoch is assumed.
pub fn get_validators<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;

    let epoch = match UrlQuery::from_request(&req) {
        // We have some parameters, so make sure it's the epoch one and parse it
        Ok(query) => query
            .only_one("epoch")?
            .parse::<u64>()
            .map(Epoch::from)
            .map_err(|e| {
                ApiError::InvalidQueryParams(format!(
                    "Invalid epoch parameter, must be a u64. {:?}",
                    e
                ))
            })?,
        // In this case, our url query did not contain any parameters, so we take the default
        Err(_) => beacon_chain.epoch().map_err(|e| {
            ApiError::ServerError(format!("Unable to determine current epoch: {:?}", e))
        })?,
    };

    let all_validators = &beacon_chain.head().beacon_state.validators;
    let active_vals: Vec<Validator> = all_validators
        .iter()
        .filter(|v| v.is_active_at(epoch))
        .cloned()
        .collect();

    ResponseBuilder::new(&req).body(&active_vals)
}

#[derive(Serialize, Encode)]
#[serde(bound = "T: EthSpec")]
pub struct StateResponse<T: EthSpec> {
    pub root: Hash256,
    pub beacon_state: BeaconState<T>,
}

/// HTTP handler to return a `BeaconState` at a given `root` or `slot`.
///
/// Will not return a state if the request slot is in the future. Will return states higher than
/// the current head by skipping slots.
pub fn get_state<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;
    let head_state = beacon_chain.head().beacon_state;

    let (key, value) = match UrlQuery::from_request(&req) {
        Ok(query) => {
            // We have *some* parameters, just check them.
            let query_params = ["root", "slot"];
            query.first_of(&query_params)?
        }
        Err(ApiError::InvalidQueryParams(_)) => {
            // No parameters provided at all, use current slot.
            (String::from("slot"), head_state.slot.to_string())
        }
        Err(e) => {
            return Err(e);
        }
    };

    let (root, state): (Hash256, BeaconState<T::EthSpec>) = match (key.as_ref(), value) {
        ("slot", value) => state_at_slot(&beacon_chain, parse_slot(&value)?)?,
        ("root", value) => {
            let root = &parse_root(&value)?;

            let state = beacon_chain
                .store
                .get(root)?
                .ok_or_else(|| ApiError::NotFound(format!("No state for root: {:?}", root)))?;

            (*root, state)
        }
        _ => return Err(ApiError::ServerError("Unexpected query parameter".into())),
    };

    let response = StateResponse {
        root,
        beacon_state: state,
    };

    ResponseBuilder::new(&req).body(&response)
}

/// HTTP handler to return a `BeaconState` root at a given `slot`.
///
/// Will not return a state if the request slot is in the future. Will return states higher than
/// the current head by skipping slots.
pub fn get_state_root<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;

    let slot_string = UrlQuery::from_request(&req)?.only_one("slot")?;
    let slot = parse_slot(&slot_string)?;

    let root = state_root_at_slot(&beacon_chain, slot)?;

    ResponseBuilder::new(&req).body(&root)
}

/// HTTP handler to return the highest finalized slot.
pub fn get_current_finalized_checkpoint<T: BeaconChainTypes + 'static>(
    req: Request<Body>,
) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;
    let head_state = beacon_chain.head().beacon_state;

    let checkpoint = head_state.finalized_checkpoint.clone();

    ResponseBuilder::new(&req).body(&checkpoint)
}

/// HTTP handler to return a `BeaconState` at the genesis block.
pub fn get_genesis_state<T: BeaconChainTypes + 'static>(req: Request<Body>) -> ApiResult {
    let beacon_chain = get_beacon_chain_from_request::<T>(&req)?;

    let (_root, state) = state_at_slot(&beacon_chain, Slot::new(0))?;

    ResponseBuilder::new(&req).body(&state)
}
<|endoftext|>


<commit_before>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::fs::File;
use std::io::{self, Seek, SeekFrom};
use std::mem::size_of;

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

/// A qcow file. Allows reading/writing clusters and appending clusters.
#[derive(Debug)]
pub struct QcowRawFile {
    file: File,
    cluster_size: u64,
    cluster_mask: u64,
}

impl QcowRawFile {
    /// Creates a `QcowRawFile` from the given `File`, `None` is returned if `cluster_size` is not
    /// a power of two.
    pub fn from(file: File, cluster_size: u64) -> Option<Self> {
        if cluster_size.count_ones() != 1 {
            return None;
        }
        Some(QcowRawFile {
            file,
            cluster_size,
            cluster_mask: cluster_size - 1,
        })
    }

    /// Reads `count` 64 bit offsets and returns them as a vector.
    /// `mask` optionally ands out some of the bits on the file.
    pub fn read_pointer_table(
        &mut self,
        offset: u64,
        count: u64,
        mask: Option<u64>,
    ) -> io::Result<Vec<u64>> {
        let mut table = Vec::with_capacity(count as usize);
        self.file.seek(SeekFrom::Start(offset))?;
        let mask = mask.unwrap_or(0xffff_ffff_ffff_ffff);
        for _ in 0..count {
            table.push(self.file.read_u64::<BigEndian>()? & mask);
        }
        Ok(table)
    }

    /// Reads a cluster's worth of 64 bit offsets and returns them as a vector.
    /// `mask` optionally ands out some of the bits on the file.
    pub fn read_pointer_cluster(&mut self, offset: u64, mask: Option<u64>) -> io::Result<Vec<u64>> {
        let count = self.cluster_size / size_of::<u64>() as u64;
        self.read_pointer_table(offset, count, mask)
    }

    /// Writes `table` of u64 pointers to `offset` in the file.
    /// `non_zero_flags` will be ORed with all non-zero values in `table`.
    /// writing.
    pub fn write_pointer_table(
        &mut self,
        offset: u64,
        table: &[u64],
        non_zero_flags: u64,
    ) -> io::Result<()> {
        self.file.seek(SeekFrom::Start(offset))?;
        for addr in table {
            let val = if *addr == 0 {
                0
            } else {
                *addr | non_zero_flags
            };
            self.file.write_u64::<BigEndian>(val)?;
        }
        Ok(())
    }

    /// Read a refcount block from the file and returns a Vec containing the block.
    /// Always returns a cluster's worth of data.
    pub fn read_refcount_block(&mut self, offset: u64) -> io::Result<Vec<u16>> {
        let count = self.cluster_size / size_of::<u16>() as u64;
        let mut table = Vec::with_capacity(count as usize);
        self.file.seek(SeekFrom::Start(offset))?;
        for _ in 0..count {
            table.push(self.file.read_u16::<BigEndian>()?);
        }
        Ok(table)
    }

    /// Writes a refcount block to the file.
    pub fn write_refcount_block(&mut self, offset: u64, table: &[u16]) -> io::Result<()> {
        self.file.seek(SeekFrom::Start(offset))?;
        for count in table {
            self.file.write_u16::<BigEndian>(*count)?;
        }
        Ok(())
    }

    /// Allocates a new cluster at the end of the current file, return the address.
    pub fn add_cluster_end(&mut self) -> io::Result<u64> {
        // Determine where the new end of the file should be and set_len, which
        // translates to truncate(2).
        let file_end: u64 = self.file.seek(SeekFrom::End(0))?;
        let new_cluster_address: u64 = (file_end + self.cluster_size - 1) & !self.cluster_mask;
        self.file.set_len(new_cluster_address + self.cluster_size)?;

        Ok(new_cluster_address)
    }

    /// Returns a reference to the underlying file.
    pub fn file(&self) -> &File {
        &self.file
    }

    /// Returns a mutable reference to the underlying file.
    pub fn file_mut(&mut self) -> &mut File {
        &mut self.file
    }

    /// Returns the size of the file's clusters.
    pub fn cluster_size(&self) -> u64 {
        self.cluster_size
    }

    /// Returns the offset of `address` within a cluster.
    pub fn cluster_offset(&self, address: u64) -> u64 {
        address & self.cluster_mask
    }
}
<commit_msg>qcow: buffer reads and writes in QcowRawFile<commit_after>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::fs::File;
use std::io::{self, BufWriter, Seek, SeekFrom};
use std::mem::size_of;

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

/// A qcow file. Allows reading/writing clusters and appending clusters.
#[derive(Debug)]
pub struct QcowRawFile {
    file: File,
    cluster_size: u64,
    cluster_mask: u64,
}

impl QcowRawFile {
    /// Creates a `QcowRawFile` from the given `File`, `None` is returned if `cluster_size` is not
    /// a power of two.
    pub fn from(file: File, cluster_size: u64) -> Option<Self> {
        if cluster_size.count_ones() != 1 {
            return None;
        }
        Some(QcowRawFile {
            file,
            cluster_size,
            cluster_mask: cluster_size - 1,
        })
    }

    /// Reads `count` 64 bit offsets and returns them as a vector.
    /// `mask` optionally ands out some of the bits on the file.
    pub fn read_pointer_table(
        &mut self,
        offset: u64,
        count: u64,
        mask: Option<u64>,
    ) -> io::Result<Vec<u64>> {
        let mut table = vec![0; count as usize];
        self.file.seek(SeekFrom::Start(offset))?;
        self.file.read_u64_into::<BigEndian>(&mut table)?;
        if let Some(m) = mask {
            for ptr in table.iter_mut() {
                *ptr &= m;
            }
        }
        Ok(table)
    }

    /// Reads a cluster's worth of 64 bit offsets and returns them as a vector.
    /// `mask` optionally ands out some of the bits on the file.
    pub fn read_pointer_cluster(&mut self, offset: u64, mask: Option<u64>) -> io::Result<Vec<u64>> {
        let count = self.cluster_size / size_of::<u64>() as u64;
        self.read_pointer_table(offset, count, mask)
    }

    /// Writes `table` of u64 pointers to `offset` in the file.
    /// `non_zero_flags` will be ORed with all non-zero values in `table`.
    /// writing.
    pub fn write_pointer_table(
        &mut self,
        offset: u64,
        table: &[u64],
        non_zero_flags: u64,
    ) -> io::Result<()> {
        self.file.seek(SeekFrom::Start(offset))?;
        let mut buffer = BufWriter::new(&self.file);
        for addr in table {
            let val = if *addr == 0 {
                0
            } else {
                *addr | non_zero_flags
            };
            buffer.write_u64::<BigEndian>(val)?;
        }
        Ok(())
    }

    /// Read a refcount block from the file and returns a Vec containing the block.
    /// Always returns a cluster's worth of data.
    pub fn read_refcount_block(&mut self, offset: u64) -> io::Result<Vec<u16>> {
        let count = self.cluster_size / size_of::<u16>() as u64;
        let mut table = vec![0; count as usize];
        self.file.seek(SeekFrom::Start(offset))?;
        self.file.read_u16_into::<BigEndian>(&mut table)?;
        Ok(table)
    }

    /// Writes a refcount block to the file.
    pub fn write_refcount_block(&mut self, offset: u64, table: &[u16]) -> io::Result<()> {
        self.file.seek(SeekFrom::Start(offset))?;
        let mut buffer = BufWriter::new(&self.file);
        for count in table {
            buffer.write_u16::<BigEndian>(*count)?;
        }
        Ok(())
    }

    /// Allocates a new cluster at the end of the current file, return the address.
    pub fn add_cluster_end(&mut self) -> io::Result<u64> {
        // Determine where the new end of the file should be and set_len, which
        // translates to truncate(2).
        let file_end: u64 = self.file.seek(SeekFrom::End(0))?;
        let new_cluster_address: u64 = (file_end + self.cluster_size - 1) & !self.cluster_mask;
        self.file.set_len(new_cluster_address + self.cluster_size)?;

        Ok(new_cluster_address)
    }

    /// Returns a reference to the underlying file.
    pub fn file(&self) -> &File {
        &self.file
    }

    /// Returns a mutable reference to the underlying file.
    pub fn file_mut(&mut self) -> &mut File {
        &mut self.file
    }

    /// Returns the size of the file's clusters.
    pub fn cluster_size(&self) -> u64 {
        self.cluster_size
    }

    /// Returns the offset of `address` within a cluster.
    pub fn cluster_offset(&self, address: u64) -> u64 {
        address & self.cluster_mask
    }
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use aarch64_esr_decoder::{decode, Decoded, FieldInfo};
use std::num::ParseIntError;
use std::ops::Deref;
use wasm_bindgen::prelude::*;
use web_sys::{Document, Element};

#[wasm_bindgen]
pub fn init() {
    #[cfg(feature = "console_error_panic_hook")]
    console_error_panic_hook::set_once();
}

#[wasm_bindgen]
pub fn decode_esr(esr: &str) -> String {
    if let Ok(esr) = parse_number(esr) {
        let decoded = decode(esr);
        if let Ok(decoded) = &decoded {
            show_decoded(esr, decoded);
        }
        format!("{:?}", decoded)
    } else {
        "Invalid ESR".to_string()
    }
}

fn make_cell(
    document: &Document,
    contents: Option<&str>,
    colspan: usize,
) -> Result<Element, JsValue> {
    let cell = document.create_element("td")?;
    cell.set_attribute("colspan", &colspan.to_string())?;
    cell.set_text_content(contents);
    Ok(cell)
}

fn show_decoded(esr: u64, decoded: &Decoded) -> Result<(), JsValue> {
    let document = web_sys::window()
        .expect("Couldn't find window")
        .document()
        .expect("Couldn't find document");
    let table = document
        .get_element_by_id("result_table")
        .expect("Couldn't find result table");
    // Remove existing contents.
    table.set_inner_html("");

    // First row has ESR in binary
    let row = document.create_element("tr")?;
    row.set_attribute("class", "value")?;
    for i in (0..u64::BITS).rev() {
        let bit = esr & (1 << i) != 0;
        let cell = make_cell(&document, Some(if bit { "1" } else { "0" }), 1)?;
        row.append_child(&cell)?;
    }
    table.append_child(&row)?;

    // Top-level description
    let row = document.create_element("tr")?;
    row.set_attribute("class", "description")?;
    let cell = make_cell(&document, decoded.description.as_deref(), 64)?;
    row.append_child(&cell)?;
    table.append_child(&row)?;

    // Top-level field names and values
    let row = document.create_element("tr")?;
    row.set_attribute("class", "name")?;
    let mut last = 64;
    add_field_cells(&document, &row, &decoded.fields, &mut last, |field| {
        Some(field.to_string())
    })?;
    table.append_child(&row)?;

    // Top-level field descriptions
    let row = document.create_element("tr")?;
    row.set_attribute("class", "description")?;
    let mut last = 64;
    add_field_cells(&document, &row, &decoded.fields, &mut last, |field| {
        field
            .decoded
            .as_ref()
            .and_then(|decoded| decoded.description.clone())
    })?;
    table.append_child(&row)?;

    // Second level field names and values
    let row = document.create_element("tr")?;
    row.set_attribute("class", "name")?;
    let mut last = 64;
    for field in &decoded.fields {
        if let Some(field_decoded) = &field.decoded {
            add_field_cells(&document, &row, &field_decoded.fields, &mut last, |field| {
                Some(field.to_string())
            })?;
        }
    }
    table.append_child(&row)?;

    // Second level field descriptions
    let row = document.create_element("tr")?;
    row.set_attribute("class", "description")?;
    let mut last = 64;
    for field in &decoded.fields {
        if let Some(field_decoded) = &field.decoded {
            add_field_cells(&document, &row, &field_decoded.fields, &mut last, |field| {
                field
                    .decoded
                    .as_ref()
                    .and_then(|decoded| decoded.description.clone())
            })?;
        }
    }
    table.append_child(&row)?;

    Ok(())
}

fn add_field_cells<F, S>(
    document: &Document,
    row: &Element,
    fields: &[FieldInfo],
    last: &mut usize,
    get_contents: F,
) -> Result<(), JsValue>
where
    F: Fn(&FieldInfo) -> Option<S>,
    S: Deref<Target = str>,
{
    for field in fields {
        if field.start + field.width != *last {
            // Add a filler
            let cell = make_cell(document, None, *last - field.start - field.width)?;
            row.append_child(&cell)?;
        }
        let cell = make_cell(document, get_contents(field).as_deref(), field.width)?;
        row.append_child(&cell)?;
        *last = field.start;
    }
    Ok(())
}

/// Parse a decimal or hexadecimal number.
fn parse_number(s: &str) -> Result<u64, ParseIntError> {
    if s.starts_with("0x") {
        u64::from_str_radix(&s[2..], 16)
    } else {
        s.parse()
    }
}
<commit_msg>Move function to more logical place.<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use aarch64_esr_decoder::{decode, Decoded, FieldInfo};
use std::num::ParseIntError;
use std::ops::Deref;
use wasm_bindgen::prelude::*;
use web_sys::{Document, Element};

#[wasm_bindgen]
pub fn init() {
    #[cfg(feature = "console_error_panic_hook")]
    console_error_panic_hook::set_once();
}

#[wasm_bindgen]
pub fn decode_esr(esr: &str) -> String {
    if let Ok(esr) = parse_number(esr) {
        let decoded = decode(esr);
        if let Ok(decoded) = &decoded {
            show_decoded(esr, decoded);
        }
        format!("{:?}", decoded)
    } else {
        "Invalid ESR".to_string()
    }
}

fn show_decoded(esr: u64, decoded: &Decoded) -> Result<(), JsValue> {
    let document = web_sys::window()
        .expect("Couldn't find window")
        .document()
        .expect("Couldn't find document");
    let table = document
        .get_element_by_id("result_table")
        .expect("Couldn't find result table");
    // Remove existing contents.
    table.set_inner_html("");

    // First row has ESR in binary
    let row = document.create_element("tr")?;
    row.set_attribute("class", "value")?;
    for i in (0..u64::BITS).rev() {
        let bit = esr & (1 << i) != 0;
        let cell = make_cell(&document, Some(if bit { "1" } else { "0" }), 1)?;
        row.append_child(&cell)?;
    }
    table.append_child(&row)?;

    // Top-level description
    let row = document.create_element("tr")?;
    row.set_attribute("class", "description")?;
    let cell = make_cell(&document, decoded.description.as_deref(), 64)?;
    row.append_child(&cell)?;
    table.append_child(&row)?;

    // Top-level field names and values
    let row = document.create_element("tr")?;
    row.set_attribute("class", "name")?;
    let mut last = 64;
    add_field_cells(&document, &row, &decoded.fields, &mut last, |field| {
        Some(field.to_string())
    })?;
    table.append_child(&row)?;

    // Top-level field descriptions
    let row = document.create_element("tr")?;
    row.set_attribute("class", "description")?;
    let mut last = 64;
    add_field_cells(&document, &row, &decoded.fields, &mut last, |field| {
        field
            .decoded
            .as_ref()
            .and_then(|decoded| decoded.description.clone())
    })?;
    table.append_child(&row)?;

    // Second level field names and values
    let row = document.create_element("tr")?;
    row.set_attribute("class", "name")?;
    let mut last = 64;
    for field in &decoded.fields {
        if let Some(field_decoded) = &field.decoded {
            add_field_cells(&document, &row, &field_decoded.fields, &mut last, |field| {
                Some(field.to_string())
            })?;
        }
    }
    table.append_child(&row)?;

    // Second level field descriptions
    let row = document.create_element("tr")?;
    row.set_attribute("class", "description")?;
    let mut last = 64;
    for field in &decoded.fields {
        if let Some(field_decoded) = &field.decoded {
            add_field_cells(&document, &row, &field_decoded.fields, &mut last, |field| {
                field
                    .decoded
                    .as_ref()
                    .and_then(|decoded| decoded.description.clone())
            })?;
        }
    }
    table.append_child(&row)?;

    Ok(())
}

fn make_cell(
    document: &Document,
    contents: Option<&str>,
    colspan: usize,
) -> Result<Element, JsValue> {
    let cell = document.create_element("td")?;
    cell.set_attribute("colspan", &colspan.to_string())?;
    cell.set_text_content(contents);
    Ok(cell)
}

fn add_field_cells<F, S>(
    document: &Document,
    row: &Element,
    fields: &[FieldInfo],
    last: &mut usize,
    get_contents: F,
) -> Result<(), JsValue>
where
    F: Fn(&FieldInfo) -> Option<S>,
    S: Deref<Target = str>,
{
    for field in fields {
        if field.start + field.width != *last {
            // Add a filler
            let cell = make_cell(document, None, *last - field.start - field.width)?;
            row.append_child(&cell)?;
        }
        let cell = make_cell(document, get_contents(field).as_deref(), field.width)?;
        row.append_child(&cell)?;
        *last = field.start;
    }
    Ok(())
}

/// Parse a decimal or hexadecimal number.
fn parse_number(s: &str) -> Result<u64, ParseIntError> {
    if s.starts_with("0x") {
        u64::from_str_radix(&s[2..], 16)
    } else {
        s.parse()
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeSet;
use std::io::Write;
use std::path::PathBuf;

use cargo_edit::{
    colorize_stderr, find, get_latest_dependency, registry_url, resolve_manifests, set_dep_version,
    shell_status, shell_warn, update_registry_index, CargoResult, Context, CrateSpec, Dependency,
    LocalManifest,
};
use clap::Args;
use indexmap::IndexMap;
use semver::{Op, VersionReq};
use termcolor::{Color, ColorSpec, StandardStream, WriteColor};

/// Upgrade dependencies as specified in the local manifest file (i.e. Cargo.toml).
#[derive(Debug, Args)]
#[clap(version)]
#[clap(after_help = "\
This command differs from `cargo update`, which updates the dependency versions recorded in the \
local lock file (Cargo.lock).

If `<dependency>`(s) are provided, only the specified dependencies will be upgraded. The version \
to upgrade to for each can be specified with e.g. `docopt@0.8.0` or `serde@>=0.9,<2.0`.

Dev, build, and all target dependencies will also be upgraded. Only dependencies from crates.io \
are supported. Git/path dependencies will be ignored.

All packages in the workspace will be upgraded if the `--workspace` flag is supplied. The \
`--workspace` flag may be supplied in the presence of a virtual manifest.

If the '--to-lockfile' flag is supplied, all dependencies will be upgraded to the currently locked \
version as recorded in the Cargo.lock file. This flag requires that the Cargo.lock file is \
up-to-date. If the lock file is missing, or it needs to be updated, cargo-upgrade will exit with \
an error. If the '--to-lockfile' flag is supplied then the network won't be accessed.")]
pub struct UpgradeArgs {
    /// Crates to be upgraded.
    dependency: Vec<String>,

    /// Path to the manifest to upgrade
    #[clap(long, value_name = "PATH", parse(from_os_str))]
    manifest_path: Option<PathBuf>,

    /// Package id of the crate to add this dependency to.
    #[clap(
        long = "package",
        short = 'p',
        value_name = "PKGID",
        conflicts_with = "all",
        conflicts_with = "workspace"
    )]
    pkgid: Vec<String>,

    /// Upgrade all packages in the workspace.
    #[clap(
        long,
        help = "[deprecated in favor of `--workspace`]",
        conflicts_with = "workspace",
        conflicts_with = "pkgid"
    )]
    all: bool,

    /// Upgrade all packages in the workspace.
    #[clap(long, conflicts_with = "all", conflicts_with = "pkgid")]
    workspace: bool,

    /// Print changes to be made without making them.
    #[clap(long)]
    dry_run: bool,

    /// Upgrade dependencies pinned in the manifest.
    #[clap(long)]
    pinned: bool,

    /// Run without accessing the network
    #[clap(long)]
    offline: bool,

    /// Upgrade all packages to the version in the lockfile.
    #[clap(long, conflicts_with = "dependency")]
    to_lockfile: bool,

    /// Crates to exclude and not upgrade.
    #[clap(long)]
    exclude: Vec<String>,

    /// Require `Cargo.toml` to be up to date
    #[clap(long)]
    locked: bool,

    /// Use verbose output
    #[clap(short, long)]
    verbose: bool,

    /// Unstable (nightly-only) flags
    #[clap(short = 'Z', value_name = "FLAG", global = true, arg_enum)]
    unstable_features: Vec<UnstableOptions>,
}

impl UpgradeArgs {
    pub fn exec(self) -> CargoResult<()> {
        exec(self)
    }

    fn workspace(&self) -> bool {
        self.all || self.workspace
    }

    fn resolve_targets(&self) -> CargoResult<Vec<cargo_metadata::Package>> {
        resolve_manifests(
            self.manifest_path.as_deref(),
            self.workspace(),
            self.pkgid.iter().map(|s| s.as_str()).collect::<Vec<_>>(),
        )
    }

    fn verbose<F>(&self, mut callback: F) -> CargoResult<()>
    where
        F: FnMut() -> CargoResult<()>,
    {
        if self.verbose {
            callback()
        } else {
            Ok(())
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, clap::ArgEnum)]
enum UnstableOptions {}

/// Main processing function. Allows us to return a `Result` so that `main` can print pretty error
/// messages.
fn exec(args: UpgradeArgs) -> CargoResult<()> {
    if args.all {
        deprecated_message("The flag `--all` has been deprecated in favor of `--workspace`")?;
    }

    if !args.offline && !args.to_lockfile && std::env::var("CARGO_IS_TEST").is_err() {
        let url = registry_url(&find(args.manifest_path.as_deref())?, None)?;
        update_registry_index(&url, false)?;
    }

    let manifests = args.resolve_targets()?;
    let locked = if std::env::var("CARGO_IS_TEST").is_err() {
        load_lockfile(&manifests)?
    } else {
        load_lockfile(&manifests).unwrap_or_default()
    };

    let selected_dependencies = args
        .dependency
        .iter()
        .map(|name| {
            let spec = CrateSpec::resolve(name)?;
            Ok((spec.name, spec.version_req))
        })
        .collect::<CargoResult<IndexMap<_, _>>>()?;
    let mut processed_keys = BTreeSet::new();

    let mut updated_registries = BTreeSet::new();
    let mut any_crate_modified = false;
    for package in manifests {
        let mut manifest = LocalManifest::try_new(package.manifest_path.as_std_path())?;
        let mut crate_modified = false;
        let manifest_path = manifest.path.clone();
        shell_status("Checking", &format!("{}'s dependencies", package.name))?;
        for dep_table in manifest.get_dependency_tables_mut() {
            for (dep_key, dep_item) in dep_table.iter_mut() {
                let dep_key = dep_key.get();
                processed_keys.insert(dep_key.to_owned());
                if !selected_dependencies.is_empty() && !selected_dependencies.contains_key(dep_key)
                {
                    args.verbose(|| {
                        shell_warn(&format!("ignoring {}, excluded by user", dep_key))
                    })?;
                    continue;
                }
                if args.exclude.contains(&dep_key.to_owned()) {
                    args.verbose(|| {
                        shell_warn(&format!("ignoring {}, excluded by user", dep_key))
                    })?;
                    continue;
                }
                let dependency = match Dependency::from_toml(&manifest_path, dep_key, dep_item) {
                    Ok(dependency) => dependency,
                    Err(err) => {
                        shell_warn(&format!("ignoring {}, invalid entry: {}", dep_key, err))?;
                        continue;
                    }
                };
                let old_version_req = match dependency.source.as_ref().and_then(|s| s.as_registry())
                {
                    Some(registry) => registry.version.clone(),
                    None => {
                        args.verbose(|| {
                            let source = dependency
                                .source()
                                .map(|s| s.to_string())
                                .unwrap_or_else(|| "unknown".to_owned());
                            shell_warn(&format!(
                                "ignoring {}, source is {}",
                                dependency.toml_key(),
                                source
                            ))
                        })?;
                        continue;
                    }
                };
                if !args.pinned {
                    if dependency.rename.is_some() {
                        args.verbose(|| {
                            shell_warn(&format!(
                                "ignoring {}, renamed dependencies are pinned",
                                dependency.toml_key(),
                            ))
                        })?;
                        continue;
                    }

                    if let Ok(version_req) = VersionReq::parse(&old_version_req) {
                        if version_req.comparators.iter().any(|comparator| {
                            matches!(comparator.op, Op::Exact | Op::Less | Op::LessEq)
                        }) {
                            args.verbose(|| {
                                shell_warn(&format!(
                                    "ignoring {}, version ({}) is pinned",
                                    dependency.toml_key(),
                                    old_version_req
                                ))
                            })?;
                            continue;
                        }
                    }
                }

                let new_version_req = if let Some(Some(new_version_req)) =
                    selected_dependencies.get(dependency.toml_key())
                {
                    new_version_req.to_owned()
                } else {
                    // Not checking `selected_dependencies.is_empty`, it was checked earlier
                    let new_version = if args.to_lockfile {
                        match find_locked_version(&dependency.name, &old_version_req, &locked) {
                            Some(new_version) => new_version,
                            None => {
                                args.verbose(|| {
                                    shell_warn(&format!(
                                        "ignoring {}, could not find package: not in lock file",
                                        dependency.toml_key(),
                                    ))
                                })?;
                                continue;
                            }
                        }
                    } else {
                        // Update indices for any alternative registries, unless
                        // we're offline.
                        let registry_url = dependency
                            .registry()
                            .map(|registry| registry_url(&manifest_path, Some(registry)))
                            .transpose()?;
                        if !args.offline && std::env::var("CARGO_IS_TEST").is_err() {
                            if let Some(registry_url) = &registry_url {
                                if updated_registries.insert(registry_url.to_owned()) {
                                    update_registry_index(registry_url, false)?;
                                }
                            }
                        }
                        let is_prerelease = old_version_req.contains('-');
                        let new_version = get_latest_dependency(
                            &dependency.name,
                            is_prerelease,
                            &manifest_path,
                            registry_url.as_ref(),
                        )
                        .map(|d| {
                            d.version()
                                .expect("registry packages always have a version")
                                .to_owned()
                        });
                        let new_version = match new_version {
                            Ok(new_version) => new_version,
                            Err(err) => {
                                args.verbose(|| {
                                    shell_warn(&format!(
                                        "ignoring {}, could not find package: {}",
                                        dependency.toml_key(),
                                        err
                                    ))
                                })?;
                                continue;
                            }
                        };
                        if old_version_compatible(&old_version_req, &new_version) {
                            args.verbose(|| {
                                shell_warn(&format!(
                                    "ignoring {}, version ({}) is compatible with {}",
                                    dependency.toml_key(),
                                    old_version_req,
                                    new_version
                                ))
                            })?;
                            continue;
                        }
                        new_version
                    };
                    let mut new_version_req = new_version;
                    let new_ver: semver::Version = new_version_req.parse()?;
                    match cargo_edit::upgrade_requirement(&old_version_req, &new_ver) {
                        Ok(Some(version)) => {
                            new_version_req = version;
                        }
                        Err(_) => {}
                        _ => {
                            new_version_req = old_version_req.clone();
                        }
                    }
                    new_version_req
                };
                if new_version_req == old_version_req {
                    args.verbose(|| {
                        shell_warn(&format!(
                            "ignoring {}, version ({}) is unchanged",
                            dependency.toml_key(),
                            new_version_req
                        ))
                    })?;
                    continue;
                }
                print_upgrade(dependency.toml_key(), &old_version_req, &new_version_req)?;
                set_dep_version(dep_item, &new_version_req)?;
                crate_modified = true;
                any_crate_modified = true;
            }
        }
        if !args.dry_run && !args.locked && crate_modified {
            manifest.write()?;
        }
    }

    if args.locked && any_crate_modified {
        anyhow::bail!("cannot upgrade due to `--locked`");
    }

    let unused = selected_dependencies
        .keys()
        .filter(|k| !processed_keys.contains(k.as_str()))
        .map(|k| k.as_str())
        .collect::<Vec<_>>();
    match unused.len() {
        0 => {}
        1 => anyhow::bail!("dependency {} doesn't exist", unused.join(", ")),
        _ => anyhow::bail!("dependencies {} don't exist", unused.join(", ")),
    }

    if args.dry_run {
        shell_warn("aborting upgrade due to dry run")?;
    }

    Ok(())
}

fn load_lockfile(targets: &[cargo_metadata::Package]) -> CargoResult<Vec<cargo_metadata::Package>> {
    // Get locked dependencies. For workspaces with multiple Cargo.toml
    // files, there is only a single lockfile, so it suffices to get
    // metadata for any one of Cargo.toml files.
    let package = targets
        .get(0)
        .ok_or_else(|| anyhow::format_err!("Invalid cargo config"))?;
    let mut cmd = cargo_metadata::MetadataCommand::new();
    cmd.manifest_path(package.manifest_path.clone());
    cmd.features(cargo_metadata::CargoOpt::AllFeatures);
    cmd.other_options(vec!["--locked".to_string()]);

    let result = cmd.exec().with_context(|| "Invalid manifest")?;

    let locked = result
        .packages
        .into_iter()
        .filter(|p| p.source.is_some()) // Source is none for local packages
        .collect::<Vec<_>>();

    Ok(locked)
}

fn find_locked_version(
    dep_name: &str,
    old_version: &str,
    locked: &[cargo_metadata::Package],
) -> Option<String> {
    let req = semver::VersionReq::parse(old_version).ok()?;
    for p in locked {
        if dep_name == p.name && req.matches(&p.version) {
            return Some(p.version.to_string());
        }
    }
    None
}

fn old_version_compatible(old_version_req: &str, new_version: &str) -> bool {
    let old_version_req = match VersionReq::parse(old_version_req) {
        Ok(req) => req,
        Err(_) => return false,
    };

    let new_version = match semver::Version::parse(new_version) {
        Ok(new_version) => new_version,
        // HACK: Skip compatibility checks on incomplete version reqs
        Err(_) => return false,
    };

    old_version_req.matches(&new_version)
}

fn deprecated_message(message: &str) -> CargoResult<()> {
    let colorchoice = colorize_stderr();
    let mut output = StandardStream::stderr(colorchoice);
    output
        .set_color(ColorSpec::new().set_fg(Some(Color::Red)).set_bold(true))
        .with_context(|| "Failed to set output colour")?;
    writeln!(output, "{}", message).with_context(|| "Failed to write deprecated message")?;
    output
        .set_color(&ColorSpec::new())
        .with_context(|| "Failed to clear output colour")?;
    Ok(())
}

/// Print a message if the new dependency version is different from the old one.
fn print_upgrade(dep_name: &str, old_version: &str, new_version: &str) -> CargoResult<()> {
    let old_version = format_version_req(old_version);
    let new_version = format_version_req(new_version);

    shell_status(
        "Upgrading",
        &format!("{dep_name}: {old_version} -> {new_version}"),
    )?;

    Ok(())
}

fn format_version_req(version: &str) -> String {
    if version.chars().next().unwrap_or('0').is_ascii_digit() {
        format!("v{}", version)
    } else {
        version.to_owned()
    }
}
<commit_msg>fix(upgrade): `--to-lockfile` changes all sources<commit_after>use std::collections::BTreeSet;
use std::io::Write;
use std::path::PathBuf;

use cargo_edit::{
    colorize_stderr, find, get_latest_dependency, registry_url, resolve_manifests, set_dep_version,
    shell_status, shell_warn, update_registry_index, CargoResult, Context, CrateSpec, Dependency,
    LocalManifest,
};
use clap::Args;
use indexmap::IndexMap;
use semver::{Op, VersionReq};
use termcolor::{Color, ColorSpec, StandardStream, WriteColor};

/// Upgrade dependencies as specified in the local manifest file (i.e. Cargo.toml).
#[derive(Debug, Args)]
#[clap(version)]
#[clap(after_help = "\
This command differs from `cargo update`, which updates the dependency versions recorded in the \
local lock file (Cargo.lock).

If `<dependency>`(s) are provided, only the specified dependencies will be upgraded. The version \
to upgrade to for each can be specified with e.g. `docopt@0.8.0` or `serde@>=0.9,<2.0`.

Dev, build, and all target dependencies will also be upgraded. Only dependencies from crates.io \
are supported. Git/path dependencies will be ignored.

All packages in the workspace will be upgraded if the `--workspace` flag is supplied. The \
`--workspace` flag may be supplied in the presence of a virtual manifest.

If the '--to-lockfile' flag is supplied, all dependencies will be upgraded to the currently locked \
version as recorded in the Cargo.lock file. This flag requires that the Cargo.lock file is \
up-to-date. If the lock file is missing, or it needs to be updated, cargo-upgrade will exit with \
an error. If the '--to-lockfile' flag is supplied then the network won't be accessed.")]
pub struct UpgradeArgs {
    /// Crates to be upgraded.
    dependency: Vec<String>,

    /// Path to the manifest to upgrade
    #[clap(long, value_name = "PATH", parse(from_os_str))]
    manifest_path: Option<PathBuf>,

    /// Package id of the crate to add this dependency to.
    #[clap(
        long = "package",
        short = 'p',
        value_name = "PKGID",
        conflicts_with = "all",
        conflicts_with = "workspace"
    )]
    pkgid: Vec<String>,

    /// Upgrade all packages in the workspace.
    #[clap(
        long,
        help = "[deprecated in favor of `--workspace`]",
        conflicts_with = "workspace",
        conflicts_with = "pkgid"
    )]
    all: bool,

    /// Upgrade all packages in the workspace.
    #[clap(long, conflicts_with = "all", conflicts_with = "pkgid")]
    workspace: bool,

    /// Print changes to be made without making them.
    #[clap(long)]
    dry_run: bool,

    /// Upgrade dependencies pinned in the manifest.
    #[clap(long)]
    pinned: bool,

    /// Run without accessing the network
    #[clap(long)]
    offline: bool,

    /// Upgrade all packages to the version in the lockfile.
    #[clap(long, conflicts_with = "dependency")]
    to_lockfile: bool,

    /// Crates to exclude and not upgrade.
    #[clap(long)]
    exclude: Vec<String>,

    /// Require `Cargo.toml` to be up to date
    #[clap(long)]
    locked: bool,

    /// Use verbose output
    #[clap(short, long)]
    verbose: bool,

    /// Unstable (nightly-only) flags
    #[clap(short = 'Z', value_name = "FLAG", global = true, arg_enum)]
    unstable_features: Vec<UnstableOptions>,
}

impl UpgradeArgs {
    pub fn exec(self) -> CargoResult<()> {
        exec(self)
    }

    fn workspace(&self) -> bool {
        self.all || self.workspace
    }

    fn resolve_targets(&self) -> CargoResult<Vec<cargo_metadata::Package>> {
        resolve_manifests(
            self.manifest_path.as_deref(),
            self.workspace(),
            self.pkgid.iter().map(|s| s.as_str()).collect::<Vec<_>>(),
        )
    }

    fn verbose<F>(&self, mut callback: F) -> CargoResult<()>
    where
        F: FnMut() -> CargoResult<()>,
    {
        if self.verbose {
            callback()
        } else {
            Ok(())
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, clap::ArgEnum)]
enum UnstableOptions {}

/// Main processing function. Allows us to return a `Result` so that `main` can print pretty error
/// messages.
fn exec(args: UpgradeArgs) -> CargoResult<()> {
    if args.all {
        deprecated_message("The flag `--all` has been deprecated in favor of `--workspace`")?;
    }

    if !args.offline && !args.to_lockfile && std::env::var("CARGO_IS_TEST").is_err() {
        let url = registry_url(&find(args.manifest_path.as_deref())?, None)?;
        update_registry_index(&url, false)?;
    }

    let manifests = args.resolve_targets()?;
    let locked = if std::env::var("CARGO_IS_TEST").is_err() {
        load_lockfile(&manifests)?
    } else {
        load_lockfile(&manifests).unwrap_or_default()
    };

    let selected_dependencies = args
        .dependency
        .iter()
        .map(|name| {
            let spec = CrateSpec::resolve(name)?;
            Ok((spec.name, spec.version_req))
        })
        .collect::<CargoResult<IndexMap<_, _>>>()?;
    let mut processed_keys = BTreeSet::new();

    let mut updated_registries = BTreeSet::new();
    let mut any_crate_modified = false;
    for package in manifests {
        let mut manifest = LocalManifest::try_new(package.manifest_path.as_std_path())?;
        let mut crate_modified = false;
        let manifest_path = manifest.path.clone();
        shell_status("Checking", &format!("{}'s dependencies", package.name))?;
        for dep_table in manifest.get_dependency_tables_mut() {
            for (dep_key, dep_item) in dep_table.iter_mut() {
                let dep_key = dep_key.get();
                processed_keys.insert(dep_key.to_owned());
                if !selected_dependencies.is_empty() && !selected_dependencies.contains_key(dep_key)
                {
                    args.verbose(|| {
                        shell_warn(&format!("ignoring {}, excluded by user", dep_key))
                    })?;
                    continue;
                }
                if args.exclude.contains(&dep_key.to_owned()) {
                    args.verbose(|| {
                        shell_warn(&format!("ignoring {}, excluded by user", dep_key))
                    })?;
                    continue;
                }
                let dependency = match Dependency::from_toml(&manifest_path, dep_key, dep_item) {
                    Ok(dependency) => dependency,
                    Err(err) => {
                        shell_warn(&format!("ignoring {}, invalid entry: {}", dep_key, err))?;
                        continue;
                    }
                };
                let old_version_req = match dependency.version() {
                    Some(version_req) => version_req.to_owned(),
                    None => {
                        args.verbose(|| {
                            let source = dependency
                                .source()
                                .map(|s| s.to_string())
                                .unwrap_or_else(|| "unknown".to_owned());
                            shell_warn(&format!(
                                "ignoring {}, source is {}",
                                dependency.toml_key(),
                                source,
                            ))
                        })?;
                        continue;
                    }
                };
                if !args.pinned {
                    if dependency.rename.is_some() {
                        args.verbose(|| {
                            shell_warn(&format!(
                                "ignoring {}, renamed dependencies are pinned",
                                dependency.toml_key(),
                            ))
                        })?;
                        continue;
                    }

                    if let Ok(version_req) = VersionReq::parse(&old_version_req) {
                        if version_req.comparators.iter().any(|comparator| {
                            matches!(comparator.op, Op::Exact | Op::Less | Op::LessEq)
                        }) {
                            args.verbose(|| {
                                shell_warn(&format!(
                                    "ignoring {}, version ({}) is pinned",
                                    dependency.toml_key(),
                                    old_version_req
                                ))
                            })?;
                            continue;
                        }
                    }
                }

                let new_version_req = if let Some(Some(new_version_req)) =
                    selected_dependencies.get(dependency.toml_key())
                {
                    new_version_req.to_owned()
                } else {
                    // Not checking `selected_dependencies.is_empty`, it was checked earlier
                    let new_version = if args.to_lockfile {
                        match find_locked_version(&dependency.name, &old_version_req, &locked) {
                            Some(new_version) => new_version,
                            None => {
                                args.verbose(|| {
                                    shell_warn(&format!(
                                        "ignoring {}, could not find package: not in lock file",
                                        dependency.toml_key(),
                                    ))
                                })?;
                                continue;
                            }
                        }
                    } else {
                        if dependency
                            .source
                            .as_ref()
                            .and_then(|s| s.as_registry())
                            .is_none()
                        {
                            args.verbose(|| {
                                let source = dependency
                                    .source()
                                    .map(|s| s.to_string())
                                    .unwrap_or_else(|| "unknown".to_owned());
                                shell_warn(&format!(
                                    "ignoring {}, source is {}",
                                    dependency.toml_key(),
                                    source
                                ))
                            })?;
                            continue;
                        }
                        // Update indices for any alternative registries, unless
                        // we're offline.
                        let registry_url = dependency
                            .registry()
                            .map(|registry| registry_url(&manifest_path, Some(registry)))
                            .transpose()?;
                        if !args.offline && std::env::var("CARGO_IS_TEST").is_err() {
                            if let Some(registry_url) = &registry_url {
                                if updated_registries.insert(registry_url.to_owned()) {
                                    update_registry_index(registry_url, false)?;
                                }
                            }
                        }
                        let is_prerelease = old_version_req.contains('-');
                        let new_version = get_latest_dependency(
                            &dependency.name,
                            is_prerelease,
                            &manifest_path,
                            registry_url.as_ref(),
                        )
                        .map(|d| {
                            d.version()
                                .expect("registry packages always have a version")
                                .to_owned()
                        });
                        let new_version = match new_version {
                            Ok(new_version) => new_version,
                            Err(err) => {
                                args.verbose(|| {
                                    shell_warn(&format!(
                                        "ignoring {}, could not find package: {}",
                                        dependency.toml_key(),
                                        err
                                    ))
                                })?;
                                continue;
                            }
                        };
                        if old_version_compatible(&old_version_req, &new_version) {
                            args.verbose(|| {
                                shell_warn(&format!(
                                    "ignoring {}, version ({}) is compatible with {}",
                                    dependency.toml_key(),
                                    old_version_req,
                                    new_version
                                ))
                            })?;
                            continue;
                        }
                        new_version
                    };
                    let mut new_version_req = new_version;
                    let new_ver: semver::Version = new_version_req.parse()?;
                    match cargo_edit::upgrade_requirement(&old_version_req, &new_ver) {
                        Ok(Some(version)) => {
                            new_version_req = version;
                        }
                        Err(_) => {}
                        _ => {
                            new_version_req = old_version_req.clone();
                        }
                    }
                    new_version_req
                };
                if new_version_req == old_version_req {
                    args.verbose(|| {
                        shell_warn(&format!(
                            "ignoring {}, version ({}) is unchanged",
                            dependency.toml_key(),
                            new_version_req
                        ))
                    })?;
                    continue;
                }
                print_upgrade(dependency.toml_key(), &old_version_req, &new_version_req)?;
                set_dep_version(dep_item, &new_version_req)?;
                crate_modified = true;
                any_crate_modified = true;
            }
        }
        if !args.dry_run && !args.locked && crate_modified {
            manifest.write()?;
        }
    }

    if args.locked && any_crate_modified {
        anyhow::bail!("cannot upgrade due to `--locked`");
    }

    let unused = selected_dependencies
        .keys()
        .filter(|k| !processed_keys.contains(k.as_str()))
        .map(|k| k.as_str())
        .collect::<Vec<_>>();
    match unused.len() {
        0 => {}
        1 => anyhow::bail!("dependency {} doesn't exist", unused.join(", ")),
        _ => anyhow::bail!("dependencies {} don't exist", unused.join(", ")),
    }

    if args.dry_run {
        shell_warn("aborting upgrade due to dry run")?;
    }

    Ok(())
}

fn load_lockfile(targets: &[cargo_metadata::Package]) -> CargoResult<Vec<cargo_metadata::Package>> {
    // Get locked dependencies. For workspaces with multiple Cargo.toml
    // files, there is only a single lockfile, so it suffices to get
    // metadata for any one of Cargo.toml files.
    let package = targets
        .get(0)
        .ok_or_else(|| anyhow::format_err!("Invalid cargo config"))?;
    let mut cmd = cargo_metadata::MetadataCommand::new();
    cmd.manifest_path(package.manifest_path.clone());
    cmd.features(cargo_metadata::CargoOpt::AllFeatures);
    cmd.other_options(vec!["--locked".to_string()]);

    let result = cmd.exec().with_context(|| "Invalid manifest")?;

    let locked = result
        .packages
        .into_iter()
        .filter(|p| p.source.is_some()) // Source is none for local packages
        .collect::<Vec<_>>();

    Ok(locked)
}

fn find_locked_version(
    dep_name: &str,
    old_version: &str,
    locked: &[cargo_metadata::Package],
) -> Option<String> {
    let req = semver::VersionReq::parse(old_version).ok()?;
    for p in locked {
        if dep_name == p.name && req.matches(&p.version) {
            return Some(p.version.to_string());
        }
    }
    None
}

fn old_version_compatible(old_version_req: &str, new_version: &str) -> bool {
    let old_version_req = match VersionReq::parse(old_version_req) {
        Ok(req) => req,
        Err(_) => return false,
    };

    let new_version = match semver::Version::parse(new_version) {
        Ok(new_version) => new_version,
        // HACK: Skip compatibility checks on incomplete version reqs
        Err(_) => return false,
    };

    old_version_req.matches(&new_version)
}

fn deprecated_message(message: &str) -> CargoResult<()> {
    let colorchoice = colorize_stderr();
    let mut output = StandardStream::stderr(colorchoice);
    output
        .set_color(ColorSpec::new().set_fg(Some(Color::Red)).set_bold(true))
        .with_context(|| "Failed to set output colour")?;
    writeln!(output, "{}", message).with_context(|| "Failed to write deprecated message")?;
    output
        .set_color(&ColorSpec::new())
        .with_context(|| "Failed to clear output colour")?;
    Ok(())
}

/// Print a message if the new dependency version is different from the old one.
fn print_upgrade(dep_name: &str, old_version: &str, new_version: &str) -> CargoResult<()> {
    let old_version = format_version_req(old_version);
    let new_version = format_version_req(new_version);

    shell_status(
        "Upgrading",
        &format!("{dep_name}: {old_version} -> {new_version}"),
    )?;

    Ok(())
}

fn format_version_req(version: &str) -> String {
    if version.chars().next().unwrap_or('0').is_ascii_digit() {
        format!("v{}", version)
    } else {
        version.to_owned()
    }
}
<|endoftext|>