<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::path::PathBuf;

use toml::Value;
use clap::App;

use error::RuntimeError as RE;
use error::RuntimeErrorKind as REK;
use error::Result;
use error::ResultExt;

/// Get a new configuration object.
///
/// The passed runtimepath is used for searching the configuration file, whereas several file
/// names are tested. If that does not work, the home directory and the XDG basedir are tested
/// with all variants.
///
/// If that doesn't work either, an error is returned.
pub fn fetch_config(searchpath: &PathBuf) -> Result<Value> {
    use std::env;
    use std::fs::File;
    use std::io::Read;
    use std::io::Write;
    use std::io::stderr;

    use xdg_basedir;
    use itertools::Itertools;

    use libimagutil::variants::generate_variants as gen_vars;
    use libimagerror::trace::trace_error;

    let variants = vec!["config", "config.toml", "imagrc", "imagrc.toml"];
    let modifier = |base: &PathBuf, v: &'static str| {
        let mut base = base.clone();
        base.push(String::from(v));
        base
    };

    vec![
        vec![searchpath.clone()],
        gen_vars(searchpath.clone(), variants.clone(), &modifier),

        env::var("HOME").map(|home| gen_vars(PathBuf::from(home), variants.clone(), &modifier))
                        .unwrap_or(vec![]),

        xdg_basedir::get_data_home().map(|data_dir| gen_vars(data_dir, variants.clone(), &modifier))
                                    .unwrap_or(vec![]),
    ].iter()
        .flatten()
        .filter(|path| path.exists() && path.is_file())
        .filter_map(|path| {
            let content = {
                let f = File::open(path);
                if f.is_err() {
                    let _ = write!(stderr(), "Error opening file: {:?}", f);
                    return None
                }
                let mut f = f.unwrap();

                let mut s = String::new();
                f.read_to_string(&mut s).ok();
                s
            };

            match ::toml::de::from_str::<::toml::Value>(&content[..]) {
                Ok(res) => Some(res),
                Err(e) => {
                    let line_col = e
                        .line_col()
                        .map(|(line, col)| {
                            format!("Line {}, Column {}", line, col)
                        })
                        .unwrap_or_else(|| String::from("Line unknown, Column unknown"));

                    let _ = write!(stderr(), "Config file parser error at {}", line_col);
                    trace_error(&e);
                    None
                }
            }
        })
        .nth(0)
        .ok_or(RE::from_kind(REK::ConfigNoConfigFileFound))
}

/// Override the configuration.
/// The `v` parameter is expected to contain 'key=value' pairs where the key is a path in the
/// TOML tree, the value to be an appropriate value.
///
/// The override fails if the configuration which is about to be overridden does not exist or
/// the `value` part cannot be converted to the type of the configuration value.
///
/// If `v` is empty, this is considered to be a successful `override_config()` call.
pub fn override_config(val: &mut Value, v: Vec<String>) -> Result<()> {
    use libimagutil::key_value_split::*;
    use toml_query::read::TomlValueReadExt;

    let iter = v.into_iter()
        .map(|s| { debug!("Trying to process '{}'", s); s })
        .filter_map(|s| s.into_kv().map(Into::into).or_else(|| {
            warn!("Could split at '=' - will be ignore override");
            None
        }))
        .map(|(k, v)| val
             .read(&k[..])
             .chain_err(|| REK::ConfigTOMLParserError)
             .map(|toml| match toml {
                Some(value) => match into_value(value, v) {
                    Some(v) => {
                        info!("Successfully overridden: {} = {}", k, v);
                        Ok(())
                    },
                    None => Err(RE::from_kind(REK::ConfigOverrideTypeNotMatching)),
                },
                None => Err(RE::from_kind(REK::ConfigOverrideKeyNotAvailable)),
            })
        );

    for elem in iter {
        let _ = try!(elem.chain_err(|| REK::ConfigOverrideError));
    }

    Ok(())
}

/// Tries to convert the String `s` into the same type as `value`.
///
/// Returns None if string cannot be converted.
///
/// Arrays and Tables are not supported and will yield `None`.
fn into_value(value: &Value, s: String) -> Option<Value> {
    use std::str::FromStr;

    match *value {
        Value::String(_)  => Some(Value::String(s)),
        Value::Integer(_) => FromStr::from_str(&s[..]).ok().map(Value::Integer),
        Value::Float(_)   => FromStr::from_str(&s[..]).ok().map(Value::Float),
        Value::Boolean(_) => {
            if s == "true" { Some(Value::Boolean(true)) }
            else if s == "false" { Some(Value::Boolean(false)) }
            else { None }
        }
        Value::Datetime(_) => Value::try_from(s).ok(),
        _ => None,
    }
}

pub trait InternalConfiguration {
    fn enable_logging(&self) -> bool {
        true
    }

    fn use_inmemory_fs(&self) -> bool {
        false
    }
}

impl<'a> InternalConfiguration for App<'a, 'a> {}

<commit_msg>Replace matching with function chaining<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::path::PathBuf;

use toml::Value;
use clap::App;

use error::RuntimeError as RE;
use error::RuntimeErrorKind as REK;
use error::Result;
use error::ResultExt;

/// Get a new configuration object.
///
/// The passed runtimepath is used for searching the configuration file, whereas several file
/// names are tested. If that does not work, the home directory and the XDG basedir are tested
/// with all variants.
///
/// If that doesn't work either, an error is returned.
pub fn fetch_config(searchpath: &PathBuf) -> Result<Value> {
    use std::env;
    use std::fs::File;
    use std::io::Read;
    use std::io::Write;
    use std::io::stderr;

    use xdg_basedir;
    use itertools::Itertools;

    use libimagutil::variants::generate_variants as gen_vars;
    use libimagerror::trace::trace_error;

    let variants = vec!["config", "config.toml", "imagrc", "imagrc.toml"];
    let modifier = |base: &PathBuf, v: &'static str| {
        let mut base = base.clone();
        base.push(String::from(v));
        base
    };

    vec![
        vec![searchpath.clone()],
        gen_vars(searchpath.clone(), variants.clone(), &modifier),

        env::var("HOME").map(|home| gen_vars(PathBuf::from(home), variants.clone(), &modifier))
                        .unwrap_or(vec![]),

        xdg_basedir::get_data_home().map(|data_dir| gen_vars(data_dir, variants.clone(), &modifier))
                                    .unwrap_or(vec![]),
    ].iter()
        .flatten()
        .filter(|path| path.exists() && path.is_file())
        .filter_map(|path| {
            let content = {
                let f = File::open(path);
                if f.is_err() {
                    let _ = write!(stderr(), "Error opening file: {:?}", f);
                    return None
                }
                let mut f = f.unwrap();

                let mut s = String::new();
                f.read_to_string(&mut s).ok();
                s
            };

            ::toml::de::from_str::<::toml::Value>(&content[..])
                .map(Some)
                .unwrap_or_else(|e| {
                    let line_col = e
                        .line_col()
                        .map(|(line, col)| format!("Line {}, Column {}", line, col))
                        .unwrap_or_else(|| String::from("Line unknown, Column unknown"));

                    let _ = write!(stderr(), "Config file parser error at {}", line_col);
                    trace_error(&e);
                    None
                })
        })
        .nth(0)
        .ok_or(RE::from_kind(REK::ConfigNoConfigFileFound))
}

/// Override the configuration.
/// The `v` parameter is expected to contain 'key=value' pairs where the key is a path in the
/// TOML tree, the value to be an appropriate value.
///
/// The override fails if the configuration which is about to be overridden does not exist or
/// the `value` part cannot be converted to the type of the configuration value.
///
/// If `v` is empty, this is considered to be a successful `override_config()` call.
pub fn override_config(val: &mut Value, v: Vec<String>) -> Result<()> {
    use libimagutil::key_value_split::*;
    use toml_query::read::TomlValueReadExt;

    let iter = v.into_iter()
        .map(|s| { debug!("Trying to process '{}'", s); s })
        .filter_map(|s| s.into_kv().map(Into::into).or_else(|| {
            warn!("Could split at '=' - will be ignore override");
            None
        }))
        .map(|(k, v)| val
             .read(&k[..])
             .chain_err(|| REK::ConfigTOMLParserError)
             .map(|toml| match toml {
                Some(value) => match into_value(value, v) {
                    Some(v) => {
                        info!("Successfully overridden: {} = {}", k, v);
                        Ok(())
                    },
                    None => Err(RE::from_kind(REK::ConfigOverrideTypeNotMatching)),
                },
                None => Err(RE::from_kind(REK::ConfigOverrideKeyNotAvailable)),
            })
        );

    for elem in iter {
        let _ = try!(elem.chain_err(|| REK::ConfigOverrideError));
    }

    Ok(())
}

/// Tries to convert the String `s` into the same type as `value`.
///
/// Returns None if string cannot be converted.
///
/// Arrays and Tables are not supported and will yield `None`.
fn into_value(value: &Value, s: String) -> Option<Value> {
    use std::str::FromStr;

    match *value {
        Value::String(_)  => Some(Value::String(s)),
        Value::Integer(_) => FromStr::from_str(&s[..]).ok().map(Value::Integer),
        Value::Float(_)   => FromStr::from_str(&s[..]).ok().map(Value::Float),
        Value::Boolean(_) => {
            if s == "true" { Some(Value::Boolean(true)) }
            else if s == "false" { Some(Value::Boolean(false)) }
            else { None }
        }
        Value::Datetime(_) => Value::try_from(s).ok(),
        _ => None,
    }
}

pub trait InternalConfiguration {
    fn enable_logging(&self) -> bool {
        true
    }

    fn use_inmemory_fs(&self) -> bool {
        false
    }
}

impl<'a> InternalConfiguration for App<'a, 'a> {}

<|endoftext|>


<commit_before>use std::{
    collections::{hash_map::Entry, HashMap},
    sync::Arc,
    time::{Duration, Instant, SystemTime},
};

use async_channel::TrySendError;
use axum::http::StatusCode;
use clap::Clap;
use futures_util::StreamExt;
use rustc_hash::FxHashMap;
use shakmaty::{
    uci::Uci, variant::VariantPosition, zobrist::Zobrist, ByColor, CastlingMode, Outcome, Position,
};
use tokio::{
    sync::{watch, RwLock},
    task::JoinHandle,
    time,
};

use crate::{
    db::Database,
    model::{
        GameInfo, GameInfoPlayer, Mode, Month, PersonalEntry, PersonalKeyBuilder, PersonalStatus,
        UserId,
    },
};

mod lila;

use lila::{Game, Lila};

const MAX_PLIES: usize = 50;

#[derive(Clap, Clone)]
pub struct IndexerOpt {
    #[clap(long = "lila", default_value = "https://lichess.org")]
    lila: String,
    #[clap(long = "bearer")]
    bearer: Option<String>,
    #[clap(long = "indexers", default_value = "16")]
    indexers: usize,
}

#[derive(Clone)]
pub struct IndexerStub {
    db: Arc<Database>,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    tx: async_channel::Sender<IndexerMessage>,
}

impl IndexerStub {
    pub fn spawn(db: Arc<Database>, opt: IndexerOpt) -> (IndexerStub, Vec<JoinHandle<()>>) {
        let indexing = Arc::new(RwLock::new(HashMap::new()));

        let (tx, rx) = async_channel::bounded(1000);
        let mut join_handles = Vec::with_capacity(opt.indexers);
        for idx in 0..opt.indexers {
            join_handles.push(tokio::spawn(
                IndexerActor {
                    idx,
                    rx: rx.clone(),
                    indexing: Arc::clone(&indexing),
                    db: Arc::clone(&db),
                    lila: Lila::new(opt.clone()),
                }
                .run(),
            ));
        }

        (IndexerStub { db, indexing, tx }, join_handles)
    }

    pub async fn num_indexing(&self) -> usize {
        let guard = self.indexing.read().await;
        guard.len()
    }

    pub async fn index_player(&self, player: &UserId) -> Option<watch::Receiver<()>> {
        // Optimization: First try subscribing to an existing indexing run,
        // without acquiring a write lock.
        {
            let guard = self.indexing.read().await;
            if let Some(sender) = guard.get(player) {
                return Some(sender.subscribe());
            }
        }

        // Check player indexing status.
        let mut status = self
            .db
            .queryable()
            .get_player_status(player)
            .expect("get player status")
            .unwrap_or_default();

        let since_created_at = match status
            .maybe_revisit_ongoing()
            .or_else(|| status.maybe_index())
        {
            Some(since) => since,
            None => return None, // Do not reindex so soon!
        };

        // Queue indexing request.
        let mut guard = self.indexing.write().await;
        let entry = match guard.entry(player.to_owned()) {
            Entry::Occupied(entry) => return Some(entry.get().subscribe()),
            Entry::Vacant(entry) => entry,
        };

        match self.tx.try_send(IndexerMessage::IndexPlayer {
            player: player.to_owned(),
            status,
            since_created_at,
        }) {
            Ok(_) => {
                let (sender, receiver) = watch::channel(());
                entry.insert(sender);
                Some(receiver)
            }
            Err(TrySendError::Full(_)) => {
                log::error!(
                    "not queuing {} because indexer queue is full",
                    player.as_str()
                );
                None
            }
            Err(TrySendError::Closed(_)) => panic!("all indexers died"),
        }
    }
}

struct IndexerActor {
    idx: usize,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    rx: async_channel::Receiver<IndexerMessage>,
    db: Arc<Database>,
    lila: Lila,
}

impl IndexerActor {
    async fn run(self) {
        while let Ok(msg) = self.rx.recv().await {
            match msg {
                IndexerMessage::IndexPlayer {
                    player,
                    status,
                    since_created_at,
                } => {
                    self.index_player(&player, status, since_created_at).await;

                    let mut guard = self.indexing.write().await;
                    guard.remove(&player);
                }
            }
        }
    }

    async fn index_player(
        &self,
        player: &UserId,
        mut status: PersonalStatus,
        since_created_at: u64,
    ) {
        let started_at = Instant::now();
        log::info!(
            "indexer {:02}: starting {} (created_at >= {})",
            self.idx,
            player.as_str(),
            since_created_at
        );

        let mut games = match self.lila.user_games(player, since_created_at).await {
            Ok(games) => games,
            Err(err) if err.status() == Some(StatusCode::NOT_FOUND) => {
                log::warn!(
                    "indexer {:02}: did not find player {}",
                    self.idx,
                    player.as_str()
                );
                return;
            }
            Err(err) => {
                log::error!("indexer {:02}: request failed: {}", self.idx, err);
                time::sleep(Duration::from_secs(5)).await;
                return;
            }
        };

        let hash = ByColor::new_with(|color| PersonalKeyBuilder::with_user_pov(player, color));
        let mut num_games = 0;
        while let Some(game) = games.next().await {
            let game = match game {
                Ok(game) => game,
                Err(err) => {
                    log::error!("indexer {:02}: {}", self.idx, err);
                    continue;
                }
            };

            self.index_game(player, &hash, game, &mut status);

            num_games += 1;
            if num_games % 1024 == 0 {
                self.db
                    .queryable()
                    .put_player_status(player, &status)
                    .expect("put player status");

                log::info!(
                    "indexer {:02}: indexed {} games for {} ...",
                    self.idx,
                    num_games,
                    player.as_str()
                );
            }
        }

        status.indexed_at = SystemTime::now();
        self.db
            .queryable()
            .put_player_status(player, &status)
            .expect("put player status");

        let elapsed = started_at.elapsed();
        log::info!(
            "indexer {:02}: finished {} games for {} in {:.3?} ({:.3?}/game, {:.1} games/s)",
            self.idx,
            num_games,
            player.as_str(),
            elapsed,
            (elapsed.as_nanos() as u64)
                .checked_div(num_games)
                .map(Duration::from_nanos)
                .unwrap_or_default(),
            (num_games as f64) / elapsed.as_secs_f64()
        );
    }

    fn index_game(
        &self,
        player: &UserId,
        hash: &ByColor<PersonalKeyBuilder>,
        game: Game,
        status: &mut PersonalStatus,
    ) {
        status.latest_created_at = game.created_at;

        if game.status.is_ongoing() {
            if status.revisit_ongoing_created_at.is_none() {
                log::debug!(
                    "indexer {:02}: will revisit ongoing game {} eventually",
                    self.idx,
                    game.id
                );
                status.revisit_ongoing_created_at = Some(game.created_at);
            }
            return;
        }

        if game.status.is_unindexable() {
            return;
        }

        if game.players.any(|p| p.user.is_none()) {
            return;
        }

        let color = match game
            .players
            .find(|p| p.user.as_ref().map_or(false, |user| user.name == *player))
        {
            Some(color) => color,
            None => {
                log::error!(
                    "indexer {:02}: {} did not play in {}",
                    self.idx,
                    player.as_str(),
                    game.id
                );
                return;
            }
        };

        // Skip game if already indexed from this side. This cannot race with
        // writes, because all writes for the same player are sequenced by
        // this actor. So making a transaction is not required.
        let queryable = self.db.queryable();
        if queryable
            .get_game_info(game.id)
            .expect("get game info")
            .map_or(false, |info| info.indexed.into_color(color))
        {
            log::debug!(
                "indexer {:02}: {}/{} already indexed",
                self.idx,
                game.id,
                color
            );
            return;
        }

        // Prepare basic information and setup initial position.
        let month = Month::from_time_saturating(game.last_move_at);
        let outcome = Outcome::from_winner(game.winner);
        let variant = game.variant.into();
        let pos = match game.initial_fen {
            Some(fen) => VariantPosition::from_setup(variant, &fen, CastlingMode::Chess960),
            None => Ok(VariantPosition::new(variant)),
        };

        let mut pos: Zobrist<_, u128> = match pos {
            Ok(pos) => Zobrist::new(pos),
            Err(err) => {
                log::warn!("indexer {:02}: not indexing {}: {}", self.idx, game.id, err);
                return;
            }
        };

        // Build an intermediate table to remove loops (due to repetitions).
        let mut table: FxHashMap<u128, Uci> =
            FxHashMap::with_capacity_and_hasher(game.moves.len(), Default::default());

        for (ply, san) in game.moves.into_iter().enumerate() {
            if ply >= MAX_PLIES {
                break;
            }

            let m = match san.to_move(&pos) {
                Ok(m) => m,
                Err(err) => {
                    log::warn!(
                        "indexer {:02}: cutting off {} at ply {}: {}: {}",
                        self.idx,
                        game.id,
                        ply,
                        err,
                        san
                    );
                    break;
                }
            };

            let uci = m.to_uci(CastlingMode::Chess960);
            table.insert(pos.zobrist_hash(), uci);

            pos.play_unchecked(&m);
        }

        // Write to database. All writes regarding this game are batched and
        // atomically committed, so the database will always be in a consistent
        // state.
        let mut batch = queryable.batch();

        batch.merge_game_info(
            game.id,
            GameInfo {
                winner: outcome.winner(),
                speed: game.speed,
                mode: Mode::from_rated(game.rated),
                month,
                players: game.players.map(|p| GameInfoPlayer {
                    name: p.user.map(|p| p.name.to_string()),
                    rating: p.rating,
                }),
                indexed: ByColor::new_with(|c| color == c),
            },
        );

        for (zobrist, uci) in table {
            batch.merge_personal(
                hash.by_color(color)
                    .with_zobrist(variant, zobrist)
                    .with_month(month),
                PersonalEntry::new_single(
                    uci.clone(),
                    game.speed,
                    Mode::from_rated(game.rated),
                    game.id,
                    outcome,
                ),
            );
        }

        batch.write().expect("atomically commit game and moves");
    }
}

enum IndexerMessage {
    IndexPlayer {
        player: UserId,
        status: PersonalStatus,
        since_created_at: u64,
    },
}
<commit_msg>apply timeouts to indexer request and stream<commit_after>use std::{
    collections::{hash_map::Entry, HashMap},
    sync::Arc,
    time::{Duration, Instant, SystemTime},
};

use async_channel::TrySendError;
use axum::http::StatusCode;
use clap::Clap;
use futures_util::StreamExt;
use rustc_hash::FxHashMap;
use shakmaty::{
    uci::Uci, variant::VariantPosition, zobrist::Zobrist, ByColor, CastlingMode, Outcome, Position,
};
use tokio::{
    sync::{watch, RwLock},
    task::JoinHandle,
    time::{sleep, timeout},
};

use crate::{
    db::Database,
    model::{
        GameInfo, GameInfoPlayer, Mode, Month, PersonalEntry, PersonalKeyBuilder, PersonalStatus,
        UserId,
    },
};

mod lila;

use lila::{Game, Lila};

const MAX_PLIES: usize = 50;

#[derive(Clap, Clone)]
pub struct IndexerOpt {
    #[clap(long = "lila", default_value = "https://lichess.org")]
    lila: String,
    #[clap(long = "bearer")]
    bearer: Option<String>,
    #[clap(long = "indexers", default_value = "16")]
    indexers: usize,
}

#[derive(Clone)]
pub struct IndexerStub {
    db: Arc<Database>,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    tx: async_channel::Sender<IndexerMessage>,
}

impl IndexerStub {
    pub fn spawn(db: Arc<Database>, opt: IndexerOpt) -> (IndexerStub, Vec<JoinHandle<()>>) {
        let indexing = Arc::new(RwLock::new(HashMap::new()));

        let (tx, rx) = async_channel::bounded(1000);
        let mut join_handles = Vec::with_capacity(opt.indexers);
        for idx in 0..opt.indexers {
            join_handles.push(tokio::spawn(
                IndexerActor {
                    idx,
                    rx: rx.clone(),
                    indexing: Arc::clone(&indexing),
                    db: Arc::clone(&db),
                    lila: Lila::new(opt.clone()),
                }
                .run(),
            ));
        }

        (IndexerStub { db, indexing, tx }, join_handles)
    }

    pub async fn num_indexing(&self) -> usize {
        let guard = self.indexing.read().await;
        guard.len()
    }

    pub async fn index_player(&self, player: &UserId) -> Option<watch::Receiver<()>> {
        // Optimization: First try subscribing to an existing indexing run,
        // without acquiring a write lock.
        {
            let guard = self.indexing.read().await;
            if let Some(sender) = guard.get(player) {
                return Some(sender.subscribe());
            }
        }

        // Check player indexing status.
        let mut status = self
            .db
            .queryable()
            .get_player_status(player)
            .expect("get player status")
            .unwrap_or_default();

        let since_created_at = match status
            .maybe_revisit_ongoing()
            .or_else(|| status.maybe_index())
        {
            Some(since) => since,
            None => return None, // Do not reindex so soon!
        };

        // Queue indexing request.
        let mut guard = self.indexing.write().await;
        let entry = match guard.entry(player.to_owned()) {
            Entry::Occupied(entry) => return Some(entry.get().subscribe()),
            Entry::Vacant(entry) => entry,
        };

        match self.tx.try_send(IndexerMessage::IndexPlayer {
            player: player.to_owned(),
            status,
            since_created_at,
        }) {
            Ok(_) => {
                let (sender, receiver) = watch::channel(());
                entry.insert(sender);
                Some(receiver)
            }
            Err(TrySendError::Full(_)) => {
                log::error!(
                    "not queuing {} because indexer queue is full",
                    player.as_str()
                );
                None
            }
            Err(TrySendError::Closed(_)) => panic!("all indexers died"),
        }
    }
}

struct IndexerActor {
    idx: usize,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    rx: async_channel::Receiver<IndexerMessage>,
    db: Arc<Database>,
    lila: Lila,
}

impl IndexerActor {
    async fn run(self) {
        while let Ok(msg) = self.rx.recv().await {
            match msg {
                IndexerMessage::IndexPlayer {
                    player,
                    status,
                    since_created_at,
                } => {
                    self.index_player(&player, status, since_created_at).await;

                    let mut guard = self.indexing.write().await;
                    guard.remove(&player);
                }
            }
        }
    }

    async fn index_player(
        &self,
        player: &UserId,
        mut status: PersonalStatus,
        since_created_at: u64,
    ) {
        let started_at = Instant::now();
        log::info!(
            "indexer {:02}: starting {} (created_at >= {})",
            self.idx,
            player.as_str(),
            since_created_at
        );

        let mut games = match timeout(
            Duration::from_secs(10),
            self.lila.user_games(player, since_created_at),
        )
        .await
        {
            Ok(Ok(games)) => games,
            Ok(Err(err)) if err.status() == Some(StatusCode::NOT_FOUND) => {
                log::warn!(
                    "indexer {:02}: did not find player {}",
                    self.idx,
                    player.as_str()
                );
                return;
            }
            Ok(Err(err)) => {
                log::error!("indexer {:02}: request failed: {}", self.idx, err);
                sleep(Duration::from_secs(5)).await;
                return;
            }
            Err(timed_out) => {
                log::error!("indexer {:02}: request to lila: {}", self.idx, timed_out);
                return;
            }
        };

        let hash = ByColor::new_with(|color| PersonalKeyBuilder::with_user_pov(player, color));

        let mut num_games = 0;
        loop {
            let game = match timeout(Duration::from_secs(10), games.next()).await {
                Ok(Some(Ok(game))) => game,
                Ok(Some(Err(err))) => {
                    log::error!("indexer {:02}: {}", self.idx, err);
                    continue;
                }
                Ok(None) => break,
                Err(timed_out) => {
                    log::error!("indexer {:02}: stream from lila: {}", self.idx, timed_out);
                    return;
                }
            };

            self.index_game(player, &hash, game, &mut status);

            num_games += 1;
            if num_games % 1024 == 0 {
                self.db
                    .queryable()
                    .put_player_status(player, &status)
                    .expect("put player status");

                log::info!(
                    "indexer {:02}: indexed {} games for {} ...",
                    self.idx,
                    num_games,
                    player.as_str()
                );
            }
        }

        status.indexed_at = SystemTime::now();
        self.db
            .queryable()
            .put_player_status(player, &status)
            .expect("put player status");

        let elapsed = started_at.elapsed();
        log::info!(
            "indexer {:02}: finished {} games for {} in {:.3?} ({:.3?}/game, {:.1} games/s)",
            self.idx,
            num_games,
            player.as_str(),
            elapsed,
            (elapsed.as_nanos() as u64)
                .checked_div(num_games)
                .map(Duration::from_nanos)
                .unwrap_or_default(),
            (num_games as f64) / elapsed.as_secs_f64()
        );
    }

    fn index_game(
        &self,
        player: &UserId,
        hash: &ByColor<PersonalKeyBuilder>,
        game: Game,
        status: &mut PersonalStatus,
    ) {
        status.latest_created_at = game.created_at;

        if game.status.is_ongoing() {
            if status.revisit_ongoing_created_at.is_none() {
                log::debug!(
                    "indexer {:02}: will revisit ongoing game {} eventually",
                    self.idx,
                    game.id
                );
                status.revisit_ongoing_created_at = Some(game.created_at);
            }
            return;
        }

        if game.status.is_unindexable() {
            return;
        }

        if game.players.any(|p| p.user.is_none()) {
            return;
        }

        let color = match game
            .players
            .find(|p| p.user.as_ref().map_or(false, |user| user.name == *player))
        {
            Some(color) => color,
            None => {
                log::error!(
                    "indexer {:02}: {} did not play in {}",
                    self.idx,
                    player.as_str(),
                    game.id
                );
                return;
            }
        };

        // Skip game if already indexed from this side. This cannot race with
        // writes, because all writes for the same player are sequenced by
        // this actor. So making a transaction is not required.
        let queryable = self.db.queryable();
        if queryable
            .get_game_info(game.id)
            .expect("get game info")
            .map_or(false, |info| info.indexed.into_color(color))
        {
            log::debug!(
                "indexer {:02}: {}/{} already indexed",
                self.idx,
                game.id,
                color
            );
            return;
        }

        // Prepare basic information and setup initial position.
        let month = Month::from_time_saturating(game.last_move_at);
        let outcome = Outcome::from_winner(game.winner);
        let variant = game.variant.into();
        let pos = match game.initial_fen {
            Some(fen) => VariantPosition::from_setup(variant, &fen, CastlingMode::Chess960),
            None => Ok(VariantPosition::new(variant)),
        };

        let mut pos: Zobrist<_, u128> = match pos {
            Ok(pos) => Zobrist::new(pos),
            Err(err) => {
                log::warn!("indexer {:02}: not indexing {}: {}", self.idx, game.id, err);
                return;
            }
        };

        // Build an intermediate table to remove loops (due to repetitions).
        let mut table: FxHashMap<u128, Uci> =
            FxHashMap::with_capacity_and_hasher(game.moves.len(), Default::default());

        for (ply, san) in game.moves.into_iter().enumerate() {
            if ply >= MAX_PLIES {
                break;
            }

            let m = match san.to_move(&pos) {
                Ok(m) => m,
                Err(err) => {
                    log::warn!(
                        "indexer {:02}: cutting off {} at ply {}: {}: {}",
                        self.idx,
                        game.id,
                        ply,
                        err,
                        san
                    );
                    break;
                }
            };

            let uci = m.to_uci(CastlingMode::Chess960);
            table.insert(pos.zobrist_hash(), uci);

            pos.play_unchecked(&m);
        }

        // Write to database. All writes regarding this game are batched and
        // atomically committed, so the database will always be in a consistent
        // state.
        let mut batch = queryable.batch();

        batch.merge_game_info(
            game.id,
            GameInfo {
                winner: outcome.winner(),
                speed: game.speed,
                mode: Mode::from_rated(game.rated),
                month,
                players: game.players.map(|p| GameInfoPlayer {
                    name: p.user.map(|p| p.name.to_string()),
                    rating: p.rating,
                }),
                indexed: ByColor::new_with(|c| color == c),
            },
        );

        for (zobrist, uci) in table {
            batch.merge_personal(
                hash.by_color(color)
                    .with_zobrist(variant, zobrist)
                    .with_month(month),
                PersonalEntry::new_single(
                    uci.clone(),
                    game.speed,
                    Mode::from_rated(game.rated),
                    game.id,
                    outcome,
                ),
            );
        }

        batch.write().expect("atomically commit game and moves");
    }
}

enum IndexerMessage {
    IndexPlayer {
        player: UserId,
        status: PersonalStatus,
        since_created_at: u64,
    },
}
<|endoftext|>


<commit_before>#[derive(Debug)]
#[allow(non_camel_case_types)]
#[allow(dead_code)] // TODO remove
pub enum Instruction {
    ALOAD(Type),
    ASTORE(Type),
    LOAD(Type, LocalVarRef),
    STORE(Type, LocalVarRef),

    ARRAYLENGTH,

    ATHROW,

    CHECKCAST(ConstPoolRef),
    INSTANCEOF(ConstPoolRef),

    ANEWARRAY(ConstPoolRef),
    MULTIANEWARRAY(ConstPoolRef, u8),
    NEW(ConstPoolRef),
    NEWARRAY(Type),

    // D2F, D2I,...
    CONVERT(Type, Type),

    ADD(Type),
    DIV(Type),
    MUL(Type),
    NEG(Type),
    REM(Type),
    SUB(Type),
    RETURN(Option<Type>),

    IINC(LocalVarRef, i16),

    AND(Type),
    OR(Type),
    SHL(Type),
    SHR(Type),
    USHR(Type),
    XOR(Type),

    DCMPG,
    DCMPL,
    FCMPG,
    FCMPL,
    LCMP,

    ACONST_NULL,
    DCONST_0,
    DCONST_1,
    FCONST_0,
    FCONST_1,
    FCONST_2,
    LCONST_0,
    LCONST_1,
    BIPUSH(i8),
    SIPUSH(i16),
    LDC(ConstPoolRef),
    LDC2_W(ConstPoolRef),

    DUP,
    DUP_X1,
    DUP_X2,
    DUP2,
    DUP2_X1,
    DUP2_X2,
    POP,
    POP2,
    SWAP,

    GETFIELD(ConstPoolRef),
    GETSTATIC(ConstPoolRef),
    PUTFIELD(ConstPoolRef),
    PUTSTATIC(ConstPoolRef),

    GOTO(i32),
    JSR(i32),
    RET(LocalVarRef),

    IF_ACMP(ComparisonEqual, i16),
    IF_ICMP(Comparison, i16),
    // comparison with zero
    IF(Comparison, i16),
    IFNULL(ComparisonEqual, i16),

    INVOKEINTERFACE(ConstPoolRef, u8),
    INVOKESPECIAL(ConstPoolRef),
    INVOKESTATIC(ConstPoolRef),
    INVOKEVIRTUAL(ConstPoolRef),

    LOOKUPSWITCH(i32, i32, Vec<(i32, i32)>),
    TABLESWITCH(i32, i32, i32, Vec<i32>),

    MONITORENTER,
    MONITOREXIT,

    NOP,
}

// index into the constant pool
pub type ConstPoolRef = u16;

// index into the local variables
pub type LocalVarRef = u16;

// true: equals, false: not equals
pub type ComparisonEqual = bool;

#[derive(Debug)]
pub enum Type {
    Reference,
    Char,
    Byte,
    Short,
    Int,
    Long,
    Float,
    Double,
}

#[derive(Debug)]
pub enum Comparison {
    EQ,
    GE,
    GT,
    LE,
    LT,
    NE,
}

impl Instruction {
    pub fn decode(bytes: &[u8]) -> Result<Vec<Instruction>, String> {
        use self::Instruction::*;
        use self::Type::*;
        use self::Comparison::*;
        fn next(index: &mut usize, bytes: &[u8]) -> Result<u8, String> {
            if *index >= bytes.len() {
                return Err("incomplete instruction".to_owned());
            }
            *index += 1;
            Ok(bytes[*index - 1])
        }

        let mut vec = Vec::new();

        let mut index = 0;
        while index < bytes.len() {
            vec.push(match bytes[index] {
                0x53 => ASTORE(Reference),
                0x01 => ACONST_NULL,
                0x19 => LOAD(Reference, next(&mut index, bytes)? as u16),
                // 0x2a => aload_0,
                // 0x2b => aload_1,
                // 0x2c => aload_2,
                // 0x2d => aload_3,
                // 0xbd => anewarray,
                // 0xb0 => areturn,
                // 0xbe => arraylength,
                // 0x3a => astore,
                // 0x4b => astore_0,
                // 0x4c => astore_1,
                // 0x4d => astore_2,
                // 0x4e => astore_3,
                // 0xbf => athrow,
                // 0x33 => baload,
                // 0x54 => bastore,
                // 0x10 => bipush,
                // 0xca => breakpoint,
                // 0x34 => caload,
                // 0x55 => castore,
                // 0xc0 => checkcast,
                // 0x90 => d2f,
                // 0x8e => d2i,
                // 0x8f => d2l,
                // 0x63 => dadd,
                // 0x31 => daload,
                // 0x52 => dastore,
                // 0x98 => dcmpg,
                // 0x97 => dcmpl,
                // 0x0e => dconst_0,
                // 0x0f => dconst_1,
                // 0x6f => ddiv,
                // 0x18 => dload,
                // 0x26 => dload_0,
                // 0x27 => dload_1,
                // 0x28 => dload_2,
                // 0x29 => dload_3,
                // 0x6b => dmul,
                // 0x77 => dneg,
                // 0x73 => drem,
                // 0xaf => dreturn,
                // 0x39 => dstore,
                // 0x47 => dstore_0,
                // 0x48 => dstore_1,
                // 0x49 => dstore_2,
                // 0x4a => dstore_3,
                // 0x67 => dsub,
                // 0x59 => dup,
                // 0x5a => dup_x1,
                // 0x5b => dup_x2,
                // 0x5c => dup2,
                // 0x5d => dup2_x1,
                // 0x5e => dup2_x2,
                // 0x8d => f2d,
                // 0x8b => f2i,
                // 0x8c => f2l,
                // 0x62 => fadd,
                // 0x30 => faload,
                // 0x51 => fastore,
                // 0x96 => fcmpg,
                // 0x95 => fcmpl,
                // 0x0b => fconst_0,
                // 0x0c => fconst_1,
                // 0x0d => fconst_2,
                // 0x6e => fdiv,
                // 0x17 => fload,
                // 0x22 => fload_0,
                // 0x23 => fload_1,
                // 0x24 => fload_2,
                // 0x25 => fload_3,
                // 0x6a => fmul,
                // 0x76 => fneg,
                // 0x72 => frem,
                // 0xae => freturn,
                // 0x38 => fstore,
                // 0x43 => fstore_0,
                // 0x44 => fstore_1,
                // 0x45 => fstore_2,
                // 0x46 => fstore_3,
                // 0x66 => fsub,
                // 0xb4 => getfield,
                // 0xb2 => getstatic,
                // 0xa7 => goto,
                // 0xc8 => goto_w,
                // 0x91 => i2b,
                // 0x92 => i2c,
                // 0x87 => i2d,
                // 0x86 => i2f,
                // 0x85 => i2l,
                // 0x93 => i2s,
                // 0x60 => iadd,
                // 0x2e => iaload,
                // 0x7e => iand,
                // 0x4f => iastore,
                // 0x02 => iconst_m1,
                // 0x03 => iconst_0,
                // 0x04 => iconst_1,
                // 0x05 => iconst_2,
                // 0x06 => iconst_3,
                // 0x07 => iconst_4,
                // 0x08 => iconst_5,
                // 0x6c => idiv,
                // 0xa5 => if_acmpeq,
                // 0xa6 => if_acmpne,
                // 0x9f => if_icmpeq,
                // 0xa2 => if_icmpge,
                // 0xa3 => if_icmpgt,
                // 0xa4 => if_icmple,
                // 0xa1 => if_icmplt,
                // 0xa0 => if_icmpne,
                // 0x99 => ifeq,
                // 0x9c => ifge,
                // 0x9d => ifgt,
                // 0x9e => ifle,
                // 0x9b => iflt,
                // 0x9a => ifne,
                // 0xc7 => ifnonnull,
                // 0xc6 => ifnull,
                // 0x84 => iinc,
                // 0x15 => iload,
                // 0x1a => iload_0,
                // 0x1b => iload_1,
                // 0x1c => iload_2,
                // 0x1d => iload_3,
                // 0xfe => impdep1,
                // 0xff => impdep2,
                // 0x68 => imul,
                // 0x74 => ineg,
                // 0xc1 => instanceof,
                // 0xba => invokedynamic,
                // 0xb9 => invokeinterface,
                // 0xb7 => invokespecial,
                // 0xb8 => invokestatic,
                // 0xb6 => invokevirtual,
                // 0x80 => ior,
                // 0x70 => irem,
                // 0xac => ireturn,
                // 0x78 => ishl,
                // 0x7a => ishr,
                // 0x36 => istore,
                // 0x3b => istore_0,
                // 0x3c => istore_1,
                // 0x3d => istore_2,
                // 0x3e => istore_3,
                // 0x64 => isub,
                // 0x7c => iushr,
                // 0x82 => ixor,
                // 0xa8 => jsr,
                // 0xc9 => jsr_w,
                // 0x8a => l2d,
                // 0x89 => l2f,
                // 0x88 => l2i,
                // 0x61 => ladd,
                // 0x2f => laload,
                // 0x7f => land,
                // 0x50 => lastore,
                // 0x94 => lcmp,
                // 0x09 => lconst_0,
                // 0x0a => lconst_1,
                // 0x12 => ldc,
                // 0x13 => ldc_w,
                // 0x14 => ldc2_w,
                // 0x6d => ldiv,
                // 0x16 => lload,
                // 0x1e => lload_0,
                // 0x1f => lload_1,
                // 0x20 => lload_2,
                // 0x21 => lload_3,
                // 0x69 => lmul,
                // 0x75 => lneg,
                // 0xab => lookupswitch,
                // 0x81 => lor,
                // 0x71 => lrem,
                // 0xad => lreturn,
                // 0x79 => lshl,
                // 0x7b => lshr,
                // 0x37 => lstore,
                // 0x3f => lstore_0,
                // 0x40 => lstore_1,
                // 0x41 => lstore_2,
                // 0x42 => lstore_3,
                // 0x65 => lsub,
                // 0x7d => lushr,
                // 0x83 => lxor,
                // 0xc2 => monitorenter,
                // 0xc3 => monitorexit,
                // 0xc5 => multianewarray,
                // 0xbb => new,
                // 0xbc => newarray,
                // 0x00 => nop,
                // 0x57 => pop,
                // 0x58 => pop2,
                // 0xb5 => putfield,
                // 0xb3 => putstatic,
                // 0xa9 => ret,
                // 0xb1 => return,
                // 0x35 => saload,
                // 0x56 => sastore,
                // 0x11 => sipush,
                // 0x5f => swap,
                // 0xaa => tableswitch,
                // 0xc4 => wide,
                op@_ => return Err(format!("Unknown Instruction {:#x}", op)),
            });
        }

        // TODO: gotos/if anpassen
        return Ok(vec);
    }
}
<commit_msg>more instructions<commit_after>#[derive(Debug)]
#[allow(non_camel_case_types)]
#[allow(dead_code)] // TODO remove
pub enum Instruction {
    ALOAD(Type),
    ASTORE(Type),
    LOAD(Type, LocalVarRef),
    STORE(Type, LocalVarRef),

    ARRAYLENGTH,

    ATHROW,

    CHECKCAST(ConstPoolRef),
    INSTANCEOF(ConstPoolRef),

    ANEWARRAY(ConstPoolRef),
    MULTIANEWARRAY(ConstPoolRef, u8),
    NEW(ConstPoolRef),
    NEWARRAY(Type),

    // D2F, D2I,...
    CONVERT(Type, Type),

    ADD(Type),
    DIV(Type),
    MUL(Type),
    NEG(Type),
    REM(Type),
    SUB(Type),
    RETURN(Option<Type>),

    IINC(LocalVarRef, i16),

    AND(Type),
    OR(Type),
    SHL(Type),
    SHR(Type),
    USHR(Type),
    XOR(Type),

    DCMPG,
    DCMPL,
    FCMPG,
    FCMPL,
    LCMP,

    ACONST_NULL,
    DCONST_0,
    DCONST_1,
    FCONST_0,
    FCONST_1,
    FCONST_2,
    LCONST_0,
    LCONST_1,
    BIPUSH(i8),
    SIPUSH(i16),
    LDC(ConstPoolRef),
    LDC2_W(ConstPoolRef),

    DUP,
    DUP_X1,
    DUP_X2,
    DUP2,
    DUP2_X1,
    DUP2_X2,
    POP,
    POP2,
    SWAP,

    GETFIELD(ConstPoolRef),
    GETSTATIC(ConstPoolRef),
    PUTFIELD(ConstPoolRef),
    PUTSTATIC(ConstPoolRef),

    GOTO(i32),
    JSR(i32),
    RET(LocalVarRef),

    IF_ACMP(ComparisonEqual, i16),
    IF_ICMP(Comparison, i16),
    // comparison with zero
    IF(Comparison, i16),
    IFNULL(ComparisonEqual, i16),

    INVOKEINTERFACE(ConstPoolRef, u8),
    INVOKESPECIAL(ConstPoolRef),
    INVOKESTATIC(ConstPoolRef),
    INVOKEVIRTUAL(ConstPoolRef),

    LOOKUPSWITCH(i32, i32, Vec<(i32, i32)>),
    TABLESWITCH(i32, i32, i32, Vec<i32>),

    MONITORENTER,
    MONITOREXIT,

    NOP,
}

// index into the constant pool
pub type ConstPoolRef = u16;

// index into the local variables
pub type LocalVarRef = u16;

// true: equals, false: not equals
pub type ComparisonEqual = bool;

#[derive(Debug)]
pub enum Type {
    Reference,
    Char,
    Byte,
    Short,
    Int,
    Long,
    Float,
    Double,
}

#[derive(Debug)]
pub enum Comparison {
    EQ,
    GE,
    GT,
    LE,
    LT,
    NE,
}

impl Instruction {
    pub fn decode(bytes: &[u8]) -> Result<Vec<Instruction>, String> {
        use self::Instruction::*;
        use self::Type::*;
        use self::Comparison::*;
        fn next(index: &mut usize, bytes: &[u8]) -> Result<u8, String> {
            if *index >= bytes.len() {
                return Err("incomplete instruction".to_owned());
            }
            *index += 1;
            Ok(bytes[*index - 1])
        }
        fn next_u16(index: &mut usize, bytes: &[u8]) -> Result<u16, String> {
            let b1 = next(index, bytes)? as u16;
            let b2 = next(index, bytes)? as u16;
            return Ok((b1 << 8) | b2);
        }

        let mut vec = Vec::new();

        let mut index = 0;
        while index < bytes.len() {
            let cur = bytes[index];
            index += 1;
            vec.push(match cur {
                0x32 => ALOAD(Reference),
                0x53 => ASTORE(Reference),
                0x01 => ACONST_NULL,
                0x19 => LOAD(Reference, next(&mut index, bytes)? as u16),
                0x2a => LOAD(Reference, 0),
                0x2b => LOAD(Reference, 1),
                0x2c => LOAD(Reference, 2),
                0x2d => LOAD(Reference, 3),
                0xbd => ANEWARRAY(next_u16(&mut index, bytes)?),
                0xb0 => RETURN(Some(Reference)),
                0xbe => ARRAYLENGTH,
                0x3a => STORE(Reference, next(&mut index, bytes)? as u16),
                0x4b => STORE(Reference, 0),
                0x4c => STORE(Reference, 1),
                0x4d => STORE(Reference, 2),
                0x4e => STORE(Reference, 3),
                0xbf => ATHROW,
                0x33 => ALOAD(Byte),
                0x54 => ASTORE(Byte),
                0x10 => BIPUSH(next(&mut index, bytes)? as i8),
                0x34 => ALOAD(Char),
                0x55 => ASTORE(Char),
                0xc0 => CHECKCAST(next_u16(&mut index, bytes)?),
                0x90 => CONVERT(Double, Float),
                0x8e => CONVERT(Double, Int),
                0x8f => CONVERT(Double, Long),
                0x63 => ADD(Double),
                0x31 => ALOAD(Double),
                0x52 => ASTORE(Double),
                0x98 => DCMPG,
                0x97 => DCMPL,
                0x0e => DCONST_0,
                0x0f => DCONST_1,
                0x6f => DIV(Double),
                0x18 => LOAD(Double, next(&mut index, bytes)? as u16),
                0x26 => LOAD(Double, 0),
                0x27 => LOAD(Double, 1),
                0x28 => LOAD(Double, 2),
                0x29 => LOAD(Double, 3),
                0x6b => MUL(Double),
                0x77 => NEG(Double),
                0x73 => REM(Double),
                0xaf => RETURN(Some(Double)),
                0x39 => STORE(Double, next(&mut index, bytes)? as u16),
                0x47 => STORE(Double, 0),
                0x48 => STORE(Double, 1),
                0x49 => STORE(Double, 2),
                0x4a => STORE(Double, 3),
                0x67 => SUB(Double),
                0x59 => DUP,
                0x5a => DUP_X1,
                0x5b => DUP_X2,
                0x5c => DUP2,
                0x5d => DUP2_X1,
                0x5e => DUP2_X2,
                0x8d => CONVERT(Float, Double),
                0x8b => CONVERT(Float, Int),
                0x8c => CONVERT(Float, Long),
                0x62 => ADD(Float),
                0x30 => ALOAD(Float),
                0x51 => ASTORE(Float),
                0x96 => FCMPG,
                0x95 => FCMPL,
                0x0b => FCONST_0,
                0x0c => FCONST_1,
                0x0d => FCONST_2,
                0x6e => fdiv,
                0x17 => fload,
                0x22 => fload_0,
                0x23 => fload_1,
                0x24 => fload_2,
                0x25 => fload_3,
                // 0x6a => fmul,
                // 0x76 => fneg,
                // 0x72 => frem,
                // 0xae => freturn,
                // 0x38 => fstore,
                // 0x43 => fstore_0,
                // 0x44 => fstore_1,
                // 0x45 => fstore_2,
                // 0x46 => fstore_3,
                // 0x66 => fsub,
                // 0xb4 => getfield,
                // 0xb2 => getstatic,
                // 0xa7 => goto,
                // 0xc8 => goto_w,
                // 0x91 => i2b,
                // 0x92 => i2c,
                // 0x87 => i2d,
                // 0x86 => i2f,
                // 0x85 => i2l,
                // 0x93 => i2s,
                // 0x60 => iadd,
                // 0x2e => iaload,
                // 0x7e => iand,
                // 0x4f => iastore,
                // 0x02 => iconst_m1,
                // 0x03 => iconst_0,
                // 0x04 => iconst_1,
                // 0x05 => iconst_2,
                // 0x06 => iconst_3,
                // 0x07 => iconst_4,
                // 0x08 => iconst_5,
                // 0x6c => idiv,
                // 0xa5 => if_acmpeq,
                // 0xa6 => if_acmpne,
                // 0x9f => if_icmpeq,
                // 0xa2 => if_icmpge,
                // 0xa3 => if_icmpgt,
                // 0xa4 => if_icmple,
                // 0xa1 => if_icmplt,
                // 0xa0 => if_icmpne,
                // 0x99 => ifeq,
                // 0x9c => ifge,
                // 0x9d => ifgt,
                // 0x9e => ifle,
                // 0x9b => iflt,
                // 0x9a => ifne,
                // 0xc7 => ifnonnull,
                // 0xc6 => ifnull,
                // 0x84 => iinc,
                // 0x15 => iload,
                // 0x1a => iload_0,
                // 0x1b => iload_1,
                // 0x1c => iload_2,
                // 0x1d => iload_3,
                // 0xfe => impdep1,
                // 0xff => impdep2,
                // 0x68 => imul,
                // 0x74 => ineg,
                // 0xc1 => instanceof,
                // 0xba => invokedynamic,
                // 0xb9 => invokeinterface,
                // 0xb7 => invokespecial,
                // 0xb8 => invokestatic,
                // 0xb6 => invokevirtual,
                // 0x80 => ior,
                // 0x70 => irem,
                // 0xac => ireturn,
                // 0x78 => ishl,
                // 0x7a => ishr,
                // 0x36 => istore,
                // 0x3b => istore_0,
                // 0x3c => istore_1,
                // 0x3d => istore_2,
                // 0x3e => istore_3,
                // 0x64 => isub,
                // 0x7c => iushr,
                // 0x82 => ixor,
                // 0xa8 => jsr,
                // 0xc9 => jsr_w,
                // 0x8a => l2d,
                // 0x89 => l2f,
                // 0x88 => l2i,
                // 0x61 => ladd,
                // 0x2f => laload,
                // 0x7f => land,
                // 0x50 => lastore,
                // 0x94 => lcmp,
                // 0x09 => lconst_0,
                // 0x0a => lconst_1,
                // 0x12 => ldc,
                // 0x13 => ldc_w,
                // 0x14 => ldc2_w,
                // 0x6d => ldiv,
                // 0x16 => lload,
                // 0x1e => lload_0,
                // 0x1f => lload_1,
                // 0x20 => lload_2,
                // 0x21 => lload_3,
                // 0x69 => lmul,
                // 0x75 => lneg,
                // 0xab => lookupswitch,
                // 0x81 => lor,
                // 0x71 => lrem,
                // 0xad => lreturn,
                // 0x79 => lshl,
                // 0x7b => lshr,
                // 0x37 => lstore,
                // 0x3f => lstore_0,
                // 0x40 => lstore_1,
                // 0x41 => lstore_2,
                // 0x42 => lstore_3,
                // 0x65 => lsub,
                // 0x7d => lushr,
                // 0x83 => lxor,
                // 0xc2 => monitorenter,
                // 0xc3 => monitorexit,
                // 0xc5 => multianewarray,
                // 0xbb => new,
                // 0xbc => newarray,
                // 0x00 => nop,
                // 0x57 => pop,
                // 0x58 => pop2,
                // 0xb5 => putfield,
                // 0xb3 => putstatic,
                // 0xa9 => ret,
                // 0xb1 => return,
                // 0x35 => saload,
                // 0x56 => sastore,
                // 0x11 => sipush,
                // 0x5f => swap,
                // 0xaa => tableswitch,
                // 0xc4 => wide,
                op@_ => return Err(format!("Unknown Instruction {:#x}", op)),
            });
        }

        // TODO: gotos/if anpassen
        return Ok(vec);
    }
}
<|endoftext|>


<commit_before>//! A lexer for the Ruby programming language.

use token::{Category, Token};
use tokenizer::{Tokenizer, StateFunction};

fn initial_state(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    for keyword in vec!["pub", "let", "mut", "match", "loop"] {
        if tokenizer.starts_with_lexeme(keyword) {
            tokenizer.tokenize_next(keyword.chars().count(), Category::Keyword);

            return Some(StateFunction(initial_state))
        }
    }

    if tokenizer.starts_with_lexeme("use") {
        tokenizer.tokenize_next(3, Category::Keyword);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("mod") {
        tokenizer.tokenize_next(3, Category::Keyword);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("extern") {
        tokenizer.tokenize_next(6, Category::Keyword);
        tokenizer.states.push(StateFunction(initial_state));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("crate") {
        tokenizer.tokenize_next(5, Category::Keyword);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("for") {
        tokenizer.tokenize_next(3, Category::Keyword);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("in") {
        tokenizer.tokenize_next(2, Category::Keyword);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("fn") {
        tokenizer.tokenize_next(2, Category::Keyword);
        tokenizer.states.push(StateFunction(function));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme(":") {
        tokenizer.advance();
        tokenizer.tokenize(Category::Key);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.has_prefix("::") {
        tokenizer.tokenize(Category::Identifier);
        tokenizer.tokenize_next(2, Category::Text);

        return Some(StateFunction(initial_state))
    } else if tokenizer.starts_with_lexeme("//") {
        return Some(StateFunction(comment))
    }

    match tokenizer.current_char() {
        Some('"') => {
            tokenizer.tokenize(Category::Text);
            tokenizer.advance();
            Some(StateFunction(inside_string))
        },
        Some('\'') => {
            tokenizer.tokenize(Category::Text);
            tokenizer.advance();
            Some(StateFunction(inside_single_quote_string))
        },
        Some('|') => {
            tokenizer.tokenize_next(1, Category::Text);
            tokenizer.states.push(StateFunction(argument));
            Some(StateFunction(whitespace))
        },
        Some('.') => {
            tokenizer.tokenize(Category::Identifier);
            tokenizer.tokenize_next(1, Category::Text);
            Some(StateFunction(initial_state))
        },
        Some('(') => {
            tokenizer.tokenize(Category::Call);
            tokenizer.tokenize_next(1, Category::Text);
            Some(StateFunction(argument))
        },
        Some('+') => {
            tokenizer.tokenize_next(1, Category::Operator);
            Some(StateFunction(initial_state))
        },
        Some(' ') | Some('\n') => {
            match tokenizer.next_non_whitespace_char() {
                Some('=') => {
                    tokenizer.tokenize(Category::Identifier);
                    tokenizer.consume_whitespace();
                    tokenizer.tokenize_next(1, Category::Text);
                },
                _ => tokenizer.consume_whitespace(),
            }
            Some(StateFunction(initial_state))
        },
        Some('=') => {
            tokenizer.tokenize(Category::Identifier);
            tokenizer.tokenize_next(1, Category::Text);
            Some(StateFunction(initial_state))
        },
        Some('[') => {
            tokenizer.tokenize(Category::Identifier);
            tokenizer.tokenize_next(1, Category::Text);
            Some(StateFunction(initial_state))
        },
        Some(']') => {
            tokenizer.tokenize_next(1, Category::Text);
            Some(StateFunction(initial_state))
        },
        Some(c) => {
            tokenizer.advance();

            if c.is_numeric() {
                Some(StateFunction(integer))
            } else {
                Some(StateFunction(initial_state))
            }
        }

        None => {
            tokenizer.tokenize(Category::Text);
            None
        }
    }
}

fn inside_string(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                '"' => {
                    tokenizer.advance();
                    tokenizer.tokenize(Category::String);
                    Some(StateFunction(initial_state))
                },
                '\\' => {
                    tokenizer.advance();
                    tokenizer.advance();
                    Some(StateFunction(inside_string))
                }
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(inside_string))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::String);
            None
        }
    }
}

fn inside_single_quote_string(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                '\'' => {
                    tokenizer.advance();
                    tokenizer.tokenize(Category::String);
                    Some(StateFunction(initial_state))
                },
                '\\' => {
                    tokenizer.advance();
                    tokenizer.advance();
                    Some(StateFunction(inside_single_quote_string))
                }
                ':' => {
                    tokenizer.advance();
                    tokenizer.tokenize(Category::Identifier);
                    Some(StateFunction(initial_state))
                },
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(inside_single_quote_string))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::String);
            None
        }
    }
}

fn whitespace(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    tokenizer.advance();
                    Some(StateFunction(whitespace))
                },
                _ => {
                    tokenizer.tokenize(Category::Whitespace);
                    match tokenizer.states.pop() {
                        Some(state) => Some(state),
                        None => Some(StateFunction(initial_state)),
                    }
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::Whitespace);
            None
        }
    }
}

fn argument(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    if tokenizer.starts_with_lexeme("true") {
        tokenizer.tokenize_next(4, Category::Boolean);
        return Some(StateFunction(argument))
    } else if tokenizer.starts_with_lexeme("false") {
        tokenizer.tokenize_next(5, Category::Boolean);
        return Some(StateFunction(argument))
    }

    match tokenizer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    tokenizer.tokenize(Category::Identifier);
                    tokenizer.states.push(StateFunction(argument));
                    Some(StateFunction(whitespace))
                },
                '|' | ')' => {
                    tokenizer.tokenize(Category::Identifier);
                    tokenizer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                },
                '=' | ',' => {
                    tokenizer.tokenize(Category::Identifier);
                    tokenizer.tokenize_next(1, Category::Text);
                    Some(StateFunction(argument))
                },
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(argument))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::Identifier);
            None
        }
    }
}

fn identifier(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    tokenizer.tokenize(Category::Identifier);
                    Some(StateFunction(initial_state))
                },
                '|' | '(' | ')' | '-' | ';' => {
                    tokenizer.tokenize(Category::Identifier);
                    tokenizer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                },
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(identifier))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::Identifier);
            None
        }
    }
}

fn function(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    tokenizer.tokenize(Category::Function);
                    Some(StateFunction(initial_state))
                },
                '(' => {
                    tokenizer.tokenize(Category::Function);
                    tokenizer.tokenize_next(1, Category::Text);
                    Some(StateFunction(argument))
                },
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(function))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::Whitespace);
            None
        }
    }
}

fn comment(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                '\n' => {
                    tokenizer.tokenize(Category::Comment);
                    Some(StateFunction(initial_state))
                },
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(comment))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::Whitespace);
            None
        }
    }
}

fn integer(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            if c.is_numeric() {
                tokenizer.advance();
                Some(StateFunction(integer))
            } else {
                tokenizer.tokenize(Category::Integer);
                Some(StateFunction(initial_state))
            }
        }

        None => {
            tokenizer.tokenize(Category::Integer);
            None
        }
    }
}

pub fn lex(data: &str) -> Vec<Token> {
    let mut tokenizer = Tokenizer::new(data);
    let mut state_function = StateFunction(initial_state);
    loop {
        let StateFunction(actual_function) = state_function;
        match actual_function(&mut tokenizer) {
            Some(f) => state_function = f,
            None => {
                match tokenizer.states.pop() {
                    Some(f) => state_function = f,
                    None => return tokenizer.tokens(),
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::lex;
    use token::Token;
    use token::Category;

    #[test]
    fn it_works() {
        let data = include_str!("../../test_data/rust.rs");
        let tokens = lex(data);
        let expected_tokens = vec![
            Token{ lexeme: "extern".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "crate".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "luthor".to_string(), category: Category::Identifier },
            Token{ lexeme: ";".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
            Token{ lexeme: "use".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "luthor".to_string(), category: Category::Identifier },
            Token{ lexeme: ";".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
            Token{ lexeme: "pub".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "fn".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "main".to_string(), category: Category::Function },
            Token{ lexeme: "(".to_string(), category: Category::Text },
            Token{ lexeme: ")".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "{".to_string(), category: Category::Text },
            Token{ lexeme: "\n    ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "let".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "mut".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "variable".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "=".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "\"string\"".to_string(), category: Category::String },
            Token{ lexeme: ";".to_string(), category: Category::Text },
            Token{ lexeme: "\n    ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "'loop_name:".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "for".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "value".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "in".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "collection".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "{}".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
            Token{ lexeme: "}".to_string(), category: Category::Text },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }
}
<commit_msg>Update Rust comment lexing to look for prefix rather than lexeme.<commit_after>//! A lexer for the Ruby programming language.

use token::{Category, Token};
use tokenizer::{Tokenizer, StateFunction};

fn initial_state(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    for keyword in vec!["pub", "let", "mut", "match", "loop"] {
        if tokenizer.starts_with_lexeme(keyword) {
            tokenizer.tokenize_next(keyword.chars().count(), Category::Keyword);

            return Some(StateFunction(initial_state))
        }
    }

    if tokenizer.starts_with_lexeme("use") {
        tokenizer.tokenize_next(3, Category::Keyword);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("mod") {
        tokenizer.tokenize_next(3, Category::Keyword);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("extern") {
        tokenizer.tokenize_next(6, Category::Keyword);
        tokenizer.states.push(StateFunction(initial_state));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("crate") {
        tokenizer.tokenize_next(5, Category::Keyword);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("for") {
        tokenizer.tokenize_next(3, Category::Keyword);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("in") {
        tokenizer.tokenize_next(2, Category::Keyword);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme("fn") {
        tokenizer.tokenize_next(2, Category::Keyword);
        tokenizer.states.push(StateFunction(function));

        return Some(StateFunction(whitespace))
    } else if tokenizer.starts_with_lexeme(":") {
        tokenizer.advance();
        tokenizer.tokenize(Category::Key);
        tokenizer.states.push(StateFunction(identifier));

        return Some(StateFunction(whitespace))
    } else if tokenizer.has_prefix("::") {
        tokenizer.tokenize(Category::Identifier);
        tokenizer.tokenize_next(2, Category::Text);

        return Some(StateFunction(initial_state))
    } else if tokenizer.has_prefix("//") {
        return Some(StateFunction(comment))
    }

    match tokenizer.current_char() {
        Some('"') => {
            tokenizer.tokenize(Category::Text);
            tokenizer.advance();
            Some(StateFunction(inside_string))
        },
        Some('\'') => {
            tokenizer.tokenize(Category::Text);
            tokenizer.advance();
            Some(StateFunction(inside_single_quote_string))
        },
        Some('|') => {
            tokenizer.tokenize_next(1, Category::Text);
            tokenizer.states.push(StateFunction(argument));
            Some(StateFunction(whitespace))
        },
        Some('.') => {
            tokenizer.tokenize(Category::Identifier);
            tokenizer.tokenize_next(1, Category::Text);
            Some(StateFunction(initial_state))
        },
        Some('(') => {
            tokenizer.tokenize(Category::Call);
            tokenizer.tokenize_next(1, Category::Text);
            Some(StateFunction(argument))
        },
        Some('+') => {
            tokenizer.tokenize_next(1, Category::Operator);
            Some(StateFunction(initial_state))
        },
        Some(' ') | Some('\n') => {
            match tokenizer.next_non_whitespace_char() {
                Some('=') => {
                    tokenizer.tokenize(Category::Identifier);
                    tokenizer.consume_whitespace();
                    tokenizer.tokenize_next(1, Category::Text);
                },
                _ => tokenizer.consume_whitespace(),
            }
            Some(StateFunction(initial_state))
        },
        Some('=') => {
            tokenizer.tokenize(Category::Identifier);
            tokenizer.tokenize_next(1, Category::Text);
            Some(StateFunction(initial_state))
        },
        Some('[') => {
            tokenizer.tokenize(Category::Identifier);
            tokenizer.tokenize_next(1, Category::Text);
            Some(StateFunction(initial_state))
        },
        Some(']') => {
            tokenizer.tokenize_next(1, Category::Text);
            Some(StateFunction(initial_state))
        },
        Some(c) => {
            tokenizer.advance();

            if c.is_numeric() {
                Some(StateFunction(integer))
            } else {
                Some(StateFunction(initial_state))
            }
        }

        None => {
            tokenizer.tokenize(Category::Text);
            None
        }
    }
}

fn inside_string(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                '"' => {
                    tokenizer.advance();
                    tokenizer.tokenize(Category::String);
                    Some(StateFunction(initial_state))
                },
                '\\' => {
                    tokenizer.advance();
                    tokenizer.advance();
                    Some(StateFunction(inside_string))
                }
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(inside_string))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::String);
            None
        }
    }
}

fn inside_single_quote_string(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                '\'' => {
                    tokenizer.advance();
                    tokenizer.tokenize(Category::String);
                    Some(StateFunction(initial_state))
                },
                '\\' => {
                    tokenizer.advance();
                    tokenizer.advance();
                    Some(StateFunction(inside_single_quote_string))
                }
                ':' => {
                    tokenizer.advance();
                    tokenizer.tokenize(Category::Identifier);
                    Some(StateFunction(initial_state))
                },
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(inside_single_quote_string))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::String);
            None
        }
    }
}

fn whitespace(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    tokenizer.advance();
                    Some(StateFunction(whitespace))
                },
                _ => {
                    tokenizer.tokenize(Category::Whitespace);
                    match tokenizer.states.pop() {
                        Some(state) => Some(state),
                        None => Some(StateFunction(initial_state)),
                    }
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::Whitespace);
            None
        }
    }
}

fn argument(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    if tokenizer.starts_with_lexeme("true") {
        tokenizer.tokenize_next(4, Category::Boolean);
        return Some(StateFunction(argument))
    } else if tokenizer.starts_with_lexeme("false") {
        tokenizer.tokenize_next(5, Category::Boolean);
        return Some(StateFunction(argument))
    }

    match tokenizer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    tokenizer.tokenize(Category::Identifier);
                    tokenizer.states.push(StateFunction(argument));
                    Some(StateFunction(whitespace))
                },
                '|' | ')' => {
                    tokenizer.tokenize(Category::Identifier);
                    tokenizer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                },
                '=' | ',' => {
                    tokenizer.tokenize(Category::Identifier);
                    tokenizer.tokenize_next(1, Category::Text);
                    Some(StateFunction(argument))
                },
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(argument))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::Identifier);
            None
        }
    }
}

fn identifier(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    tokenizer.tokenize(Category::Identifier);
                    Some(StateFunction(initial_state))
                },
                '|' | '(' | ')' | '-' | ';' => {
                    tokenizer.tokenize(Category::Identifier);
                    tokenizer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                },
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(identifier))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::Identifier);
            None
        }
    }
}

fn function(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    tokenizer.tokenize(Category::Function);
                    Some(StateFunction(initial_state))
                },
                '(' => {
                    tokenizer.tokenize(Category::Function);
                    tokenizer.tokenize_next(1, Category::Text);
                    Some(StateFunction(argument))
                },
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(function))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::Whitespace);
            None
        }
    }
}

fn comment(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            match c {
                '\n' => {
                    tokenizer.tokenize(Category::Comment);
                    Some(StateFunction(initial_state))
                },
                _ => {
                    tokenizer.advance();
                    Some(StateFunction(comment))
                }
            }
        }

        None => {
            tokenizer.tokenize(Category::Whitespace);
            None
        }
    }
}

fn integer(tokenizer: &mut Tokenizer) -> Option<StateFunction> {
    match tokenizer.current_char() {
        Some(c) => {
            if c.is_numeric() {
                tokenizer.advance();
                Some(StateFunction(integer))
            } else {
                tokenizer.tokenize(Category::Integer);
                Some(StateFunction(initial_state))
            }
        }

        None => {
            tokenizer.tokenize(Category::Integer);
            None
        }
    }
}

pub fn lex(data: &str) -> Vec<Token> {
    let mut tokenizer = Tokenizer::new(data);
    let mut state_function = StateFunction(initial_state);
    loop {
        let StateFunction(actual_function) = state_function;
        match actual_function(&mut tokenizer) {
            Some(f) => state_function = f,
            None => {
                match tokenizer.states.pop() {
                    Some(f) => state_function = f,
                    None => return tokenizer.tokens(),
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::lex;
    use token::Token;
    use token::Category;

    #[test]
    fn it_works() {
        let data = include_str!("../../test_data/rust.rs");
        let tokens = lex(data);
        let expected_tokens = vec![
            Token{ lexeme: "extern".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "crate".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "luthor".to_string(), category: Category::Identifier },
            Token{ lexeme: ";".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
            Token{ lexeme: "use".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "luthor".to_string(), category: Category::Identifier },
            Token{ lexeme: ";".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
            Token{ lexeme: "pub".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "fn".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "main".to_string(), category: Category::Function },
            Token{ lexeme: "(".to_string(), category: Category::Text },
            Token{ lexeme: ")".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "{".to_string(), category: Category::Text },
            Token{ lexeme: "\n    ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "let".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "mut".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "variable".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "=".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "\"string\"".to_string(), category: Category::String },
            Token{ lexeme: ";".to_string(), category: Category::Text },
            Token{ lexeme: "\n    ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "'loop_name:".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "for".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "value".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "in".to_string(), category: Category::Keyword },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "collection".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "{}".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
            Token{ lexeme: "}".to_string(), category: Category::Text },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }
}
<|endoftext|>