<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.
// spell-checker:ignore TODO canonicalizes direntry pathbuf symlinked
//! Recursively copy the contents of a directory.
//!
//! See the [`copy_directory`] function for more information.
#[cfg(windows)]
use std::borrow::Cow;
use std::collections::HashSet;
use std::env;
use std::fs;
use std::io;
use std::path::Path;

use uucore::display::Quotable;
use uucore::error::UIoError;
use uucore::fs::{canonicalize, FileInformation, MissingHandling, ResolveMode};
use walkdir::WalkDir;

use crate::{
    copy_attributes, copy_file, copy_link, preserve_hardlinks, CopyResult, Error, Options,
    TargetSlice,
};

/// Continue next iteration of loop if result of expression is error
macro_rules! or_continue(
    ($expr:expr) => (match $expr {
        Ok(temp) => temp,
        Err(error) => {
            show_error!("{}", error);
            continue
        },
    })
);

#[cfg(target_os = "windows")]
fn adjust_canonicalization(p: &Path) -> Cow<Path> {
    // In some cases, \\? can be missing on some Windows paths.  Add it at the
    // beginning unless the path is prefixed with a device namespace.
    const VERBATIM_PREFIX: &str = r#"\\?"#;
    const DEVICE_NS_PREFIX: &str = r#"\\."#;

    let has_prefix = p
        .components()
        .next()
        .and_then(|comp| comp.as_os_str().to_str())
        .map(|p_str| p_str.starts_with(VERBATIM_PREFIX) || p_str.starts_with(DEVICE_NS_PREFIX))
        .unwrap_or_default();

    if has_prefix {
        p.into()
    } else {
        Path::new(VERBATIM_PREFIX).join(p).into()
    }
}

/// Read the contents of the directory `root` and recursively copy the
/// contents to `target`.
///
/// Any errors encountered copying files in the tree will be logged but
/// will not cause a short-circuit.
pub(crate) fn copy_directory(
    root: &Path,
    target: &TargetSlice,
    options: &Options,
    symlinked_files: &mut HashSet<FileInformation>,
    source_in_command_line: bool,
) -> CopyResult<()> {
    if !options.recursive {
        return Err(format!("omitting directory {}", root.quote()).into());
    }

    // if no-dereference is enabled and this is a symlink, copy it as a file
    if !options.dereference(source_in_command_line) && root.is_symlink() {
        return copy_file(
            root,
            target,
            options,
            symlinked_files,
            source_in_command_line,
        );
    }

    // check if root is a prefix of target
    if path_has_prefix(target, root)? {
        return Err(format!(
            "cannot copy a directory, {}, into itself, {}",
            root.quote(),
            target.join(root.file_name().unwrap()).quote()
        )
        .into());
    }

    let current_dir =
        env::current_dir().unwrap_or_else(|e| crash!(1, "failed to get current directory {}", e));

    let root_path = current_dir.join(root);

    let root_parent = if target.exists() {
        root_path.parent()
    } else {
        Some(root_path.as_path())
    };

    #[cfg(unix)]
    let mut hard_links: Vec<(String, u64)> = vec![];
    let preserve_hard_links = options.preserve_hard_links();

    // This should be changed once Redox supports hardlinks
    #[cfg(any(windows, target_os = "redox"))]
    let mut hard_links: Vec<(String, u64)> = vec![];

    for path in WalkDir::new(root)
        .same_file_system(options.one_file_system)
        .follow_links(options.dereference)
    {
        let p = or_continue!(path);
        let path = current_dir.join(p.path());

        let local_to_root_parent = match root_parent {
            Some(parent) => {
                #[cfg(windows)]
                {
                    // On Windows, some paths are starting with \\?
                    // but not always, so, make sure that we are consistent for strip_prefix
                    // See https://docs.microsoft.com/en-us/windows/win32/fileio/naming-a-file for more info
                    let parent_can = adjust_canonicalization(parent);
                    let path_can = adjust_canonicalization(&path);

                    or_continue!(&path_can.strip_prefix(&parent_can)).to_path_buf()
                }
                #[cfg(not(windows))]
                {
                    or_continue!(path.strip_prefix(parent)).to_path_buf()
                }
            }
            None => path.clone(),
        };

        let local_to_target = target.join(&local_to_root_parent);
        if p.path().is_symlink() && !options.dereference {
            copy_link(&path, &local_to_target, symlinked_files)?;
        } else if path.is_dir() && !local_to_target.exists() {
            if target.is_file() {
                return Err("cannot overwrite non-directory with directory".into());
            }
            fs::create_dir_all(local_to_target)?;
        } else if !path.is_dir() {
            if preserve_hard_links {
                let mut found_hard_link = false;
                let source = path.to_path_buf();
                let dest = local_to_target.as_path().to_path_buf();
                preserve_hardlinks(&mut hard_links, &source, &dest, &mut found_hard_link)?;
                if !found_hard_link {
                    match copy_file(
                        path.as_path(),
                        local_to_target.as_path(),
                        options,
                        symlinked_files,
                        false,
                    ) {
                        Ok(_) => Ok(()),
                        Err(err) => {
                            if source.is_symlink() {
                                // silent the error with a symlink
                                // In case we do --archive, we might copy the symlink
                                // before the file itself
                                Ok(())
                            } else {
                                Err(err)
                            }
                        }
                    }?;
                }
            } else {
                // At this point, `path` is just a plain old file.
                // Terminate this function immediately if there is any
                // kind of error *except* a "permission denied" error.
                //
                // TODO What other kinds of errors, if any, should
                // cause us to continue walking the directory?
                match copy_file(
                    path.as_path(),
                    local_to_target.as_path(),
                    options,
                    symlinked_files,
                    false,
                ) {
                    Ok(_) => {}
                    Err(Error::IoErrContext(e, _))
                        if e.kind() == std::io::ErrorKind::PermissionDenied =>
                    {
                        show!(uio_error!(
                            e,
                            "cannot open {} for reading",
                            p.path().quote()
                        ));
                    }
                    Err(e) => return Err(e),
                }
            }
        }
    }
    // Copy the attributes from the root directory to the target directory.
    copy_attributes(root, target, &options.preserve_attributes)?;
    Ok(())
}

/// Decide whether the second path is a prefix of the first.
///
/// This function canonicalizes the paths via
/// [`uucore::fs::canonicalize`] before comparing.
///
/// # Errors
///
/// If there is an error determining the canonical, absolute form of
/// either path.
///
/// # Examples
///
/// ```rust,ignore
/// assert!(path_has_prefix(Path::new("/usr/bin"), Path::new("/usr")))
/// assert!(!path_has_prefix(Path::new("/usr"), Path::new("/usr/bin")))
/// assert!(!path_has_prefix(Path::new("/usr/bin"), Path::new("/var/log")))
/// ```
pub fn path_has_prefix(p1: &Path, p2: &Path) -> io::Result<bool> {
    let pathbuf1 = canonicalize(p1, MissingHandling::Normal, ResolveMode::Logical)?;
    let pathbuf2 = canonicalize(p2, MissingHandling::Normal, ResolveMode::Logical)?;

    Ok(pathbuf1.starts_with(pathbuf2))
}
<commit_msg>cp: break up copy_directory() into helper funcs<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.
// spell-checker:ignore TODO canonicalizes direntry pathbuf symlinked
//! Recursively copy the contents of a directory.
//!
//! See the [`copy_directory`] function for more information.
#[cfg(windows)]
use std::borrow::Cow;
use std::collections::HashSet;
use std::env;
use std::fs;
use std::io;
use std::path::{Path, PathBuf, StripPrefixError};

use uucore::display::Quotable;
use uucore::error::UIoError;
use uucore::fs::{canonicalize, FileInformation, MissingHandling, ResolveMode};
use walkdir::{DirEntry, WalkDir};

use crate::{
    copy_attributes, copy_file, copy_link, preserve_hardlinks, CopyResult, Error, Options,
    TargetSlice,
};

/// Ensure a Windows path starts with a `\\?`.
#[cfg(target_os = "windows")]
fn adjust_canonicalization(p: &Path) -> Cow<Path> {
    // In some cases, \\? can be missing on some Windows paths.  Add it at the
    // beginning unless the path is prefixed with a device namespace.
    const VERBATIM_PREFIX: &str = r#"\\?"#;
    const DEVICE_NS_PREFIX: &str = r#"\\."#;

    let has_prefix = p
        .components()
        .next()
        .and_then(|comp| comp.as_os_str().to_str())
        .map(|p_str| p_str.starts_with(VERBATIM_PREFIX) || p_str.starts_with(DEVICE_NS_PREFIX))
        .unwrap_or_default();

    if has_prefix {
        p.into()
    } else {
        Path::new(VERBATIM_PREFIX).join(p).into()
    }
}

/// Get a descendant path relative to the given parent directory.
///
/// If `root_parent` is `None`, then this just returns the `path`
/// itself. Otherwise, this function strips the parent prefix from the
/// given `path`, leaving only the portion of the path relative to the
/// parent.
fn get_local_to_root_parent(
    path: &Path,
    root_parent: Option<&Path>,
) -> Result<PathBuf, StripPrefixError> {
    match root_parent {
        Some(parent) => {
            // On Windows, some paths are starting with \\?
            // but not always, so, make sure that we are consistent for strip_prefix
            // See https://docs.microsoft.com/en-us/windows/win32/fileio/naming-a-file for more info
            #[cfg(windows)]
            let (path, parent) = (
                adjust_canonicalization(path),
                adjust_canonicalization(parent),
            );
            let path = path.strip_prefix(&parent)?;
            Ok(path.to_path_buf())
        }
        None => Ok(path.to_path_buf()),
    }
}

/// Paths that are invariant throughout the traversal when copying a directory.
struct Context<'a> {
    /// The current working directory at the time of starting the traversal.
    current_dir: PathBuf,

    /// The path to the parent of the source directory, if any.
    root_parent: Option<PathBuf>,

    /// The target path to which the directory will be copied.
    target: &'a Path,
}

impl<'a> Context<'a> {
    fn new(root: &'a Path, target: &'a Path) -> std::io::Result<Self> {
        let current_dir = env::current_dir()?;
        let root_path = current_dir.join(root);
        let root_parent = if target.exists() {
            root_path.parent().map(|p| p.to_path_buf())
        } else {
            Some(root_path)
        };
        Ok(Self {
            current_dir,
            root_parent,
            target,
        })
    }
}

/// Data needed to perform a single copy operation while traversing a directory.
///
/// For convenience while traversing a directory, the [`Entry::new`]
/// function allows creating an entry from a [`Context`] and a
/// [`walkdir::DirEntry`].
///
/// # Examples
///
/// For example, if the source directory structure is `a/b/c`, the
/// target is `d/`, a directory that already exists, and the copy
/// command is `cp -r a/b/c d`, then the overall set of copy
/// operations could be represented as three entries,
///
/// ```rust,ignore
/// let operations = [
///     Entry {
///         source_absolute: "/tmp/a".into(),
///         source_relative: "a".into(),
///         local_to_target: "d/a".into(),
///         target_is_file: false,
///     }
///     Entry {
///         source_absolute: "/tmp/a/b".into(),
///         source_relative: "a/b".into(),
///         local_to_target: "d/a/b".into(),
///         target_is_file: false,
///     }
///     Entry {
///         source_absolute: "/tmp/a/b/c".into(),
///         source_relative: "a/b/c".into(),
///         local_to_target: "d/a/b/c".into(),
///         target_is_file: false,
///     }
/// ];
/// ```
struct Entry {
    /// The absolute path to file or directory to copy.
    source_absolute: PathBuf,

    /// The relative path to file or directory to copy.
    source_relative: PathBuf,

    /// The path to the destination, relative to the target.
    local_to_target: PathBuf,

    /// Whether the destination is a file.
    target_is_file: bool,
}

impl Entry {
    fn new(context: &Context, direntry: &DirEntry) -> Result<Self, StripPrefixError> {
        let source_relative = direntry.path().to_path_buf();
        let source_absolute = context.current_dir.join(&source_relative);
        let descendant =
            get_local_to_root_parent(&source_absolute, context.root_parent.as_deref())?;
        let local_to_target = context.target.join(descendant);
        let target_is_file = context.target.is_file();
        Ok(Self {
            source_absolute,
            source_relative,
            local_to_target,
            target_is_file,
        })
    }
}

/// Copy a single entry during a directory traversal.
fn copy_direntry(
    entry: Entry,
    options: &Options,
    symlinked_files: &mut HashSet<FileInformation>,
    preserve_hard_links: bool,
    hard_links: &mut Vec<(String, u64)>,
) -> CopyResult<()> {
    let Entry {
        source_absolute,
        source_relative,
        local_to_target,
        target_is_file,
    } = entry;

    // If the source is a symbolic link and the options tell us not to
    // dereference the link, then copy the link object itself.
    if source_absolute.is_symlink() && !options.dereference {
        return copy_link(&source_absolute, &local_to_target, symlinked_files);
    }

    // If the source is a directory and the destination does not
    // exist, ...
    if source_absolute.is_dir() && !local_to_target.exists() {
        if target_is_file {
            return Err("cannot overwrite non-directory with directory".into());
        } else {
            // TODO Since the calling code is traversing from the root
            // of the directory structure, I don't think
            // `create_dir_all()` will have any benefit over
            // `create_dir()`, since all the ancestor directories
            // should have already been created.
            fs::create_dir_all(local_to_target)?;
            return Ok(());
        }
    }

    // If the source is not a directory, then we need to copy the file.
    if !source_absolute.is_dir() {
        if preserve_hard_links {
            let mut found_hard_link = false;
            let dest = local_to_target.as_path().to_path_buf();
            preserve_hardlinks(hard_links, &source_absolute, &dest, &mut found_hard_link)?;
            if !found_hard_link {
                match copy_file(
                    &source_absolute,
                    local_to_target.as_path(),
                    options,
                    symlinked_files,
                    false,
                ) {
                    Ok(_) => Ok(()),
                    Err(err) => {
                        if source_absolute.is_symlink() {
                            // silent the error with a symlink
                            // In case we do --archive, we might copy the symlink
                            // before the file itself
                            Ok(())
                        } else {
                            Err(err)
                        }
                    }
                }?;
            }
        } else {
            // At this point, `path` is just a plain old file.
            // Terminate this function immediately if there is any
            // kind of error *except* a "permission denied" error.
            //
            // TODO What other kinds of errors, if any, should
            // cause us to continue walking the directory?
            match copy_file(
                &source_absolute,
                local_to_target.as_path(),
                options,
                symlinked_files,
                false,
            ) {
                Ok(_) => {}
                Err(Error::IoErrContext(e, _))
                    if e.kind() == std::io::ErrorKind::PermissionDenied =>
                {
                    show!(uio_error!(
                        e,
                        "cannot open {} for reading",
                        source_relative.quote(),
                    ));
                }
                Err(e) => return Err(e),
            }
        }
    }

    // In any other case, there is nothing to do, so we just return to
    // continue the traversal.
    Ok(())
}

/// Read the contents of the directory `root` and recursively copy the
/// contents to `target`.
///
/// Any errors encountered copying files in the tree will be logged but
/// will not cause a short-circuit.
pub(crate) fn copy_directory(
    root: &Path,
    target: &TargetSlice,
    options: &Options,
    symlinked_files: &mut HashSet<FileInformation>,
    source_in_command_line: bool,
) -> CopyResult<()> {
    if !options.recursive {
        return Err(format!("omitting directory {}", root.quote()).into());
    }

    // if no-dereference is enabled and this is a symlink, copy it as a file
    if !options.dereference(source_in_command_line) && root.is_symlink() {
        return copy_file(
            root,
            target,
            options,
            symlinked_files,
            source_in_command_line,
        );
    }

    // check if root is a prefix of target
    if path_has_prefix(target, root)? {
        return Err(format!(
            "cannot copy a directory, {}, into itself, {}",
            root.quote(),
            target.join(root.file_name().unwrap()).quote()
        )
        .into());
    }

    let mut hard_links: Vec<(String, u64)> = vec![];
    let preserve_hard_links = options.preserve_hard_links();

    // Collect some paths here that are invariant during the traversal
    // of the given directory, like the current working directory and
    // the target directory.
    let context = match Context::new(root, target) {
        Ok(c) => c,
        Err(e) => return Err(format!("failed to get current directory {}", e).into()),
    };

    // Traverse the contents of the directory, copying each one.
    for direntry_result in WalkDir::new(root)
        .same_file_system(options.one_file_system)
        .follow_links(options.dereference)
    {
        match direntry_result {
            Ok(direntry) => {
                let entry = Entry::new(&context, &direntry)?;
                copy_direntry(
                    entry,
                    options,
                    symlinked_files,
                    preserve_hard_links,
                    &mut hard_links,
                )?;
            }
            // Print an error message, but continue traversing the directory.
            Err(e) => show_error!("{}", e),
        }
    }
    // Copy the attributes from the root directory to the target directory.
    copy_attributes(root, target, &options.preserve_attributes)?;
    Ok(())
}

/// Decide whether the second path is a prefix of the first.
///
/// This function canonicalizes the paths via
/// [`uucore::fs::canonicalize`] before comparing.
///
/// # Errors
///
/// If there is an error determining the canonical, absolute form of
/// either path.
///
/// # Examples
///
/// ```rust,ignore
/// assert!(path_has_prefix(Path::new("/usr/bin"), Path::new("/usr")))
/// assert!(!path_has_prefix(Path::new("/usr"), Path::new("/usr/bin")))
/// assert!(!path_has_prefix(Path::new("/usr/bin"), Path::new("/var/log")))
/// ```
pub fn path_has_prefix(p1: &Path, p2: &Path) -> io::Result<bool> {
    let pathbuf1 = canonicalize(p1, MissingHandling::Normal, ResolveMode::Logical)?;
    let pathbuf2 = canonicalize(p2, MissingHandling::Normal, ResolveMode::Logical)?;

    Ok(pathbuf1.starts_with(pathbuf2))
}
<|endoftext|>


<commit_before>use error::{Error, ErrorCode};
use query::{Condition, SortRule, Range, Query};
use schema::Type;
use value::{Value, ValueIter};

/// The driver trait which all drivers will implement. Designed to be
/// interoperable with any data source, however the driver also assumes a
/// collection based data model.
///
/// ## Collection Based Data Model
/// Any value the driver can access is assumed to have a specific “type”
/// associated with it. This way the driver can map to structures like
/// collections in MongoDB, tables or views in a SQL database like PostgreSQL,
/// or labels in Neo4j.
///
/// In addition, operations on the driver are done in a CRUD fashion for a
/// specific type. Looking at the specific methods, like `read`, heavy
/// inspiration is taken from SQL databases and MongoDB which do all their CRUD
/// on collections. The unit for CRUD in an Ardite driver is a type.
///
/// Originally the driver was designed to work strictly like Falcor or GraphQL
/// by assuming a “graph” like structure for *all* data. This turns out to be a
/// nice abstraction, however it is inneficient for creating performant and
/// flexible systems. Abandoning the graph structure of drivers also helps the
/// driver implementors who might find it incredibly difficult and repetitive
/// to copy the same graph interface on their relational database.
///
/// ## Relationship Between the Schema and the Driver
/// The driver and the schema should never interact with each other. The schema
/// should be managed completely by Ardite and used to validate queries and
/// values before being sent to the driver. The driver should only be a
/// low-level consistent interface to a data source and should not make any
/// intelligent decisions about how an Ardite program should run.
///
/// Some unique driver-specific features are encouraged to be included in the
/// driver, such as PostgreSQL computed columns. However, these features should
/// *not* be derived from the schema, these features *should* be consistent,
/// and the driver must consistently perform all other expected functionality
/// *before* implementing new features.
///
/// This sharp seperation is most primarily for a seperation of concerns. It
/// allows developers working with the driver higher up in Ardite to make basic
/// assumptions about the driver. The seperation also prevents code reuse. The
/// `validate_query` or `validate_value` methods on `Schema` may be tempting to
/// use in the driver. There should never be any confusion over where in the
/// algorithm these methods should be called.
pub trait Driver {
  /// Connects to a driver and returns a driver instance. After calling this
  /// the driver is ready to roll!
  ///
  /// No schema definition is provided to the driver in its construction step.
  fn connect(uri: &str) -> Result<Self, Error> where Self: Sized;

  /// Lazily read some values from the driver.
  ///
  /// Designed against a couple of database specifications. Including the
  /// following:
  ///
  /// - [SQL `SELECT` statement][1].
  /// - [MongoDB `find` command][2].
  ///
  /// [1]: http://www.postgresql.org/docs/current/static/sql-select.html
  /// [2]: https://docs.mongodb.org/manual/reference/command/find/
  fn read(
    &self,
    type_: &Type,
    condition: Condition,
    sort: Vec<SortRule>,
    range: Range,
    query: Query
  ) -> Result<ValueIter, Error>;

  /// Read a single value from the driver. The default implementation uses the
  /// driver read method with a range of one.
  ///
  /// If a condition matches more than one value (while not recommended for
  /// this method) the first of these values, using the default sorting
  /// algorithm of the database, is returned.
  ///
  /// This method may be optionally optimized by the driver.
  fn read_one(
    &self,
    type_: &Type,
    condition: Condition,
    query: Query
  ) -> Result<Value, Error> {
    let mut values: Vec<_> = try!(self.read(
      type_,
      condition,
      Default::default(),
      Range::new(None, Some(1)),
      query
    )).collect();

    if values.len() > 1 {
      Err(Error::internal("Read with a limit of one returned more than one value."))
    } else if let Some(value) = values.pop() {
      Ok(value)
    } else {
      Err(Error::new(ErrorCode::NotFound, "No value was found for the condition.", None))
    }
  }
}
<commit_msg>doc(driver): bolster inline documentation<commit_after>use error::{Error, ErrorCode};
use query::{Condition, SortRule, Range, Query};
use schema::Type;
use value::{Value, ValueIter};

/// The driver trait which all drivers will implement. Designed to be
/// interoperable with any data source, however the driver also assumes a
/// collection based data model.
pub trait Driver {
  /// Connects to a driver and returns a driver instance. After calling this
  /// the driver is ready to roll!
  ///
  /// No schema definition is provided to the driver in its construction step.
  fn connect(uri: &str) -> Result<Self, Error> where Self: Sized;

  /// Lazily read some values from the driver.
  ///
  /// Designed against a couple of database specifications. Including the
  /// following:
  ///
  /// - [SQL `SELECT` statement][1].
  /// - [MongoDB `find` command][2].
  ///
  /// [1]: http://www.postgresql.org/docs/current/static/sql-select.html
  /// [2]: https://docs.mongodb.org/manual/reference/command/find/
  fn read(
    &self,
    type_: &Type,
    condition: Condition,
    sort: Vec<SortRule>,
    range: Range,
    query: Query
  ) -> Result<ValueIter, Error>;

  /// Read a single value from the driver. The default implementation uses the
  /// driver read method with a range of one.
  ///
  /// If a condition matches more than one value (while not recommended for
  /// this method) the first of these values, using the default sorting
  /// algorithm of the database, is returned.
  ///
  /// This method may be optionally optimized by the driver.
  fn read_one(
    &self,
    type_: &Type,
    condition: Condition,
    query: Query
  ) -> Result<Value, Error> {
    let mut values: Vec<_> = try!(self.read(
      type_,
      condition,
      Default::default(),
      Range::new(None, Some(1)),
      query
    )).collect();

    if values.len() > 1 {
      Err(Error::internal("Read with a limit of one returned more than one value."))
    } else if let Some(value) = values.pop() {
      Ok(value)
    } else {
      Err(Error::new(ErrorCode::NotFound, "No value was found for the condition.", None))
    }
  }
}
<|endoftext|>


<commit_before>extern crate walkdir;
extern crate chrono;

use std::path;
use self::walkdir::{WalkDir, DirEntry, Error as WalkDirError};
use std::fs::File;
use std::io::BufReader;
use std::io::BufRead;
use self::chrono::{DateTime, UTC};
use std::error::Error;
use std::str::FromStr;

pub struct ELBLogEntry {
    timestamp: DateTime<UTC>,
    elb_name: String,
    client_address: String,
    backend_address: String,
    request_processing_time: f32,
    backend_processing_time: f32,
    response_processing_time: f32,
    elb_status_code: u16,
    backend_status_code: u16,
    received_bytes: String,
    sent_bytes: String,
    request_method: String,
    request_url: String,
    request_http_version: String
}

pub fn file_list(dir: &path::Path, filenames: &mut Vec<DirEntry>) -> Result<usize, WalkDirError> {
    for entry in WalkDir::new(dir).min_depth(1) {
        match entry {
            Err(err) => return Err(err),
            Ok(entry) => filenames.push(entry),
        }
    }
    Ok(filenames.len())
}

pub fn process_files(runtime_context: &::RuntimeContext, filenames: Vec<walkdir::DirEntry>) -> usize {
    let debug = runtime_context.debug;
    let mut record_count = 0;
    for filename in filenames {
        debug!(debug, "Processing file {}.", filename.path().display());
        match File::open(filename.path()) {
            Ok(file) => {
                let buffered_file = BufReader::new(&file);
                let recs: Vec<_> = buffered_file.lines()
                    .map(|x| {
                        parse_line(&(x.unwrap()))
                    })
                    .collect();
                record_count += recs.len();
                debug!(debug, "Found {} records in file {}.", recs.len(), filename.path().display());
            },
            Err(e) => {
                println!("ERROR: {}", e);
            }
        }
    }

    record_count
}

#[derive(Debug)]
pub struct ParsingError{
    property: &'static str,
    inner_description: Box<Error>,
}

#[derive(Debug)]
pub struct ParsingErrors {
    record: String,
    errors: Vec<ParsingError>,
}

const TIMESTAMP: &'static str = "timestamp";
const REQUEST_PROCESSING_TIME: &'static str = "request processing time";
const BACKEND_PROCESSING_TIME: &'static str = "backend processing time";
const RESPONSE_PROCESSING_TIME: &'static str = "response processing time";
const ELB_STATUS_CODE: &'static str = "ELB status code";
const BE_STATUS_CODE: &'static str = "backend status code";

pub fn parse_line(line: &String) -> Result<Box<ELBLogEntry>, Box<ParsingErrors>> {
    let split_line: Vec<_> = line.split(" ").collect();
    let mut errors: Vec<ParsingError> = Vec::new();

    let ts = parse_property::<DateTime<UTC>>(split_line[0], TIMESTAMP, &mut errors);
    let req_proc_time = parse_property::<f32>(split_line[4], REQUEST_PROCESSING_TIME, &mut errors);
    let be_proc_time = parse_property::<f32>(split_line[5], BACKEND_PROCESSING_TIME, &mut errors);
    let res_proc_time = parse_property::<f32>(split_line[6], RESPONSE_PROCESSING_TIME, &mut errors);
    let elb_sc = parse_property::<u16>(split_line[7], ELB_STATUS_CODE, &mut errors);
    let be_sc = parse_property::<u16>(split_line[8], BE_STATUS_CODE, &mut errors);

    if errors.is_empty() {
        Ok(Box::new(
            ELBLogEntry {
                timestamp: ts.unwrap(),
                elb_name: split_line[1].to_string(),
                client_address: split_line[2].to_string(),
                backend_address: split_line[3].to_string(),
                request_processing_time: req_proc_time.unwrap(),
                backend_processing_time: be_proc_time.unwrap(),
                response_processing_time: res_proc_time.unwrap(),
                elb_status_code: elb_sc.unwrap(),
                backend_status_code: be_sc.unwrap(),
                received_bytes: split_line[9].to_string(),
                sent_bytes: split_line[10].to_string(),
                request_method: split_line[11].trim_matches('"').to_string(),
                request_url: split_line[12].to_string(),
                request_http_version: split_line[13].trim_matches('"').to_string()
            }
        ))
    } else {
        Err(Box::new(
            ParsingErrors {
                record: line.clone(),
                errors: errors
            }
        ))
    }
}

fn parse_property<T>(raw_prop: &str, prop_name: &'static str, errors: &mut Vec<ParsingError>) -> Option<T>
    where T: FromStr,
    T::Err: Error + 'static,
{
    match raw_prop.parse::<T>() {
        Ok(parsed) => Some(parsed),

        Err(e) => {
            errors.push(
                ParsingError {
                    property: prop_name,
                    inner_description: Box::new(e),
                }
            );
            None
        }
    }
}

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

    const TEST_LINE: &'static str = "2015-08-15T23:43:05.302180Z elb-name 172.16.1.6:54814 \
    172.16.1.5:9000 0.000039 0.145507 0.00003 200 200 0 7582 \
    \"GET http://some.domain.com:80/path0/path1?param0=p0&param1=p1 HTTP/1.1\" \
    ";

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_http_version() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_http_version, "HTTP/1.1")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_url() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_url, "http://some.domain.com:80/path0/path1?param0=p0&param1=p1")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_method() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_method, "GET")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_sent_bytes() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.sent_bytes, "7582")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_received_bytes() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.received_bytes, "0")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_status_code() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_status_code, 200)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_elb_status_code() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.elb_status_code, 200)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_response_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.response_processing_time, 0.00003)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_processing_time, 0.145507)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_processing_time, 0.000039)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_address() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_address, "172.16.1.5:9000")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_client_address() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.client_address, "172.16.1.6:54814")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_timestamp() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(format!("{:?}", elb_log_entry.timestamp), "2015-08-15T23:43:05.302180Z")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_elb_name() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.elb_name, "elb-name")
	}
}
<commit_msg>Converted the bytes sent and received to u64.<commit_after>extern crate walkdir;
extern crate chrono;

use std::path;
use self::walkdir::{WalkDir, DirEntry, Error as WalkDirError};
use std::fs::File;
use std::io::BufReader;
use std::io::BufRead;
use self::chrono::{DateTime, UTC};
use std::error::Error;
use std::str::FromStr;

pub struct ELBLogEntry {
    timestamp: DateTime<UTC>,
    elb_name: String,
    client_address: String,
    backend_address: String,
    request_processing_time: f32,
    backend_processing_time: f32,
    response_processing_time: f32,
    elb_status_code: u16,
    backend_status_code: u16,
    received_bytes: u64,
    sent_bytes: u64,
    request_method: String,
    request_url: String,
    request_http_version: String
}

pub fn file_list(dir: &path::Path, filenames: &mut Vec<DirEntry>) -> Result<usize, WalkDirError> {
    for entry in WalkDir::new(dir).min_depth(1) {
        match entry {
            Err(err) => return Err(err),
            Ok(entry) => filenames.push(entry),
        }
    }
    Ok(filenames.len())
}

pub fn process_files(runtime_context: &::RuntimeContext, filenames: Vec<walkdir::DirEntry>) -> usize {
    let debug = runtime_context.debug;
    let mut record_count = 0;
    for filename in filenames {
        debug!(debug, "Processing file {}.", filename.path().display());
        match File::open(filename.path()) {
            Ok(file) => {
                let buffered_file = BufReader::new(&file);
                let recs: Vec<_> = buffered_file.lines()
                    .map(|x| {
                        parse_line(&(x.unwrap()))
                    })
                    .collect();
                record_count += recs.len();
                debug!(debug, "Found {} records in file {}.", recs.len(), filename.path().display());
            },
            Err(e) => {
                println!("ERROR: {}", e);
            }
        }
    }

    record_count
}

#[derive(Debug)]
pub struct ParsingError{
    property: &'static str,
    inner_description: Box<Error>,
}

#[derive(Debug)]
pub struct ParsingErrors {
    record: String,
    errors: Vec<ParsingError>,
}

const TIMESTAMP: &'static str = "timestamp";
const REQUEST_PROCESSING_TIME: &'static str = "request processing time";
const BACKEND_PROCESSING_TIME: &'static str = "backend processing time";
const RESPONSE_PROCESSING_TIME: &'static str = "response processing time";
const ELB_STATUS_CODE: &'static str = "ELB status code";
const BE_STATUS_CODE: &'static str = "backend status code";
const RECEIVED_BYTES: &'static str = "received bytes";
const SENT_BYTES: &'static str = "sent bytes";

pub fn parse_line(line: &String) -> Result<Box<ELBLogEntry>, Box<ParsingErrors>> {
    let split_line: Vec<_> = line.split(" ").collect();
    let mut errors: Vec<ParsingError> = Vec::new();

    let ts = parse_property::<DateTime<UTC>>(split_line[0], TIMESTAMP, &mut errors);
    let req_proc_time = parse_property::<f32>(split_line[4], REQUEST_PROCESSING_TIME, &mut errors);
    let be_proc_time = parse_property::<f32>(split_line[5], BACKEND_PROCESSING_TIME, &mut errors);
    let res_proc_time = parse_property::<f32>(split_line[6], RESPONSE_PROCESSING_TIME, &mut errors);
    let elb_sc = parse_property::<u16>(split_line[7], ELB_STATUS_CODE, &mut errors);
    let be_sc = parse_property::<u16>(split_line[8], BE_STATUS_CODE, &mut errors);
    let bytes_received = parse_property::<u64>(split_line[9], RECEIVED_BYTES, &mut errors);
    let bytes_sent = parse_property::<u64>(split_line[10], SENT_BYTES, &mut errors);

    if errors.is_empty() {
        Ok(Box::new(
            ELBLogEntry {
                timestamp: ts.unwrap(),
                elb_name: split_line[1].to_string(),
                client_address: split_line[2].to_string(),
                backend_address: split_line[3].to_string(),
                request_processing_time: req_proc_time.unwrap(),
                backend_processing_time: be_proc_time.unwrap(),
                response_processing_time: res_proc_time.unwrap(),
                elb_status_code: elb_sc.unwrap(),
                backend_status_code: be_sc.unwrap(),
                received_bytes: bytes_received.unwrap(),
                sent_bytes: bytes_sent.unwrap(),
                request_method: split_line[11].trim_matches('"').to_string(),
                request_url: split_line[12].to_string(),
                request_http_version: split_line[13].trim_matches('"').to_string()
            }
        ))
    } else {
        Err(Box::new(
            ParsingErrors {
                record: line.clone(),
                errors: errors
            }
        ))
    }
}

fn parse_property<T>(raw_prop: &str, prop_name: &'static str, errors: &mut Vec<ParsingError>) -> Option<T>
    where T: FromStr,
    T::Err: Error + 'static,
{
    match raw_prop.parse::<T>() {
        Ok(parsed) => Some(parsed),

        Err(e) => {
            errors.push(
                ParsingError {
                    property: prop_name,
                    inner_description: Box::new(e),
                }
            );
            None
        }
    }
}

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

    const TEST_LINE: &'static str = "2015-08-15T23:43:05.302180Z elb-name 172.16.1.6:54814 \
    172.16.1.5:9000 0.000039 0.145507 0.00003 200 200 0 7582 \
    \"GET http://some.domain.com:80/path0/path1?param0=p0&param1=p1 HTTP/1.1\" \
    ";

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_http_version() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_http_version, "HTTP/1.1")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_url() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_url, "http://some.domain.com:80/path0/path1?param0=p0&param1=p1")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_method() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_method, "GET")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_sent_bytes() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.sent_bytes, 7582)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_received_bytes() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.received_bytes, 0)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_status_code() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_status_code, 200)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_elb_status_code() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.elb_status_code, 200)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_response_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.response_processing_time, 0.00003)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_processing_time, 0.145507)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_processing_time, 0.000039)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_address() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_address, "172.16.1.5:9000")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_client_address() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.client_address, "172.16.1.6:54814")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_timestamp() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(format!("{:?}", elb_log_entry.timestamp), "2015-08-15T23:43:05.302180Z")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_elb_name() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.elb_name, "elb-name")
	}
}
<|endoftext|>


<commit_before>use std::thread;
use std::cell::UnsafeCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, RecvError};
use basetypes::*;
use chess_move::MoveStack;
use tt::*;
use position::Position;

pub enum Command {
    Search {
        search_id: usize,
        position: Position,
        depth: u8,
        lower_bound: Value,
        upper_bound: Value,
    },
    Stop,
    Exit,
}

pub enum Report {
    Progress {
        search_id: usize,
        searched_nodes: NodeCount,
    },
    Done {
        search_id: usize,
        searched_nodes: NodeCount,
        value: Option<Value>,
    },
}

pub fn run_search(tt: Arc<TranspositionTable>,
                  commands: Receiver<Command>,
                  reports: Sender<Report>) {
    thread_local!(
        static MOVE_STACK: UnsafeCell<MoveStack> = UnsafeCell::new(MoveStack::new())
    );
    MOVE_STACK.with(|s| {
        let mut move_stack = unsafe { &mut *s.get() };
        let mut pending_command = None;
        loop {
            // If there is a pending command, we take it, otherwise we
            // block and wait to receive a new one.
            let command = pending_command.take()
                                         .unwrap_or(commands.recv()
                                                            .or::<RecvError>(Ok(Command::Exit))
                                                            .unwrap());
            match command {
                Command::Search { search_id, mut position, depth, lower_bound, upper_bound } => {
                    let mut reported_nodes = 0;
                    let mut unreported_nodes = 0;
                    let value = search(&tt,
                                       &mut position,
                                       move_stack,
                                       &mut unreported_nodes,
                                       &mut |n| {
                                           reported_nodes += n;
                                           reports.send(Report::Progress {
                                                      search_id: search_id,
                                                      searched_nodes: reported_nodes,
                                                  })
                                                  .ok();
                                           match commands.try_recv() {
                                               Ok(x) => {
                                                   pending_command = Some(x);
                                                   Err(TerminatedSearch)
                                               }
                                               _ => Ok(()),
                                           }
                                       },
                                       lower_bound,
                                       upper_bound,
                                       depth)
                                    .ok();
                    reported_nodes += unreported_nodes;
                    reports.send(Report::Progress {
                               search_id: search_id,
                               searched_nodes: reported_nodes,
                           })
                           .ok();
                    reports.send(Report::Done {
                               search_id: search_id,
                               searched_nodes: reported_nodes,
                               value: value,
                           })
                           .ok();
                    move_stack.clear();
                }
                Command::Stop => continue,
                Command::Exit => break,
            }
        }
    })
}

pub fn run_deepening(tt: Arc<TranspositionTable>,
                     commands: Receiver<Command>,
                     reports: Sender<Report>) {
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        run_search(tt, slave_commands_rx, slave_reports_tx);
    });
    let mut pending_command = None;
    loop {
        // If there is a pending command, we take it, otherwise we
        // block and wait to receive a new one.
        match pending_command.take().unwrap_or(commands.recv()
                                                       .or::<RecvError>(Ok(Command::Exit))
                                                       .unwrap()) {
            Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                let mut searched_nodes_final = 0;
                let mut value_final = None;
                // TODO: May be use MIN_PLY..(depth+1).
                'depthloop: for n in 1..(depth + 1) {
                    slave_commands_tx.send(Command::Search {
                                         search_id: n as usize,
                                         position: position.clone(),
                                         depth: n,
                                         lower_bound: lower_bound,
                                         upper_bound: upper_bound,
                                     })
                                     .unwrap();
                    loop {
                        match slave_reports_rx.recv().unwrap() {
                            Report::Progress { searched_nodes, .. } => {
                                reports.send(Report::Progress {
                                           search_id: search_id,
                                           searched_nodes: searched_nodes_final + searched_nodes,
                                       })
                                       .ok();
                                if pending_command.is_none() {
                                    pending_command = match commands.try_recv() {
                                        Ok(cmd) => {
                                            slave_commands_tx.send(Command::Stop).unwrap();
                                            Some(cmd)
                                        }
                                        _ => None,
                                    }
                                }
                            }
                            Report::Done { searched_nodes, value, .. } => {
                                searched_nodes_final += searched_nodes;
                                if n == depth {
                                    value_final = value;
                                }
                                if pending_command.is_none() {
                                    break;
                                } else {
                                    break 'depthloop;
                                }
                            }
                        }
                    }
                }
                reports.send(Report::Done {
                           search_id: search_id,
                           searched_nodes: searched_nodes_final,
                           value: value_final,
                       })
                       .ok();
            }
            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }
            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}

// Represents a terminated search condition.
struct TerminatedSearch;

// Helper function for `run()`. It implements the principal variation
// search algorithm. When returning `Err(TerminatedSearch)`, this
// function may leave un-restored move lists in `moves`.
fn search(tt: &TranspositionTable,
          p: &mut Position,
          moves: &mut MoveStack,
          nc: &mut NodeCount,
          report: &mut FnMut(NodeCount) -> Result<(), TerminatedSearch>,
          mut alpha: Value, // lower bound
          beta: Value, // upper bound
          depth: u8)
          -> Result<Value, TerminatedSearch> {
    assert!(alpha < beta);
    if depth == 0 {
        // On leaf nodes, do quiescence search.
        let (value, nodes) = p.evaluate_quiescence(alpha, beta);
        *nc += nodes;
        Ok(value)
    } else {
        moves.save();
        p.generate_moves(moves);
        // TODO: Consult `tt` here.

        let mut bound_type = BOUND_UPPER;
        let mut move16 = 0;
        let mut no_moves_yet = true;
        while let Some(m) = moves.remove_best_move() {
            if p.do_move(m) {
                // From time to time, we report how many nodes had
                // been searched since the last report. This also
                // gives an opportunity for the search to be
                // terminated.
                *nc += 1;
                if *nc > NODE_COUNT_REPORT_INTERVAL {
                    try!(report(*nc));
                    *nc = 0;
                }

                let value = if no_moves_yet {
                    // The first move we analyze with a fully open
                    // window (alpha, beta).
                    -try!(search(tt, p, moves, nc, report, -beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that
                    // they are not better than our current best
                    // move. For this purpose we analyze them with a
                    // null window (alpha, alpha + 1). This is faster
                    // than a full window search. Only when we are
                    // certain that the move is better than our
                    // current best move, we do a full-window search.
                    match -try!(search(tt, p, moves, nc, report, -alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(search(tt, p, moves, nc, report, -beta, -alpha, depth - 1)),
                    }
                };
                no_moves_yet = false;
                p.undo_move();

                if value >= beta {
                    // This move is too good, so that the opponent
                    // will not allow this line of play to
                    // happen. Therefore we can stop here.
                    alpha = beta;
                    bound_type = BOUND_LOWER;
                    move16 = m.move16();
                    break;
                }
                if value > alpha {
                    // We found ourselves a new best move.
                    alpha = value;
                    bound_type = BOUND_EXACT;
                    move16 = m.move16();
                }
            }
        }
        moves.restore();
        if no_moves_yet {
            // No legal moves -- this is a final position.
            alpha = p.evaluate_final();
        }
        // TODO: Do not store values with too small `depth`.
        tt.store(p.hash(),
                 EntryData::new(alpha, bound_type, depth, move16, 0));
        Ok(alpha)
    }
}

const NODE_COUNT_REPORT_INTERVAL: NodeCount = 10000;

#[cfg(test)]
mod tests {
    use super::search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let mut p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let value = search(&TranspositionTable::new(),
                           &mut p,
                           &mut MoveStack::new(),
                           &mut 0,
                           &mut |_| Ok(()),
                           -20000,
                           20000,
                           2)
                        .ok()
                        .unwrap();
        assert!(value < -300);
        let value = search(&TranspositionTable::new(),
                           &mut p,
                           &mut MoveStack::new(),
                           &mut 0,
                           &mut |_| Ok(()),
                           -20000,
                           20000,
                           3)
                        .ok()
                        .unwrap();
        assert!(value >= 20000);
    }
}
<commit_msg>changed name run_search -> run<commit_after>use std::thread;
use std::cell::UnsafeCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, RecvError};
use basetypes::*;
use chess_move::MoveStack;
use tt::*;
use position::Position;

pub enum Command {
    Search {
        search_id: usize,
        position: Position,
        depth: u8,
        lower_bound: Value,
        upper_bound: Value,
    },
    Stop,
    Exit,
}

pub enum Report {
    Progress {
        search_id: usize,
        searched_nodes: NodeCount,
    },
    Done {
        search_id: usize,
        searched_nodes: NodeCount,
        value: Option<Value>,
    },
}

pub fn run(tt: Arc<TranspositionTable>, commands: Receiver<Command>, reports: Sender<Report>) {
    thread_local!(
        static MOVE_STACK: UnsafeCell<MoveStack> = UnsafeCell::new(MoveStack::new())
    );
    MOVE_STACK.with(|s| {
        let mut move_stack = unsafe { &mut *s.get() };
        let mut pending_command = None;
        loop {
            // If there is a pending command, we take it, otherwise we
            // block and wait to receive a new one.
            let command = pending_command.take()
                                         .unwrap_or(commands.recv()
                                                            .or::<RecvError>(Ok(Command::Exit))
                                                            .unwrap());
            match command {
                Command::Search { search_id, mut position, depth, lower_bound, upper_bound } => {
                    let mut reported_nodes = 0;
                    let mut unreported_nodes = 0;
                    let value = search(&tt,
                                       &mut position,
                                       move_stack,
                                       &mut unreported_nodes,
                                       &mut |n| {
                                           reported_nodes += n;
                                           reports.send(Report::Progress {
                                                      search_id: search_id,
                                                      searched_nodes: reported_nodes,
                                                  })
                                                  .ok();
                                           match commands.try_recv() {
                                               Ok(x) => {
                                                   pending_command = Some(x);
                                                   Err(TerminatedSearch)
                                               }
                                               _ => Ok(()),
                                           }
                                       },
                                       lower_bound,
                                       upper_bound,
                                       depth)
                                    .ok();
                    reported_nodes += unreported_nodes;
                    reports.send(Report::Progress {
                               search_id: search_id,
                               searched_nodes: reported_nodes,
                           })
                           .ok();
                    reports.send(Report::Done {
                               search_id: search_id,
                               searched_nodes: reported_nodes,
                               value: value,
                           })
                           .ok();
                    move_stack.clear();
                }
                Command::Stop => continue,
                Command::Exit => break,
            }
        }
    })
}

pub fn run_deepening(tt: Arc<TranspositionTable>,
                     commands: Receiver<Command>,
                     reports: Sender<Report>) {
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        run(tt, slave_commands_rx, slave_reports_tx);
    });
    let mut pending_command = None;
    loop {
        // If there is a pending command, we take it, otherwise we
        // block and wait to receive a new one.
        match pending_command.take().unwrap_or(commands.recv()
                                                       .or::<RecvError>(Ok(Command::Exit))
                                                       .unwrap()) {
            Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                let mut searched_nodes_final = 0;
                let mut value_final = None;
                // TODO: May be use MIN_PLY..(depth+1).
                'depthloop: for n in 1..(depth + 1) {
                    slave_commands_tx.send(Command::Search {
                                         search_id: n as usize,
                                         position: position.clone(),
                                         depth: n,
                                         lower_bound: lower_bound,
                                         upper_bound: upper_bound,
                                     })
                                     .unwrap();
                    loop {
                        match slave_reports_rx.recv().unwrap() {
                            Report::Progress { searched_nodes, .. } => {
                                reports.send(Report::Progress {
                                           search_id: search_id,
                                           searched_nodes: searched_nodes_final + searched_nodes,
                                       })
                                       .ok();
                                if pending_command.is_none() {
                                    pending_command = match commands.try_recv() {
                                        Ok(cmd) => {
                                            slave_commands_tx.send(Command::Stop).unwrap();
                                            Some(cmd)
                                        }
                                        _ => None,
                                    }
                                }
                            }
                            Report::Done { searched_nodes, value, .. } => {
                                searched_nodes_final += searched_nodes;
                                if n == depth {
                                    value_final = value;
                                }
                                if pending_command.is_none() {
                                    break;
                                } else {
                                    break 'depthloop;
                                }
                            }
                        }
                    }
                }
                reports.send(Report::Done {
                           search_id: search_id,
                           searched_nodes: searched_nodes_final,
                           value: value_final,
                       })
                       .ok();
            }
            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }
            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}

// Represents a terminated search condition.
struct TerminatedSearch;

// Helper function for `run()`. It implements the principal variation
// search algorithm. When returning `Err(TerminatedSearch)`, this
// function may leave un-restored move lists in `moves`.
fn search(tt: &TranspositionTable,
          p: &mut Position,
          moves: &mut MoveStack,
          nc: &mut NodeCount,
          report: &mut FnMut(NodeCount) -> Result<(), TerminatedSearch>,
          mut alpha: Value, // lower bound
          beta: Value, // upper bound
          depth: u8)
          -> Result<Value, TerminatedSearch> {
    assert!(alpha < beta);
    if depth == 0 {
        // On leaf nodes, do quiescence search.
        let (value, nodes) = p.evaluate_quiescence(alpha, beta);
        *nc += nodes;
        Ok(value)
    } else {
        moves.save();
        p.generate_moves(moves);
        // TODO: Consult `tt` here.

        let mut bound_type = BOUND_UPPER;
        let mut move16 = 0;
        let mut no_moves_yet = true;
        while let Some(m) = moves.remove_best_move() {
            if p.do_move(m) {
                // From time to time, we report how many nodes had
                // been searched since the last report. This also
                // gives an opportunity for the search to be
                // terminated.
                *nc += 1;
                if *nc > NODE_COUNT_REPORT_INTERVAL {
                    try!(report(*nc));
                    *nc = 0;
                }

                let value = if no_moves_yet {
                    // The first move we analyze with a fully open
                    // window (alpha, beta).
                    -try!(search(tt, p, moves, nc, report, -beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that
                    // they are not better than our current best
                    // move. For this purpose we analyze them with a
                    // null window (alpha, alpha + 1). This is faster
                    // than a full window search. Only when we are
                    // certain that the move is better than our
                    // current best move, we do a full-window search.
                    match -try!(search(tt, p, moves, nc, report, -alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(search(tt, p, moves, nc, report, -beta, -alpha, depth - 1)),
                    }
                };
                no_moves_yet = false;
                p.undo_move();

                if value >= beta {
                    // This move is too good, so that the opponent
                    // will not allow this line of play to
                    // happen. Therefore we can stop here.
                    alpha = beta;
                    bound_type = BOUND_LOWER;
                    move16 = m.move16();
                    break;
                }
                if value > alpha {
                    // We found ourselves a new best move.
                    alpha = value;
                    bound_type = BOUND_EXACT;
                    move16 = m.move16();
                }
            }
        }
        moves.restore();
        if no_moves_yet {
            // No legal moves -- this is a final position.
            alpha = p.evaluate_final();
        }
        // TODO: Do not store values with too small `depth`.
        tt.store(p.hash(),
                 EntryData::new(alpha, bound_type, depth, move16, 0));
        Ok(alpha)
    }
}

const NODE_COUNT_REPORT_INTERVAL: NodeCount = 10000;

#[cfg(test)]
mod tests {
    use super::search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let mut p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let value = search(&TranspositionTable::new(),
                           &mut p,
                           &mut MoveStack::new(),
                           &mut 0,
                           &mut |_| Ok(()),
                           -20000,
                           20000,
                           2)
                        .ok()
                        .unwrap();
        assert!(value < -300);
        let value = search(&TranspositionTable::new(),
                           &mut p,
                           &mut MoveStack::new(),
                           &mut 0,
                           &mut |_| Ok(()),
                           -20000,
                           20000,
                           3)
                        .ok()
                        .unwrap();
        assert!(value >= 20000);
    }
}
<|endoftext|>


<commit_before>use shakmaty::fen::Fen;
use shakmaty::uci::Uci;
use shakmaty::Color;

use serde::Deserialize;
use serde_with::{serde_as, CommaSeparator, DisplayFromStr, FromInto, StringWithSeparator};

use crate::lila::LilaVariant;
use crate::model::{Mode, Speed};

#[serde_as]
#[derive(Deserialize)]
pub struct PersonalQuery {
    variant: LilaVariant,
    #[serde_as(as = "DisplayFromStr")]
    fen: Fen,
    #[serde_as(as = "StringWithSeparator<CommaSeparator, Uci>")]
    #[serde(default)]
    play: Vec<Uci>,
    modes: Option<Vec<Mode>>,
    speeds: Option<Vec<Speed>>,
    player: String,
    #[serde_as(as = "FromInto<ColorProxy>")]
    color: Color,
    //since: _SinceYear,
}

struct _SinceYear(u8); // since 2000 or so

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
enum ColorProxy {
    White,
    Black,
}

impl From<Color> for ColorProxy {
    fn from(color: Color) -> ColorProxy {
        match color {
            Color::White => ColorProxy::White,
            Color::Black => ColorProxy::Black,
        }
    }
}

impl From<ColorProxy> for Color {
    fn from(color: ColorProxy) -> Color {
        match color {
            ColorProxy::White => Color::White,
            ColorProxy::Black => Color::Black,
        }
    }
}
<commit_msg>allow underscores in fen<commit_after>use crate::lila::LilaVariant;
use crate::model::{Mode, Speed};
use serde::Deserialize;
use serde_with::{serde_as, CommaSeparator, DisplayFromStr, FromInto, StringWithSeparator};
use shakmaty::fen::{Fen, ParseFenError};
use shakmaty::uci::Uci;
use shakmaty::Color;
use std::fmt;
use std::str::FromStr;

#[serde_as]
#[derive(Deserialize)]
pub struct PersonalQuery {
    variant: LilaVariant,
    #[serde_as(as = "DisplayFromStr")]
    fen: LaxFen,
    #[serde_as(as = "StringWithSeparator<CommaSeparator, Uci>")]
    #[serde(default)]
    play: Vec<Uci>,
    modes: Option<Vec<Mode>>,
    speeds: Option<Vec<Speed>>,
    player: String,
    #[serde_as(as = "FromInto<ColorProxy>")]
    color: Color,
    //since: _SinceYear,
}

struct _SinceYear(u8); // since 2000 or so

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
enum ColorProxy {
    White,
    Black,
}

impl From<Color> for ColorProxy {
    fn from(color: Color) -> ColorProxy {
        match color {
            Color::White => ColorProxy::White,
            Color::Black => ColorProxy::Black,
        }
    }
}

impl From<ColorProxy> for Color {
    fn from(color: ColorProxy) -> Color {
        match color {
            ColorProxy::White => Color::White,
            ColorProxy::Black => Color::Black,
        }
    }
}

pub struct LaxFen(Fen);

impl From<Fen> for LaxFen {
    fn from(fen: Fen) -> LaxFen {
        LaxFen(fen)
    }
}

impl From<LaxFen> for Fen {
    fn from(LaxFen(fen): LaxFen) -> Fen {
        fen
    }
}

impl FromStr for LaxFen {
    type Err = ParseFenError;

    fn from_str(s: &str) -> Result<LaxFen, ParseFenError> {
        s.replace("_", " ").parse().map(LaxFen)
    }
}
<|endoftext|>


<commit_before>use sa::{insert, suffix_array};

// Generate the BWT of input data (calls the given function with the BWT data as it's generated)
pub fn bwt(input: &[u8]) -> Vec<u8> {
    suffix_array(input).into_iter().map(|i| {
        // BWT[i] = S[SA[i] - 1]
        if i == 0 { 0 } else { input[(i - 1) as usize] }
    }).collect()
}

// Takes a frequency map of bytes and generates the index of first occurrence
// of each byte.
fn generate_occurrence_index(map: &mut Vec<u32>) {
    let mut idx = 0;
    for i in 0..map.len() {
        let c = map[i];
        map[i] = idx;
        idx += c;
    }
}

// Invert the BWT data (generate the original data)
pub fn ibwt(input: &[u8]) -> Vec<u8> {
    // get the byte distribution
    let mut map = Vec::new();
    for i in input {
        insert(&mut map, *i);
    }

    generate_occurrence_index(&mut map);

    // generate the LF vector
    let mut lf = vec![0; input.len()];
    for (i, c) in input.iter().enumerate() {
        let byte = *c as usize;
        let val = map[byte];
        lf[i] = val;
        map[byte] = val + 1;
    }

    let mut idx = 0;
    // construct the sequence by traversing through the LF vector
    let mut output = vec![0; input.len()];
    for i in (0..(input.len() - 1)).rev() {
        output[i] = input[idx];
        idx = lf[idx] as usize;
    }

    output.pop();
    output
}

#[derive(Clone, Debug, RustcEncodable, RustcDecodable)]
pub struct FMIndex {
    // BW-transformed data
    data: Vec<u8>,
    // forward frequency of each character in the BWT data
    cache: Vec<u32>,
    // incremental character frequencies
    occ_map: Vec<u32>,
    // LF-mapping for backward search
    lf_vec: Vec<u32>,
}

impl FMIndex {
    pub fn new(data: &[u8]) -> FMIndex {
        FMIndex::new_from_bwt(bwt(data))
    }

    pub fn new_from_bwt(bwt_data: Vec<u8>) -> FMIndex {
        let mut map = Vec::new();
        let mut count = vec![0u32; bwt_data.len()];
        let mut idx = 0;
        // generate the frequency map and forward frequency vector from BWT
        for i in &bwt_data {
            let value = insert(&mut map, *i);
            count[idx] = value;
            idx += 1;
        }

        generate_occurrence_index(&mut map);

        let mut lf_vec = count.clone();
        let mut lf_occ_map = map.clone();
        // generate the LF vector (just like inverting the BWT)
        for (i, c) in bwt_data.iter().enumerate() {
            let idx = *c as usize;
            let val = lf_occ_map[idx];
            lf_vec[i] = val;
            lf_occ_map[idx] = val + 1;
        }

        let mut i = 0;
        let mut counter = bwt_data.len() as u32;
        // Only difference is that we replace the LF indices with the lengths of prefix
        // from a particular position (in other words, the number of times
        // it would take us to get to the start of string).
        for _ in 0..bwt_data.len() {
            let next = lf_vec[i];
            lf_vec[i] = counter;
            i = next as usize;
            counter -= 1;
        }

        FMIndex {
            data: bwt_data,
            cache: count,
            occ_map: map,
            lf_vec: lf_vec,
        }
    }

    // Get the index of the nearest occurrence of a character in the BWT data
    fn nearest(&self, idx: usize, ch: u8) -> usize {
        match self.occ_map.get(ch as usize) {
            Some(res) => {
                *res as usize + (0..idx).rev()
                                        .find(|&i| self.data[i] == ch)
                                        .map(|i| self.cache[i] as usize)
                                        .unwrap_or(0)
            },
            None => 0,
        }
    }

    // Find the positions of occurrences of sub-string in the original data.
    pub fn search(&self, query: &str) -> Vec<usize> {
        let mut top = 0;
        let mut bottom = self.data.len();
        for ch in query.as_bytes().iter().rev() {
            top = self.nearest(top, *ch);
            bottom = self.nearest(bottom, *ch);
        }

        (top..bottom).map(|idx| {
            let i = self.nearest(idx, self.data[idx]);
            // wrap around on overflow, which usually occurs only for the
            // last index of LF vector (or the first index of original string)
            self.lf_vec[i] as usize % self.data.len()
        }).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::{FMIndex, bwt, ibwt};

    #[test]
    fn test_bwt_and_ibwt() {
        let text = String::from("ATCTAGGAGATCTGAATCTAGTTCAACTAGCTAGATCTAGAGACAGCTAA");
        let bw = bwt(&text.as_bytes().to_owned());
        let ibw = ibwt(&bw);
        assert_eq!(String::from("AATCGGAGTTGCTTTG\u{0}AGTAGTGATTTTAAGAAAAAACCCCCCTAAAACG"),
                   String::from_utf8(bw).unwrap());
        assert_eq!(text, String::from_utf8(ibw).unwrap());
    }

    #[test]
    fn test_fm_index() {
        let text = String::from("GCGTGCCCAGGGCACTGCCGCTGCAGGCGTAGGCATCGCATCACACGCGT");
        let index = FMIndex::new(&text.as_bytes().to_owned());
        let mut result = index.search("TG");
        result.sort();
        assert_eq!(result, vec![3, 15, 21]);
        let mut result = index.search("GCGT");
        result.sort();
        assert_eq!(result, vec![0, 26, 46]);
        assert_eq!(vec![1], index.search("CGTGCCC"));
    }
}
<commit_msg>Perf bug: Bail out when character is not present<commit_after>use sa::{insert, suffix_array};

// Generate the BWT of input data (calls the given function with the BWT data as it's generated)
pub fn bwt(input: &[u8]) -> Vec<u8> {
    suffix_array(input).into_iter().map(|i| {
        // BWT[i] = S[SA[i] - 1]
        if i == 0 { 0 } else { input[(i - 1) as usize] }
    }).collect()
}

// Takes a frequency map of bytes and generates the index of first occurrence
// of each byte.
fn generate_occurrence_index(map: &mut Vec<u32>) {
    let mut idx = 0;
    for i in 0..map.len() {
        let c = map[i];
        map[i] = idx;
        idx += c;
    }
}

// Invert the BWT data (generate the original data)
pub fn ibwt(input: &[u8]) -> Vec<u8> {
    // get the byte distribution
    let mut map = Vec::new();
    for i in input {
        insert(&mut map, *i);
    }

    generate_occurrence_index(&mut map);

    // generate the LF vector
    let mut lf = vec![0; input.len()];
    for (i, c) in input.iter().enumerate() {
        let byte = *c as usize;
        let val = map[byte];
        lf[i] = val;
        map[byte] = val + 1;
    }

    let mut idx = 0;
    // construct the sequence by traversing through the LF vector
    let mut output = vec![0; input.len()];
    for i in (0..(input.len() - 1)).rev() {
        output[i] = input[idx];
        idx = lf[idx] as usize;
    }

    output.pop();
    output
}

#[derive(Clone, Debug, RustcEncodable, RustcDecodable)]
pub struct FMIndex {
    // BW-transformed data
    data: Vec<u8>,
    // forward frequency of each character in the BWT data
    cache: Vec<u32>,
    // incremental character frequencies
    occ_map: Vec<u32>,
    // LF-mapping for backward search
    lf_vec: Vec<u32>,
}

impl FMIndex {
    pub fn new(data: &[u8]) -> FMIndex {
        FMIndex::new_from_bwt(bwt(data))
    }

    pub fn new_from_bwt(bwt_data: Vec<u8>) -> FMIndex {
        let mut map = Vec::new();
        let mut count = vec![0u32; bwt_data.len()];
        let mut idx = 0;
        // generate the frequency map and forward frequency vector from BWT
        for i in &bwt_data {
            let value = insert(&mut map, *i);
            count[idx] = value;
            idx += 1;
        }

        generate_occurrence_index(&mut map);

        let mut lf_vec = count.clone();
        let mut lf_occ_map = map.clone();
        // generate the LF vector (just like inverting the BWT)
        for (i, c) in bwt_data.iter().enumerate() {
            let idx = *c as usize;
            let val = lf_occ_map[idx];
            lf_vec[i] = val;
            lf_occ_map[idx] = val + 1;
        }

        let mut i = 0;
        let mut counter = bwt_data.len() as u32;
        // Only difference is that we replace the LF indices with the lengths of prefix
        // from a particular position (in other words, the number of times
        // it would take us to get to the start of string).
        for _ in 0..bwt_data.len() {
            let next = lf_vec[i];
            lf_vec[i] = counter;
            i = next as usize;
            counter -= 1;
        }

        FMIndex {
            data: bwt_data,
            cache: count,
            occ_map: map,
            lf_vec: lf_vec,
        }
    }

    // Get the index of the nearest occurrence of a character in the BWT data
    fn nearest(&self, idx: usize, ch: u8) -> usize {
        match self.occ_map.get(ch as usize) {
            Some(res) if *res > 0 => {
                *res as usize + (0..idx).rev()
                                        .find(|&i| self.data[i] == ch)
                                        .map(|i| self.cache[i] as usize)
                                        .unwrap_or(0)
            },
            _ => 0,
        }
    }

    // Find the positions of occurrences of sub-string in the original data.
    pub fn search(&self, query: &str) -> Vec<usize> {
        let mut top = 0;
        let mut bottom = self.data.len();
        for ch in query.as_bytes().iter().rev() {
            top = self.nearest(top, *ch);
            bottom = self.nearest(bottom, *ch);
        }

        (top..bottom).map(|idx| {
            let i = self.nearest(idx, self.data[idx]);
            // wrap around on overflow, which usually occurs only for the
            // last index of LF vector (or the first index of original string)
            self.lf_vec[i] as usize % self.data.len()
        }).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::{FMIndex, bwt, ibwt};

    #[test]
    fn test_bwt_and_ibwt() {
        let text = String::from("ATCTAGGAGATCTGAATCTAGTTCAACTAGCTAGATCTAGAGACAGCTAA");
        let bw = bwt(&text.as_bytes().to_owned());
        let ibw = ibwt(&bw);
        assert_eq!(String::from("AATCGGAGTTGCTTTG\u{0}AGTAGTGATTTTAAGAAAAAACCCCCCTAAAACG"),
                   String::from_utf8(bw).unwrap());
        assert_eq!(text, String::from_utf8(ibw).unwrap());
    }

    #[test]
    fn test_fm_index() {
        let text = String::from("GCGTGCCCAGGGCACTGCCGCTGCAGGCGTAGGCATCGCATCACACGCGT");
        let index = FMIndex::new(&text.as_bytes().to_owned());
        let mut result = index.search("TG");
        result.sort();
        assert_eq!(result, vec![3, 15, 21]);
        let mut result = index.search("GCGT");
        result.sort();
        assert_eq!(result, vec![0, 26, 46]);
        assert_eq!(vec![1], index.search("CGTGCCC"));
    }
}
<|endoftext|>


<commit_before>extern crate rand;

use atty::Stream;
use std;
use rand::Rng;
use std::thread::sleep;
use std::time::Duration;

// A struct to contain info we need to print with every character
pub struct Control {
    pub seed: f64,
    pub spread: f64,
    pub frequency: f64,
    pub background_mode: bool,
    pub dialup_mode: bool,
}

pub fn print_lines_lol<I : Iterator<Item=S>, S : AsRef<str>>(lines: I, c: &mut Control) {
    for line in lines {
        print_line_lol(line.as_ref(), c);

        if c.dialup_mode {
            let stall = Duration::from_millis(rand::thread_rng().gen_range(30, 700));
            sleep(stall);
        }
    }

    // Reset the foreground color, since we are at the end of the output
    print!("\x1b[39m");
}

// Takes in s a string of characters terminated by an implicit newline
// prints them using colored_print
// TODO Ignores most escape sequences
// TODO Adds the color to a color escape sequence
fn print_line_lol(s: &str, c: &mut Control) {
    let original_seed = c.seed;
    let mut skipping = false;
    let mut whitespace_after_newline = true;

    if !atty::is(Stream::Stdout) {
        println!("{}", s);
        return;
    }

    for character in s.chars() {
        // Strip out any color chars
        if character == '\x1b' {
            skipping = true;
            continue;
        }
        if skipping && character == 'm' {
            skipping = false;
            continue;
        }
        if skipping {
            continue;
        }

        if !character.is_whitespace() {
            whitespace_after_newline = false;
        }

        if whitespace_after_newline {
            print!("{}", character);
            continue;
        }

        c.seed += 1.0;

        if c.background_mode {
            let bg = get_color_tuple(c);
            let fg = calc_fg_color(bg);
            colored_print_with_background(fg, bg, character);
        } else {
            let fg = get_color_tuple(c);
            colored_print(fg, character);
        }
    }

    // Reset the background color if we should, since otherwise it will bleed all the way to the
    // end of the terminal
    if c.background_mode {
        print!("\x1b[49m");
    }

    // Print a newline since one is implied
    println!();
    c.seed = original_seed + 1.0; // Reset the seed, but bump it a bit
}

fn calc_fg_color(bg: (u8, u8, u8)) -> (u8, u8, u8) {
    // Currently, it only computes the forground clolor based on some threshold
    // on grayscale value.
    // TODO: Add a better algorithm for computing forground color
    if conv_grayscale(bg) > 0xA0_u8 {
        (0u8, 0u8, 0u8)
    } else {
        (0xffu8, 0xffu8, 0xffu8)
    }
}

fn linear_to_srgb(intensity: f64) -> f64 {
    if intensity <= 0.003_130_8 {
        12.92 * intensity
    } else {
        1.055 * intensity.powf(1.0 / 2.4) - 0.055
    }
}

fn srgb_to_linear(intensity: f64) -> f64 {
    if intensity < 0.04045 {
        intensity / 12.92
    } else {
        ((intensity + 0.055) / 1.055).powf(2.4)
    }
}

fn conv_grayscale(color: (u8, u8, u8)) -> u8 {
    // See https://en.wikipedia.org/wiki/Grayscale#Converting_color_to_grayscale
    const SCALE: f64 = 256.0;

    // Changing SRGB to Linear for gamma correction
    let red = srgb_to_linear(f64::from(color.0) / SCALE);
    let green = srgb_to_linear(f64::from(color.1) / SCALE);
    let blue = srgb_to_linear(f64::from(color.2) / SCALE);

    // Converting to grayscale
    let gray_linear = red * 0.299 + green * 0.587 + blue * 0.114;

    // Gamma correction
    let gray_srgb = linear_to_srgb(gray_linear);

    (gray_srgb * SCALE) as u8
}

fn colored_print(fg: (u8, u8, u8), c: char) {
    print!("\x1b[38;2;{};{};{}m{}", fg.0, fg.1, fg.2, c);
}

fn colored_print_with_background(fg: (u8, u8, u8), bg: (u8, u8, u8), c: char) {
    print!(
        "\x1b[38;2;{};{};{};48;2;{};{};{}m{}",
        fg.0, fg.1, fg.2, bg.0, bg.1, bg.2, c
    );
}

fn get_color_tuple(c: &Control) -> (u8, u8, u8) {
    let i = c.frequency * c.seed / c.spread;
    let red = i.sin() * 127.00 + 128.00;
    let green = (i + (std::f64::consts::PI * 2.00 / 3.00)).sin() * 127.00 + 128.00;
    let blue = (i + (std::f64::consts::PI * 4.00 / 3.00)).sin() * 127.00 + 128.00;

    (red as u8, green as u8, blue as u8)
}
<commit_msg>Add hacky(ish) support for recognizing when we are actually in an escape sequence. Does not render `sl` correctly yet, though, becasue it buffers input incorrectly<commit_after>extern crate rand;

use atty::Stream;
use std;
use rand::Rng;
use std::thread::sleep;
use std::time::Duration;

// A struct to contain info we need to print with every character
pub struct Control {
    pub seed: f64,
    pub spread: f64,
    pub frequency: f64,
    pub background_mode: bool,
    pub dialup_mode: bool,
}

pub fn print_lines_lol<I : Iterator<Item=S>, S : AsRef<str>>(lines: I, c: &mut Control) {
    for line in lines {
        print_line_lol(line.as_ref(), c);

        if c.dialup_mode {
            let stall = Duration::from_millis(rand::thread_rng().gen_range(30, 700));
            sleep(stall);
        }
    }
}

// Takes in s a string of characters
// prints them using colored_print, and a newline at the end
// Ignores most escape sequences
// TODO Adds the color to a color escape sequence
fn print_line_lol(s: &str, c: &mut Control) {
    let original_seed = c.seed;
    let mut ignore_whitespace = c.background_mode;

    if !atty::is(Stream::Stdout) {
        println!("{}", s);
        return;
    }

    let mut iter = s.chars();
    while let Some(character) = iter.next() {
        // Consume escape sequences
        if character == '\x1b' {
            // Escape sequences seem to be one of many different categories
            // CSI sequences are \e \[ [bytes in 0x30-0x3F] [bytes in 0x20-0x2F] [final byte in 0x40-0x7E]
            // nF Escape seq are \e [bytes in 0x20-0x2F] [byte in 0x30-0x7E]
            // Fp Escape seq are \e [byte in 0x30-0x3F] [I have no idea, but `sl` creates one where the next byte is the end of the escape sequence, so assume that]
            // Fe Escape seq are \e [byte in 0x40-0x5F] [I have no idea, '' though sl doesn't make one]
            // Fs Escape seq are \e [byte in 0x60-0x7E] [I have no idea, '' though sl doesn't make one]
            // Otherwise the next byte is the whole escape sequence (maybe? I can't exactly tell, but I will go with it)
            // We will consume up to, but not through, the next printable character
            // In addition, we print everything in the escape sequence, even if it is a color (that will be overriden)
            print!("\x1b");
            let mut escape_sequence_character = iter.next().expect("Escape character with no escape sequence after it");
            print!("{}", escape_sequence_character);
            match escape_sequence_character {
            '[' => {
                loop {
                    escape_sequence_character = iter.next().expect("CSI escape sequence did not terminate");
                    print!("{}", escape_sequence_character);
                    match escape_sequence_character {
                    '\x30' ..= '\x3F' => continue,
                    '\x20' ..= '\x2F' => {
                        loop {
                            escape_sequence_character = iter.next().expect("CSI escape sequence did not terminate");
                            print!("{}", escape_sequence_character);
                            match escape_sequence_character {
                            '\x20' ..= '\x2F' => continue,
                            '\x40' ..= '\x7E' => break,
                            _ => panic!("CSI escape sequence terminated with an incorrect value"),
                            }
                        }
                        break;
                    },
                    '\x40' ..= '\x7E' => break,
                    _ => panic!("CSI escape sequence terminated with an incorrect value"),
                    }
                }
            },
            '\x20' ..= '\x2F' => {
                loop {
                    escape_sequence_character = iter.next().expect("nF escape sequence did not terminate");
                    print!("{}", escape_sequence_character);
                    match escape_sequence_character {
                    '\x20' ..= '\x2F' => continue,
                    '\x30' ..= '\x7E' => break,
                    _ => panic!("nF escape sequence terminated with an incorrect value"),
                    }
                }
            },
//            '\x30' ..= '\x3F' => panic!("Fp escape sequences are not supported"),
//            '\x40' ..= '\x5F' => panic!("Fe escape sequences are not supported"),
//            '\x60' ..= '\x7E' => panic!("Fs escape sequences are not supported"),
            // be lazy and assume in all other cases we consume exactly 1 byte
            _ => (),
            }
        }

        // In background mode, don't print colorful whitespace until the first printable character
        if ignore_whitespace && character.is_whitespace() {
            print!("{}", character);
            continue;
        } else {
            ignore_whitespace = false;
        }

        colored_print(character, c);
        c.seed += 1.0;
    }

    // Reset the background color if we should, since otherwise it will bleed all the way to the
    // end of the terminal
    if c.background_mode {
        print!("\x1b[49m");
    }

    // Reset the foreground color, since we are at the end of the output, and print the last newline
    // This has to happen here, rather than only at the end of the program, because otherwise
    // newlines become weird
    println!("\x1b[39m");

    c.seed = original_seed + 1.0; // Reset the seed, but bump it a bit
}

fn calc_fg_color(bg: (u8, u8, u8)) -> (u8, u8, u8) {
    // Currently, it only computes the forground clolor based on some threshold
    // on grayscale value.
    // TODO: Add a better algorithm for computing forground color
    if conv_grayscale(bg) > 0xA0_u8 {
        (0u8, 0u8, 0u8)
    } else {
        (0xffu8, 0xffu8, 0xffu8)
    }
}

fn linear_to_srgb(intensity: f64) -> f64 {
    if intensity <= 0.003_130_8 {
        12.92 * intensity
    } else {
        1.055 * intensity.powf(1.0 / 2.4) - 0.055
    }
}

fn srgb_to_linear(intensity: f64) -> f64 {
    if intensity < 0.04045 {
        intensity / 12.92
    } else {
        ((intensity + 0.055) / 1.055).powf(2.4)
    }
}

fn conv_grayscale(color: (u8, u8, u8)) -> u8 {
    // See https://en.wikipedia.org/wiki/Grayscale#Converting_color_to_grayscale
    const SCALE: f64 = 256.0;

    // Changing SRGB to Linear for gamma correction
    let red = srgb_to_linear(f64::from(color.0) / SCALE);
    let green = srgb_to_linear(f64::from(color.1) / SCALE);
    let blue = srgb_to_linear(f64::from(color.2) / SCALE);

    // Converting to grayscale
    let gray_linear = red * 0.299 + green * 0.587 + blue * 0.114;

    // Gamma correction
    let gray_srgb = linear_to_srgb(gray_linear);

    (gray_srgb * SCALE) as u8
}

fn colored_print(character: char, c: &mut Control) {
    if c.background_mode {
        let bg = get_color_tuple(c);
        let fg = calc_fg_color(bg);
        print!(
            "\x1b[38;2;{};{};{};48;2;{};{};{}m{}",
            fg.0, fg.1, fg.2, bg.0, bg.1, bg.2, character
        );
    } else {
        let fg = get_color_tuple(c);
        print!("\x1b[38;2;{};{};{}m{}", fg.0, fg.1, fg.2, character);
    }
}

fn get_color_tuple(c: &Control) -> (u8, u8, u8) {
    let i = c.frequency * c.seed / c.spread;
    let red = i.sin() * 127.00 + 128.00;
    let green = (i + (std::f64::consts::PI * 2.00 / 3.00)).sin() * 127.00 + 128.00;
    let blue = (i + (std::f64::consts::PI * 4.00 / 3.00)).sin() * 127.00 + 128.00;

    (red as u8, green as u8, blue as u8)
}
<|endoftext|>