<commit_before>//! Implements a simple parser and extractor for a Unix Archive.
//!
//! There are two "common" formats: BSD and SysV
//!
//! This crate currently only implements the SysV version, which essentially postfixes all
//! names in the archive with a / as a sigil for the end of the name, and uses a special symbol
//! index for looking up symbols faster.

use scroll::{self, Pread};

use strtab;
use error::{Result, Error};

use std::usize;
use std::collections::HashMap;

pub const SIZEOF_MAGIC: usize = 8;
/// The magic number of a Unix Archive
pub const MAGIC: &'static [u8; SIZEOF_MAGIC] = b"!<arch>\x0A";

const SIZEOF_FILE_IDENTIFER: usize = 16;
const SIZEOF_FILE_SIZE: usize = 10;

#[repr(C)]
#[derive(Debug, Clone, PartialEq, Pread, Pwrite, SizeWith)]
/// A Unix Archive Header - meta data for the file/byte blob/whatever that follows exactly after.
/// All data is right-padded with spaces ASCII `0x20`. The Binary layout is as follows:
///
/// |Offset|Length|Name                       |Format     |
/// |:-----|:-----|:--------------------------|:----------|
/// |0     |16    |File identifier            |ASCII      |
/// |16    |12    |File modification timestamp|Decimal    |
/// |28    |6     |Owner ID                   |Decimal    |
/// |34    |6     |Group ID                   |Decimal    |
/// |40    |8     |File mode                  |Octal      |
/// |48    |10    |Filesize in bytes          |Decimal    |
/// |58    |2     |Ending characters          |`0x60 0x0A`|
///
/// Byte alignment is according to the following:
/// > Each archive file member begins on an even byte boundary; a newline is inserted between files
/// > if necessary. Nevertheless, the size given reflects the actual size of the file exclusive
/// > of padding.
pub struct MemberHeader {
    /// The identifier, or name for this file/whatever.
    pub identifier: [u8; 16],
    /// The timestamp for when this file was last modified. Base 10 number
    pub timestamp: [u8; 12],
    /// The file's owner's id. Base 10 string number
    pub owner_id: [u8; 6],
    /// The file's group id. Base 10 string number
    pub group_id: [u8; 6],
    /// The file's permissions mode. Base 8 number number
    pub mode: [u8; 8],
    /// The size of this file. Base 10 string number
    pub file_size: [u8; 10],
    /// The file header's terminator, always `0x60 0x0A`
    pub terminator: [u8; 2],
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Header<'a> {
    pub name: &'a str,
    pub size: usize,
}

pub const SIZEOF_HEADER: usize = SIZEOF_FILE_IDENTIFER + 12 + 6 + 6 + 8 + SIZEOF_FILE_SIZE + 2;

impl MemberHeader {
    pub fn name(&self) -> Result<&str> {
        Ok(self.identifier.pread_with::<&str>(0, ::scroll::ctx::StrCtx::Length(SIZEOF_FILE_IDENTIFER))?)
    }
    pub fn size(&self) -> Result<usize> {
        match usize::from_str_radix(self.file_size.pread_with::<&str>(0, ::scroll::ctx::StrCtx::Length(self.file_size.len()))?.trim_right(), 10) {
            Ok(file_size) => Ok(file_size),
            Err(err) => Err(Error::Malformed(format!("{:?} Bad file_size in header: {:?}", err, self)))
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
/// Represents a single entry in the archive
pub struct Member<'a> {
    /// The entry header
    pub header: Header<'a>,
    /// File offset from the start of the archive to where the file begins
    pub offset: u64,
}

impl<'a> Member<'a> {
    /// Tries to parse the header in `R`, as well as the offset in `R.
    /// **NOTE** the Seek will be pointing at the first byte of whatever the file is, skipping padding.
    /// This is because just like members in the archive, the data section is 2-byte aligned.
    pub fn parse(buffer: &'a [u8], offset: &mut usize) -> Result<Member<'a>> {
        let name = buffer.pread_with::<&str>(*offset, ::scroll::ctx::StrCtx::Length(SIZEOF_FILE_IDENTIFER))?;
        let archive_header = buffer.gread::<MemberHeader>(offset)?;
        let header = Header { name: name, size: archive_header.size()? };
        // skip newline padding if we're on an uneven byte boundary
        if *offset & 1 == 1 {
            *offset += 1;
        }
        Ok(Member {
            header: header,
            offset: *offset as u64,
        })
    }

    /// The size of the Member's content, in bytes. Does **not** include newline padding,
    /// nor the size of the file header.
    pub fn size(&self) -> usize {
        self.header.size
    }

    fn trim(name: &str) -> &str {
        name.trim_right_matches(' ').trim_right_matches('/')
    }

    /// The untrimmed raw member name, i.e., includes right-aligned space padding and `'/'` end-of-string
    /// identifier
    pub fn name(&self) -> &'a str {
        self.header.name
    }

}

#[derive(Debug, Default)]
/// The special index member signified by the name `'/'`.
/// The data element contains a list of symbol indexes and symbol names, giving their offsets
/// into the archive for a given name.
pub struct Index {
    /// Big Endian number of symbol_indexes and strings
    pub size: usize,
    /// Big Endian u32 index into the archive for this symbol (index in array is the index into the string table)
    pub symbol_indexes: Vec<u32>,
    /// Set of zero-terminated strings indexed by above. Number of strings = `self.size`
    pub strtab: Vec<String>,
}

/// SysV Archive Variant Symbol Lookup Table "Magic" Name
const INDEX_NAME: &'static str = "/               ";
/// SysV Archive Variant Extended Filename String Table Name
const NAME_INDEX_NAME: &'static str = "//              ";

impl Index {
    /// Parses the given byte buffer into an Index. NB: the buffer must be the start of the index
    pub fn parse(buffer: &[u8], size: usize) -> Result<Index> {
        let mut offset = &mut 0;
        let sizeof_table = buffer.gread_with::<u32>(offset, scroll::BE)? as usize;
        let mut indexes = Vec::with_capacity(sizeof_table);
        for _ in 0..sizeof_table {
            indexes.push(buffer.gread_with::<u32>(offset, scroll::BE)?);
        }
        let sizeof_strtab = size - ((sizeof_table * 4) + 4);
        let strtab = strtab::Strtab::parse(buffer, *offset, sizeof_strtab, 0x0)?;
        Ok (Index {
            size: sizeof_table,
            symbol_indexes: indexes,
            strtab: strtab.to_vec()?, // because i'm lazy
        })
    }
}

/// Member names greater than 16 bytes are indirectly referenced using a `/<idx` schema,
/// where `idx` is an offset into a newline delimited string table directly following the `//` member
/// of the archive.
#[derive(Debug, Default)]
struct NameIndex<'a> {
    strtab: strtab::Strtab<'a>
}

impl<'a> NameIndex<'a> {
    pub fn parse(buffer: &'a [u8], offset: &mut usize, size: usize) -> Result<NameIndex<'a>> {
        // This is a total hack, because strtab returns "" if idx == 0, need to change
        // but previous behavior might rely on this, as ELF strtab's have "" at 0th index...
        let hacked_size = size + 1;
        let strtab = strtab::Strtab::parse(buffer, *offset-1, hacked_size, '\n' as u8)?;
        // precious time was lost when refactoring because strtab::parse doesn't update the mutable seek...
        *offset += hacked_size - 2;
        Ok (NameIndex {
            strtab: strtab
        })
    }

    pub fn get(&self, name: &str) -> Result<&str> {
        let idx = name.trim_left_matches('/').trim_right();
        match usize::from_str_radix(idx, 10) {
            Ok(idx) => {
                let name = &self.strtab[idx+1];
                if name != "" {
                    Ok(name.trim_right_matches('/'))
                }  else {
                    return Err(Error::Malformed(format!("Could not find {:?} in index", name).into()));
                }
            },
            Err (_) => {
                return Err(Error::Malformed(format!("Bad name index {:?} in index", name).into()));
            }
        }
    }
}

// TODO: add pretty printer fmt::Display with number of members, and names of members, along with
// the values of the index symbols once implemented
#[derive(Debug)]
/// An in-memory representation of a parsed Unix Archive
pub struct Archive<'a> {
    // we can chuck this because the symbol index is a better representation, but we keep for
    // debugging
    index: Index,
    extended_names: NameIndex<'a>,
    // the array of members, which are indexed by the members hash and symbol index
    member_array: Vec<Member<'a>>,
    members: HashMap<String, usize>,
    // symbol -> member
    symbol_index: HashMap<String, usize>
}

impl<'a> Archive<'a> {
    pub fn parse(buffer: &'a [u8]) -> Result<Archive<'a>> {
        let mut magic = [0u8; SIZEOF_MAGIC];
        let mut offset = &mut 0usize;
        buffer.gread_inout(offset, &mut magic)?;
        if &magic != MAGIC {
            use scroll::Pread;
            return Err(Error::BadMagic(magic.pread(0)?).into());
        }
        let mut member_array = Vec::new();
        let size = buffer.as_ref().len() as u64;
        let mut pos = 0u64;
        let mut index = Index::default();
        let mut extended_names = NameIndex::default();
        loop {
            if pos >= size { break }
            // if the member is on an uneven byte boundary, we bump the buffer
            if pos & 1 == 1 {
                *offset += 1;
            }
            let member = Member::parse(buffer, offset)?;
            let name = member.name();
            let size = member.size();
            if name == INDEX_NAME {
                *offset = member.offset as usize;
                // get the member data (the index in this case)
                // FIXME, TODO: gread_slice needs to return
                //let data: &[u8] = buffer.gread_slice(offset, size)?;
                let data: &[u8] = &buffer[*offset..];
                *offset += size;
                // parse it
                index = Index::parse(&data, size)?;
                pos = *offset as u64;
            } else if name == NAME_INDEX_NAME {
                *offset = member.offset as usize;
                extended_names = NameIndex::parse(buffer, offset, size)?;
                pos = *offset as u64;
            } else {
                // we move the buffer past the file blob
                *offset += size;
                pos = *offset as u64;
                member_array.push(member);
            }
        }

        // this preprocesses the member names so they are searchable by their canonical versions
        let mut members = HashMap::new();
        for (i, member) in member_array.iter().enumerate() {
            let key = {
                let name = member.name();
                if name.starts_with("/") {
                    try!(extended_names.get(name))
                } else {
                    Member::trim(name)
            }}.to_owned();
            members.insert(key, i);
        }

        let mut symbol_index = HashMap::new();
        let mut last_symidx = 0u32;
        let mut last_member = 0usize;
        for (i, symidx) in index.symbol_indexes.iter().enumerate() {
            let name = index.strtab[i].to_owned();
            if *symidx == last_symidx {
                symbol_index.insert(name, last_member);
            } else {
                for (memidx, member) in member_array.iter().enumerate() {
                    if *symidx == (member.offset - SIZEOF_HEADER as u64) as u32 {
                        symbol_index.insert(name, memidx);
                        last_symidx = *symidx;
                        last_member = memidx;
                        break
                    }
                }
            }
        }

        let archive = Archive {
            index: index,
            member_array: member_array,
            extended_names: extended_names,
            members: members,
            symbol_index: symbol_index,
        };
        Ok(archive)
    }

    pub fn get (&self, member: &str) -> Option<&Member> {
        if let Some(idx) = self.members.get(member) {
            Some(&self.member_array[*idx])
        } else {
            None
        }
    }

    /// Returns a slice of the raw bytes for the given `member` in the scrollable `buffer`
    pub fn extract<'b>(&self, member: &str, buffer: &'b [u8]) -> Result<&'b [u8]> {
        if let Some(member) = self.get(member) {
            let bytes = buffer.pread_with(member.offset as usize, member.size())?;
            Ok(bytes)
        } else {
            Err(Error::Malformed(format!("Cannot extract member {:?}", member).into()))
        }
    }

    pub fn members(&self) -> Vec<&String> {
        self.members.keys().collect()
    }

    /// Returns the member's name which contains the given `symbol`, if it is in the archive
    pub fn member_of_symbol (&self, symbol: &str) -> Option<&str> {
        if let Some(idx) = self.symbol_index.get(symbol) {
            let name = (self.member_array[*idx]).name();
            if name.starts_with("/") {
                Some(self.extended_names.get(name).unwrap())
            } else {
                Some(Member::trim(name))
            }
        } else {
            None
        }
    }
}
<commit_msg>archive: add summarize method for interacting with internals of the archive. add some docs<commit_after>//! Implements a simple parser and extractor for a Unix Archive.
//!
//! There are two "common" formats: BSD and SysV
//!
//! This crate currently only implements the SysV version, which essentially postfixes all
//! names in the archive with a / as a sigil for the end of the name, and uses a special symbol
//! index for looking up symbols faster.

use scroll::{self, Pread};

use strtab;
use error::{Result, Error};

use std::usize;
use std::collections::HashMap;

pub const SIZEOF_MAGIC: usize = 8;
/// The magic number of a Unix Archive
pub const MAGIC: &'static [u8; SIZEOF_MAGIC] = b"!<arch>\x0A";

const SIZEOF_FILE_IDENTIFER: usize = 16;
const SIZEOF_FILE_SIZE: usize = 10;

#[repr(C)]
#[derive(Debug, Clone, PartialEq, Pread, Pwrite, SizeWith)]
/// A Unix Archive Header - meta data for the file/byte blob/whatever that follows exactly after.
/// All data is right-padded with spaces ASCII `0x20`. The Binary layout is as follows:
///
/// |Offset|Length|Name                       |Format     |
/// |:-----|:-----|:--------------------------|:----------|
/// |0     |16    |File identifier            |ASCII      |
/// |16    |12    |File modification timestamp|Decimal    |
/// |28    |6     |Owner ID                   |Decimal    |
/// |34    |6     |Group ID                   |Decimal    |
/// |40    |8     |File mode                  |Octal      |
/// |48    |10    |Filesize in bytes          |Decimal    |
/// |58    |2     |Ending characters          |`0x60 0x0A`|
///
/// Byte alignment is according to the following:
/// > Each archive file member begins on an even byte boundary; a newline is inserted between files
/// > if necessary. Nevertheless, the size given reflects the actual size of the file exclusive
/// > of padding.
pub struct MemberHeader {
    /// The identifier, or name for this file/whatever.
    pub identifier: [u8; 16],
    /// The timestamp for when this file was last modified. Base 10 number
    pub timestamp: [u8; 12],
    /// The file's owner's id. Base 10 string number
    pub owner_id: [u8; 6],
    /// The file's group id. Base 10 string number
    pub group_id: [u8; 6],
    /// The file's permissions mode. Base 8 number number
    pub mode: [u8; 8],
    /// The size of this file. Base 10 string number
    pub file_size: [u8; 10],
    /// The file header's terminator, always `0x60 0x0A`
    pub terminator: [u8; 2],
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Header<'a> {
    pub name: &'a str,
    pub size: usize,
}

pub const SIZEOF_HEADER: usize = SIZEOF_FILE_IDENTIFER + 12 + 6 + 6 + 8 + SIZEOF_FILE_SIZE + 2;

impl MemberHeader {
    pub fn name(&self) -> Result<&str> {
        Ok(self.identifier.pread_with::<&str>(0, ::scroll::ctx::StrCtx::Length(SIZEOF_FILE_IDENTIFER))?)
    }
    pub fn size(&self) -> Result<usize> {
        match usize::from_str_radix(self.file_size.pread_with::<&str>(0, ::scroll::ctx::StrCtx::Length(self.file_size.len()))?.trim_right(), 10) {
            Ok(file_size) => Ok(file_size),
            Err(err) => Err(Error::Malformed(format!("{:?} Bad file_size in header: {:?}", err, self)))
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
/// Represents a single entry in the archive
pub struct Member<'a> {
    /// The entry header
    pub header: Header<'a>,
    /// File offset from the start of the archive to where the file begins
    pub offset: u64,
}

impl<'a> Member<'a> {
    /// Tries to parse the header in `R`, as well as the offset in `R.
    /// **NOTE** the Seek will be pointing at the first byte of whatever the file is, skipping padding.
    /// This is because just like members in the archive, the data section is 2-byte aligned.
    pub fn parse(buffer: &'a [u8], offset: &mut usize) -> Result<Member<'a>> {
        let name = buffer.pread_with::<&str>(*offset, ::scroll::ctx::StrCtx::Length(SIZEOF_FILE_IDENTIFER))?;
        let archive_header = buffer.gread::<MemberHeader>(offset)?;
        let header = Header { name: name, size: archive_header.size()? };
        // skip newline padding if we're on an uneven byte boundary
        if *offset & 1 == 1 {
            *offset += 1;
        }
        Ok(Member {
            header: header,
            offset: *offset as u64,
        })
    }

    /// The size of the Member's content, in bytes. Does **not** include newline padding,
    /// nor the size of the file header.
    pub fn size(&self) -> usize {
        self.header.size
    }

    fn trim(name: &str) -> &str {
        name.trim_right_matches(' ').trim_right_matches('/')
    }

    /// The untrimmed raw member name, i.e., includes right-aligned space padding and `'/'` end-of-string
    /// identifier
    pub fn name(&self) -> &'a str {
        self.header.name
    }

}

#[derive(Debug, Default)]
/// The special index member signified by the name `'/'`.
/// The data element contains a list of symbol indexes and symbol names, giving their offsets
/// into the archive for a given name.
pub struct Index {
    /// Big Endian number of symbol_indexes and strings
    pub size: usize,
    /// Big Endian u32 index into the archive for this symbol (index in array is the index into the string table)
    pub symbol_indexes: Vec<u32>,
    /// Set of zero-terminated strings indexed by above. Number of strings = `self.size`
    pub strtab: Vec<String>,
}

/// SysV Archive Variant Symbol Lookup Table "Magic" Name
const INDEX_NAME: &'static str = "/               ";
/// SysV Archive Variant Extended Filename String Table Name
const NAME_INDEX_NAME: &'static str = "//              ";

impl Index {
    /// Parses the given byte buffer into an Index. NB: the buffer must be the start of the index
    pub fn parse(buffer: &[u8], size: usize) -> Result<Index> {
        let mut offset = &mut 0;
        let sizeof_table = buffer.gread_with::<u32>(offset, scroll::BE)? as usize;
        let mut indexes = Vec::with_capacity(sizeof_table);
        for _ in 0..sizeof_table {
            indexes.push(buffer.gread_with::<u32>(offset, scroll::BE)?);
        }
        let sizeof_strtab = size - ((sizeof_table * 4) + 4);
        let strtab = strtab::Strtab::parse(buffer, *offset, sizeof_strtab, 0x0)?;
        Ok (Index {
            size: sizeof_table,
            symbol_indexes: indexes,
            strtab: strtab.to_vec()?, // because i'm lazy
        })
    }
}

/// Member names greater than 16 bytes are indirectly referenced using a `/<idx` schema,
/// where `idx` is an offset into a newline delimited string table directly following the `//` member
/// of the archive.
#[derive(Debug, Default)]
struct NameIndex<'a> {
    strtab: strtab::Strtab<'a>
}

impl<'a> NameIndex<'a> {
    pub fn parse(buffer: &'a [u8], offset: &mut usize, size: usize) -> Result<NameIndex<'a>> {
        // This is a total hack, because strtab returns "" if idx == 0, need to change
        // but previous behavior might rely on this, as ELF strtab's have "" at 0th index...
        let hacked_size = size + 1;
        let strtab = strtab::Strtab::parse(buffer, *offset-1, hacked_size, '\n' as u8)?;
        // precious time was lost when refactoring because strtab::parse doesn't update the mutable seek...
        *offset += hacked_size - 2;
        Ok (NameIndex {
            strtab: strtab
        })
    }

    pub fn get(&self, name: &str) -> Result<&str> {
        let idx = name.trim_left_matches('/').trim_right();
        match usize::from_str_radix(idx, 10) {
            Ok(idx) => {
                let name = &self.strtab[idx+1];
                if name != "" {
                    Ok(name.trim_right_matches('/'))
                }  else {
                    return Err(Error::Malformed(format!("Could not find {:?} in index", name).into()));
                }
            },
            Err (_) => {
                return Err(Error::Malformed(format!("Bad name index {:?} in index", name).into()));
            }
        }
    }
}

// TODO: add pretty printer fmt::Display with number of members, and names of members, along with
// the values of the index symbols once implemented
#[derive(Debug)]
/// An in-memory representation of a parsed Unix Archive
pub struct Archive<'a> {
    // we can chuck this because the symbol index is a better representation, but we keep for
    // debugging
    index: Index,
    extended_names: NameIndex<'a>,
    // the array of members, which are indexed by the members hash and symbol index
    member_array: Vec<Member<'a>>,
    members: HashMap<String, usize>,
    // symbol -> member
    symbol_index: HashMap<String, usize>
}

impl<'a> Archive<'a> {
    pub fn parse(buffer: &'a [u8]) -> Result<Archive<'a>> {
        let mut magic = [0u8; SIZEOF_MAGIC];
        let mut offset = &mut 0usize;
        buffer.gread_inout(offset, &mut magic)?;
        if &magic != MAGIC {
            use scroll::Pread;
            return Err(Error::BadMagic(magic.pread(0)?).into());
        }
        let mut member_array = Vec::new();
        let size = buffer.as_ref().len() as u64;
        let mut pos = 0u64;
        let mut index = Index::default();
        let mut extended_names = NameIndex::default();
        loop {
            if pos >= size { break }
            // if the member is on an uneven byte boundary, we bump the buffer
            if pos & 1 == 1 {
                *offset += 1;
            }
            let member = Member::parse(buffer, offset)?;
            let name = member.name();
            let size = member.size();
            if name == INDEX_NAME {
                *offset = member.offset as usize;
                // get the member data (the index in this case)
                // FIXME, TODO: gread_slice needs to return
                //let data: &[u8] = buffer.gread_slice(offset, size)?;
                let data: &[u8] = &buffer[*offset..];
                *offset += size;
                // parse it
                index = Index::parse(&data, size)?;
                pos = *offset as u64;
            } else if name == NAME_INDEX_NAME {
                *offset = member.offset as usize;
                extended_names = NameIndex::parse(buffer, offset, size)?;
                pos = *offset as u64;
            } else {
                // we move the buffer past the file blob
                *offset += size;
                pos = *offset as u64;
                member_array.push(member);
            }
        }

        // this preprocesses the member names so they are searchable by their canonical versions
        let mut members = HashMap::new();
        for (i, member) in member_array.iter().enumerate() {
            let key = {
                let name = member.name();
                if name.starts_with("/") {
                    try!(extended_names.get(name))
                } else {
                    Member::trim(name)
            }}.to_owned();
            members.insert(key, i);
        }

        let mut symbol_index = HashMap::new();
        let mut last_symidx = 0u32;
        let mut last_member = 0usize;
        for (i, symidx) in index.symbol_indexes.iter().enumerate() {
            let name = index.strtab[i].to_owned();
            if *symidx == last_symidx {
                symbol_index.insert(name, last_member);
            } else {
                for (memidx, member) in member_array.iter().enumerate() {
                    if *symidx == (member.offset - SIZEOF_HEADER as u64) as u32 {
                        symbol_index.insert(name, memidx);
                        last_symidx = *symidx;
                        last_member = memidx;
                        break
                    }
                }
            }
        }

        let archive = Archive {
            index: index,
            member_array: member_array,
            extended_names: extended_names,
            members: members,
            symbol_index: symbol_index,
        };
        Ok(archive)
    }

    /// Get the member named `member` in this archive, if any
    pub fn get (&self, member: &str) -> Option<&Member> {
        if let Some(idx) = self.members.get(member) {
            Some(&self.member_array[*idx])
        } else {
            None
        }
    }

    /// Returns a slice of the raw bytes for the given `member` in the scrollable `buffer`
    pub fn extract<'b>(&self, member: &str, buffer: &'b [u8]) -> Result<&'b [u8]> {
        if let Some(member) = self.get(member) {
            let bytes = buffer.pread_with(member.offset as usize, member.size())?;
            Ok(bytes)
        } else {
            Err(Error::Malformed(format!("Cannot extract member {:?}", member).into()))
        }
    }

    /// Gets a summary of this archive, returning a list of membername, the member, and the list of symbols the member contains
    pub fn summarize(&self) -> Vec<(&String, &Member, Vec<&String>)> {
        let symbols = self.symbol_index.iter().collect::<Vec<_>>();
        let mut res = Vec::new();
        for (member_name, idx) in self.members.iter() {
            let member = &self.member_array[*idx];
            let symbols = symbols.iter().filter_map(|&(ref symbol, ref mem_idx)| if *mem_idx == idx { Some(*symbol) } else { None }).collect::<Vec<_>>();
            res.push((member_name, member, symbols));
        }
        res
    }

    /// Get the list of member names in this archive
    pub fn members(&self) -> Vec<&String> {
        self.members.keys().collect()
    }

    /// Returns the member's name which contains the given `symbol`, if it is in the archive
    pub fn member_of_symbol (&self, symbol: &str) -> Option<&str> {
        if let Some(idx) = self.symbol_index.get(symbol) {
            let name = (self.member_array[*idx]).name();
            if name.starts_with("/") {
                Some(self.extended_names.get(name).unwrap())
            } else {
                Some(Member::trim(name))
            }
        } else {
            None
        }
    }
}
<|endoftext|>


<commit_before>use crate::wheel;
use crate::streaming::{self, StreamingSieve};

#[cfg(target_pointer_width = "32")]
const SQRT: usize = 1 << 16;
#[cfg(target_pointer_width = "64")]
const SQRT: usize = 1 << 32;

enum Early {
    Two,
    Three,
    Five,
    Done,
}

/// An iterator over all primes.
///
/// This will yield primes indefinitely (bits in `usize`
/// permitting). If there is an known upper bound, sieving first with
/// `Sieve` and using its `primes_from` method may be more efficient,
/// especially if the bound is small.
///
/// This requires *O(sqrt(p))* memory to yield prime `p`, where `X` is
/// the maximum value of `usize`.
///
/// # Examples
///
/// ```rust
/// let count = primal::Primes::all().take_while(|p| *p < 1_000_000).count();
/// println!("{}", count);
/// ```
pub struct Primes {
    early: Early,
    base: usize,
    ones: primal_bit::IntoOnes,
    streaming: StreamingSieve,
    sieving_primes: Option<Box<Primes>>,
    left_over: Option<usize>,
}

impl Primes {
    /// The sequence `2, 3, 5, 7, 11, ...`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// // print the first 20 primes
    /// for p in primal::Primes::all().take(20) {
    ///     println!("{}", p);
    /// }
    /// ```
    pub fn all() -> Primes {
        Primes::sqrt(SQRT)
    }

    fn sqrt(sqrt: usize) -> Primes {
        let (sieving, limit) = if sqrt < streaming::isqrt(streaming::SEG_LEN) {
            (None, sqrt * sqrt)
        } else {
            (Some(Box::new(Primes::sqrt(streaming::isqrt(sqrt)))),
             streaming::SEG_LEN)
        };
        let mut streaming = StreamingSieve::new(limit);

        let ones = {
            let (_, bits) = streaming.next().unwrap();
            bits.clone().into_ones()
        };

        // we manually add the primes
        streaming.small = None;
        // go to the end.
        streaming.limit = !0;

        Primes {
            early: Early::Two,
            base: 0,
            ones,
            streaming,
            sieving_primes: sieving,
            left_over: None,
        }
    }

    #[inline]
    fn from_bit_index(&self, i: usize) -> Option<usize> {
        let w = wheel::from_bit_index(i);
        self.base.checked_add(w)
    }

    fn advance_ones(&mut self) -> bool {
        let low = self.streaming.low;
        let high = low.saturating_add(streaming::SEG_LEN);

        for q in self.left_over.into_iter().chain(self.sieving_primes.as_mut().unwrap()) {
            if q.saturating_mul(q) >= high {
                self.left_over = Some(q);
                break
            }
            if q > streaming::isqrt(streaming::SEG_LEN) {
                self.streaming.add_sieving_prime(q, low);
            }
        }

        match self.streaming.next() {
            Some((_, bits)) => {
                self.base = low;
                self.ones = bits.clone().into_ones();
                true
            },
            None => false,
        }
    }
}

// See also `Iterator for SievePrimes` with nearly identical code.
impl Iterator for Primes {
    type Item = usize;

    #[inline]
    fn next(&mut self) -> Option<usize> {
        match self.early {
            Early::Done => {}
            Early::Two => {
                self.early = Early::Three;
                return Some(2)
            }
            Early::Three => {
                self.early = Early::Five;
                return Some(3)
            }
            Early::Five => {
                self.early = Early::Done;
                return Some(5)
            }
        }

        loop {
            if let Some(i) = self.ones.next() {
                return self.from_bit_index(i);
            }
            if !self.advance_ones() {
                return None;
            }
        }
    }

    fn fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc
    where
        F: FnMut(Acc, Self::Item) -> Acc
    {
        match self.early {
            Early::Done => {}
            Early::Two => {
                acc = f(acc, 2);
                acc = f(acc, 3);
                acc = f(acc, 5);
            }
            Early::Three => {
                acc = f(acc, 3);
                acc = f(acc, 5);
            }
            Early::Five => {
                acc = f(acc, 5);
            }
        }
        loop {
            while let Some(i) = self.ones.next() {
                match self.from_bit_index(i) {
                    Some(p) => acc = f(acc, p),
                    None => return acc,
                }
            }
            if !self.advance_ones() {
                return acc;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::Sieve;
    use super::Primes;

    fn check_equality(limit: usize) {
        let sieve = Sieve::new(limit);

        let real = sieve.primes_from(0).take_while(|x| *x < limit);
        let computed = Primes::all().take_while(|x| *x < limit);

        let mut i = 0;
        for (r, c) in real.zip(computed) {
            assert_eq!(c, r);
            i += 1;
        }
        assert_eq!(sieve.prime_pi(limit), i);
    }

    #[test]
    fn equality() {
        check_equality(20_000_000);
    }

    #[test]
    fn equality_huge() {
        let limit = if cfg!(feature = "slow_tests") {
            // This takes a minute or so in debug mode, but it does work!
            ::std::u32::MAX as usize
        } else {
            100_000_000
        };
        check_equality(limit);
    }

    #[test]
    #[should_panic = "123456791"]
    fn fold() {
        // There's no termination until we exceed `usize::MAX`, which
        // will take too long, but we can cut it short by unwinding.
        Primes::all().fold((), |(), p| {
            if p > 123456789 {
                panic!(format!("{}", p));
            }
        })
    }
}
<commit_msg>Use panic formatting in a test to avoid a deprecation<commit_after>use crate::wheel;
use crate::streaming::{self, StreamingSieve};

#[cfg(target_pointer_width = "32")]
const SQRT: usize = 1 << 16;
#[cfg(target_pointer_width = "64")]
const SQRT: usize = 1 << 32;

enum Early {
    Two,
    Three,
    Five,
    Done,
}

/// An iterator over all primes.
///
/// This will yield primes indefinitely (bits in `usize`
/// permitting). If there is an known upper bound, sieving first with
/// `Sieve` and using its `primes_from` method may be more efficient,
/// especially if the bound is small.
///
/// This requires *O(sqrt(p))* memory to yield prime `p`, where `X` is
/// the maximum value of `usize`.
///
/// # Examples
///
/// ```rust
/// let count = primal::Primes::all().take_while(|p| *p < 1_000_000).count();
/// println!("{}", count);
/// ```
pub struct Primes {
    early: Early,
    base: usize,
    ones: primal_bit::IntoOnes,
    streaming: StreamingSieve,
    sieving_primes: Option<Box<Primes>>,
    left_over: Option<usize>,
}

impl Primes {
    /// The sequence `2, 3, 5, 7, 11, ...`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// // print the first 20 primes
    /// for p in primal::Primes::all().take(20) {
    ///     println!("{}", p);
    /// }
    /// ```
    pub fn all() -> Primes {
        Primes::sqrt(SQRT)
    }

    fn sqrt(sqrt: usize) -> Primes {
        let (sieving, limit) = if sqrt < streaming::isqrt(streaming::SEG_LEN) {
            (None, sqrt * sqrt)
        } else {
            (Some(Box::new(Primes::sqrt(streaming::isqrt(sqrt)))),
             streaming::SEG_LEN)
        };
        let mut streaming = StreamingSieve::new(limit);

        let ones = {
            let (_, bits) = streaming.next().unwrap();
            bits.clone().into_ones()
        };

        // we manually add the primes
        streaming.small = None;
        // go to the end.
        streaming.limit = !0;

        Primes {
            early: Early::Two,
            base: 0,
            ones,
            streaming,
            sieving_primes: sieving,
            left_over: None,
        }
    }

    #[inline]
    fn from_bit_index(&self, i: usize) -> Option<usize> {
        let w = wheel::from_bit_index(i);
        self.base.checked_add(w)
    }

    fn advance_ones(&mut self) -> bool {
        let low = self.streaming.low;
        let high = low.saturating_add(streaming::SEG_LEN);

        for q in self.left_over.into_iter().chain(self.sieving_primes.as_mut().unwrap()) {
            if q.saturating_mul(q) >= high {
                self.left_over = Some(q);
                break
            }
            if q > streaming::isqrt(streaming::SEG_LEN) {
                self.streaming.add_sieving_prime(q, low);
            }
        }

        match self.streaming.next() {
            Some((_, bits)) => {
                self.base = low;
                self.ones = bits.clone().into_ones();
                true
            },
            None => false,
        }
    }
}

// See also `Iterator for SievePrimes` with nearly identical code.
impl Iterator for Primes {
    type Item = usize;

    #[inline]
    fn next(&mut self) -> Option<usize> {
        match self.early {
            Early::Done => {}
            Early::Two => {
                self.early = Early::Three;
                return Some(2)
            }
            Early::Three => {
                self.early = Early::Five;
                return Some(3)
            }
            Early::Five => {
                self.early = Early::Done;
                return Some(5)
            }
        }

        loop {
            if let Some(i) = self.ones.next() {
                return self.from_bit_index(i);
            }
            if !self.advance_ones() {
                return None;
            }
        }
    }

    fn fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc
    where
        F: FnMut(Acc, Self::Item) -> Acc
    {
        match self.early {
            Early::Done => {}
            Early::Two => {
                acc = f(acc, 2);
                acc = f(acc, 3);
                acc = f(acc, 5);
            }
            Early::Three => {
                acc = f(acc, 3);
                acc = f(acc, 5);
            }
            Early::Five => {
                acc = f(acc, 5);
            }
        }
        loop {
            while let Some(i) = self.ones.next() {
                match self.from_bit_index(i) {
                    Some(p) => acc = f(acc, p),
                    None => return acc,
                }
            }
            if !self.advance_ones() {
                return acc;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::Sieve;
    use super::Primes;

    fn check_equality(limit: usize) {
        let sieve = Sieve::new(limit);

        let real = sieve.primes_from(0).take_while(|x| *x < limit);
        let computed = Primes::all().take_while(|x| *x < limit);

        let mut i = 0;
        for (r, c) in real.zip(computed) {
            assert_eq!(c, r);
            i += 1;
        }
        assert_eq!(sieve.prime_pi(limit), i);
    }

    #[test]
    fn equality() {
        check_equality(20_000_000);
    }

    #[test]
    fn equality_huge() {
        let limit = if cfg!(feature = "slow_tests") {
            // This takes a minute or so in debug mode, but it does work!
            ::std::u32::MAX as usize
        } else {
            100_000_000
        };
        check_equality(limit);
    }

    #[test]
    #[should_panic = "123456791"]
    fn fold() {
        // There's no termination until we exceed `usize::MAX`, which
        // will take too long, but we can cut it short by unwinding.
        Primes::all().fold((), |(), p| {
            if p > 123456789 {
                panic!("{}", p);
            }
        })
    }
}
<|endoftext|>


<commit_before>//Implements http://rosettacode.org/wiki/Sorting_algorithms/Bubble_sort
fn bubble_sort(v: &mut[int]) {
    let mut done = v.len()<1;
    while !done {
        done = true;
        for index in range(0, v.len()-1) {
            if v[index] > v[index+1] {
                done = false;
                v.swap(index, index+1);
            }
        }
    }
}

fn check_sort(v: &[int]) {
    if v.len() > 1 {
        for i in range(0, v.len()-1) {
            assert!(v[i] <= v[i+1]);
        }
    }
}

#[test]
fn test_rosetta_vector() {
    let mut numbers = [4, 65, 2, -31, 0, 99, 2, 83, 782, 1];
    bubble_sort(numbers);
    check_sort(numbers);
}

#[test]
fn test_empty_vector() {
    let mut numbers = [];
    bubble_sort(numbers);
    check_sort(numbers);
}

#[test]
fn test_one_element_vector() {
    let mut numbers = [0];
    bubble_sort(numbers);
    check_sort(numbers);
}

#[test]
fn test_repeat_vector() {
    let mut numbers = [1, 1, 1, 1, 1];
    bubble_sort(numbers);
    check_sort(numbers);
}

#[test]
fn test_worst_case_vector() {
    let mut numbers = [20, 10, 0, -1, -5];
    bubble_sort(numbers);
    check_sort(numbers);
}

#[test]
fn test_already_sorted_vector() {
    let mut numbers = [-1, 0, 3, 6, 99];
    bubble_sort(numbers);
    check_sort(numbers);
}
<commit_msg>optim : doesn't have to check the whole range each time<commit_after>//Implements http://rosettacode.org/wiki/Sorting_algorithms/Bubble_sort
fn bubble_sort(v: &mut[int]) {
    let mut done = v.len()<1;
    let mut length = v.len();
    while !done {
        done = true;
        length = length-1;
        for index in range(0, length) {
            if v[index] > v[index+1] {
                done = false;
                v.swap(index, index+1);
            }
        }
    }
}

fn check_sort(v: &[int]) {
    if v.len() > 1 {
        for i in range(0, v.len()-1) {
            assert!(v[i] <= v[i+1]);
        }
    }
}

#[test]
fn test_rosetta_vector() {
    let mut numbers = [4, 65, 2, -31, 0, 99, 2, 83, 782, 1];
    bubble_sort(numbers);
    check_sort(numbers);
}

#[test]
fn test_empty_vector() {
    let mut numbers = [];
    bubble_sort(numbers);
    check_sort(numbers);
}

#[test]
fn test_one_element_vector() {
    let mut numbers = [0];
    bubble_sort(numbers);
    check_sort(numbers);
}

#[test]
fn test_repeat_vector() {
    let mut numbers = [1, 1, 1, 1, 1];
    bubble_sort(numbers);
    check_sort(numbers);
}

#[test]
fn test_worst_case_vector() {
    let mut numbers = [20, 10, 0, -1, -5];
    bubble_sort(numbers);
    check_sort(numbers);
}

#[test]
fn test_already_sorted_vector() {
    let mut numbers = [-1, 0, 3, 6, 99];
    bubble_sort(numbers);
    check_sort(numbers);
}
<|endoftext|>


<commit_before>/*
 This Source Code Form is subject to the terms of the Mozilla Public
 License, v. 2.0. If a copy of the MPL was not distributed with this
 file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/

use rusqlite::{Connection, Row};

use crate::content::Content;
use crate::fromdb::FromDb;
use crate::lrobject::LrId;

/// A collection as defined in Lightroom
pub struct Collection {
    /// Local id of the collection
    id: LrId,
    /// Name of the collection (displayed in the UI)
    pub name: String,
    /// Parent of the `Collection`
    pub parent: LrId,
    /// is system only (seems to be the Quick Pick collection)
    pub system_only: bool,
    /// content definition of the collection
    pub content: Option<Content>,
}

impl FromDb for Collection {
    fn read_from(row: &Row) -> rusqlite::Result<Self> {
        Ok(Collection {
            id: row.get(0)?,
            name: row.get(2)?,
            parent: row.get(3).unwrap_or(0),
            system_only: matches!(row.get::<usize, f64>(4)? as i64, 0),
            content: None,
        })
    }
    fn read_db_tables() -> &'static str {
        "AgLibraryCollection"
    }
    fn read_db_columns() -> &'static str {
        "id_local,genealogy,name,parent,systemOnly"
    }
}

impl Collection {
    /// Return the local_id of the collection.
    pub fn id(&self) -> LrId {
        self.id
    }

    /// Read the `content` for this collection from the database.
    pub fn read_content(&self, conn: &Connection) -> Content {
        Content::from_db(conn, "AgLibraryCollectionContent", "collection", self.id)
    }
}
<commit_msg>Fix the system_only mapping for collection<commit_after>/*
 This Source Code Form is subject to the terms of the Mozilla Public
 License, v. 2.0. If a copy of the MPL was not distributed with this
 file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/

use rusqlite::{Connection, Row};

use crate::content::Content;
use crate::fromdb::FromDb;
use crate::lrobject::LrId;

/// A collection as defined in Lightroom
pub struct Collection {
    /// Local id of the collection
    id: LrId,
    /// Name of the collection (displayed in the UI)
    pub name: String,
    /// Parent of the `Collection`
    pub parent: LrId,
    /// is system only (seems to be the Quick Pick collection)
    pub system_only: bool,
    /// content definition of the collection
    pub content: Option<Content>,
}

impl FromDb for Collection {
    fn read_from(row: &Row) -> rusqlite::Result<Self> {
        Ok(Collection {
            id: row.get(0)?,
            name: row.get(2)?,
            parent: row.get(3).unwrap_or(0),
            system_only: !matches!(row.get::<usize, f64>(4)? as i64, 0),
            content: None,
        })
    }
    fn read_db_tables() -> &'static str {
        "AgLibraryCollection"
    }
    fn read_db_columns() -> &'static str {
        "id_local,genealogy,name,parent,systemOnly"
    }
}

impl Collection {
    /// Return the local_id of the collection.
    pub fn id(&self) -> LrId {
        self.id
    }

    /// Read the `content` for this collection from the database.
    pub fn read_content(&self, conn: &Connection) -> Content {
        Content::from_db(conn, "AgLibraryCollectionContent", "collection", self.id)
    }
}
<|endoftext|>


<commit_before>//! Full implementation of HMAC in accordance with [RFC 2104](http://tools.ietf.org/html/rfc2104)
//!
//!# Examples
//!
//!```
//!use rust_oauth::crypto::{sha1, hmac};
//!let key = "key".as_bytes();
//!let msg = "The quick brown fox jumps over the lazy dog".as_bytes();
//!let hmac = hmac::hmac(|x|{sha1::sha1(x)}, msg, key);
//!
//!
//!

use std::ops::{BitXor};

const BLOCKSIZE : uint = 64;
const IPAD : [u8, ..BLOCKSIZE] = [0x36u8, ..BLOCKSIZE];
const OPAD : [u8, ..BLOCKSIZE] = [0x5cu8, ..BLOCKSIZE];

struct U8BLOCK([u8, ..BLOCKSIZE]);

/// Generate the hmac using the hashing function, message, and key provided.
pub fn hmac(hash : |&[u8]|->[u8, ..20], msg : &[u8], key : &[u8]) -> [u8, ..20]{
    let mut key_new : [u8, ..BLOCKSIZE] = [0u8, ..BLOCKSIZE];

    if key.len() > BLOCKSIZE{
        let hash = hash(key);
        for x in range(0, hash.len()){
            key_new[x] = hash[x];
        }
    } else {
        for x in range(0, key.len()){
            key_new[x] = key[x];
        }
    }
    let mut v = Vec::new();
    v.push_all(&(U8BLOCK(key_new) ^ IPAD));
    v.push_all(msg);
    let temp : [u8, ..20] = hash(v.as_slice());
    println!("");
    for x in v.iter(){
        print!("{}", *x as char);
    }; println!("");
    for x in temp.iter(){
        print!("{0:X} ", *x);
    }; println!("");
    v = Vec::new();
    v.push_all(&(U8BLOCK(key_new) ^ OPAD));
    v.push_all(&temp);
    for x in v.iter(){
        print!("{}", *x as char);
        }; println!("");
    for x in hash(v.as_slice()).iter(){
        print!("{} ", *x);
        }; println!("");
    hash(v.as_slice())

}

impl BitXor<[u8, ..BLOCKSIZE], [u8, ..BLOCKSIZE]> for U8BLOCK{
    fn bitxor(self, _rhs : [u8, ..BLOCKSIZE]) -> [u8, ..BLOCKSIZE]{
        let U8BLOCK(t) = self;
        let mut x = t;
        for i in range(0, BLOCKSIZE){
            x[i] = t[i] ^ _rhs[i];
        };
        x
    }
}

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

    #[test]
    fn hamc_test1(){
        let key = "key".as_bytes();
        let msg = "The quick brown fox jumps over the lazy dog".as_bytes();
        let h = hmac(|x|{sha1(x)}, msg, key);
        println!("");
        assert!(h ==
            [0xdeu8, 0x7cu8, 0x9bu8, 0x85u8, 0xb8u8,
             0xb7u8, 0x8au8, 0xa6u8, 0xbcu8, 0x8au8,
             0x7au8, 0x36u8, 0xf7u8, 0x0au8, 0x90u8,
             0x70u8, 0x1cu8, 0x9du8, 0xb4u8, 0xd9u8]);
    }
}
<commit_msg>Fixed hmac function signature<commit_after>//! Full implementation of HMAC in accordance with [RFC 2104](http://tools.ietf.org/html/rfc2104)
//!
//!# Examples
//!
//!```
//!use rust_oauth::crypto::hmac;
//!let key = "key".as_bytes();
//!let msg = "The quick brown fox jumps over the lazy dog".as_bytes();
//!let hmac = hmac::hmac_sha1(msg, key);
//!
//!
//!

use std::ops::{BitXor};
use super::sha1;

// HMAC constants defined in RFC
const BLOCKSIZE : uint = 64;
const IPAD : [u8, ..BLOCKSIZE] = [0x36u8, ..BLOCKSIZE];
const OPAD : [u8, ..BLOCKSIZE] = [0x5cu8, ..BLOCKSIZE];

struct U8BLOCK([u8, ..BLOCKSIZE]);

impl BitXor<[u8, ..BLOCKSIZE], [u8, ..BLOCKSIZE]> for U8BLOCK{
    fn bitxor(self, _rhs : [u8, ..BLOCKSIZE]) -> [u8, ..BLOCKSIZE]{
        let U8BLOCK(t) = self;
        let mut x = t;
        for i in range(0, BLOCKSIZE){
            x[i] = t[i] ^ _rhs[i];
        };
        x
    }
}

/// Generate the hmac using the hashing function, message, and key provided.
pub fn hmac_sha1(msg : &[u8], key : &[u8]) -> [u8, ..20]{
    let mut key_new : [u8, ..BLOCKSIZE] = [0u8, ..BLOCKSIZE];

    if key.len() > BLOCKSIZE{
        let hash = sha1::sha1(key);
        for x in range(0, hash.len()){
            key_new[x] = hash[x];
        }
    } else {
        for x in range(0, key.len()){
            key_new[x] = key[x];
        }
    }
    let mut v = Vec::new();
    v.push_all(&(U8BLOCK(key_new) ^ IPAD));
    v.push_all(msg);
    let temp : [u8, ..20] = sha1::sha1(v.as_slice());
    v = Vec::new();
    v.push_all(&(U8BLOCK(key_new) ^ OPAD));
    v.push_all(&temp);
    sha1::sha1(v.as_slice())

}

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

    #[test]
    fn hamc_test1(){
        let key = "key".as_bytes();
        let msg = "The quick brown fox jumps over the lazy dog".as_bytes();
        let h = hmac_sha1(msg, key);
        println!("");
        assert!(h ==
            [0xdeu8, 0x7cu8, 0x9bu8, 0x85u8, 0xb8u8,
             0xb7u8, 0x8au8, 0xa6u8, 0xbcu8, 0x8au8,
             0x7au8, 0x36u8, 0xf7u8, 0x0au8, 0x90u8,
             0x70u8, 0x1cu8, 0x9du8, 0xb4u8, 0xd9u8]);
    }
}
<|endoftext|>


<commit_before>//! Learning benchmark domains module.
use crate::geometry::Space;

macro_rules! impl_into {
    (Transition < S, $type:ty > => Transition < S,() >) => {
        impl<S> Into<Transition<S, ()>> for Transition<S, $type> {
            fn into(self) -> Transition<S, ()> {
                self.drop_action()
            }
        }
    };
}

/// Container class for data associated with a domain observation.
#[derive(Clone)]
pub enum Observation<S> {
    /// Fully observed state of the environment.
    Full(S),

    /// Partially observed state of the environment.
    Partial(S),

    /// Terminal state of the environment.
    Terminal(S),
}

impl<S> Observation<S> {
    /// Helper function returning a reference to the state values for the given
    /// observation.
    pub fn state(&self) -> &S {
        use self::Observation::*;

        match self {
            Full(ref state) | Partial(ref state) | Terminal(ref state) => state,
        }
    }

    /// Returns true if the state was fully observed, otherwise false.
    pub fn is_full(&self) -> bool {
        match self {
            Observation::Full(_) => true,
            _ => false,
        }
    }

    /// Returns true if the state was only partially observed, otherwise false.
    pub fn is_partial(&self) -> bool {
        match self {
            Observation::Partial(_) => true,
            _ => false,
        }
    }

    /// Returns true if the observation is the terminal state, otherwise false.
    pub fn is_terminal(&self) -> bool {
        match self {
            Observation::Terminal(_) => true,
            _ => false,
        }
    }
}

/// Container class for data associated with a domain transition.
#[derive(Clone)]
pub struct Transition<S, A> {
    /// State transitioned _from_, `s`.
    pub from: Observation<S>,

    /// Action taken to initiate the transition (control tasks).
    pub action: A,

    /// Reward obtained from the transition.
    pub reward: f64,

    /// State transitioned _to_, `s'`.
    pub to: Observation<S>,
}

impl<S, A> Transition<S, A> {
    /// Return references to the `from` and `to` states associated with this transition.
    pub fn states(&self) -> (&S, &S) {
        (self.from.state(), self.to.state())
    }

    /// Apply a closure to the `from` and `to` states associated with this transition.
    pub fn map_states<O>(&self, f: impl Fn(&S) -> O) -> (O, O) {
        (f(self.from.state()), f(self.to.state()))
    }

    /// Returns true if the transition ends in a terminal state.
    pub fn terminated(&self) -> bool {
        self.to.is_terminal()
    }

    /// Replace the action associated with this transition and return a new instance.
    pub fn replace_action<T>(self, action: T) -> Transition<S, T> {
        Transition {
            from: self.from,
            action: action,
            reward: self.reward,
            to: self.to,
        }
    }

    /// Drop the action associated with this transition and return a new instance.
    pub fn drop_action(self) -> Transition<S, ()> {
        self.replace_action(())
    }
}

impl_into!(Transition<S, u8> => Transition<S, ()>);
impl_into!(Transition<S, u16> => Transition<S, ()>);
impl_into!(Transition<S, u32> => Transition<S, ()>);
impl_into!(Transition<S, u64> => Transition<S, ()>);
impl_into!(Transition<S, usize> => Transition<S, ()>);
impl_into!(Transition<S, i8> => Transition<S, ()>);
impl_into!(Transition<S, i16> => Transition<S, ()>);
impl_into!(Transition<S, i32> => Transition<S, ()>);
impl_into!(Transition<S, i64> => Transition<S, ()>);
impl_into!(Transition<S, isize> => Transition<S, ()>);
impl_into!(Transition<S, f32> => Transition<S, ()>);
impl_into!(Transition<S, f64> => Transition<S, ()>);

/// An interface for constructing reinforcement learning problem domains.
pub trait Domain {
    /// State space representation type class.
    type StateSpace: Space;

    /// Action space representation type class.
    type ActionSpace: Space;

    /// Emit an observation of the current state of the environment.
    fn emit(&self) -> Observation<<Self::StateSpace as Space>::Value>;

    /// Transition the environment forward a single step given an action, `a`.
    fn step(
        &mut self,
        a: <Self::ActionSpace as Space>::Value,
    ) -> Transition<<Self::StateSpace as Space>::Value, <Self::ActionSpace as Space>::Value>;

    /// Returns true if the current state is terminal.
    fn is_terminal(&self) -> bool;

    /// Compute the reward associated with a transition from one state to
    /// another.
    fn reward(
        &self,
        from: &Observation<<Self::StateSpace as Space>::Value>,
        to: &Observation<<Self::StateSpace as Space>::Value>,
    ) -> f64;

    /// Returns an instance of the state space type class.
    fn state_space(&self) -> Self::StateSpace;

    /// Returns an instance of the action space type class.
    fn action_space(&self) -> Self::ActionSpace;
}

mod ode;
use self::ode::*;

mod grid_world;

import_all!(mountain_car);
import_all!(cart_pole);
import_all!(acrobat);
import_all!(hiv);
import_all!(cliff_walk);

#[cfg(feature = "openai")]
import_all!(openai);
<commit_msg>feat(domains): add Debug implementaiton to Observation and Transition<commit_after>//! Learning benchmark domains module.
use crate::geometry::Space;

macro_rules! impl_into {
    (Transition < S, $type:ty > => Transition < S,() >) => {
        impl<S> Into<Transition<S, ()>> for Transition<S, $type> {
            fn into(self) -> Transition<S, ()> {
                self.drop_action()
            }
        }
    };
}

/// Container class for data associated with a domain observation.
#[derive(Clone, Debug)]
pub enum Observation<S> {
    /// Fully observed state of the environment.
    Full(S),

    /// Partially observed state of the environment.
    Partial(S),

    /// Terminal state of the environment.
    Terminal(S),
}

impl<S> Observation<S> {
    /// Helper function returning a reference to the state values for the given
    /// observation.
    pub fn state(&self) -> &S {
        use self::Observation::*;

        match self {
            Full(ref state) | Partial(ref state) | Terminal(ref state) => state,
        }
    }

    /// Returns true if the state was fully observed, otherwise false.
    pub fn is_full(&self) -> bool {
        match self {
            Observation::Full(_) => true,
            _ => false,
        }
    }

    /// Returns true if the state was only partially observed, otherwise false.
    pub fn is_partial(&self) -> bool {
        match self {
            Observation::Partial(_) => true,
            _ => false,
        }
    }

    /// Returns true if the observation is the terminal state, otherwise false.
    pub fn is_terminal(&self) -> bool {
        match self {
            Observation::Terminal(_) => true,
            _ => false,
        }
    }
}

/// Container class for data associated with a domain transition.
#[derive(Clone, Debug)]
pub struct Transition<S, A> {
    /// State transitioned _from_, `s`.
    pub from: Observation<S>,

    /// Action taken to initiate the transition (control tasks).
    pub action: A,

    /// Reward obtained from the transition.
    pub reward: f64,

    /// State transitioned _to_, `s'`.
    pub to: Observation<S>,
}

impl<S, A> Transition<S, A> {
    /// Return references to the `from` and `to` states associated with this transition.
    pub fn states(&self) -> (&S, &S) {
        (self.from.state(), self.to.state())
    }

    /// Apply a closure to the `from` and `to` states associated with this transition.
    pub fn map_states<O>(&self, f: impl Fn(&S) -> O) -> (O, O) {
        (f(self.from.state()), f(self.to.state()))
    }

    /// Returns true if the transition ends in a terminal state.
    pub fn terminated(&self) -> bool {
        self.to.is_terminal()
    }

    /// Replace the action associated with this transition and return a new instance.
    pub fn replace_action<T>(self, action: T) -> Transition<S, T> {
        Transition {
            from: self.from,
            action: action,
            reward: self.reward,
            to: self.to,
        }
    }

    /// Drop the action associated with this transition and return a new instance.
    pub fn drop_action(self) -> Transition<S, ()> {
        self.replace_action(())
    }
}

impl_into!(Transition<S, u8> => Transition<S, ()>);
impl_into!(Transition<S, u16> => Transition<S, ()>);
impl_into!(Transition<S, u32> => Transition<S, ()>);
impl_into!(Transition<S, u64> => Transition<S, ()>);
impl_into!(Transition<S, usize> => Transition<S, ()>);
impl_into!(Transition<S, i8> => Transition<S, ()>);
impl_into!(Transition<S, i16> => Transition<S, ()>);
impl_into!(Transition<S, i32> => Transition<S, ()>);
impl_into!(Transition<S, i64> => Transition<S, ()>);
impl_into!(Transition<S, isize> => Transition<S, ()>);
impl_into!(Transition<S, f32> => Transition<S, ()>);
impl_into!(Transition<S, f64> => Transition<S, ()>);

/// An interface for constructing reinforcement learning problem domains.
pub trait Domain {
    /// State space representation type class.
    type StateSpace: Space;

    /// Action space representation type class.
    type ActionSpace: Space;

    /// Emit an observation of the current state of the environment.
    fn emit(&self) -> Observation<<Self::StateSpace as Space>::Value>;

    /// Transition the environment forward a single step given an action, `a`.
    fn step(
        &mut self,
        a: <Self::ActionSpace as Space>::Value,
    ) -> Transition<<Self::StateSpace as Space>::Value, <Self::ActionSpace as Space>::Value>;

    /// Returns true if the current state is terminal.
    fn is_terminal(&self) -> bool;

    /// Compute the reward associated with a transition from one state to
    /// another.
    fn reward(
        &self,
        from: &Observation<<Self::StateSpace as Space>::Value>,
        to: &Observation<<Self::StateSpace as Space>::Value>,
    ) -> f64;

    /// Returns an instance of the state space type class.
    fn state_space(&self) -> Self::StateSpace;

    /// Returns an instance of the action space type class.
    fn action_space(&self) -> Self::ActionSpace;
}

mod ode;
use self::ode::*;

mod grid_world;

import_all!(mountain_car);
import_all!(cart_pole);
import_all!(acrobat);
import_all!(hiv);
import_all!(cliff_walk);

#[cfg(feature = "openai")]
import_all!(openai);
<|endoftext|>


<commit_before>// A port of selecta's scoring algorithm.
// selecta (c) 2014 John Hawthorn
// Licensed under the MIT license

use std::cmp::Ordering;
use super::eq;
use super::mat::Mat;
use super::bonus::compute_bonus;
use super::consts::*;

/// A score derived from comparing needle and haystack strings, with optional
/// match position information.
#[derive(Debug)]
pub struct Score {
    /// The computed score value
    pub value: f64,

    /// Optional vector of match positions
    pub positions: Option<Vec<usize>>
}

impl Score {
    /// Creates a new Score from the provided needle and haystack
    ///
    /// # Examples
    ///
    /// ```
    /// use rff::fuzzy::Score;
    /// let score = Score::new("abc", "abc");
    /// assert_eq!(score.value, std::f64::INFINITY);
    /// ```
    pub fn new(needle: &str, haystack: &str) -> Score {
        Scorer::new(needle, haystack).calculate().as_score()
    }

    /// Creates a new Score from the provided needle and haystack, calculating
    /// match positions.
    ///
    /// # Examples
    ///
    /// ```
    /// use rff::fuzzy::Score;
    /// let score = Score::with_positions("abc", "abc");
    /// assert_eq!(score.value, std::f64::INFINITY);
    /// assert_eq!(score.positions, Some(vec![0, 1, 2]));
    /// ```
    pub fn with_positions(needle: &str, haystack: &str) -> Score {
        Scorer::new(needle, haystack).calculate().with_positions().as_score()
    }
}

impl PartialOrd for Score {
    #[inline]
    fn partial_cmp(&self, other: &Score) -> Option<Ordering> {
        self.value.partial_cmp(&other.value)
    }
}

impl PartialEq for Score {
    #[inline]
    fn eq(&self, other: &Score) -> bool {
        self.value == other.value
    }
}

struct Scorer<'a> {
    needle: &'a str,
    haystack: &'a str,
    len_n: usize,
    len_h: usize,
    d: Mat,
    m: Mat,
    positions: Option<Vec<usize>>
}

impl<'a> Scorer<'a> {
    #[inline]
    fn new(needle: &'a str, haystack: &'a str) -> Scorer<'a> {
        let len_n = needle.chars().count();
        let len_h = haystack.chars().count();

        let m = Mat::new(len_n, len_h);
        let d = Mat::new(len_n, len_h);

        Scorer {
            needle: needle,
            haystack: haystack,
            len_n: len_n,
            len_h: len_h,
            m: m,
            d: d,
            positions: None
        }
    }

    #[inline]
    fn calculate(mut self) -> Scorer<'a> {
        if self.len_n == 0 || self.len_n == self.len_h {
            return self
        }

        let bonus = compute_bonus(self.haystack);

        for (i, n) in self.needle.chars().enumerate() {
            let mut prev_score = SCORE_MIN;
            let gap_score = if i == self.len_n - 1 { SCORE_GAP_TRAILING } else { SCORE_GAP_INNER };

            for (j, h) in self.haystack.chars().enumerate() {
                if eq(n, h) {
                    let mut score = SCORE_MIN;

                    let bonus_score = bonus[j];

                    if i == 0 {
                        score = ((j as f64) * SCORE_GAP_LEADING) + bonus_score;
                    } else if j > 0 {
                        let m = self.m.get(i - 1, j - 1).unwrap();
                        let d = self.d.get(i - 1, j - 1).unwrap();

                        score = (m + bonus_score).max(d + SCORE_MATCH_CONSECUTIVE);
                    }

                    prev_score = score.max(prev_score + gap_score);

                    self.d.set(i, j, score);
                    self.m.set(i, j, prev_score);
                } else {
                    self.d.set(i, j, SCORE_MIN);
                    self.m.set(i, j, prev_score + gap_score);
                    prev_score += gap_score;
                }
            }
        }

        self
    }

    #[inline]
    fn with_positions(mut self) -> Scorer<'a> {
        if self.len_n == 0 {
            return self;
        } else if self.len_n == self.len_h {
            self.positions = Some((0..self.len_n).collect());
            return self;
        }

        let mut positions = vec![0 as usize; self.len_n];
        let mut match_required = false;

        let mut j = self.len_h - 1;

        for i in (0..self.len_n).rev() {
            while j > (0 as usize) {
                let last = if i > 0 && j > 0 { self.d.get(i - 1, j - 1).unwrap() } else { 0.0 };

                let d = self.d.get(i, j).unwrap();
                let m = self.m.get(i, j).unwrap();

                if d != SCORE_MIN && (match_required || d == m) {
                    if i > 0 && j > 0 && m == last + SCORE_MATCH_CONSECUTIVE {
                        match_required = true;
                    }

                    positions[i] = j;

                    break;
                }

                j -= 1
            }
        }

        self.positions = Some(positions);
        self
    }

    #[inline]
    fn as_score(self) -> Score {
        let value = match self.len_n {
            0 => SCORE_MIN,
            _ if self.len_n == self.len_h => SCORE_MAX,
            _ => self.m.get(self.len_n - 1, self.len_h - 1).unwrap_or(SCORE_MIN)
        };

        Score {
            value: value,
            positions: self.positions,
        }
    }
}
#[cfg(test)]
mod tests {
    use super::*;

    fn score(needle: &str, haystack: &str) -> Score {
        Score::new(needle, haystack)
    }

    #[test]
    fn test_eq() {
        let a = Score { value: 1.0, positions: None };
        let b = Score { value: 1.0, positions: None };
        assert_eq!(a, b);
    }

    #[test]
    fn test_cmp() {
        let a = Score { value: 2.0, positions: None };
        let b = Score { value: 1.0, positions: None };
        assert!(a > b);
        assert!(b < a);

        let b = Score { value: 2.0, positions: None };
        assert!(a == b);
    }

    #[test]
    fn relative_scores() {
        // App/Models/Order is better than App/MOdels/zRder
        assert!(score("amor", "app/models/order") > score("amor", "app/models/zrder"));

        // App/MOdels/foo is better than App/M/fOo
        assert!(score("amo", "app/m/foo") < score("amo", "app/models/foo"));

        // GEMFIle.Lock < GEMFILe
        assert!(score("gemfil", "Gemfile.lock") < score("gemfil", "Gemfile"));

        // GEMFIle.Lock < GEMFILe
        assert!(score("gemfil", "Gemfile.lock") < score("gemfil", "Gemfile"));

        // Prefer shorter scorees
        assert!(score("abce", "abcdef") > score("abce", "abc de"));

        // Prefer shorter candidates
        assert!(score("test", "tests") > score("test", "testing"));

        // Scores first letter highly
        assert!(score("test", "testing") > score("test", "/testing"));

        // Prefer shorter scorees
        assert!(score("abc", "    a b c ") > score("abc", " a  b  c "));
        assert!(score("abc", " a b c    ") > score("abc", " a  b  c "));
    }

    #[test]
    fn score_utf8() {
        assert_eq!(score("ß", "öäßéè").value, -0.02);
    }

    #[test]
    fn positions() {
        macro_rules! test_positions {
            ($needle:expr, $haystack:expr, $result:expr) => {
                let score = Score::with_positions($needle, $haystack);
                assert_eq!(score.positions, Some($result));
            }
        }

        test_positions!("amo", "app/models/foo", vec![0, 4, 5]);
        test_positions!("amor", "app/models/order", vec![0, 4, 11, 12]);
        test_positions!("as", "tags", vec![1, 3]);
        test_positions!("abc", "a/a/b/c/c", vec![2, 4, 6]);
        test_positions!("foo", "foo", vec![0, 1, 2]);
        test_positions!("drivers", "/path/to/drivers/file.txt", vec![9, 10, 11, 12, 13, 14, 15]);
    }
}
<commit_msg>Do not calculate score for strings over 1024 chars<commit_after>// A port of selecta's scoring algorithm.
// selecta (c) 2014 John Hawthorn
// Licensed under the MIT license

use std::cmp::Ordering;
use super::eq;
use super::mat::Mat;
use super::bonus::compute_bonus;
use super::consts::*;

/// A score derived from comparing needle and haystack strings, with optional
/// match position information.
#[derive(Debug)]
pub struct Score {
    /// The computed score value
    pub value: f64,

    /// Optional vector of match positions
    pub positions: Option<Vec<usize>>
}

impl Score {
    /// Creates a new Score from the provided needle and haystack
    ///
    /// # Examples
    ///
    /// ```
    /// use rff::fuzzy::Score;
    /// let score = Score::new("abc", "abc");
    /// assert_eq!(score.value, std::f64::INFINITY);
    /// ```
    pub fn new(needle: &str, haystack: &str) -> Score {
        Scorer::new(needle, haystack).calculate().as_score()
    }

    /// Creates a new Score from the provided needle and haystack, calculating
    /// match positions.
    ///
    /// # Examples
    ///
    /// ```
    /// use rff::fuzzy::Score;
    /// let score = Score::with_positions("abc", "abc");
    /// assert_eq!(score.value, std::f64::INFINITY);
    /// assert_eq!(score.positions, Some(vec![0, 1, 2]));
    /// ```
    pub fn with_positions(needle: &str, haystack: &str) -> Score {
        Scorer::new(needle, haystack).calculate().with_positions().as_score()
    }
}

impl PartialOrd for Score {
    #[inline]
    fn partial_cmp(&self, other: &Score) -> Option<Ordering> {
        self.value.partial_cmp(&other.value)
    }
}

impl PartialEq for Score {
    #[inline]
    fn eq(&self, other: &Score) -> bool {
        self.value == other.value
    }
}

struct Scorer<'a> {
    needle: &'a str,
    haystack: &'a str,
    len_n: usize,
    len_h: usize,
    d: Mat,
    m: Mat,
    score: f64,
    positions: Option<Vec<usize>>
}

impl<'a> Scorer<'a> {
    #[inline]
    fn new(needle: &'a str, haystack: &'a str) -> Scorer<'a> {
        let len_n = needle.chars().count();
        let len_h = haystack.chars().count();

        let m = Mat::new(len_n, len_h);
        let d = Mat::new(len_n, len_h);

        Scorer {
            needle: needle,
            haystack: haystack,
            len_n: len_n,
            len_h: len_h,
            m: m,
            d: d,
            score: 0.0,
            positions: None
        }
    }

    #[inline]
    fn calculate(mut self) -> Scorer<'a> {
        if self.len_n == 0 {
            self.score = SCORE_MIN;
            return self;
        }

        if self.len_n == self.len_h {
            self.score = SCORE_MAX;
            return self;
        }

        if self.len_h > 1024 {
            // unreasonably large candidate
            self.score = SCORE_MIN;
            return self;
        }

        let bonus = compute_bonus(self.haystack);

        for (i, n) in self.needle.chars().enumerate() {
            let mut prev_score = SCORE_MIN;
            let gap_score = if i == self.len_n - 1 { SCORE_GAP_TRAILING } else { SCORE_GAP_INNER };

            for (j, h) in self.haystack.chars().enumerate() {
                if eq(n, h) {
                    let mut score = SCORE_MIN;

                    let bonus_score = bonus[j];

                    if i == 0 {
                        score = ((j as f64) * SCORE_GAP_LEADING) + bonus_score;
                    } else if j > 0 {
                        let m = self.m.get(i - 1, j - 1).unwrap();
                        let d = self.d.get(i - 1, j - 1).unwrap();

                        score = (m + bonus_score).max(d + SCORE_MATCH_CONSECUTIVE);
                    }

                    prev_score = score.max(prev_score + gap_score);

                    self.d.set(i, j, score);
                    self.m.set(i, j, prev_score);
                } else {
                    self.d.set(i, j, SCORE_MIN);
                    self.m.set(i, j, prev_score + gap_score);
                    prev_score += gap_score;
                }
            }
        }

        self.score = self.m.get(self.len_n - 1, self.len_h - 1).unwrap_or(SCORE_MIN);
        self
    }

    #[inline]
    fn with_positions(mut self) -> Scorer<'a> {
        if self.len_n == 0 {
            return self;
        } else if self.len_n == self.len_h {
            self.positions = Some((0..self.len_n).collect());
            return self;
        }

        let mut positions = vec![0 as usize; self.len_n];
        let mut match_required = false;

        let mut j = self.len_h - 1;

        for i in (0..self.len_n).rev() {
            while j > (0 as usize) {
                let last = if i > 0 && j > 0 { self.d.get(i - 1, j - 1).unwrap() } else { 0.0 };

                let d = self.d.get(i, j).unwrap();
                let m = self.m.get(i, j).unwrap();

                if d != SCORE_MIN && (match_required || d == m) {
                    if i > 0 && j > 0 && m == last + SCORE_MATCH_CONSECUTIVE {
                        match_required = true;
                    }

                    positions[i] = j;

                    break;
                }

                j -= 1
            }
        }

        self.positions = Some(positions);
        self
    }

    #[inline]
    fn as_score(self) -> Score {
        Score {
            value: self.score,
            positions: self.positions,
        }
    }
}
#[cfg(test)]
mod tests {
    use super::*;

    fn score(needle: &str, haystack: &str) -> Score {
        Score::new(needle, haystack)
    }

    #[test]
    fn test_eq() {
        let a = Score { value: 1.0, positions: None };
        let b = Score { value: 1.0, positions: None };
        assert_eq!(a, b);
    }

    #[test]
    fn test_cmp() {
        let a = Score { value: 2.0, positions: None };
        let b = Score { value: 1.0, positions: None };
        assert!(a > b);
        assert!(b < a);

        let b = Score { value: 2.0, positions: None };
        assert!(a == b);
    }

    #[test]
    fn relative_scores() {
        // App/Models/Order is better than App/MOdels/zRder
        assert!(score("amor", "app/models/order") > score("amor", "app/models/zrder"));

        // App/MOdels/foo is better than App/M/fOo
        assert!(score("amo", "app/m/foo") < score("amo", "app/models/foo"));

        // GEMFIle.Lock < GEMFILe
        assert!(score("gemfil", "Gemfile.lock") < score("gemfil", "Gemfile"));

        // GEMFIle.Lock < GEMFILe
        assert!(score("gemfil", "Gemfile.lock") < score("gemfil", "Gemfile"));

        // Prefer shorter scorees
        assert!(score("abce", "abcdef") > score("abce", "abc de"));

        // Prefer shorter candidates
        assert!(score("test", "tests") > score("test", "testing"));

        // Scores first letter highly
        assert!(score("test", "testing") > score("test", "/testing"));

        // Prefer shorter scorees
        assert!(score("abc", "    a b c ") > score("abc", " a  b  c "));
        assert!(score("abc", " a b c    ") > score("abc", " a  b  c "));
    }

    #[test]
    fn score_utf8() {
        assert_eq!(score("ß", "öäßéè").value, -0.02);
    }

    #[test]
    fn positions() {
        macro_rules! test_positions {
            ($needle:expr, $haystack:expr, $result:expr) => {
                let score = Score::with_positions($needle, $haystack);
                assert_eq!(score.positions, Some($result));
            }
        }

        test_positions!("amo", "app/models/foo", vec![0, 4, 5]);
        test_positions!("amor", "app/models/order", vec![0, 4, 11, 12]);
        test_positions!("as", "tags", vec![1, 3]);
        test_positions!("abc", "a/a/b/c/c", vec![2, 4, 6]);
        test_positions!("foo", "foo", vec![0, 1, 2]);
        test_positions!("drivers", "/path/to/drivers/file.txt", vec![9, 10, 11, 12, 13, 14, 15]);
    }
}
<|endoftext|>


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

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

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

mod lila;

use lila::{Game, Lila};

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            num_games += 1;
            if num_games % 1024 == 0 {
                log::info!(
                    "indexer {:02}: indexed {} games for {} ...",
                    self.idx,
                    num_games,
                    player.as_str()
                );
            }
        }

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

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

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

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

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

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

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

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

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

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

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

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

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

            pos.play_unchecked(&m);
        }

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

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

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

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

enum IndexerMessage {
    IndexPlayer {
        player: UserId,
        status: PersonalStatus,
        since_created_at: u64,
    },
}
<commit_msg>limit to 50 plies per game<commit_after>use std::{
    collections::{hash_map::Entry, HashMap},
    sync::Arc,
    time::{Duration, Instant, SystemTime},
};

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

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

mod lila;

use lila::{Game, Lila};

const MAX_PLIES: usize = 50;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            num_games += 1;
            if num_games % 1024 == 0 {
                log::info!(
                    "indexer {:02}: indexed {} games for {} ...",
                    self.idx,
                    num_games,
                    player.as_str()
                );
            }
        }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            pos.play_unchecked(&m);
        }

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

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

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

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

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