<commit_before>use fixedbitset::FixedBitSet;

use crate::core::{Card, Hand};

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Round {
    Starting,
    Preflop,
    Flop,
    Turn,
    River,
    Showdown,
}

impl Round {
    pub fn advance(&self) -> Option<Self> {
        match *self {
            Round::Starting => Some(Round::Preflop),
            Round::Preflop => Some(Round::Flop),
            Round::Flop => Some(Round::Turn),
            Round::Turn => Some(Round::River),
            Round::River => Some(Round::Showdown),
            _ => None,
        }
    }
}

pub struct RoundData {
    pub player_active: FixedBitSet,
    // The minimum allowed bet.
    pub min_bet: usize,
    // The value to be called.
    pub bet: usize,
    // How much each player has put in so far.
    pub player_bet: Vec<usize>,
    // The number of times the player has put in money.
    pub bet_count: Vec<usize>,
    // The index of the next player to act.
    pub to_act_idx: usize,
}

impl RoundData {
    pub fn advance(&mut self) {
        loop {
            self.to_act_idx = (self.to_act_idx + 1) % self.player_active.len();
            if self.player_active.count_ones(..) == 0
                || self.player_active.contains(self.to_act_idx)
            {
                break;
            }
        }
    }
}

pub struct GameState {
    pub num_players: usize,
    pub player_active: FixedBitSet,
    total_pot: usize,
    pub stacks: Vec<usize>,
    pub big_blind: usize,
    pub small_blind: usize,
    pub hands: Vec<Hand>,
    pub dealer_idx: usize,
    pub last_raise: usize,
    pub round: Round,
    pub round_data: Vec<RoundData>,
    pub board: Vec<Card>,
}

impl GameState {
    pub fn new(stacks: Vec<usize>, big_blind: usize, small_blind: usize) -> Self {
        let num_players = stacks.len();

        // Everyone's active right now
        let mut active_mask = FixedBitSet::with_capacity(num_players);
        active_mask.set_range(.., true);
        GameState {
            num_players,
            stacks,
            big_blind,
            small_blind,
            player_active: active_mask,
            dealer_idx: 0,
            last_raise: 0,
            total_pot: 0,
            hands: vec![Hand::default(); num_players],
            round: Round::Starting,
            board: vec![],
            round_data: vec![],
        }
    }

    pub fn advance_round(&mut self) {
        match self.round {
            Round::Starting => self.advance_preflop(),
            _ => self.advance_normal(),
        }
    }

    fn advance_preflop(&mut self) {
        self.round = Round::Preflop;
        self.round_data.push(self.round_data());
        self.do_bet(self.small_blind, true);
        self.do_bet(self.big_blind, true);
    }

    fn advance_normal(&mut self) {
        self.round = self.round.advance().unwrap();
        self.round_data.push(self.round_data())
    }

    fn round_data(&self) -> RoundData {
        RoundData {
            player_bet: vec![0; self.num_players],
            bet_count: vec![0; self.num_players],
            bet: 0,
            min_bet: self.big_blind,
            player_active: self.player_active.clone(),
            to_act_idx: self.left_of_dealer(),
        }
    }

    pub fn do_bet(&mut self, ammount: usize, is_forced: bool) -> Option<usize> {
        if let Some(round_idx) = self.round_index() {
            // Grab the current round data.
            let rd = &mut self.round_data[round_idx];

            // Which player is next to act
            let idx = rd.to_act_idx;
            // Can't bet what's not there
            let bet_ammount = ammount.min(self.stacks[idx]);
            // Take the money out.
            self.stacks[idx] -= bet_ammount;
            // If they put money into the pot then they are done this turn.
            rd.player_active.set(idx, is_forced);

            rd.player_bet[idx] += bet_ammount;
            rd.bet_count[idx] += 1;

            let previous_bet = rd.bet;

            // The amount to be called is
            // the maximum anyone has wagered.
            rd.bet = rd.bet.max(rd.player_bet[idx]);

            // Keep the maximum bet ammount. Anything
            // smaller should be due to going all in.
            rd.min_bet = rd.min_bet.max(rd.bet - previous_bet);

            self.total_pot += bet_ammount;

            // We're out and can't continue.
            if self.stacks[idx] == 0 {
                self.player_active.set(idx, false);
                rd.player_active.set(idx, false);
            }

            // Advance the next to act.
            rd.advance();

            Some(bet_ammount)
        } else {
            None
        }
    }

    pub fn left_of_dealer(&self) -> usize {
        (self.dealer_idx + 1) % self.num_players
    }

    fn round_index(&self) -> Option<usize> {
        match self.round {
            Round::Preflop => Some(0),
            Round::Flop => Some(1),
            Round::Turn => Some(2),
            Round::River => Some(3),
            _ => None,
        }
    }
}
<commit_msg>Clean up game state a little more<commit_after>use fixedbitset::FixedBitSet;

use crate::core::{Card, Hand};

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Round {
    Starting,
    Preflop,
    Flop,
    Turn,
    River,
    Showdown,
}

impl Round {
    pub fn advance(&self) -> Option<Self> {
        match *self {
            Round::Starting => Some(Round::Preflop),
            Round::Preflop => Some(Round::Flop),
            Round::Flop => Some(Round::Turn),
            Round::Turn => Some(Round::River),
            Round::River => Some(Round::Showdown),
            _ => None,
        }
    }
}

pub struct RoundData {
    pub player_active: FixedBitSet,
    // The minimum allowed bet.
    pub min_bet: usize,
    // The value to be called.
    pub bet: usize,
    // How much each player has put in so far.
    pub player_bet: Vec<usize>,
    // The number of times the player has put in money.
    pub bet_count: Vec<usize>,
    // The index of the next player to act.
    pub to_act_idx: usize,
}

impl RoundData {
    pub fn advance(&mut self) {
        loop {
            self.to_act_idx = (self.to_act_idx + 1) % self.player_active.len();
            if self.player_active.count_ones(..) == 0
                || self.player_active.contains(self.to_act_idx)
            {
                break;
            }
        }
    }
}

pub struct GameState {
    /// The number of players that started
    pub num_players: usize,
    /// Which players are still active in the game.
    pub player_active: FixedBitSet,
    /// The total ammount in all pots
    total_pot: usize,
    /// How much is left in each player's stack
    pub stacks: Vec<usize>,
    /// The big blind size
    pub big_blind: usize,
    /// The small blind size
    pub small_blind: usize,
    /// The hands for each player. We keep hands
    /// even if the player is not currently active.
    pub hands: Vec<Hand>,
    /// The index of the player who's the dealer
    pub dealer_idx: usize,
    pub round: Round,
    pub round_data: Vec<RoundData>,
    pub board: Vec<Card>,
}

impl GameState {
    pub fn new(stacks: Vec<usize>, big_blind: usize, small_blind: usize) -> Self {
        let num_players = stacks.len();

        // Everyone's active right now
        let mut active_mask = FixedBitSet::with_capacity(num_players);
        active_mask.set_range(.., true);
        GameState {
            num_players,
            stacks,
            big_blind,
            small_blind,
            player_active: active_mask,
            dealer_idx: 0,
            total_pot: 0,
            hands: vec![Hand::default(); num_players],
            round: Round::Starting,
            board: vec![],
            round_data: vec![],
        }
    }

    pub fn num_active_players(&self) -> usize {
        self.player_active.count_ones(..)
    }

    pub fn is_complete(&self) -> bool {
        self.num_active_players() == 1 || self.round == Round::Showdown
    }

    pub fn advance_round(&mut self) {
        match self.round {
            Round::Starting => self.advance_preflop(),
            _ => self.advance_normal(),
        }
    }

    fn advance_preflop(&mut self) {
        self.round = Round::Preflop;
        self.round_data.push(self.round_data());
        self.do_bet(self.small_blind, true);
        self.do_bet(self.big_blind, true);
    }

    fn advance_normal(&mut self) {
        self.round = self.round.advance().unwrap();
        self.round_data.push(self.round_data())
    }

    fn round_data(&self) -> RoundData {
        RoundData {
            player_bet: vec![0; self.num_players],
            bet_count: vec![0; self.num_players],
            bet: 0,
            min_bet: self.big_blind,
            player_active: self.player_active.clone(),
            to_act_idx: self.left_of_dealer(),
        }
    }

    pub fn do_bet(&mut self, ammount: usize, is_forced: bool) -> Option<usize> {
        if let Some(round_idx) = self.round_index() {
            // Grab the current round data.
            let rd = &mut self.round_data[round_idx];

            // Which player is next to act
            let idx = rd.to_act_idx;
            // Can't bet what's not there
            let bet_ammount = ammount.min(self.stacks[idx]);
            // Take the money out.
            self.stacks[idx] -= bet_ammount;
            // If they put money into the pot then they are done this turn.
            rd.player_active.set(idx, is_forced);

            rd.player_bet[idx] += bet_ammount;
            rd.bet_count[idx] += 1;

            let previous_bet = rd.bet;

            // The amount to be called is
            // the maximum anyone has wagered.
            rd.bet = rd.bet.max(rd.player_bet[idx]);

            // Keep the maximum bet ammount. Anything
            // smaller should be due to going all in.
            rd.min_bet = rd.min_bet.max(rd.bet - previous_bet);

            self.total_pot += bet_ammount;

            // We're out and can't continue.
            if self.stacks[idx] == 0 {
                self.player_active.set(idx, false);
                rd.player_active.set(idx, false);
            }

            // Advance the next to act.
            rd.advance();

            Some(bet_ammount)
        } else {
            None
        }
    }

    pub fn left_of_dealer(&self) -> usize {
        (self.dealer_idx + 1) % self.num_players
    }

    fn round_index(&self) -> Option<usize> {
        match self.round {
            Round::Preflop => Some(0),
            Round::Flop => Some(1),
            Round::Turn => Some(2),
            Round::River => Some(3),
            _ => None,
        }
    }
}
<|endoftext|>


<commit_before>#![feature(core)]
#![feature(rand)]
/**
 * Simulate Bell's quantum entanglement experiments
 */

use std::rand::random;

use Direction::{SpinUp, SpinDown, SpinSuper};
use Plan::{Trivial, OddBall};

#[derive(Copy)]
#[derive(Debug)]
#[derive(PartialEq)]
enum Direction {
    SpinUp,
    SpinDown,
    SpinSuper,
}

/*
impl fmt::String for Direction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let r = match self {
            &SpinUp => "SpinUp",
            &SpinDown => "SpinDown",
            &SpinSuper => "SpinSuper"
        };
        write!(f, "{}", r)
    }
}

impl fmt::String for Detector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Detector::{D12, D3, D9};
        let r = match self {
            &D12 => "12 oclock",
            &D3 => "3 oclock",
            &D9 => "9 oclock"
        };
        write!(f, "{}", r)
    }
}
*/

struct Pair<T> {
    lhs: T,
    rhs: T
}

struct Particle {
    spin: Direction
}

enum Plan {
    Trivial, // up-up-up -> down-down-down
    OddBall  // up-down-up -> down-up-down
}

#[derive(Debug)]
#[derive(PartialEq)]
enum Detector {
    D12, // 12 o'clock
    D3, // 3 o'clock
    D9, // 9 o' clock
}

impl Particle {
    fn new_pair () -> Pair<Particle> {
        let d1 = SpinSuper;
        let d2 = SpinSuper;
        let p1 = Particle{spin: d1};
        let p2 = Particle{spin: d2};

        return Pair{lhs: p1, rhs: p2};
    }

    pub fn measure (&mut self, detector: &Detector) {
        self.spin = match detector {
            &Detector::D12 => SpinUp,
            &Detector::D3 => match rand32() {
                0.0  ... 0.25 => SpinDown,
                0.25 ... 1.00 => SpinUp,
                _ => unreachable!()
            },
            &Detector::D9 => match rand32() {
                0.0  ... 0.25 => SpinDown,
                0.25 ... 1.00 => SpinUp,
                _ => unreachable!()
            },
        };
    }

    // NOTE on spooky and premeditated only if measured in the same
    // direction must the spins be opposite

    // measure with with a message
    pub fn spooky (&mut self, friend: &mut Particle) -> Pair<Direction> {

        let detector1 = Detector::D12;
        let detector2 = Particle::get_detector_direction();

        self.measure (&detector1);

        if detector1 == detector2 { // 1/3
            friend.spin = match self.spin {
                SpinUp => SpinDown,
                SpinDown => SpinUp,
                _ => unreachable!()
            };
        } else { // 2/3
            friend.measure(&detector2);
        }

        return Pair{lhs: self.spin, rhs: friend.spin};
    }

    fn get_detector_direction () -> Detector {

        let detector = match rand32() {
            0.0       ... 0.3333333 => Detector::D12,
            0.3333333 ... 0.6666666 => Detector::D3,
            0.6666666 ... 1.00 => Detector::D9,
            _ => unreachable!()
        };

        return detector;
    }

    pub fn hidden_information (&mut self, friend: &mut Particle, plan: Plan) -> Pair<Direction> {

        let detector = Particle::get_detector_direction();
        let spin = match plan {
            Trivial => SpinUp,
            OddBall => match detector {
                Detector::D3 => SpinDown,
                _ => SpinUp
            }
        };
        friend.spin = spin;

        let detector = Particle::get_detector_direction();
        let spin = match plan {
            Trivial => SpinDown,
            OddBall => match detector {
                Detector::D3 => SpinUp,
                _ => SpinDown
            }
        };
        self.spin = spin;

        return Pair{lhs: self.spin, rhs: friend.spin};
    }
}

fn rand32 () -> f32 {
    return random::<f32>();
}

/**
 * run spooky for trials number of trials
 */
fn run_spooky (trials: f64) {
    let mut num_different: f64 = 0f64;

    let mut lhs;
    let mut rhs;

    for _ in 0 .. trials as usize {
        let particles = Particle::new_pair();
        lhs = particles.lhs;
        rhs = particles.rhs;

        lhs.spooky(&mut rhs);

        if lhs.spin != rhs.spin {
            num_different += 1.0;
        }
    }

    println!("Percent different for spooky {}%", 100f64 * (num_different) / trials as f64);
    println!("      Should be about 1/2 or 50%");
}

/**
 * Run test for hidden information, set the plan_probability to choose
 * OddBall vs Trivial
 */
fn run_hidden(trials: f64, plan_probability: f32) {

    let mut num_different = 0f64;
    let mut lhs;
    let mut rhs;

    for _ in 0 .. trials as usize {
        let particles = Particle::new_pair();
        lhs = particles.lhs;
        rhs = particles.rhs;

        if rand32() < plan_probability {
            lhs.hidden_information(&mut rhs, OddBall);
        } else {
            lhs.hidden_information(&mut rhs, Trivial);
        }

        //println!("lhs.spin {}, rhs.spin {}", lhs.spin, rhs.spin);
        if lhs.spin != rhs.spin {
            num_different += 1.0;
        }
    }

    println!("Percent different for hidden info {}%", 100f64 * (num_different) / trials as f64);
    println!("  Should be greater than 5/9th or {}%", 100.0 * 5.0/9.0);
    println!("With OddBall choosen {}% of the time and Trivial {}% of the time",
             100.0 * plan_probability,
             100.0 * (1.0 - plan_probability));

}

// hidden_information should give +55.6% difference, spooky would give 50%
fn main () {

    let trials: f64 = 1000000f64;

    run_spooky(trials);
    run_hidden(trials, 0.5);
    run_hidden(trials, 1.0);
    run_hidden(trials, 0.0);
}
<commit_msg>update to fmt::Display<commit_after>#![feature(core)]
#![feature(rand)]
/**
 * Simulate Bell's quantum entanglement experiments
 */

use std::rand::random;
use std::fmt;

use Direction::{SpinUp, SpinDown, SpinSuper};
use Plan::{Trivial, OddBall};

#[derive(Copy)]
#[derive(Debug)]
#[derive(PartialEq)]
enum Direction {
    SpinUp,
    SpinDown,
    SpinSuper,
}

impl fmt::Display for Direction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let r = match self {
            &SpinUp => "SpinUp",
            &SpinDown => "SpinDown",
            &SpinSuper => "SpinSuper"
        };
        write!(f, "{}", r)
    }
}

impl fmt::Display for Detector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Detector::{D12, D3, D9};
        let r = match self {
            &D12 => "12 oclock",
            &D3 => "3 oclock",
            &D9 => "9 oclock"
        };
        write!(f, "{}", r)
    }
}

struct Pair<T> {
    lhs: T,
    rhs: T
}

struct Particle {
    spin: Direction
}

enum Plan {
    Trivial, // up-up-up -> down-down-down
    OddBall  // up-down-up -> down-up-down
}

#[derive(Debug)]
#[derive(PartialEq)]
enum Detector {
    D12, // 12 o'clock
    D3, // 3 o'clock
    D9, // 9 o' clock
}

impl Particle {
    fn new_pair () -> Pair<Particle> {
        let d1 = SpinSuper;
        let d2 = SpinSuper;
        let p1 = Particle{spin: d1};
        let p2 = Particle{spin: d2};

        return Pair{lhs: p1, rhs: p2};
    }

    pub fn measure (&mut self, detector: &Detector) {
        self.spin = match detector {
            &Detector::D12 => SpinUp,
            &Detector::D3 => match rand32() {
                0.0  ... 0.25 => SpinDown,
                0.25 ... 1.00 => SpinUp,
                _ => unreachable!()
            },
            &Detector::D9 => match rand32() {
                0.0  ... 0.25 => SpinDown,
                0.25 ... 1.00 => SpinUp,
                _ => unreachable!()
            },
        };
    }

    // NOTE on spooky and premeditated only if measured in the same
    // direction must the spins be opposite

    // measure with with a message
    pub fn spooky (&mut self, friend: &mut Particle) -> Pair<Direction> {

        let detector1 = Detector::D12;
        let detector2 = Particle::get_detector_direction();

        self.measure (&detector1);

        if detector1 == detector2 { // 1/3
            friend.spin = match self.spin {
                SpinUp => SpinDown,
                SpinDown => SpinUp,
                _ => unreachable!()
            };
        } else { // 2/3
            friend.measure(&detector2);
        }

        return Pair{lhs: self.spin, rhs: friend.spin};
    }

    fn get_detector_direction () -> Detector {

        let detector = match rand32() {
            0.0       ... 0.3333333 => Detector::D12,
            0.3333333 ... 0.6666666 => Detector::D3,
            0.6666666 ... 1.00 => Detector::D9,
            _ => unreachable!()
        };

        return detector;
    }

    pub fn hidden_information (&mut self, friend: &mut Particle, plan: Plan) -> Pair<Direction> {

        let detector = Particle::get_detector_direction();
        let spin = match plan {
            Trivial => SpinUp,
            OddBall => match detector {
                Detector::D3 => SpinDown,
                _ => SpinUp
            }
        };
        friend.spin = spin;

        let detector = Particle::get_detector_direction();
        let spin = match plan {
            Trivial => SpinDown,
            OddBall => match detector {
                Detector::D3 => SpinUp,
                _ => SpinDown
            }
        };
        self.spin = spin;

        return Pair{lhs: self.spin, rhs: friend.spin};
    }
}

fn rand32 () -> f32 {
    return random::<f32>();
}

/**
 * run spooky for trials number of trials
 */
fn run_spooky (trials: f64) {
    let mut num_different: f64 = 0f64;

    let mut lhs;
    let mut rhs;

    for _ in 0 .. trials as usize {
        let particles = Particle::new_pair();
        lhs = particles.lhs;
        rhs = particles.rhs;

        lhs.spooky(&mut rhs);

        if lhs.spin != rhs.spin {
            num_different += 1.0;
        }
    }

    println!("Percent different for spooky {}%", 100f64 * (num_different) / trials as f64);
    println!("      Should be about 1/2 or 50%");
}

/**
 * Run test for hidden information, set the plan_probability to choose
 * OddBall vs Trivial
 */
fn run_hidden(trials: f64, plan_probability: f32) {

    let mut num_different = 0f64;
    let mut lhs;
    let mut rhs;

    for _ in 0 .. trials as usize {
        let particles = Particle::new_pair();
        lhs = particles.lhs;
        rhs = particles.rhs;

        if rand32() < plan_probability {
            lhs.hidden_information(&mut rhs, OddBall);
        } else {
            lhs.hidden_information(&mut rhs, Trivial);
        }

        //println!("lhs.spin {}, rhs.spin {}", lhs.spin, rhs.spin);
        if lhs.spin != rhs.spin {
            num_different += 1.0;
        }
    }

    println!("Percent different for hidden info {}%", 100f64 * (num_different) / trials as f64);
    println!("  Should be greater than 5/9th or {}%", 100.0 * 5.0/9.0);
    println!("With OddBall choosen {}% of the time and Trivial {}% of the time",
             100.0 * plan_probability,
             100.0 * (1.0 - plan_probability));

}

// hidden_information should give +55.6% difference, spooky would give 50%
fn main () {

    let trials: f64 = 1000000f64;

    run_spooky(trials);
    run_hidden(trials, 0.5);
    run_hidden(trials, 1.0);
    run_hidden(trials, 0.0);
}
<|endoftext|>


<commit_before>use std::mem;
use std::sync::mpsc::Sender;
use std::rc::Rc;
use std::cell::RefCell;

use communication::Pushable;

// TODO : Using an Observer requires a &mut reference, and should have the "No races!" property:
// TODO : If you hold a &mut ref, no one else can call open/push/shut. Don't let go of that &mut!
// TODO : Probably a good place to insist on RAII... (see ObserverSession)

// observer trait
pub trait Observer {
    type Time;
    type Data;
    fn open(&mut self, time: &Self::Time);   // new punctuation, essentially ...
    fn push(&mut self, time: &Self::Data);   // reveals push data to the observer.
    fn shut(&mut self, time: &Self::Time);   // indicates that we are done for now.
}

// extension trait for creating an RAII observer session from any observer
pub trait ObserverSessionExt : Observer {
    fn session<'a>(&'a mut self, time: &'a <Self as Observer>::Time) -> ObserverSession<'a, Self>;
}

impl<O: Observer> ObserverSessionExt for O {
    #[inline(always)] fn session<'a>(&'a mut self, time: &'a <O as Observer>::Time) -> ObserverSession<'a, O> {
        ObserverSession::new(self, time)
    }
}

// Attempt at RAII for observers. Intended to prevent mis-sequencing of open/push/shut.
pub struct ObserverSession<'a, O:Observer+'a> where <O as Observer>::Time: 'a {
    observer:   &'a mut O,
    time:       &'a O::Time,
}

impl<'a, O:Observer+'a> Drop for ObserverSession<'a, O> where <O as Observer>::Time: 'a {
    #[inline(always)] fn drop(&mut self) { self.observer.shut(self.time); }
}

impl<'a, O:Observer+'a> ObserverSession<'a, O> where <O as Observer>::Time : 'a {
    #[inline(always)] fn new(obs: &'a mut O, time: &'a O::Time) -> ObserverSession<'a, O> {
        obs.open(time);
        ObserverSession { observer: obs, time: time }
    }
    #[inline(always)] pub fn push(&mut self, data: &O::Data) {
        self.observer.push(data);
    }
}

// implementation for inter-thread queues
impl<T:Clone+'static, D:Clone+'static> Observer for (Vec<D>, Rc<RefCell<Vec<(T, Vec<D>)>>>) {
    type Time = T;
    type Data = D;
    #[inline(always)] fn open(&mut self,_time: &T) { }
    #[inline(always)] fn push(&mut self, data: &D) { self.0.push(data.clone()); }
    #[inline(always)] fn shut(&mut self, time: &T) { let vec = mem::replace(&mut self.0, Vec::new()); self.1.borrow_mut().push((time.clone(), vec)); }
}

// implementation for inter-thread queues
impl<T:Clone+Send+'static, D:Clone+Send+'static> Observer for (Vec<D>, Sender<(T, Vec<D>)>) {
    type Time = T;
    type Data = D;
    #[inline(always)] fn open(&mut self,_time: &T) { }
    #[inline(always)] fn push(&mut self, data: &D) { self.0.push(data.clone()); }
    #[inline(always)] fn shut(&mut self, time: &T) { let vec = mem::replace(&mut self.0, Vec::new()); self.1.send((time.clone(), vec)).ok().expect("send error"); }
}

// an observer broadcasting to many observers
pub struct BroadcastObserver<O: Observer> {
    pub observers:  Vec<O>,
}

impl<O: Observer> Observer for BroadcastObserver<O> {
    type Time = O::Time;
    type Data = O::Data;
    #[inline(always)] fn open(&mut self, time: &O::Time) { for observer in self.observers.iter_mut() { observer.open(time); } }
    #[inline(always)] fn push(&mut self, data: &O::Data) { for observer in self.observers.iter_mut() { observer.push(data); } }
    #[inline(always)] fn shut(&mut self, time: &O::Time) { for observer in self.observers.iter_mut() { observer.shut(time); } }
}

// an observer routing between many observers
pub struct ExchangeObserver<O: Observer, H: Fn(&O::Data) -> u64> {
    pub observers:  Vec<O>,
    pub hash_func:  H,
}

impl<O: Observer, H: Fn(&O::Data) -> u64+'static> Observer for ExchangeObserver<O, H> where O::Data : Clone {
    type Time = O::Time;
    type Data = O::Data;
    #[inline(always)] fn open(&mut self, time: &O::Time) -> () { for observer in self.observers.iter_mut() { observer.open(time); } }
    #[inline(always)] fn push(&mut self, data: &O::Data) -> () {
        let dst = (self.hash_func)(data) % self.observers.len() as u64;
        self.observers[dst as usize].push(data);
    }
    #[inline(always)] fn shut(&mut self, time: &O::Time) -> () { for observer in self.observers.iter_mut() { observer.shut(time); } }
}

// an observer buffering records before sending
pub struct BufferedObserver<D, O: Observer> {
    limit:      u64,
    buffer:     Vec<D>,
    observer:   O,
}

impl<D: Clone+'static, O: Observer<Data = Vec<D>>> Observer for BufferedObserver<D, O> where O::Time : Clone + 'static {
    type Time = O::Time;
    type Data = D;
    #[inline(always)] fn open(&mut self, time: &O::Time) { self.observer.open(time); }
    #[inline(always)] fn push(&mut self, data: &D) -> () {
        self.buffer.push(data.clone());
        if self.buffer.len() as u64 > self.limit {
            self.observer.push(&mut self.buffer);
            self.buffer.clear();
        }
    }
    #[inline(always)] fn shut(&mut self, time: &O::Time) -> () {
        self.observer.push(&self.buffer);
        self.observer.shut(time);
        self.buffer.clear();
    }
}

// dual to BufferedObserver, flattens out buffers
pub struct FlattenedObserver<O: Observer> {
    observer:   O,
}

impl<O: Observer> Observer for FlattenedObserver<O> {
    type Time = O::Time;
    type Data = Vec<O::Data>;
    #[inline(always)] fn open(&mut self, time: &O::Time) -> () { self.observer.open(time); }
    #[inline(always)] fn push(&mut self, data: &Vec<O::Data>) -> () { for datum in data.iter() { self.observer.push(datum); } }
    #[inline(always)] fn shut(&mut self, time: &O::Time) -> () { self.observer.shut(time); }
}

// discriminated union of two observers
pub enum ObserverPair<O1: Observer, O2: Observer> {
    Type1(O1),
    Type2(O2),
}

impl<T, D, O1: Observer<Time=T, Data=D>, O2: Observer<Time=T, Data=D>> Observer for ObserverPair<O1, O2> {
    type Time = T;
    type Data = D;
    #[inline(always)]
    fn open(&mut self, time: &T) {
        match *self {
            ObserverPair::Type1(ref mut observer) => observer.open(time),
            ObserverPair::Type2(ref mut observer) => observer.open(time),
        }
    }
    #[inline(always)]
    fn push(&mut self, data: &D) {
        match *self {
            ObserverPair::Type1(ref mut observer) => observer.push(data),
            ObserverPair::Type2(ref mut observer) => observer.push(data),
        }
    }
    #[inline(always)]
    fn shut(&mut self, time: &T) {
        match *self {
            ObserverPair::Type1(ref mut observer) => observer.shut(time),
            ObserverPair::Type2(ref mut observer) => observer.shut(time),
        }
    }
}
<commit_msg>simplified session creation<commit_after>use std::mem;
use std::sync::mpsc::Sender;
use std::rc::Rc;
use std::cell::RefCell;

use communication::Pushable;

// TODO : Using an Observer requires a &mut reference, and should have the "No races!" property:
// TODO : If you hold a &mut ref, no one else can call open/push/shut. Don't let go of that &mut!
// TODO : Probably a good place to insist on RAII... (see ObserverSession)

// observer trait
pub trait Observer {
    type Time;
    type Data;
    fn open(&mut self, time: &Self::Time);   // new punctuation, essentially ...
    fn push(&mut self, time: &Self::Data);   // reveals push data to the observer.
    fn shut(&mut self, time: &Self::Time);   // indicates that we are done for now.
}

// extension trait for creating an RAII observer session from any observer
pub trait ObserverSessionExt : Observer {
    fn session<'a>(&'a mut self, time: &'a <Self as Observer>::Time) -> ObserverSession<'a, Self>;
}

impl<O: Observer> ObserverSessionExt for O {
    #[inline(always)] fn session<'a>(&'a mut self, time: &'a <O as Observer>::Time) -> ObserverSession<'a, O> {
        self.open(time);
        ObserverSession { observer: self, time: time }
    }
}

// Attempt at RAII for observers. Intended to prevent mis-sequencing of open/push/shut.
pub struct ObserverSession<'a, O:Observer+'a> where <O as Observer>::Time: 'a {
    observer:   &'a mut O,
    time:       &'a O::Time,
}

impl<'a, O:Observer+'a> Drop for ObserverSession<'a, O> where <O as Observer>::Time: 'a {
    #[inline(always)] fn drop(&mut self) { self.observer.shut(self.time); }
}

impl<'a, O:Observer+'a> ObserverSession<'a, O> where <O as Observer>::Time : 'a {
    #[inline(always)] pub fn push(&mut self, data: &O::Data) { self.observer.push(data); }
}

// implementation for intra-thread queues (a Vec<D>)
impl<T:Clone+'static, D:Clone+'static> Observer for (Vec<D>, Rc<RefCell<Vec<(T, Vec<D>)>>>) {
    type Time = T;
    type Data = D;
    #[inline(always)] fn open(&mut self,_time: &T) { }
    #[inline(always)] fn push(&mut self, data: &D) { self.0.push(data.clone()); }
    #[inline(always)] fn shut(&mut self, time: &T) { let vec = mem::replace(&mut self.0, Vec::new()); self.1.borrow_mut().push((time.clone(), vec)); }
}

// implementation for inter-thread queues
impl<T:Clone+Send+'static, D:Clone+Send+'static> Observer for (Vec<D>, Sender<(T, Vec<D>)>) {
    type Time = T;
    type Data = D;
    #[inline(always)] fn open(&mut self,_time: &T) { }
    #[inline(always)] fn push(&mut self, data: &D) { self.0.push(data.clone()); }
    #[inline(always)] fn shut(&mut self, time: &T) { let vec = mem::replace(&mut self.0, Vec::new()); self.1.send((time.clone(), vec)).ok().expect("send error"); }
}

// an observer broadcasting to many observers
pub struct BroadcastObserver<O: Observer> {
    pub observers:  Vec<O>,
}

impl<O: Observer> Observer for BroadcastObserver<O> {
    type Time = O::Time;
    type Data = O::Data;
    #[inline(always)] fn open(&mut self, time: &O::Time) { for observer in self.observers.iter_mut() { observer.open(time); } }
    #[inline(always)] fn push(&mut self, data: &O::Data) { for observer in self.observers.iter_mut() { observer.push(data); } }
    #[inline(always)] fn shut(&mut self, time: &O::Time) { for observer in self.observers.iter_mut() { observer.shut(time); } }
}

// an observer routing between many observers
pub struct ExchangeObserver<O: Observer, H: Fn(&O::Data) -> u64> {
    pub observers:  Vec<O>,
    pub hash_func:  H,
}

impl<O: Observer, H: Fn(&O::Data) -> u64+'static> Observer for ExchangeObserver<O, H> where O::Data : Clone {
    type Time = O::Time;
    type Data = O::Data;
    #[inline(always)] fn open(&mut self, time: &O::Time) -> () { for observer in self.observers.iter_mut() { observer.open(time); } }
    #[inline(always)] fn push(&mut self, data: &O::Data) -> () {
        let dst = (self.hash_func)(data) % self.observers.len() as u64;
        self.observers[dst as usize].push(data);
    }
    #[inline(always)] fn shut(&mut self, time: &O::Time) -> () { for observer in self.observers.iter_mut() { observer.shut(time); } }
}

// an observer buffering records before sending
pub struct BufferedObserver<D, O: Observer> {
    limit:      u64,
    buffer:     Vec<D>,
    observer:   O,
}

impl<D: Clone+'static, O: Observer<Data = Vec<D>>> Observer for BufferedObserver<D, O> where O::Time : Clone + 'static {
    type Time = O::Time;
    type Data = D;
    #[inline(always)] fn open(&mut self, time: &O::Time) { self.observer.open(time); }
    #[inline(always)] fn push(&mut self, data: &D) -> () {
        self.buffer.push(data.clone());
        if self.buffer.len() as u64 > self.limit {
            self.observer.push(&mut self.buffer);
            self.buffer.clear();
        }
    }
    #[inline(always)] fn shut(&mut self, time: &O::Time) -> () {
        self.observer.push(&self.buffer);
        self.observer.shut(time);
        self.buffer.clear();
    }
}

// dual to BufferedObserver, flattens out buffers
pub struct FlattenedObserver<O: Observer> {
    observer:   O,
}

impl<O: Observer> Observer for FlattenedObserver<O> {
    type Time = O::Time;
    type Data = Vec<O::Data>;
    #[inline(always)] fn open(&mut self, time: &O::Time) -> () { self.observer.open(time); }
    #[inline(always)] fn push(&mut self, data: &Vec<O::Data>) -> () { for datum in data.iter() { self.observer.push(datum); } }
    #[inline(always)] fn shut(&mut self, time: &O::Time) -> () { self.observer.shut(time); }
}

// discriminated union of two observers
pub enum ObserverPair<O1: Observer, O2: Observer> {
    Type1(O1),
    Type2(O2),
}

impl<T, D, O1: Observer<Time=T, Data=D>, O2: Observer<Time=T, Data=D>> Observer for ObserverPair<O1, O2> {
    type Time = T;
    type Data = D;
    #[inline(always)]
    fn open(&mut self, time: &T) {
        match *self {
            ObserverPair::Type1(ref mut observer) => observer.open(time),
            ObserverPair::Type2(ref mut observer) => observer.open(time),
        }
    }
    #[inline(always)]
    fn push(&mut self, data: &D) {
        match *self {
            ObserverPair::Type1(ref mut observer) => observer.push(data),
            ObserverPair::Type2(ref mut observer) => observer.push(data),
        }
    }
    #[inline(always)]
    fn shut(&mut self, time: &T) {
        match *self {
            ObserverPair::Type1(ref mut observer) => observer.shut(time),
            ObserverPair::Type2(ref mut observer) => observer.shut(time),
        }
    }
}
<|endoftext|>


<commit_before>use std::{cmp::min, env, fs::File, io, mem, time::Duration};

use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, rngs::SmallRng, Rng, SeedableRng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 50;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    filename: String,
    client: reqwest::blocking::Client,
    rng: SmallRng,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(filename: String) -> Importer {
        Importer {
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            rng: SmallRng::from_seed([
                0x19, 0x29, 0xab, 0x17, 0xc6, 0xfa, 0xb0, 0xe9, 0x4b, 0x44, 0xd8, 0x07, 0x09, 0xbf,
                0x1d, 0x87, 0xbd, 0xd8, 0xb3, 0x2f, 0xe1, 0xe2, 0xa0, 0x1a, 0x9e, 0x30, 0x98, 0xd7,
                0xef, 0xd5, 0x7a, 0x1d,
            ]),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put("http://127.0.0.1:9001/import/lichess")
            .json(&self.batch)
            .send()
            .expect("send batch");

        println!(
            "{}: {}: {} - {}",
            self.filename,
            self.batch
                .last()
                .and_then(|g| g.date.as_ref())
                .unwrap_or(&String::new()),
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            if value.as_bytes() == b"rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1" {
                // https://github.com/ornicar/lichess-db/issues/40
                self.current.fen = None;
            } else {
                self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
            }
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(false, |name| name != "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.0,

                _ if rating >= 2500 => 1.0,

                Speed::Rapid if rating >= 2200 => 1.0,
                Speed::Rapid if rating >= 2000 => 0.30,
                Speed::Rapid if rating >= 1800 => 0.18,
                Speed::Rapid if rating >= 1600 => 0.16,

                Speed::Blitz if rating >= 2200 => 0.16,
                Speed::Blitz if rating >= 2000 => 0.07,
                Speed::Blitz if rating >= 1800 => 0.05,
                Speed::Blitz if rating >= 1600 => 0.05,

                Speed::Bullet if rating >= 2200 => 0.20,
                Speed::Bullet if rating >= 2000 => 0.11,
                Speed::Bullet if rating >= 1800 => 0.07,
                Speed::Bullet if rating >= 1600 => 0.07,

                Speed::UltraBullet => 1.0,

                _ => 0.02,
            }
        } else {
            // variant games
            if rating >= 1600 {
                1.0
            } else {
                0.5
            }
        };

        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1501
            && probability >= self.rng.sample(OpenClosed01)
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

fn main() -> Result<(), io::Error> {
    for arg in env::args().skip(1) {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.ends_with(".bz2") {
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(arg);
        reader.read_all(&mut importer)?;
    }

    Ok(())
}
<commit_msg>fix standard bool flipped in pgn indexer<commit_after>use std::{cmp::min, env, fs::File, io, mem, time::Duration};

use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, rngs::SmallRng, Rng, SeedableRng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 50;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    filename: String,
    client: reqwest::blocking::Client,
    rng: SmallRng,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(filename: String) -> Importer {
        Importer {
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            rng: SmallRng::from_seed([
                0x19, 0x29, 0xab, 0x17, 0xc6, 0xfa, 0xb0, 0xe9, 0x4b, 0x44, 0xd8, 0x07, 0x09, 0xbf,
                0x1d, 0x87, 0xbd, 0xd8, 0xb3, 0x2f, 0xe1, 0xe2, 0xa0, 0x1a, 0x9e, 0x30, 0x98, 0xd7,
                0xef, 0xd5, 0x7a, 0x1d,
            ]),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put("http://127.0.0.1:9001/import/lichess")
            .json(&self.batch)
            .send()
            .expect("send batch");

        println!(
            "{}: {}: {} - {}",
            self.filename,
            self.batch
                .last()
                .and_then(|g| g.date.as_ref())
                .unwrap_or(&String::new()),
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            if value.as_bytes() == b"rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1" {
                // https://github.com/ornicar/lichess-db/issues/40
                self.current.fen = None;
            } else {
                self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
            }
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(true, |name| name == "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.0,

                _ if rating >= 2500 => 1.0,

                Speed::Rapid if rating >= 2200 => 1.0,
                Speed::Rapid if rating >= 2000 => 0.30,
                Speed::Rapid if rating >= 1800 => 0.18,
                Speed::Rapid if rating >= 1600 => 0.16,

                Speed::Blitz if rating >= 2200 => 0.16,
                Speed::Blitz if rating >= 2000 => 0.07,
                Speed::Blitz if rating >= 1800 => 0.05,
                Speed::Blitz if rating >= 1600 => 0.05,

                Speed::Bullet if rating >= 2200 => 0.20,
                Speed::Bullet if rating >= 2000 => 0.11,
                Speed::Bullet if rating >= 1800 => 0.07,
                Speed::Bullet if rating >= 1600 => 0.07,

                Speed::UltraBullet => 1.0,

                _ => 0.02,
            }
        } else {
            // variant games
            if rating >= 1600 {
                1.0
            } else {
                0.5
            }
        };

        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1501
            && probability >= self.rng.sample(OpenClosed01)
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

fn main() -> Result<(), io::Error> {
    for arg in env::args().skip(1) {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.ends_with(".bz2") {
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(arg);
        reader.read_all(&mut importer)?;
    }

    Ok(())
}
<|endoftext|>


<commit_before><commit_msg>make sure GRAB_JOB is handled after CAN_DO body<commit_after><|endoftext|>


<commit_before>use std::error;

#[derive(Debug)]
pub struct RootEntry {
    pub filename:  [u8; 8],
    pub extension: [u8; 3],
    attrs: u8,
    pub reserved: u16,
    pub creation_time: u16,
    pub creation_date: u16,
    pub last_access_date: u16,
    _ignore: u16,
    pub last_write_time: u16,
    pub last_write_date: u16,
    pub first_logical_cluster: u16,
    pub file_size: u32, // in bytes
}

#[allow(dead_code)]
impl RootEntry {
    pub fn filename(&self) -> Result<String, Box<error::Error>> {
        let mut my_fn = self.filename.to_vec();
        let mut name = my_fn
            .drain(..)
            .take_while(|&c| c != ' ' as u8)
            .collect::<Vec<u8>>();
        name.push('.' as u8);
        name.extend(self.extension.iter());

        match String::from_utf8(name) {
            Ok(s) => Ok(s),
            Err(err) => Err(From::from(err)),
        }
    }

    pub fn is_read_only(&self)    -> bool {
        self.attrs & 0x01 == 0x01
    }
    pub fn is_hidden(&self)       -> bool {
        self.attrs & 0x02 == 0x02
    }
    pub fn is_system(&self)       -> bool {
        self.attrs & 0x04 == 0x04
    }
    pub fn is_volume_label(&self) -> bool {
        self.attrs & 0x08 == 0x08
    }
    pub fn is_subdir(&self)       -> bool {
        self.attrs & 0x10 == 0x10
    }
    pub fn is_archive(&self)      -> bool {
        self.attrs & 0x20 == 0x20
    }

    pub fn is_free(&self) -> bool {
        self.filename[0] == 0 || self.filename[0] == 0xe5
    }

    pub fn rest_are_free(&self) -> bool {
        self.filename[0] == 0
    }

    pub fn filename_full(&self) -> String {
        let filename = String::from_utf8(
            Vec::from(&self.filename[..])
        );
        let extension = String::from_utf8(
            Vec::from(&self.extension[..])
        );

        if filename.is_ok() && extension.is_ok() {
            format!(
                "{}.{}",
                filename.unwrap(),
                extension.unwrap()
            )
        } else {
            "BAD FILENAME".to_string()
        }
    }
}
<commit_msg>Improve RootEntry pretty-printing<commit_after>use std::error;
use std::fmt;
use std::fmt::Debug;

pub struct RootEntry {
    pub filename:  [u8; 8],
    pub extension: [u8; 3],
    attrs: u8,
    reserved: u16,
    pub creation_time: u16,
    pub creation_date: u16,
    pub last_access_date: u16,
    _ignore: u16,
    pub last_write_time: u16,
    pub last_write_date: u16,
    pub first_logical_cluster: u16,
    pub file_size: u32, // in bytes
}

#[allow(dead_code)]
impl RootEntry {
    pub fn filename(&self) -> Result<String, Box<error::Error>> {
        let mut my_fn = self.filename.to_vec();
        let mut name = my_fn
            .drain(..)
            .take_while(|&c| c != ' ' as u8)
            .collect::<Vec<u8>>();
        name.push('.' as u8);
        name.extend(self.extension.iter());

        match String::from_utf8(name) {
            Ok(s) => Ok(s),
            Err(err) => Err(From::from(err)),
        }
    }

    pub fn is_read_only(&self)    -> bool {
        self.attrs & 0x01 == 0x01
    }
    pub fn is_hidden(&self)       -> bool {
        self.attrs & 0x02 == 0x02
    }
    pub fn is_system(&self)       -> bool {
        self.attrs & 0x04 == 0x04
    }
    pub fn is_volume_label(&self) -> bool {
        self.attrs & 0x08 == 0x08
    }
    pub fn is_subdir(&self)       -> bool {
        self.attrs & 0x10 == 0x10
    }
    pub fn is_archive(&self)      -> bool {
        self.attrs & 0x20 == 0x20
    }

    pub fn is_free(&self) -> bool {
        self.filename[0] == 0 || self.filename[0] == 0xe5
    }

    pub fn rest_are_free(&self) -> bool {
        self.filename[0] == 0
    }

    pub fn filename_full(&self) -> String {
        let filename = String::from_utf8(
            Vec::from(&self.filename[..])
        );
        let extension = String::from_utf8(
            Vec::from(&self.extension[..])
        );

        if filename.is_ok() && extension.is_ok() {
            format!(
                "{}.{}",
                filename.unwrap(),
                extension.unwrap()
            )
        } else {
            "BAD FILENAME".to_string()
        }
    }
}

impl Debug for RootEntry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("RootEntry")
            .field(
                "filename",
                &self.filename().unwrap_or("bad filename".to_string())
            )
            .field("attrs",                 &self.attrs)
            .field("creation_time",         &self.creation_time)
            .field("creation_date",         &self.creation_date)
            .field("last_access_date",      &self.last_access_date)
            .field("last_write_time",       &self.last_write_time)
            .field("last_write_date",       &self.last_write_date)
            .field("first_logical_cluster", &self.first_logical_cluster)
            .field("file_size",             &format!("{:#x}", self.file_size))
            .finish()
    }
}
<|endoftext|>


<commit_before>use auto_enums::auto_enum;
use glob::{glob_with, MatchOptions};
use ion_shell::{
    expansion::{unescape, Expander},
    Shell,
};
use liner::{BasicCompleter, Completer, CursorPosition, Event, EventKind};
use std::{env, iter, path::PathBuf, str};

pub struct IonCompleter<'a, 'b> {
    shell:             &'b Shell<'a>,
    history_completer: Option<BasicCompleter>,
}

/// Escapes filenames from the completer so that special characters will be properly escaped.
///
/// NOTE: Perhaps we should submit a PR to Liner to add a &'static [u8] field to
/// `FilenameCompleter` so that we don't have to perform the escaping ourselves?
fn escape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    for character in input.bytes() {
        match character {
            b'(' | b')' | b'[' | b']' | b'&' | b'$' | b'@' | b'{' | b'}' | b'<' | b'>' | b';'
            | b'"' | b'\'' | b'#' | b'^' | b'*' | b' ' => output.push(b'\\'),
            _ => (),
        }
        output.push(character);
    }
    unsafe { String::from_utf8_unchecked(output) }
}

impl<'a, 'b> IonCompleter<'a, 'b> {
    pub fn new(shell: &'b Shell<'a>) -> Self { IonCompleter { shell, history_completer: None } }
}

impl<'a, 'b> Completer for IonCompleter<'a, 'b> {
    fn completions(&mut self, start: &str) -> Vec<String> {
        let mut completions = IonFileCompleter::new(None, &self.shell).completions(start);

        if let Some(ref mut history) = &mut self.history_completer {
            let vars = self.shell.variables();

            completions.extend(history.completions(start));
            // Initialize a new completer from the definitions collected.
            // Creates a list of definitions from the shell environment that
            // will be used
            // in the creation of a custom completer.
            completions.extend(
                self.shell
                    .builtins()
                    .keys()
                    // Add built-in commands to the completer's definitions.
                    .map(ToString::to_string)
                    // Add the aliases to the completer's definitions.
                    .chain(vars.aliases().map(|(key, _)| key.to_string()))
                    // Add the list of available functions to the completer's
                    // definitions.
                    .chain(vars.functions().map(|(key, _)| key.to_string()))
                    // Add the list of available variables to the completer's
                    // definitions. TODO: We should make
                    // it free to do String->SmallString
                    //       and mostly free to go back (free if allocated)
                    .chain(vars.string_vars().map(|(s, _)| ["$", &s].concat()))
                    .filter(|s| s.starts_with(start)),
            );
            // Creates completers containing definitions from all directories
            // listed
            // in the environment's **$PATH** variable.
            let file_completers: Vec<_> = if let Some(paths) = env::var_os("PATH") {
                env::split_paths(&paths)
                    .map(|s| IonFileCompleter::new(Some(s), &self.shell))
                    .collect()
            } else {
                vec![IonFileCompleter::new(Some("/bin/".into()), &self.shell)]
            };
            // Merge the collected definitions with the file path definitions.
            completions.extend(MultiCompleter::new(file_completers).completions(start));
        }
        completions
    }

    fn on_event<W: std::io::Write>(&mut self, event: Event<'_, '_, W>) {
        if let EventKind::BeforeComplete = event.kind {
            let (words, pos) = event.editor.get_words_and_cursor_position();

            let filename = match pos {
                CursorPosition::InWord(index) => index > 0,
                CursorPosition::InSpace(Some(_), _) => true,
                CursorPosition::InSpace(None, _) => false,
                CursorPosition::OnWordLeftEdge(index) => index >= 1,
                CursorPosition::OnWordRightEdge(index) => words
                    .into_iter()
                    .nth(index)
                    .map(|(start, end)| event.editor.current_buffer().range(start, end))
                    .and_then(|filename| {
                        Some(complete_as_file(&env::current_dir().ok()?, &filename, index))
                    })
                    .filter(|&x| x)
                    .is_some(),
            };

            // Add the history list to the completer's definitions.
            // Map each underlying `liner::Buffer` into a `String`.
            self.history_completer = if filename {
                Some(BasicCompleter::new(
                    event
                        .editor
                        .context()
                        .history
                        .buffers
                        .iter()
                        .map(ToString::to_string)
                        .collect(),
                ))
            } else {
                None
            };
        }
    }
}

/// Performs escaping to an inner `FilenameCompleter` to enable a handful of special cases
/// needed by the shell, such as expanding '~' to a home directory, or adding a backslash
/// when a special character is contained within an expanded filename.
pub struct IonFileCompleter<'a, 'b> {
    shell: &'b Shell<'a>,
    /// The directory the expansion takes place in
    path: PathBuf,
}

impl<'a, 'b> IonFileCompleter<'a, 'b> {
    pub fn new(path: Option<PathBuf>, shell: &'b Shell<'a>) -> Self {
        let path = path.unwrap_or_default();
        IonFileCompleter { shell, path }
    }
}

impl<'a, 'b> Completer for IonFileCompleter<'a, 'b> {
    /// When the tab key is pressed, **Liner** will use this method to perform completions of
    /// filenames. As our `IonFileCompleter` is a wrapper around **Liner**'s
    /// `FilenameCompleter`,
    /// the purpose of our custom `Completer` is to expand possible `~` characters in the
    /// `start`
    /// value that we receive from the prompt, grab completions from the inner
    /// `FilenameCompleter`,
    /// and then escape the resulting filenames, as well as remove the expanded form of the `~`
    /// character and re-add the `~` character in it's place.
    fn completions(&mut self, start: &str) -> Vec<String> {
        // Dereferencing the raw pointers here should be entirely safe, theoretically,
        // because no changes will occur to either of the underlying references in the
        // duration between creation of the completers and execution of their
        // completions.
        let expanded = match self.shell.tilde(start) {
            Ok(expanded) => expanded,
            Err(why) => {
                eprintln!("ion: {}", why);
                return vec![start.into()];
            }
        };
        // Now we obtain completions for the `expanded` form of the `start` value.
        let completions = filename_completion(&expanded, &self.path);
        if expanded == start {
            return completions.collect();
        }
        // We can do that by obtaining the index position where the tilde character
        // ends. We don't search with `~` because we also want to
        // handle other tilde variants.
        let t_index = start.find('/').unwrap_or(1);
        // `tilde` is the tilde pattern, and `search` is the pattern that follows.
        let (tilde, search) = start.split_at(t_index);

        if search.len() < 2 {
            // If the length of the search pattern is less than 2, the search pattern is
            // empty, and thus the completions actually contain files and directories in
            // the home directory.

            // The tilde pattern will actually be our `start` command in itself,
            // and the completed form will be all of the characters beyond the length of
            // the expanded form of the tilde pattern.
            completions.map(|completion| [start, &completion[expanded.len()..]].concat()).collect()
        // To save processing time, we should get obtain the index position where our
        // search pattern begins, and re-use that index to slice the completions so
        // that we may re-add the tilde character with the completion that follows.
        } else if let Some(e_index) = expanded.rfind(search) {
            // And then we will need to take those completions and remove the expanded form
            // of the tilde pattern and replace it with that pattern yet again.
            completions
                .map(|completion| {
                    println!("'{}' : {}", completion, e_index);
                    escape(&[tilde, &completion[e_index..]].concat())
                })
                .collect()
        } else {
            Vec::new()
        }
    }
}

#[auto_enum]
fn filename_completion<'a>(start: &'a str, path: &'a PathBuf) -> impl Iterator<Item = String> + 'a {
    let unescaped_start = unescape(start);

    let mut split_start = unescaped_start.split('/');
    let mut string = String::with_capacity(128);

    // When 'start' is an absolute path, "/..." gets split to ["", "..."]
    // So we skip the first element and add "/" to the start of the string
    if unescaped_start.starts_with('/') {
        split_start.next();
        string.push('/');
    } else {
        string.push_str(&path.to_string_lossy());
    }

    for element in split_start {
        string.push_str(element);
        if element != "." && element != ".." {
            string.push('*');
        }
        string.push('/');
    }

    string.pop(); // pop out the last '/' character
    if string.ends_with('.') {
        string.push('*')
    }

    let globs = glob_with(
        &string,
        MatchOptions {
            case_sensitive:              true,
            require_literal_separator:   true,
            require_literal_leading_dot: false,
        },
    )
    .ok()
    .map(|completions| {
        completions.filter_map(Result::ok).filter_map(move |file| {
            let out = file.to_str()?;
            let mut joined = String::with_capacity(out.len() + 3); // worst case senario
            if unescaped_start.starts_with("./") {
                joined.push_str("./");
            }
            joined.push_str(out);
            if file.is_dir() {
                joined.push('/');
            }
            Some(escape(&joined))
        })
    });

    #[auto_enum(Iterator)]
    match globs {
        Some(iter) => iter,
        None => iter::once(escape(start)),
    }
}

/// A completer that combines suggestions from multiple completers.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MultiCompleter<A>(Vec<A>);

impl<A> MultiCompleter<A> {
    pub fn new(completions: Vec<A>) -> Self { MultiCompleter(completions) }
}

impl<A> Completer for MultiCompleter<A>
where
    A: Completer,
{
    fn completions(&mut self, start: &str) -> Vec<String> {
        self.0.iter_mut().flat_map(|comp| comp.completions(start)).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn filename_completion() {
        let shell = Shell::default();
        let mut completer = IonFileCompleter::new(None, &shell);
        assert_eq!(completer.completions("testing"), vec!["testing/"]);
        assert_eq!(completer.completions("testing/file"), vec!["testing/file_with_text"]);
        assert_eq!(completer.completions("~"), vec!["~/"]);
        assert_eq!(completer.completions("tes/fil"), vec!["testing/file_with_text"]);
    }
}

/// Infer if the given filename is actually a partial filename
fn complete_as_file(current_dir: &PathBuf, filename: &str, index: usize) -> bool {
    let filename = filename.trim();
    let mut file = current_dir.clone();
    file.push(&filename);
    // If the user explicitly requests a file through this syntax then complete as
    // a file
    filename.starts_with('.') ||
    // If the file starts with a dollar sign, it's a variable, not a file
    (!filename.starts_with('$') &&
    // Once we are beyond the first string, assume its a file
    (index > 0 ||
    // If we are referencing a file that exists then just complete to that file
    file.exists() ||
    // If we have a partial file inside an existing directory, e.g. /foo/b when
    // /foo/bar exists, then treat it as file as long as `foo` isn't the
    // current directory, otherwise this would apply to any string `foo`
    file.parent().filter(|parent| parent.exists() && parent != current_dir).is_some()))
    // By default assume its not a file
}
<commit_msg>Remove a left out comment in completion<commit_after>use auto_enums::auto_enum;
use glob::{glob_with, MatchOptions};
use ion_shell::{
    expansion::{unescape, Expander},
    Shell,
};
use liner::{BasicCompleter, Completer, CursorPosition, Event, EventKind};
use std::{env, iter, path::PathBuf, str};

pub struct IonCompleter<'a, 'b> {
    shell:             &'b Shell<'a>,
    history_completer: Option<BasicCompleter>,
}

/// Escapes filenames from the completer so that special characters will be properly escaped.
///
/// NOTE: Perhaps we should submit a PR to Liner to add a &'static [u8] field to
/// `FilenameCompleter` so that we don't have to perform the escaping ourselves?
fn escape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    for character in input.bytes() {
        match character {
            b'(' | b')' | b'[' | b']' | b'&' | b'$' | b'@' | b'{' | b'}' | b'<' | b'>' | b';'
            | b'"' | b'\'' | b'#' | b'^' | b'*' | b' ' => output.push(b'\\'),
            _ => (),
        }
        output.push(character);
    }
    unsafe { String::from_utf8_unchecked(output) }
}

impl<'a, 'b> IonCompleter<'a, 'b> {
    pub fn new(shell: &'b Shell<'a>) -> Self { IonCompleter { shell, history_completer: None } }
}

impl<'a, 'b> Completer for IonCompleter<'a, 'b> {
    fn completions(&mut self, start: &str) -> Vec<String> {
        let mut completions = IonFileCompleter::new(None, &self.shell).completions(start);

        if let Some(ref mut history) = &mut self.history_completer {
            let vars = self.shell.variables();

            completions.extend(history.completions(start));
            // Initialize a new completer from the definitions collected.
            // Creates a list of definitions from the shell environment that
            // will be used
            // in the creation of a custom completer.
            completions.extend(
                self.shell
                    .builtins()
                    .keys()
                    // Add built-in commands to the completer's definitions.
                    .map(ToString::to_string)
                    // Add the aliases to the completer's definitions.
                    .chain(vars.aliases().map(|(key, _)| key.to_string()))
                    // Add the list of available functions to the completer's
                    // definitions.
                    .chain(vars.functions().map(|(key, _)| key.to_string()))
                    // Add the list of available variables to the completer's
                    // definitions. TODO: We should make
                    // it free to do String->SmallString
                    //       and mostly free to go back (free if allocated)
                    .chain(vars.string_vars().map(|(s, _)| ["$", &s].concat()))
                    .filter(|s| s.starts_with(start)),
            );
            // Creates completers containing definitions from all directories
            // listed
            // in the environment's **$PATH** variable.
            let file_completers: Vec<_> = if let Some(paths) = env::var_os("PATH") {
                env::split_paths(&paths)
                    .map(|s| IonFileCompleter::new(Some(s), &self.shell))
                    .collect()
            } else {
                vec![IonFileCompleter::new(Some("/bin/".into()), &self.shell)]
            };
            // Merge the collected definitions with the file path definitions.
            completions.extend(MultiCompleter::new(file_completers).completions(start));
        }
        completions
    }

    fn on_event<W: std::io::Write>(&mut self, event: Event<'_, '_, W>) {
        if let EventKind::BeforeComplete = event.kind {
            let (words, pos) = event.editor.get_words_and_cursor_position();

            let filename = match pos {
                CursorPosition::InWord(index) => index > 0,
                CursorPosition::InSpace(Some(_), _) => true,
                CursorPosition::InSpace(None, _) => false,
                CursorPosition::OnWordLeftEdge(index) => index >= 1,
                CursorPosition::OnWordRightEdge(index) => words
                    .into_iter()
                    .nth(index)
                    .map(|(start, end)| event.editor.current_buffer().range(start, end))
                    .and_then(|filename| {
                        Some(complete_as_file(&env::current_dir().ok()?, &filename, index))
                    })
                    .filter(|&x| x)
                    .is_some(),
            };

            // Add the history list to the completer's definitions.
            // Map each underlying `liner::Buffer` into a `String`.
            self.history_completer = if filename {
                Some(BasicCompleter::new(
                    event
                        .editor
                        .context()
                        .history
                        .buffers
                        .iter()
                        .map(ToString::to_string)
                        .collect(),
                ))
            } else {
                None
            };
        }
    }
}

/// Performs escaping to an inner `FilenameCompleter` to enable a handful of special cases
/// needed by the shell, such as expanding '~' to a home directory, or adding a backslash
/// when a special character is contained within an expanded filename.
pub struct IonFileCompleter<'a, 'b> {
    shell: &'b Shell<'a>,
    /// The directory the expansion takes place in
    path: PathBuf,
}

impl<'a, 'b> IonFileCompleter<'a, 'b> {
    pub fn new(path: Option<PathBuf>, shell: &'b Shell<'a>) -> Self {
        let path = path.unwrap_or_default();
        IonFileCompleter { shell, path }
    }
}

impl<'a, 'b> Completer for IonFileCompleter<'a, 'b> {
    /// When the tab key is pressed, **Liner** will use this method to perform completions of
    /// filenames. As our `IonFileCompleter` is a wrapper around **Liner**'s
    /// `FilenameCompleter`,
    /// the purpose of our custom `Completer` is to expand possible `~` characters in the
    /// `start`
    /// value that we receive from the prompt, grab completions from the inner
    /// `FilenameCompleter`,
    /// and then escape the resulting filenames, as well as remove the expanded form of the `~`
    /// character and re-add the `~` character in it's place.
    fn completions(&mut self, start: &str) -> Vec<String> {
        // Dereferencing the raw pointers here should be entirely safe, theoretically,
        // because no changes will occur to either of the underlying references in the
        // duration between creation of the completers and execution of their
        // completions.
        let expanded = match self.shell.tilde(start) {
            Ok(expanded) => expanded,
            Err(why) => {
                eprintln!("ion: {}", why);
                return vec![start.into()];
            }
        };
        // Now we obtain completions for the `expanded` form of the `start` value.
        let completions = filename_completion(&expanded, &self.path);
        if expanded == start {
            return completions.collect();
        }
        // We can do that by obtaining the index position where the tilde character
        // ends. We don't search with `~` because we also want to
        // handle other tilde variants.
        let t_index = start.find('/').unwrap_or(1);
        // `tilde` is the tilde pattern, and `search` is the pattern that follows.
        let (tilde, search) = start.split_at(t_index);

        if search.len() < 2 {
            // If the length of the search pattern is less than 2, the search pattern is
            // empty, and thus the completions actually contain files and directories in
            // the home directory.

            // The tilde pattern will actually be our `start` command in itself,
            // and the completed form will be all of the characters beyond the length of
            // the expanded form of the tilde pattern.
            completions.map(|completion| [start, &completion[expanded.len()..]].concat()).collect()
        // To save processing time, we should get obtain the index position where our
        // search pattern begins, and re-use that index to slice the completions so
        // that we may re-add the tilde character with the completion that follows.
        } else if let Some(e_index) = expanded.rfind(search) {
            // And then we will need to take those completions and remove the expanded form
            // of the tilde pattern and replace it with that pattern yet again.
            completions
                .map(|completion| escape(&[tilde, &completion[e_index..]].concat()))
                .collect()
        } else {
            Vec::new()
        }
    }
}

#[auto_enum]
fn filename_completion<'a>(start: &'a str, path: &'a PathBuf) -> impl Iterator<Item = String> + 'a {
    let unescaped_start = unescape(start);

    let mut split_start = unescaped_start.split('/');
    let mut string = String::with_capacity(128);

    // When 'start' is an absolute path, "/..." gets split to ["", "..."]
    // So we skip the first element and add "/" to the start of the string
    if unescaped_start.starts_with('/') {
        split_start.next();
        string.push('/');
    } else {
        string.push_str(&path.to_string_lossy());
    }

    for element in split_start {
        string.push_str(element);
        if element != "." && element != ".." {
            string.push('*');
        }
        string.push('/');
    }

    string.pop(); // pop out the last '/' character
    if string.ends_with('.') {
        string.push('*')
    }

    let globs = glob_with(
        &string,
        MatchOptions {
            case_sensitive:              true,
            require_literal_separator:   true,
            require_literal_leading_dot: false,
        },
    )
    .ok()
    .map(|completions| {
        completions.filter_map(Result::ok).filter_map(move |file| {
            let out = file.to_str()?;
            let mut joined = String::with_capacity(out.len() + 3); // worst case senario
            if unescaped_start.starts_with("./") {
                joined.push_str("./");
            }
            joined.push_str(out);
            if file.is_dir() {
                joined.push('/');
            }
            Some(escape(&joined))
        })
    });

    #[auto_enum(Iterator)]
    match globs {
        Some(iter) => iter,
        None => iter::once(escape(start)),
    }
}

/// A completer that combines suggestions from multiple completers.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MultiCompleter<A>(Vec<A>);

impl<A> MultiCompleter<A> {
    pub fn new(completions: Vec<A>) -> Self { MultiCompleter(completions) }
}

impl<A> Completer for MultiCompleter<A>
where
    A: Completer,
{
    fn completions(&mut self, start: &str) -> Vec<String> {
        self.0.iter_mut().flat_map(|comp| comp.completions(start)).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn filename_completion() {
        let shell = Shell::default();
        let mut completer = IonFileCompleter::new(None, &shell);
        assert_eq!(completer.completions("testing"), vec!["testing/"]);
        assert_eq!(completer.completions("testing/file"), vec!["testing/file_with_text"]);
        assert_eq!(completer.completions("~"), vec!["~/"]);
        assert_eq!(completer.completions("tes/fil"), vec!["testing/file_with_text"]);
    }
}

/// Infer if the given filename is actually a partial filename
fn complete_as_file(current_dir: &PathBuf, filename: &str, index: usize) -> bool {
    let filename = filename.trim();
    let mut file = current_dir.clone();
    file.push(&filename);
    // If the user explicitly requests a file through this syntax then complete as
    // a file
    filename.starts_with('.') ||
    // If the file starts with a dollar sign, it's a variable, not a file
    (!filename.starts_with('$') &&
    // Once we are beyond the first string, assume its a file
    (index > 0 ||
    // If we are referencing a file that exists then just complete to that file
    file.exists() ||
    // If we have a partial file inside an existing directory, e.g. /foo/b when
    // /foo/bar exists, then treat it as file as long as `foo` isn't the
    // current directory, otherwise this would apply to any string `foo`
    file.parent().filter(|parent| parent.exists() && parent != current_dir).is_some()))
    // By default assume its not a file
}
<|endoftext|>