<commit_before>use core::SegmentId;
use error::TantivyError;
use schema::Schema;
use serde_json;
use std::borrow::BorrowMut;
use std::fmt;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use Result;
use indexer::LockType;
use super::pool::LeasedItem;
use super::pool::Pool;
use super::segment::create_segment;
use super::segment::Segment;
use core::searcher::Searcher;
use core::IndexMeta;
use core::SegmentMeta;
use core::SegmentReader;
use core::META_FILEPATH;
#[cfg(feature = "mmap")]
use directory::MmapDirectory;
use directory::{Directory, RAMDirectory};
use directory::{ManagedDirectory};
use indexer::index_writer::open_index_writer;
use indexer::index_writer::HEAP_SIZE_MIN;
use indexer::segment_updater::save_new_metas;
use num_cpus;
use std::path::Path;
use tokenizer::TokenizerManager;
use IndexWriter;
use schema::FieldType;
use schema::Field;
use tokenizer::BoxedTokenizer;

fn load_metas(directory: &Directory) -> Result<IndexMeta> {
    let meta_data = directory.atomic_read(&META_FILEPATH)?;
    let meta_string = String::from_utf8_lossy(&meta_data);
    serde_json::from_str(&meta_string)
        .map_err(|_| TantivyError::CorruptedFile(META_FILEPATH.clone()))
}

/// Search Index
pub struct Index {
    directory: ManagedDirectory,
    schema: Schema,
    num_searchers: Arc<AtomicUsize>,
    searcher_pool: Arc<Pool<Searcher>>,
    tokenizers: TokenizerManager,
}

impl Index {
    /// Creates a new index using the `RAMDirectory`.
    ///
    /// The index will be allocated in anonymous memory.
    /// This should only be used for unit tests.
    pub fn create_in_ram(schema: Schema) -> Index {
        let ram_directory = RAMDirectory::create();
        Index::create(ram_directory, schema).expect("Creating a RAMDirectory should never fail")
    }

    /// Creates a new index in a given filepath.
    /// The index will use the `MMapDirectory`.
    ///
    /// If a previous index was in this directory, then its meta file will be destroyed.
    #[cfg(feature = "mmap")]
    pub fn create_in_dir<P: AsRef<Path>>(directory_path: P, schema: Schema) -> Result<Index> {
        let mmap_directory = MmapDirectory::open(directory_path)?;
        Index::create(mmap_directory, schema)
    }

    /// Creates a new index in a temp directory.
    ///
    /// The index will use the `MMapDirectory` in a newly created directory.
    /// The temp directory will be destroyed automatically when the `Index` object
    /// is destroyed.
    ///
    /// The temp directory is only used for testing the `MmapDirectory`.
    /// For other unit tests, prefer the `RAMDirectory`, see: `create_in_ram`.
    #[cfg(feature = "mmap")]
    pub fn create_from_tempdir(schema: Schema) -> Result<Index> {
        let mmap_directory = MmapDirectory::create_from_tempdir()?;
        Index::create(mmap_directory, schema)
    }

    /// Creates a new index given an implementation of the trait `Directory`
    pub fn create<Dir: Directory>(dir: Dir, schema: Schema) -> Result<Index> {
        let directory = ManagedDirectory::new(dir)?;
        Index::from_directory(directory, schema)
    }

    /// Create a new index from a directory.
    fn from_directory(mut directory: ManagedDirectory, schema: Schema) -> Result<Index> {
        save_new_metas(schema.clone(), 0, directory.borrow_mut())?;
        let metas = IndexMeta::with_schema(schema);
        Index::create_from_metas(directory, &metas)
    }

    /// Creates a new index given a directory and an `IndexMeta`.
    fn create_from_metas(directory: ManagedDirectory, metas: &IndexMeta) -> Result<Index> {
        let schema = metas.schema.clone();
        let n_cpus = num_cpus::get();
        let index = Index {
            directory,
            schema,
            num_searchers: Arc::new(AtomicUsize::new(n_cpus)),
            searcher_pool: Arc::new(Pool::new()),
            tokenizers: TokenizerManager::default(),
        };
        index.load_searchers()?;
        Ok(index)
    }

    /// Accessor for the tokenizer manager.
    pub fn tokenizers(&self) -> &TokenizerManager {
        &self.tokenizers
    }

    /// Helper to access the tokenizer associated to a specific field.
    pub fn tokenizer_for_field(&self, field: Field) -> Result<Box<BoxedTokenizer>> {
        let field_entry = self.schema.get_field_entry(field);
        let field_type = field_entry.field_type();
        let tokenizer_manager: &TokenizerManager = self.tokenizers();
        let tokenizer_name_opt: Option<Box<BoxedTokenizer>> =
            match field_type {
                FieldType::Str(text_options) => {
                    text_options
                        .get_indexing_options()
                        .map(|text_indexing_options| text_indexing_options.tokenizer().to_string())
                        .and_then(|tokenizer_name| tokenizer_manager.get(&tokenizer_name))
                },
                _ => {
                    None
                }
            };
        match tokenizer_name_opt {
            Some(tokenizer) => {
                Ok(tokenizer)
            }
            None => {
                Err(TantivyError:: SchemaError(format!("{:?} is not a text field.", field_entry.name())))
            }
        }
    }

    /// Opens a new directory from an index path.
    #[cfg(feature = "mmap")]
    pub fn open_in_dir<P: AsRef<Path>>(directory_path: P) -> Result<Index> {
        let mmap_directory = MmapDirectory::open(directory_path)?;
        Index::open(mmap_directory)
    }

    /// Open the index using the provided directory
    pub fn open<D: Directory>(directory: D) -> Result<Index> {
        let directory = ManagedDirectory::new(directory)?;
        let metas = load_metas(&directory)?;
        Index::create_from_metas(directory, &metas)
    }

    /// Reads the index meta file from the directory.
    pub fn load_metas(&self) -> Result<IndexMeta> {
        load_metas(self.directory())
    }

    /// Open a new index writer. Attempts to acquire a lockfile.
    ///
    /// The lockfile should be deleted on drop, but it is possible
    /// that due to a panic or other error, a stale lockfile will be
    /// left in the index directory. If you are sure that no other
    /// `IndexWriter` on the system is accessing the index directory,
    /// it is safe to manually delete the lockfile.
    ///
    /// - `num_threads` defines the number of indexing workers that
    /// should work at the same time.
    ///
    /// - `overall_heap_size_in_bytes` sets the amount of memory
    /// allocated for all indexing thread.
    /// Each thread will receive a budget of  `overall_heap_size_in_bytes / num_threads`.
    ///
    /// # Errors
    /// If the lockfile already exists, returns `Error::FileAlreadyExists`.
    /// # Panics
    /// If the heap size per thread is too small, panics.
    pub fn writer_with_num_threads(
        &self,
        num_threads: usize,
        overall_heap_size_in_bytes: usize,
    ) -> Result<IndexWriter> {

        let directory_lock = LockType::IndexWriterLock.acquire_lock(&self.directory)?;
        let heap_size_in_bytes_per_thread = overall_heap_size_in_bytes / num_threads;
        open_index_writer(
            self,
            num_threads,
            heap_size_in_bytes_per_thread,
            directory_lock,
        )
    }

    /// Creates a multithreaded writer
    ///
    /// Tantivy will automatically define the number of threads to use.
    /// `overall_heap_size_in_bytes` is the total target memory usage that will be split
    /// between a given number of threads.
    ///
    /// # Errors
    /// If the lockfile already exists, returns `Error::FileAlreadyExists`.
    /// # Panics
    /// If the heap size per thread is too small, panics.
    pub fn writer(&self, overall_heap_size_in_bytes: usize) -> Result<IndexWriter> {
        let mut num_threads = num_cpus::get();
        let heap_size_in_bytes_per_thread = overall_heap_size_in_bytes / num_threads;
        if heap_size_in_bytes_per_thread < HEAP_SIZE_MIN {
            num_threads = (overall_heap_size_in_bytes / HEAP_SIZE_MIN).max(1);
        }
        self.writer_with_num_threads(num_threads, overall_heap_size_in_bytes)
    }

    /// Accessor to the index schema
    ///
    /// The schema is actually cloned.
    pub fn schema(&self) -> Schema {
        self.schema.clone()
    }

    /// Returns the list of segments that are searchable
    pub fn searchable_segments(&self) -> Result<Vec<Segment>> {
        Ok(self.searchable_segment_metas()?
            .into_iter()
            .map(|segment_meta| self.segment(segment_meta))
            .collect())
    }

    #[doc(hidden)]
    pub fn segment(&self, segment_meta: SegmentMeta) -> Segment {
        create_segment(self.clone(), segment_meta)
    }

    /// Creates a new segment.
    pub fn new_segment(&self) -> Segment {
        let segment_meta = SegmentMeta::new(SegmentId::generate_random(), 0);
        self.segment(segment_meta)
    }

    /// Return a reference to the index directory.
    pub fn directory(&self) -> &ManagedDirectory {
        &self.directory
    }

    /// Return a mutable reference to the index directory.
    pub fn directory_mut(&mut self) -> &mut ManagedDirectory {
        &mut self.directory
    }

    /// Reads the meta.json and returns the list of
    /// `SegmentMeta` from the last commit.
    pub fn searchable_segment_metas(&self) -> Result<Vec<SegmentMeta>> {
        Ok(self.load_metas()?.segments)
    }

    /// Returns the list of segment ids that are searchable.
    pub fn searchable_segment_ids(&self) -> Result<Vec<SegmentId>> {
        Ok(self.searchable_segment_metas()?
            .iter()
            .map(|segment_meta| segment_meta.id())
            .collect())
    }

    /// Sets the number of searchers to use
    ///
    /// Only works after the next call to `load_searchers`
    pub fn set_num_searchers(&mut self, num_searchers: usize) {
        self.num_searchers.store(num_searchers, Ordering::Release);
    }

    /// Creates a new generation of searchers after

    /// a change of the set of searchable indexes.
    ///
    /// This needs to be called when a new segment has been
    /// published or after a merge.
    pub fn load_searchers(&self) -> Result<()> {
        let _meta_lock = LockType::MetaLock.acquire_lock(self.directory())?;
        let searchable_segments = self.searchable_segments()?;
        let segment_readers: Vec<SegmentReader> = searchable_segments
            .iter()
            .map(SegmentReader::open)
            .collect::<Result<_>>()?;
        let schema = self.schema();
        let num_searchers: usize = self.num_searchers.load(Ordering::Acquire);
        let searchers = (0..num_searchers)
            .map(|_| Searcher::new(schema.clone(), self.clone(), segment_readers.clone()))
            .collect();
        self.searcher_pool.publish_new_generation(searchers);
        Ok(())
    }

    /// Returns a searcher
    ///
    /// This method should be called every single time a search
    /// query is performed.
    /// The searchers are taken from a pool of `num_searchers` searchers.
    /// If no searcher is available
    /// this may block.
    ///
    /// The same searcher must be used for a given query, as it ensures
    /// the use of a consistent segment set.
    pub fn searcher(&self) -> LeasedItem<Searcher> {
        self.searcher_pool.acquire()
    }
}

impl fmt::Debug for Index {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Index({:?})", self.directory)
    }
}

impl Clone for Index {
    fn clone(&self) -> Index {
        Index {
            directory: self.directory.clone(),
            schema: self.schema.clone(),
            num_searchers: Arc::clone(&self.num_searchers),
            searcher_pool: Arc::clone(&self.searcher_pool),
            tokenizers: self.tokenizers.clone(),
        }
    }
}

#[cfg(test)]
mod tests {
    use Index;
    use schema::{SchemaBuilder, TEXT, INT_INDEXED};

    #[test]
    fn test_indexer_for_field() {
        let mut schema_builder = SchemaBuilder::default();
        let num_likes_field = schema_builder.add_u64_field("num_likes", INT_INDEXED);
        let body_field = schema_builder.add_text_field("body", TEXT);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        assert!(index.tokenizer_for_field(body_field).is_ok());
        assert_eq!(
            format!("{:?}", index.tokenizer_for_field(num_likes_field).err()),
            "Some(SchemaError(\"\\\"num_likes\\\" is not a text field.\"))"
        );
    }

}<commit_msg>Closes #414. (#417)<commit_after>use core::SegmentId;
use error::TantivyError;
use schema::Schema;
use serde_json;
use std::borrow::BorrowMut;
use std::fmt;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use Result;
use indexer::LockType;
use super::pool::LeasedItem;
use super::pool::Pool;
use super::segment::create_segment;
use super::segment::Segment;
use core::searcher::Searcher;
use core::IndexMeta;
use core::SegmentMeta;
use core::SegmentReader;
use core::META_FILEPATH;
#[cfg(feature = "mmap")]
use directory::MmapDirectory;
use directory::{Directory, RAMDirectory};
use directory::{ManagedDirectory};
use indexer::index_writer::open_index_writer;
use indexer::index_writer::HEAP_SIZE_MIN;
use indexer::segment_updater::save_new_metas;
use num_cpus;
use std::path::Path;
use tokenizer::TokenizerManager;
use IndexWriter;
use schema::FieldType;
use schema::Field;
use tokenizer::BoxedTokenizer;

fn load_metas(directory: &Directory) -> Result<IndexMeta> {
    let meta_data = directory.atomic_read(&META_FILEPATH)?;
    let meta_string = String::from_utf8_lossy(&meta_data);
    serde_json::from_str(&meta_string)
        .map_err(|_| TantivyError::CorruptedFile(META_FILEPATH.clone()))
}

/// Search Index
pub struct Index {
    directory: ManagedDirectory,
    schema: Schema,
    num_searchers: Arc<AtomicUsize>,
    searcher_pool: Arc<Pool<Searcher>>,
    tokenizers: TokenizerManager,
}

impl Index {
    /// Creates a new index using the `RAMDirectory`.
    ///
    /// The index will be allocated in anonymous memory.
    /// This should only be used for unit tests.
    pub fn create_in_ram(schema: Schema) -> Index {
        let ram_directory = RAMDirectory::create();
        Index::create(ram_directory, schema).expect("Creating a RAMDirectory should never fail")
    }

    /// Creates a new index in a given filepath.
    /// The index will use the `MMapDirectory`.
    ///
    /// If a previous index was in this directory, then its meta file will be destroyed.
    #[cfg(feature = "mmap")]
    pub fn create_in_dir<P: AsRef<Path>>(directory_path: P, schema: Schema) -> Result<Index> {
        let mmap_directory = MmapDirectory::open(directory_path)?;
        Index::create(mmap_directory, schema)
    }

    /// Creates a new index in a temp directory.
    ///
    /// The index will use the `MMapDirectory` in a newly created directory.
    /// The temp directory will be destroyed automatically when the `Index` object
    /// is destroyed.
    ///
    /// The temp directory is only used for testing the `MmapDirectory`.
    /// For other unit tests, prefer the `RAMDirectory`, see: `create_in_ram`.
    #[cfg(feature = "mmap")]
    pub fn create_from_tempdir(schema: Schema) -> Result<Index> {
        let mmap_directory = MmapDirectory::create_from_tempdir()?;
        Index::create(mmap_directory, schema)
    }

    /// Creates a new index given an implementation of the trait `Directory`
    pub fn create<Dir: Directory>(dir: Dir, schema: Schema) -> Result<Index> {
        let directory = ManagedDirectory::new(dir)?;
        Index::from_directory(directory, schema)
    }

    /// Create a new index from a directory.
    fn from_directory(mut directory: ManagedDirectory, schema: Schema) -> Result<Index> {
        save_new_metas(schema.clone(), 0, directory.borrow_mut())?;
        let metas = IndexMeta::with_schema(schema);
        Index::create_from_metas(directory, &metas)
    }

    /// Creates a new index given a directory and an `IndexMeta`.
    fn create_from_metas(directory: ManagedDirectory, metas: &IndexMeta) -> Result<Index> {
        let schema = metas.schema.clone();
        let n_cpus = num_cpus::get();
        let index = Index {
            directory,
            schema,
            num_searchers: Arc::new(AtomicUsize::new(n_cpus)),
            searcher_pool: Arc::new(Pool::new()),
            tokenizers: TokenizerManager::default(),
        };
        index.load_searchers()?;
        Ok(index)
    }

    /// Accessor for the tokenizer manager.
    pub fn tokenizers(&self) -> &TokenizerManager {
        &self.tokenizers
    }

    /// Helper to access the tokenizer associated to a specific field.
    pub fn tokenizer_for_field(&self, field: Field) -> Result<Box<BoxedTokenizer>> {
        let field_entry = self.schema.get_field_entry(field);
        let field_type = field_entry.field_type();
        let tokenizer_manager: &TokenizerManager = self.tokenizers();
        let tokenizer_name_opt: Option<Box<BoxedTokenizer>> =
            match field_type {
                FieldType::Str(text_options) => {
                    text_options
                        .get_indexing_options()
                        .map(|text_indexing_options| text_indexing_options.tokenizer().to_string())
                        .and_then(|tokenizer_name| tokenizer_manager.get(&tokenizer_name))
                },
                _ => {
                    None
                }
            };
        match tokenizer_name_opt {
            Some(tokenizer) => {
                Ok(tokenizer)
            }
            None => {
                Err(TantivyError:: SchemaError(format!("{:?} is not a text field.", field_entry.name())))
            }
        }
    }

    /// Opens a new directory from an index path.
    #[cfg(feature = "mmap")]
    pub fn open_in_dir<P: AsRef<Path>>(directory_path: P) -> Result<Index> {
        let mmap_directory = MmapDirectory::open(directory_path)?;
        Index::open(mmap_directory)
    }

    /// Open the index using the provided directory
    pub fn open<D: Directory>(directory: D) -> Result<Index> {
        let directory = ManagedDirectory::new(directory)?;
        let metas = load_metas(&directory)?;
        Index::create_from_metas(directory, &metas)
    }

    /// Reads the index meta file from the directory.
    pub fn load_metas(&self) -> Result<IndexMeta> {
        load_metas(self.directory())
    }

    /// Open a new index writer. Attempts to acquire a lockfile.
    ///
    /// The lockfile should be deleted on drop, but it is possible
    /// that due to a panic or other error, a stale lockfile will be
    /// left in the index directory. If you are sure that no other
    /// `IndexWriter` on the system is accessing the index directory,
    /// it is safe to manually delete the lockfile.
    ///
    /// - `num_threads` defines the number of indexing workers that
    /// should work at the same time.
    ///
    /// - `overall_heap_size_in_bytes` sets the amount of memory
    /// allocated for all indexing thread.
    /// Each thread will receive a budget of  `overall_heap_size_in_bytes / num_threads`.
    ///
    /// # Errors
    /// If the lockfile already exists, returns `Error::FileAlreadyExists`.
    /// # Panics
    /// If the heap size per thread is too small, panics.
    pub fn writer_with_num_threads(
        &self,
        num_threads: usize,
        overall_heap_size_in_bytes: usize,
    ) -> Result<IndexWriter> {

        let directory_lock = LockType::IndexWriterLock.acquire_lock(&self.directory)?;
        let heap_size_in_bytes_per_thread = overall_heap_size_in_bytes / num_threads;
        open_index_writer(
            self,
            num_threads,
            heap_size_in_bytes_per_thread,
            directory_lock,
        )
    }

    /// Creates a multithreaded writer
    ///
    /// Tantivy will automatically define the number of threads to use.
    /// `overall_heap_size_in_bytes` is the total target memory usage that will be split
    /// between a given number of threads.
    ///
    /// # Errors
    /// If the lockfile already exists, returns `Error::FileAlreadyExists`.
    /// # Panics
    /// If the heap size per thread is too small, panics.
    pub fn writer(&self, overall_heap_size_in_bytes: usize) -> Result<IndexWriter> {
        let mut num_threads = num_cpus::get();
        let heap_size_in_bytes_per_thread = overall_heap_size_in_bytes / num_threads;
        if heap_size_in_bytes_per_thread < HEAP_SIZE_MIN {
            num_threads = (overall_heap_size_in_bytes / HEAP_SIZE_MIN).max(1);
        }
        self.writer_with_num_threads(num_threads, overall_heap_size_in_bytes)
    }

    /// Accessor to the index schema
    ///
    /// The schema is actually cloned.
    pub fn schema(&self) -> Schema {
        self.schema.clone()
    }

    /// Returns the list of segments that are searchable
    pub fn searchable_segments(&self) -> Result<Vec<Segment>> {
        Ok(self.searchable_segment_metas()?
            .into_iter()
            .map(|segment_meta| self.segment(segment_meta))
            .collect())
    }

    #[doc(hidden)]
    pub fn segment(&self, segment_meta: SegmentMeta) -> Segment {
        create_segment(self.clone(), segment_meta)
    }

    /// Creates a new segment.
    pub fn new_segment(&self) -> Segment {
        let segment_meta = SegmentMeta::new(SegmentId::generate_random(), 0);
        self.segment(segment_meta)
    }

    /// Return a reference to the index directory.
    pub fn directory(&self) -> &ManagedDirectory {
        &self.directory
    }

    /// Return a mutable reference to the index directory.
    pub fn directory_mut(&mut self) -> &mut ManagedDirectory {
        &mut self.directory
    }

    /// Reads the meta.json and returns the list of
    /// `SegmentMeta` from the last commit.
    pub fn searchable_segment_metas(&self) -> Result<Vec<SegmentMeta>> {
        Ok(self.load_metas()?.segments)
    }

    /// Returns the list of segment ids that are searchable.
    pub fn searchable_segment_ids(&self) -> Result<Vec<SegmentId>> {
        Ok(self.searchable_segment_metas()?
            .iter()
            .map(|segment_meta| segment_meta.id())
            .collect())
    }

    /// Sets the number of searchers to use
    ///
    /// Only works after the next call to `load_searchers`
    pub fn set_num_searchers(&mut self, num_searchers: usize) {
        self.num_searchers.store(num_searchers, Ordering::Release);
    }

    /// Update searchers so that they reflect the state of the last
    /// `.commit()`.
    ///
    /// If indexing happens in the same process as searching,
    /// you most likely want to call `.load_searchers()` right after each
    /// successful call to `.commit()`.
    ///
    /// If indexing and searching happen in different processes, the way to
    /// get the freshest `index` at all time, is to watch `meta.json` and
    /// call `load_searchers` whenever a changes happen.
    pub fn load_searchers(&self) -> Result<()> {
        let _meta_lock = LockType::MetaLock.acquire_lock(self.directory())?;
        let searchable_segments = self.searchable_segments()?;
        let segment_readers: Vec<SegmentReader> = searchable_segments
            .iter()
            .map(SegmentReader::open)
            .collect::<Result<_>>()?;
        let schema = self.schema();
        let num_searchers: usize = self.num_searchers.load(Ordering::Acquire);
        let searchers = (0..num_searchers)
            .map(|_| Searcher::new(schema.clone(), self.clone(), segment_readers.clone()))
            .collect();
        self.searcher_pool.publish_new_generation(searchers);
        Ok(())
    }

    /// Returns a searcher
    ///
    /// This method should be called every single time a search
    /// query is performed.
    /// The searchers are taken from a pool of `num_searchers` searchers.
    /// If no searcher is available
    /// this may block.
    ///
    /// The same searcher must be used for a given query, as it ensures
    /// the use of a consistent segment set.
    pub fn searcher(&self) -> LeasedItem<Searcher> {
        self.searcher_pool.acquire()
    }
}

impl fmt::Debug for Index {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Index({:?})", self.directory)
    }
}

impl Clone for Index {
    fn clone(&self) -> Index {
        Index {
            directory: self.directory.clone(),
            schema: self.schema.clone(),
            num_searchers: Arc::clone(&self.num_searchers),
            searcher_pool: Arc::clone(&self.searcher_pool),
            tokenizers: self.tokenizers.clone(),
        }
    }
}

#[cfg(test)]
mod tests {
    use Index;
    use schema::{SchemaBuilder, TEXT, INT_INDEXED};

    #[test]
    fn test_indexer_for_field() {
        let mut schema_builder = SchemaBuilder::default();
        let num_likes_field = schema_builder.add_u64_field("num_likes", INT_INDEXED);
        let body_field = schema_builder.add_text_field("body", TEXT);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        assert!(index.tokenizer_for_field(body_field).is_ok());
        assert_eq!(
            format!("{:?}", index.tokenizer_for_field(num_likes_field).err()),
            "Some(SchemaError(\"\\\"num_likes\\\" is not a text field.\"))"
        );
    }

}<|endoftext|>


<commit_before>use sodiumoxide::crypto::stream;
use blob::{ChunkRef, Kind};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

pub struct PlainText(Vec<u8>);
pub struct CipherText(Vec<u8>);
pub struct CipherTextRef<'a>(&'a [u8]);

impl PlainText {
    pub fn new(bytes: Vec<u8>) -> PlainText {
        PlainText(bytes)
    }
    pub fn from_vec(mut bytes: &mut Vec<u8>) -> PlainText {
        let mut pt = PlainText(Vec::with_capacity(bytes.len()));
        pt.0.append(&mut bytes);
        pt
    }
    pub fn append(&mut self, mut other: PlainText) {
        self.0.append(&mut other.0);
    }
    pub fn len(&self) -> usize {
        self.0.len()
    }
    pub fn into_vec(self) -> Vec<u8> {
        self.0
    }
    pub fn as_ref(&self) -> &[u8] {
        &self.0[..]
    }
}

impl CipherText {
    pub fn new(ct: Vec<u8>) -> CipherText {
        CipherText(ct)
    }
    pub fn from(mut other_ct: &mut CipherText) -> CipherText {
        let mut ct = CipherText::new(Vec::with_capacity(other_ct.len()));
        ct.append(&mut other_ct);
        ct
    }
    pub fn append(&mut self, other: &mut CipherText) {
        self.0.append(&mut other.0);
    }
    pub fn len(&self) -> usize {
        self.0.len()
    }
    pub fn empty_into(&mut self, out: &mut CipherText) {
        out.0.append(&mut self.0)
    }
    pub fn into_vec(self) -> Vec<u8> {
        self.0
    }
    pub fn random_pad_upto(&mut self, final_size: usize) {
        let size = self.len();
        if final_size > size {
            self.0.resize(final_size, 0);
            let key = stream::salsa20::gen_key();
            let nonce = stream::salsa20::gen_nonce();
            stream::salsa20::stream_xor_inplace(&mut self.0[size..], &nonce, &key);
        }
    }
    pub fn as_ref(&self) -> CipherTextRef {
        CipherTextRef(self.0.as_ref())
    }
}

impl<'a> CipherTextRef<'a> {
    pub fn new(bytes: &'a [u8]) -> CipherTextRef<'a> {
        CipherTextRef(bytes)
    }
    pub fn slice(&self, from: usize, to: usize) -> CipherTextRef<'a> {
        CipherTextRef(&self.0[from..to])
    }
    pub fn len(&self) -> usize {
        self.0.len()
    }
    pub fn split_from_right(&self, len: usize) -> (CipherTextRef<'a>, CipherTextRef<'a>) {
        assert!(self.len() >= len);
        (self.slice(0, self.len() - len), self.slice(self.len() - len, self.len()))
    }
}

pub struct RefKey {}

impl RefKey {
    pub fn add_sealed_size(size: usize) -> usize {
        size + 0  // TODO: should be len of seal(.., vec[0; size])
    }
    pub fn seal(cref: &mut ChunkRef, pt: PlainText) -> CipherText {
        // TODO(jos): WIP: Plug in encryption/crypto here.
        // Update cref with key.
        cref.length = pt.len();
        CipherText(pt.0)
    }

    pub fn unseal(cref: &ChunkRef, ct: CipherTextRef) -> PlainText {
        PlainText(ct.slice(cref.offset, cref.offset + cref.length).0.to_vec())
    }
}

pub struct FixedKey {
    key: Vec<u8>,
}

impl FixedKey {
    pub fn new(key: Vec<u8>) -> FixedKey {
        FixedKey { key: key }
    }

    pub fn seal(&self, pt: PlainText) -> CipherText {
        // TODO(jos): WIP: Plug in encryption/crypto here.
        // Seal with our fixed key.
        let nonce = stream::xsalsa20::gen_nonce();
        let mut ct = CipherText(pt.0);

        // Serialize ciphertext length.
        let mut ct_len = vec![];
        ct_len.write_u64::<LittleEndian>(ct.len() as u64).unwrap();
        assert_eq!(ct_len.len(), 8);

        // Add nonce and ciphertext length.
        ct.0.extend_from_slice(nonce.as_ref());
        ct.0.extend_from_slice(&ct_len[..]);
        // Add version.
        ct.0.extend_from_slice(&[1]);
        ct
    }

    fn version<'a>(ct: CipherTextRef<'a>) -> (CipherTextRef<'a>, CipherTextRef<'a>) {
        let (rest, version) = ct.split_from_right(1);
        assert_eq!(1, version.0[0]);
        (rest, version)
    }

    fn ciphertext_and_nonce<'a>(ct: CipherTextRef<'a>) -> (CipherTextRef<'a>, CipherTextRef<'a>, &'a [u8]) {
        let (rest, mut ct_len_ref) = ct.split_from_right(8);
        let ct_len = ct_len_ref.0.read_u64::<LittleEndian>().unwrap();
        let (rest, nonce) = rest.split_from_right(stream::xsalsa20::NONCEBYTES);
        let (rest, ct) = rest.split_from_right(ct_len as usize);
        (rest, ct, nonce.0)
    }

    pub fn unseal<'a, 'b>(&'a self, ct: CipherTextRef<'b>) -> (CipherTextRef<'b>, PlainText) {
        // TODO(jos): WIP: Plug in encryption/crypto here.
        // Unseal with our fixed key.

        // Recover version with size of ciphertext.
        let (rest, _) = Self::version(ct);
        let (rest, ct, nonce) = Self::ciphertext_and_nonce(rest);
        let nonce = stream::xsalsa20::Nonce::from_slice(nonce);
        // TODO(jos): WIP: Use nonce to decrypt ct.
        drop(nonce);

        (rest, PlainText(ct.0.to_vec()))
    }
}
<commit_msg>Rustfmt.<commit_after>use sodiumoxide::crypto::stream;
use blob::{ChunkRef, Kind};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

pub struct PlainText(Vec<u8>);
pub struct CipherText(Vec<u8>);
pub struct CipherTextRef<'a>(&'a [u8]);

impl PlainText {
    pub fn new(bytes: Vec<u8>) -> PlainText {
        PlainText(bytes)
    }
    pub fn from_vec(mut bytes: &mut Vec<u8>) -> PlainText {
        let mut pt = PlainText(Vec::with_capacity(bytes.len()));
        pt.0.append(&mut bytes);
        pt
    }
    pub fn append(&mut self, mut other: PlainText) {
        self.0.append(&mut other.0);
    }
    pub fn len(&self) -> usize {
        self.0.len()
    }
    pub fn into_vec(self) -> Vec<u8> {
        self.0
    }
    pub fn as_ref(&self) -> &[u8] {
        &self.0[..]
    }
}

impl CipherText {
    pub fn new(ct: Vec<u8>) -> CipherText {
        CipherText(ct)
    }
    pub fn from(mut other_ct: &mut CipherText) -> CipherText {
        let mut ct = CipherText::new(Vec::with_capacity(other_ct.len()));
        ct.append(&mut other_ct);
        ct
    }
    pub fn append(&mut self, other: &mut CipherText) {
        self.0.append(&mut other.0);
    }
    pub fn len(&self) -> usize {
        self.0.len()
    }
    pub fn empty_into(&mut self, out: &mut CipherText) {
        out.0.append(&mut self.0)
    }
    pub fn into_vec(self) -> Vec<u8> {
        self.0
    }
    pub fn random_pad_upto(&mut self, final_size: usize) {
        let size = self.len();
        if final_size > size {
            self.0.resize(final_size, 0);
            let key = stream::salsa20::gen_key();
            let nonce = stream::salsa20::gen_nonce();
            stream::salsa20::stream_xor_inplace(&mut self.0[size..], &nonce, &key);
        }
    }
    pub fn as_ref(&self) -> CipherTextRef {
        CipherTextRef(self.0.as_ref())
    }
}

impl<'a> CipherTextRef<'a> {
    pub fn new(bytes: &'a [u8]) -> CipherTextRef<'a> {
        CipherTextRef(bytes)
    }
    pub fn slice(&self, from: usize, to: usize) -> CipherTextRef<'a> {
        CipherTextRef(&self.0[from..to])
    }
    pub fn len(&self) -> usize {
        self.0.len()
    }
    pub fn split_from_right(&self, len: usize) -> (CipherTextRef<'a>, CipherTextRef<'a>) {
        assert!(self.len() >= len);
        (self.slice(0, self.len() - len), self.slice(self.len() - len, self.len()))
    }
}

pub struct RefKey {}

impl RefKey {
    pub fn add_sealed_size(size: usize) -> usize {
        size + 0  // TODO: should be len of seal(.., vec[0; size])
    }
    pub fn seal(cref: &mut ChunkRef, pt: PlainText) -> CipherText {
        // TODO(jos): WIP: Plug in encryption/crypto here.
        // Update cref with key.
        cref.length = pt.len();
        CipherText(pt.0)
    }

    pub fn unseal(cref: &ChunkRef, ct: CipherTextRef) -> PlainText {
        PlainText(ct.slice(cref.offset, cref.offset + cref.length).0.to_vec())
    }
}

pub struct FixedKey {
    key: Vec<u8>,
}

impl FixedKey {
    pub fn new(key: Vec<u8>) -> FixedKey {
        FixedKey { key: key }
    }

    pub fn seal(&self, pt: PlainText) -> CipherText {
        // TODO(jos): WIP: Plug in encryption/crypto here.
        // Seal with our fixed key.
        let nonce = stream::xsalsa20::gen_nonce();
        let mut ct = CipherText(pt.0);

        // Serialize ciphertext length.
        let mut ct_len = vec![];
        ct_len.write_u64::<LittleEndian>(ct.len() as u64).unwrap();
        assert_eq!(ct_len.len(), 8);

        // Add nonce and ciphertext length.
        ct.0.extend_from_slice(nonce.as_ref());
        ct.0.extend_from_slice(&ct_len[..]);
        // Add version.
        ct.0.extend_from_slice(&[1]);
        ct
    }

    fn version<'a>(ct: CipherTextRef<'a>) -> (CipherTextRef<'a>, CipherTextRef<'a>) {
        let (rest, version) = ct.split_from_right(1);
        assert_eq!(1, version.0[0]);
        (rest, version)
    }

    fn ciphertext_and_nonce<'a>(ct: CipherTextRef<'a>)
                                -> (CipherTextRef<'a>, CipherTextRef<'a>, &'a [u8]) {
        let (rest, mut ct_len_ref) = ct.split_from_right(8);
        let ct_len = ct_len_ref.0.read_u64::<LittleEndian>().unwrap();
        let (rest, nonce) = rest.split_from_right(stream::xsalsa20::NONCEBYTES);
        let (rest, ct) = rest.split_from_right(ct_len as usize);
        (rest, ct, nonce.0)
    }

    pub fn unseal<'a, 'b>(&'a self, ct: CipherTextRef<'b>) -> (CipherTextRef<'b>, PlainText) {
        // TODO(jos): WIP: Plug in encryption/crypto here.
        // Unseal with our fixed key.

        // Recover version with size of ciphertext.
        let (rest, _) = Self::version(ct);
        let (rest, ct, nonce) = Self::ciphertext_and_nonce(rest);
        let nonce = stream::xsalsa20::Nonce::from_slice(nonce);
        // TODO(jos): WIP: Use nonce to decrypt ct.
        drop(nonce);

        (rest, PlainText(ct.0.to_vec()))
    }
}
<|endoftext|>


<commit_before>// TODO: get num cols of terminal dynamically
// TODO: care about the given print options

use super::{Report, ReportKind, RemarkKind, Snippet};
use code::{FileMap, LineIdx};
use term_painter::ToStyle;
use term_painter::Color::*;
use std::default::Default;

/// Options for printing on the terminal. By `default()` everything is enabled.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct PrintOptions {
    /// Use of unicode allowed?
    pub unicode: bool,
    /// Use of colors allowed?
    pub color: bool,
    /// Is line wrapping allowed?
    pub line_wrap: bool,
}

impl Default for PrintOptions {
    fn default() -> Self {
        PrintOptions {
            unicode: true,
            color: true,
            line_wrap: true,
        }
    }
}

/// Pretty prints a report
///
/// **Note**: right now, the `PrintOptions` are ignored.
pub fn print(rep: &Report, src: &FileMap, _: PrintOptions) {
    trace!("Printing report: {:#?}", rep);
    trace!("Printing with filemap: {:#?}", src);

    // print header
    let title = match rep.kind {
        ReportKind::Error => White.bold().bg(Red).paint("ERROR"),
        ReportKind::Warning => White.bold().bg(Yellow).paint("WARNING"),
    };

    let (sep, line) = if let Some(span) = rep.span {
        (" : ", if span.is_dummy() {
            "<dummy-span>".into()
        } else {
            let start = src.get_loc(span.lo);
            let end = src.get_loc(span.hi);
            trace!("Span is from {:?} to {:?}", start, end);

            if start.line != end.line {
                format!("{}-{}", start.line, end.line)
            } else {
                start.line.to_string()
            }
        })
    } else {
        ("", "".into())
    };

    println!("+---- {} in {}{}{} ----+",
        title,
        src.filename(),
        sep,
        Magenta.bold().paint(line)
    );

    for rem in &rep.remarks {
        trace!("Handling Remark {:?}", rem);

        // print message
        let (title, title_len) = match rem.kind {
            RemarkKind::Error => (Red.paint("error:"), 6),
            RemarkKind::Warning => (Yellow.paint("warning:"), 8),
            RemarkKind::Note => (Green.paint("note:"), 5),
        };

        print!("      =====>  {} ", title);
        // spaces + big arrow + spaces + title + space
        let indent = 6 + 6 + 2 + title_len + 1;
        let block_width = 80 - indent;

        let mut col = 0;
        for word in rem.desc.split_whitespace() {
            let word_len = word.chars().count();
            if col + word_len >= block_width && col != 0 {
                println!("");
                print!("           >  {0:>1$} ", " ", title_len);
                col = 0;
            }
            print!("{} ", White.bold().paint(word));
            col += word_len + 1;
        }
        println!("");

        // print code snippet
        match rem.snippet {
            Snippet::Orig(_) | Snippet::Replace { .. } => {
                print_snippet(src, &rem.snippet);
                println!("");
            },
            _ => {},
        }
    }
    println!("");
}

fn print_snippet(src: &FileMap, snippet: &Snippet) {
    let span = match *snippet {
        Snippet::Orig(span) => span,
        Snippet::Replace { span, .. } => span,
        _ => unreachable!(),
    };

    let start = src.get_loc(span.lo);
    let end = src.get_loc(span.hi);
    trace!("Span is from {:?} to {:?}", start, end);

    // ----- Dummyspan -----
    if span.is_dummy() {
        println!("   {} {} ! no snippet due to <dummy-span>, this is a bug !",
            Magenta.bold().paint("?"),
            Magenta.bold().paint("|"),
        );
    }

    // ----- Singleline -----
    else if start.line == end.line {
        let line_orig = expect_line(src, start.line);
        trace!("Printing single line span. Orig line: {:?}", line_orig);

        // let mut line = line_orig;
        // // replace string if requested
        // if let Snippet::Replace { replace_span, replace_with, .. } = snippet {

        // }

        // replace tabs
        let line = line_orig.replace("\t", "    ");
        let num_tabs = line_orig[..start.col.0 as usize]
            .chars()
            .filter(|&c| c == '\t')
            .count();

        // adjust cols in case of replaced tabs
        let startcol = start.col.0 as usize + 3*num_tabs;
        let endcol = end.col.0 as usize + 3*num_tabs;

        print!(
            "{:>#4} {} ",
            Magenta.bold().paint(start.line),
            Magenta.bold().paint("|")
        );
        if let &Snippet::Replace { span, ref with } = snippet {
            println!("{}{}{}",
                &line[..startcol],
                Green.paint(with),
                &line[endcol..],
            );

            Green.with(|| {
                println!("      {: <2$}{:^<3$}",
                    " ", "^",
                    startcol + 1,
                    with.len(),
                );
            });
        } else {
            // print line (with marked span section)
            println!("{}{}{}",
                &line[..startcol],
                Yellow.paint(&line[startcol..endcol]),
                &line[endcol..],
            );

            Yellow.with(|| {
                println!("      {: <2$}{:^<3$}",
                    " ", "^",
                    startcol + 1,
                    endcol - startcol,
                );
            });
        }
    }

    // ----- Multiline -----
    else {
        // print first line
        if let Some(line) = src.get_line(start.line) {
            let startcol = start.col.0 as usize;
            println!("{:>#4} {} {}{}",
                Magenta.bold().paint(start.line),
                Magenta.bold().paint("|"),
                &line[..startcol],
                Yellow.paint(&line[startcol..]),
            );
        }

        // print all lines that are completely in the span
        for line_idx in (start.line.0 + 1)..end.line.0 {
            let line_idx = LineIdx(line_idx);
            if let Some(line) = src.get_line(line_idx) {
                println!("{:>#4} {} {}",
                    Magenta.bold().paint(line_idx),
                    Magenta.bold().paint("|"),
                    line
                );
            }
        }

        // print last line
        if let Some(line) = src.get_line(end.line) {
            let endcol = end.col.0 as usize;
            println!("{:>#4} {} {}{}",
                Magenta.bold().paint(end.line),
                Magenta.bold().paint("|"),
                Yellow.paint(&line[..endcol]),
                &line[endcol..],
            );
        }
    }
}

fn expect_line(src: &FileMap, loc: Loc) -> &str {
    src.get_line(start.line)
        .expect("`Loc` from FileMap should return a valid line")
}
<commit_msg>Cleanup single-line print code<commit_after>// TODO: get num cols of terminal dynamically
// TODO: care about the given print options

use super::{Report, ReportKind, RemarkKind, Snippet};
use code::{FileMap, LineIdx};
use term_painter::ToStyle;
use term_painter::Color::*;
use std::default::Default;

/// Options for printing on the terminal. By `default()` everything is enabled.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct PrintOptions {
    /// Use of unicode allowed?
    pub unicode: bool,
    /// Use of colors allowed?
    pub color: bool,
    /// Is line wrapping allowed?
    pub line_wrap: bool,
}

impl Default for PrintOptions {
    fn default() -> Self {
        PrintOptions {
            unicode: true,
            color: true,
            line_wrap: true,
        }
    }
}

/// Pretty prints a report
///
/// **Note**: right now, the `PrintOptions` are ignored.
pub fn print(rep: &Report, src: &FileMap, _: PrintOptions) {
    trace!("Printing report: {:#?}", rep);
    trace!("Printing with filemap: {:#?}", src);

    // print header
    let title = match rep.kind {
        ReportKind::Error => White.bold().bg(Red).paint("ERROR"),
        ReportKind::Warning => White.bold().bg(Yellow).paint("WARNING"),
    };

    let (sep, line) = if let Some(span) = rep.span {
        (" : ", if span.is_dummy() {
            "<dummy-span>".into()
        } else {
            let start = src.get_loc(span.lo);
            let end = src.get_loc(span.hi);
            trace!("Span is from {:?} to {:?}", start, end);

            if start.line != end.line {
                format!("{}-{}", start.line, end.line)
            } else {
                start.line.to_string()
            }
        })
    } else {
        ("", "".into())
    };

    println!("+---- {} in {}{}{} ----+",
        title,
        src.filename(),
        sep,
        Magenta.bold().paint(line)
    );

    for rem in &rep.remarks {
        trace!("Handling Remark {:?}", rem);

        // print message
        let (title, title_len) = match rem.kind {
            RemarkKind::Error => (Red.paint("error:"), 6),
            RemarkKind::Warning => (Yellow.paint("warning:"), 8),
            RemarkKind::Note => (Green.paint("note:"), 5),
        };

        print!("      =====>  {} ", title);
        // spaces + big arrow + spaces + title + space
        let indent = 6 + 6 + 2 + title_len + 1;
        let block_width = 80 - indent;

        let mut col = 0;
        for word in rem.desc.split_whitespace() {
            let word_len = word.chars().count();
            if col + word_len >= block_width && col != 0 {
                println!("");
                print!("           >  {0:>1$} ", " ", title_len);
                col = 0;
            }
            print!("{} ", White.bold().paint(word));
            col += word_len + 1;
        }
        println!("");

        // print code snippet
        match rem.snippet {
            Snippet::Orig(_) | Snippet::Replace { .. } => {
                print_snippet(src, &rem.snippet);
                println!("");
            },
            _ => {},
        }
    }
    println!("");
}

fn print_snippet(src: &FileMap, snippet: &Snippet) {
    let span = match *snippet {
        Snippet::Orig(span) => span,
        Snippet::Replace { span, .. } => span,
        _ => unreachable!(),
    };

    let start = src.get_loc(span.lo);
    let end = src.get_loc(span.hi);
    trace!("Span is from {:?} to {:?}", start, end);

    // ----- Dummyspan -----
    if span.is_dummy() {
        println!("   {} {} ! no snippet due to <dummy-span>, this is a bug !",
            Magenta.bold().paint("?"),
            Magenta.bold().paint("|"),
        );
    }

    // ----- Singleline -----
    else if start.line == end.line {
        let line_orig = expect_line(src, start.line);
        trace!("Printing single line span. Orig line: {:?}", line_orig);

        // replace tabs
        let line = line_orig.replace("\t", "    ");
        let num_tabs = line_orig[..start.col.0 as usize]
            .chars()
            .filter(|&c| c == '\t')
            .count();

        // adjust cols in case of replaced tabs
        let startcol = start.col.0 as usize + 3*num_tabs;
        let endcol = end.col.0 as usize + 3*num_tabs;

        let (middle, underline_len, color) = match *snippet {
            Snippet::Replace { ref with, .. } => {
                (&with[..], with.len(), Green)
            },
            Snippet::Orig(_) =>  {
                (&line[startcol..endcol], endcol - startcol, Yellow)
            },
            _ => unreachable!(),
        };

        // print the line
        println!("{:>#4} {} {}{}{}",
            Magenta.bold().paint(start.line),
            Magenta.bold().paint("|"),
            &line[..startcol],
            color.paint(middle),
            &line[endcol..],
        );

        // print the underline
        color.with(|| {
            println!("      {: <2$}{:^<3$}",
                " ", "^",
                startcol + 1,
                underline_len,
            );
        });
    }

    // ----- Multiline -----
    else {
        // print first line
        if let Some(line) = src.get_line(start.line) {
            let startcol = start.col.0 as usize;
            println!("{:>#4} {} {}{}",
                Magenta.bold().paint(start.line),
                Magenta.bold().paint("|"),
                &line[..startcol],
                Yellow.paint(&line[startcol..]),
            );
        }

        // print all lines that are completely in the span
        for line_idx in (start.line.0 + 1)..end.line.0 {
            let line_idx = LineIdx(line_idx);
            if let Some(line) = src.get_line(line_idx) {
                println!("{:>#4} {} {}",
                    Magenta.bold().paint(line_idx),
                    Magenta.bold().paint("|"),
                    line
                );
            }
        }

        // print last line
        if let Some(line) = src.get_line(end.line) {
            let endcol = end.col.0 as usize;
            println!("{:>#4} {} {}{}",
                Magenta.bold().paint(end.line),
                Magenta.bold().paint("|"),
                Yellow.paint(&line[..endcol]),
                &line[endcol..],
            );
        }
    }
}

fn expect_line(src: &FileMap, line: LineIdx) -> &str {
    src.get_line(line).expect("`Loc` from FileMap should return a valid line")
}
<|endoftext|>


<commit_before>use std::sync::mpsc::{Receiver, Sender};
use std::rc::Rc;

use action::ExecResult;
use super::{config, Condition, Context, Event, Message, TimerEvent};
use reactor;

#[derive(Debug)]
pub enum Request {
    Event(Event),
    Exit
}

#[derive(Debug)]
pub enum Response {
    Event(ExecResult),
    Exit
}

pub struct Dispatcher {
    contexts: Vec<Context>,
    output_channel: Sender<Response>,
    exits_received: u32
}

impl Dispatcher {
    pub fn new(contexts: Vec<config::Context>, action_output_channel: Sender<Response>) -> Dispatcher {
        let contexts = contexts.into_iter().map(|ctx| Context::from(ctx)).collect::<Vec<Context>>();
        Dispatcher {
            contexts: contexts,
            output_channel: action_output_channel,
            exits_received: 0
        }
    }

    pub fn start_loop(&mut self, channel: Receiver<Request>) {
        for i in channel.iter() {
            match i {
                Request::Event(event) => self.dispatch(event),
                Request::Exit => {
                    if self.on_exit() {
                        break;
                    }
                }
            }
        }
    }

    fn on_exit(&mut self) -> bool {
        self.exits_received += 1;
        let _ = self.output_channel.send(Response::Exit);
        self.exits_received >= 2
    }

    pub fn dispatch(&mut self, event: Event) {
        match event {
            Event::Message(event) => {
                let event = Rc::new(event);
                self.on_message(event);
            },
            Event::Timer(ref event) => {
                self.on_timer(event);
            }
        };
    }

    fn on_message(&mut self, event: Rc<Message>) {
        for context in self.contexts.iter_mut() {
            if let Some(result) = context.on_message(event.clone()) {
                for i in result.into_iter() {
                    let _ = self.output_channel.send(i.into());
                }
            }
        }
    }

    fn on_timer(&mut self, event: &TimerEvent) {
        for context in self.contexts.iter_mut() {
            if let Some(result) = context.on_timer(event) {
                for i in result.into_iter() {
                    let _ = self.output_channel.send(i.into());
                }
            }
        }
    }
}

mod handlers {
    mod exit {
        use dispatcher::Request;
        use condition::Condition;
        use reactor::EventHandler;

        struct ExitHandler{
            condition: Condition,
            stops: u32
        }

        impl ExitHandler {
            fn new(condition: Condition) -> ExitHandler {
                ExitHandler {
                    condition: condition,
                    stops: 0
                }
            }
        }

        impl EventHandler<Request> for ExitHandler {
            fn handle_event(&mut self, event: Request) {
                if let Request::Exit = event {
                    self.stops += 1;

                    if self.stops >= 2 {
                        self.condition.activate();
                    }
                } else {
                    unreachable!("An ExitHandler should only receive Exit events");
                }
            }
        }
    }

    pub mod event {
        use dispatcher::Request;
        use reactor;

        pub struct EventHandler;

        impl reactor::EventHandler<Request> for EventHandler {
            fn handle_event(&mut self, event: Request) {
                if let Request::Event(_) = event {
                    println!("Event recvd");
                } else {
                    unreachable!("An EventHandler should only receive Event events");
                }
            }
        }
    }
}
<commit_msg>dispatcher: make ExitHandler public<commit_after>use std::sync::mpsc::{Receiver, Sender};
use std::rc::Rc;

use action::ExecResult;
use super::{config, Condition, Context, Event, Message, TimerEvent};
use reactor;

#[derive(Debug)]
pub enum Request {
    Event(Event),
    Exit
}

#[derive(Debug)]
pub enum Response {
    Event(ExecResult),
    Exit
}

pub struct Dispatcher {
    contexts: Vec<Context>,
    output_channel: Sender<Response>,
    exits_received: u32
}

impl Dispatcher {
    pub fn new(contexts: Vec<config::Context>, action_output_channel: Sender<Response>) -> Dispatcher {
        let contexts = contexts.into_iter().map(|ctx| Context::from(ctx)).collect::<Vec<Context>>();
        Dispatcher {
            contexts: contexts,
            output_channel: action_output_channel,
            exits_received: 0
        }
    }

    pub fn start_loop(&mut self, channel: Receiver<Request>) {
        for i in channel.iter() {
            match i {
                Request::Event(event) => self.dispatch(event),
                Request::Exit => {
                    if self.on_exit() {
                        break;
                    }
                }
            }
        }
    }

    fn on_exit(&mut self) -> bool {
        self.exits_received += 1;
        let _ = self.output_channel.send(Response::Exit);
        self.exits_received >= 2
    }

    pub fn dispatch(&mut self, event: Event) {
        match event {
            Event::Message(event) => {
                let event = Rc::new(event);
                self.on_message(event);
            },
            Event::Timer(ref event) => {
                self.on_timer(event);
            }
        };
    }

    fn on_message(&mut self, event: Rc<Message>) {
        for context in self.contexts.iter_mut() {
            if let Some(result) = context.on_message(event.clone()) {
                for i in result.into_iter() {
                    let _ = self.output_channel.send(i.into());
                }
            }
        }
    }

    fn on_timer(&mut self, event: &TimerEvent) {
        for context in self.contexts.iter_mut() {
            if let Some(result) = context.on_timer(event) {
                for i in result.into_iter() {
                    let _ = self.output_channel.send(i.into());
                }
            }
        }
    }
}

mod handlers {
    pub mod exit {
        use dispatcher::Request;
        use condition::Condition;
        use reactor::EventHandler;

        pub struct ExitHandler{
            condition: Condition,
            stops: u32
        }

        impl ExitHandler {
            pub fn new(condition: Condition) -> ExitHandler {
                ExitHandler {
                    condition: condition,
                    stops: 0
                }
            }
        }

        impl EventHandler<Request> for ExitHandler {
            fn handle_event(&mut self, event: Request) {
                if let Request::Exit = event {
                    self.stops += 1;

                    if self.stops >= 2 {
                        self.condition.activate();
                    }
                } else {
                    unreachable!("An ExitHandler should only receive Exit events");
                }
            }
        }
    }

    pub mod event {
        use dispatcher::Request;
        use reactor;

        pub struct EventHandler;

        impl reactor::EventHandler<Request> for EventHandler {
            fn handle_event(&mut self, event: Request) {
                if let Request::Event(_) = event {
                    println!("Event recvd");
                } else {
                    unreachable!("An EventHandler should only receive Event events");
                }
            }
        }
    }
}
<|endoftext|>


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

use ::engine::AppInfo;
use ::resource::ResourceManager;
use ::window::WindowManager;
use ::util::Logger;

use std::cell::RefCell;
use std::rc::Rc;
use std::boxed::Box;
use std::process;

/*===============================================================================================*/
/*------STATIC VARIABLES-------------------------------------------------------------------------*/
/*===============================================================================================*/

static mut APP_POINTER: Option <*mut App> = None;

/*===============================================================================================*/
/*------APP STRUCT-------------------------------------------------------------------------------*/
/*===============================================================================================*/

/// The app
///
/// This is the main control center of ionCore.
/// It is in charge of initialization, updating, and shutdown of all modules,
/// as well as the handing of any inter-module communication.
pub struct App {

    // Public
    /// The application info.
    pub app_info: AppInfo,
    /// The resource manager.
    pub resource_mgr: Rc<RefCell<ResourceManager>>,
    /// The window manager.
    pub window_mgr: Rc<RefCell<WindowManager>>,

    // Private
    _is_in_main_loop: bool,
    _should_exit: bool,
}

/*===============================================================================================*/
/*------APP PUBLIC METHODS-----------------------------------------------------------------------*/
/*===============================================================================================*/

impl App {

    /// The main app loop.
    pub fn run (&mut self) {

        self._is_in_main_loop = true;

        loop {

            let should_exit = self._should_exit;

            if !should_exit {

                self._on_pre_render ();
                self._on_render ();
                self._on_post_render ();
            }

            else {

                self._is_in_main_loop = false;
                return;
            }
        }
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Releases all resources, and exits the application.
    pub fn exit (&mut self) {

        // Check if in main loop
        if self._is_in_main_loop {

            info! ("Exiting main loop.");
            self._should_exit = true;
        }

        else {

            info! ("Shutting down ion Core.");

            // Release the managers, and release self
            self._release_managers ();
            App::_terminate ();
        }
    }

/*===============================================================================================*/
/*------APP PUBLIC STATIC METHODS----------------------------------------------------------------*/
/*===============================================================================================*/

    /// Initializes the app
    pub fn init (app_info: AppInfo) {

        // Check if not already initialized
        if !App::is_initialized () {

            Logger::init ("./ionCore.log", true).unwrap ();
            info! ("Initializing ionCore | Version: {}", env! ("CARGO_PKG_VERSION"));

            let ab = Box::new (App {

                app_info: app_info,
                resource_mgr: Rc::new (RefCell::new (ResourceManager::new ())),
                window_mgr: Rc::new (RefCell::new (WindowManager::new ())),
                _is_in_main_loop: false,
                _should_exit: false,
            });

            unsafe {APP_POINTER = Some (Box::into_raw (ab))};

            // Init the managers
            App::get_instance ().unwrap ()._init_managers ();
        }
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Checks if the app has been initialized.
    ///
    /// # Return value
    /// A bool returning whether the app has been initialized.
    pub fn is_initialized () -> bool {

        unsafe {

            match APP_POINTER {

                Some (pointer) => !pointer.is_null (),
                None => false
            }
        }
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Returns a reference to the app instance.
    pub fn get_instance () -> Result<&'static App, ()> {

        if App::is_initialized () {
            return Ok (unsafe {&*APP_POINTER.unwrap ()});
        }

        Err (())
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Returns a mutable reference to the app instance.
    pub fn get_instance_mut () -> Result<&'static mut App, ()> {

        if App::is_initialized () {
            return Ok (unsafe {&mut *APP_POINTER.unwrap ()});
        }

        Err (())
    }

/*===============================================================================================*/
/*------APP PRIVATE METHODS----------------------------------------------------------------------*/
/*===============================================================================================*/

    // Initializes the various managers.
    fn _init_managers (&self) {

        self.resource_mgr.borrow_mut ().init ();
        self.window_mgr.borrow_mut   ().init ();
    }

/*-----------------------------------------------------------------------------------------------*/

    // TODO: Finish me
    // On pre render
    fn _on_pre_render (&self) {
        self.window_mgr.borrow_mut ().on_pre_render ();
    }

/*-----------------------------------------------------------------------------------------------*/

    // TODO: Finish me
    // On render
    fn _on_render (&self) {
        self.window_mgr.borrow_mut ().on_render ();
    }

/*-----------------------------------------------------------------------------------------------*/

    // TODO: Finish me
    // On post render
    fn _on_post_render (&self) {
        self.window_mgr.borrow_mut ().on_post_render ();
    }

/*-----------------------------------------------------------------------------------------------*/

    // Releases the managers
    fn _release_managers (&self) {
        self.window_mgr.borrow_mut ().release ();
    }

/*===============================================================================================*/
/*------APP PRIVATE STATIC METHODS---------------------------------------------------------------*/
/*===============================================================================================*/

    // Release the app pointer, and terminate the application.
    fn _terminate () {

        unsafe {

            drop (Box::from_raw (APP_POINTER.unwrap ()));
            APP_POINTER = None;
        };

        // Release the logger, and shutdown the application
        info! ("Terminating the application.");
        Logger::release ();

        process::exit (0);
    }
}
<commit_msg>Converted app creation to use builder method<commit_after>/*===============================================================================================*/
// Copyright 2016 Kyle Finlay
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/*===============================================================================================*/

use ::engine::AppInfo;
use ::resource::ResourceManager;
use ::window::WindowManager;
use ::util::{Version, Logger};

use std::cell::RefCell;
use std::rc::Rc;
use std::boxed::Box;
use std::process;

/*===============================================================================================*/
/*------STATIC VARIABLES-------------------------------------------------------------------------*/
/*===============================================================================================*/

pub static mut APP_POINTER: Option <*mut App> = None;

/*===============================================================================================*/
/*------APP STRUCT-------------------------------------------------------------------------------*/
/*===============================================================================================*/

/// The app
///
/// This is the main control center of ionCore.
/// It is in charge of initialization, updating, and shutdown of all modules,
/// as well as the handing of any inter-module communication.
pub struct App {

    // Public
    /// The application info.
    pub app_info: AppInfo,
    /// The resource manager.
    pub resource_mgr: Rc<RefCell<ResourceManager>>,
    /// The window manager.
    pub window_mgr: Rc<RefCell<WindowManager>>,

    // Private
    _is_in_main_loop: bool,
    _should_exit: bool,
}

/*===============================================================================================*/
/*------APP PUBLIC METHODS-----------------------------------------------------------------------*/
/*===============================================================================================*/

impl App {

    /// The main app loop.
    pub fn run (&mut self) {

        self._is_in_main_loop = true;

        loop {

            let should_exit = self._should_exit;

            if !should_exit {

                self._on_pre_render ();
                self._on_render ();
                self._on_post_render ();
            }

            else {

                self._is_in_main_loop = false;
                return;
            }
        }
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Releases all resources, and exits the application.
    pub fn exit (&mut self) {

        // Check if in main loop
        if self._is_in_main_loop {

            info! ("Exiting main loop.");
            self._should_exit = true;
        }

        else {

            info! ("Shutting down ion Core.");

            // Release the managers, and release self
            self._release_managers ();
            App::_terminate ();
        }
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Initializes the app
    pub fn init (&self) {

        Logger::init ("./ionCore.log", true).unwrap ();
        info! ("Initializing ionCore | Version: {}", env! ("CARGO_PKG_VERSION"));

        // Init the managers
        self._init_managers ();
    }

/*===============================================================================================*/
/*------APP PUBLIC STATIC METHODS----------------------------------------------------------------*/
/*===============================================================================================*/

    /// Returns a new App Builder.
    pub fn builder () -> AppBuilder {
        AppBuilder::new ()
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Checks if the app has been initialized.
    ///
    /// # Return value
    /// A bool returning whether the app has been initialized.
    pub fn is_initialized () -> bool {

        unsafe {

            match APP_POINTER {

                Some (pointer) => !pointer.is_null (),
                None => false
            }
        }
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Returns a reference to the app instance.
    pub fn get_instance () -> Result<&'static App, ()> {

        if App::is_initialized () {
            return Ok (unsafe {&*APP_POINTER.unwrap ()});
        }

        Err (())
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Returns a mutable reference to the app instance.
    pub fn get_instance_mut () -> Result<&'static mut App, ()> {

        if App::is_initialized () {
            return Ok (unsafe {&mut *APP_POINTER.unwrap ()});
        }

        Err (())
    }

/*===============================================================================================*/
/*------APP PRIVATE METHODS----------------------------------------------------------------------*/
/*===============================================================================================*/

    // Initializes the various managers.
    fn _init_managers (&self) {

        self.resource_mgr.borrow_mut ().init ();
        self.window_mgr.borrow_mut   ().init ();
    }

/*-----------------------------------------------------------------------------------------------*/

    // TODO: Finish me
    // On pre render
    fn _on_pre_render (&self) {
        self.window_mgr.borrow_mut ().on_pre_render ();
    }

/*-----------------------------------------------------------------------------------------------*/

    // TODO: Finish me
    // On render
    fn _on_render (&self) {
        self.window_mgr.borrow_mut ().on_render ();
    }

/*-----------------------------------------------------------------------------------------------*/

    // TODO: Finish me
    // On post render
    fn _on_post_render (&self) {
        self.window_mgr.borrow_mut ().on_post_render ();
    }

/*-----------------------------------------------------------------------------------------------*/

    // Releases the managers
    fn _release_managers (&self) {
        self.window_mgr.borrow_mut ().release ();
    }

/*===============================================================================================*/
/*------APP PRIVATE STATIC METHODS---------------------------------------------------------------*/
/*===============================================================================================*/

    // Release the app pointer, and terminate the application.
    fn _terminate () {

        unsafe {

            drop (Box::from_raw (APP_POINTER.unwrap ()));
            APP_POINTER = None;
        };

        // Release the logger, and shutdown the application
        info! ("Terminating the application.");
        Logger::release ();

        process::exit (0);
    }
}

/*===============================================================================================*/
/*------APP BUILDER STRUCT-----------------------------------------------------------------------*/
/*===============================================================================================*/

/// Used for building a new, static app instance.
pub struct AppBuilder {

    // Private
    _app_name: String,
    _app_developer: String,
    _app_publisher: String,
    _app_version: Version,
}

/*===============================================================================================*/
/*------APP BUILDER PUBLIC METHODS---------------------------------------------------------------*/
/*===============================================================================================*/

impl AppBuilder {

    /// Sets the app name.
    pub fn name (&mut self, name: &str) -> &mut Self {

        self._app_name = name.to_string ();
        self
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Sets the app developer.
    pub fn developer (&mut self, developer: &str) -> &mut Self {

        self._app_developer = developer.to_string ();
        self
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Sets the app publisher
    pub fn publisher (&mut self, publisher: &str) -> &mut Self {

        self._app_publisher = publisher.to_string ();
        self
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Sets the app version
    pub fn version (&mut self, version: Version) -> &mut Self {

        self._app_version = version;
        self
    }

/*-----------------------------------------------------------------------------------------------*/

    /// Builds the app.
    pub fn build (&self) -> Result<&'static mut App, ()> {

        // Check if initialized
        if !App::is_initialized () {

            let app_info = AppInfo {

                app_name: self._app_name.clone (),
                app_developer: self._app_developer.clone (),
                app_publisher: self._app_publisher.clone (),
                app_version: self._app_version
            };

            let ab = Box::new (App {

                app_info: app_info,
                resource_mgr: Rc::new (RefCell::new (ResourceManager::new ())),
                window_mgr: Rc::new (RefCell::new (WindowManager::new ())),
                _is_in_main_loop: false,
                _should_exit: false,
            });

            unsafe {APP_POINTER = Some (Box::into_raw (ab))};
        }

        App::get_instance_mut ()
    }

/*===============================================================================================*/
/*------APP BUILDER PUBLIC STATIC METHODS--------------------------------------------------------*/
/*===============================================================================================*/

    /// Returns a new app builder instance.
    pub fn new () -> AppBuilder {

        AppBuilder {

            _app_name: "Unknown".to_string (),
            _app_developer: "Unknown".to_string (),
            _app_publisher: "Unknown".to_string (),
            _app_version: Version::new ()
        }
    }
}

/*-----------------------------------------------------------------------------------------------*/

impl Default for AppBuilder {

    fn default () -> AppBuilder {
        AppBuilder::new ()
    }
}
<|endoftext|>


<commit_before>pub mod search;

use std::thread;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
use basetypes::*;
use uci::{UciEngine, UciEngineFactory, EngineReply, OptionName, OptionDescription};
use position::Position;
use chess_move::*;
use tt::TranspositionTable;
use rand;
use rand::distributions::{Sample, Range};

const VERSION: &'static str = "0.1";

/// Implements `UciEngine` trait.
pub struct Engine {
    position: Position,
    replies: Vec<EngineReply>,
    is_thinking: bool,
    infinite: bool,
    ponder: bool,
    best_move: String,
    move_stack: Vec<Move>,
    tt: Arc<TranspositionTable>,
    commands: Sender<search::Command>,
    reports: Receiver<search::Progress>,
    results: Receiver<search::Done>,
    searcher: thread::JoinHandle<()>,
}

impl Engine {
    /// Creates a new instance.
    ///
    /// `tt_size_mb` is the preferred size of the transposition
    /// table in Mbytes.
    pub fn new(tt_size_mb: usize) -> Engine {
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        let (results_tx, results_rx) = channel();
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt = Arc::new(tt);
        Engine {
            position: Position::from_history("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk \
                                              - 0 1",
                                             &mut vec![].into_iter())
                          .ok()
                          .unwrap(),
            replies: vec![],
            is_thinking: false,
            infinite: false,
            ponder: false,
            best_move: "0000".to_string(),
            move_stack: Vec::new(),
            tt: tt.clone(),
            commands: commands_tx,
            reports: reports_rx,
            results: results_rx,
            searcher: thread::spawn(move || {
                search::run(&tt, commands_rx, reports_tx, results_tx);
            }),
        }
    }
}

impl UciEngine for Engine {
    #[allow(unused_variables)]
    fn set_option(&mut self, name: &str, value: &str) {
        self.replies.push(EngineReply::Info(vec![("string".to_string(),
                                                  format!("{} -> {}", name, value))]));
    }

    fn new_game(&mut self) {}

    #[allow(unused_variables)]
    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if self.is_thinking {
            return;
        }
        self.position = Position::from_history(fen, moves).ok().unwrap();
    }

    #[allow(unused_variables)]
    fn go(&mut self,
          searchmoves: Option<Vec<String>>,
          ponder: bool,
          wtime: Option<u64>,
          btime: Option<u64>,
          winc: Option<u64>,
          binc: Option<u64>,
          movestogo: Option<u64>,
          depth: Option<u64>,
          nodes: Option<u64>,
          mate: Option<u64>,
          movetime: Option<u64>,
          infinite: bool) {
        if self.is_thinking {
            return;
        }

        // Get some legal move
        let s = &mut self.move_stack;
        let p = &mut self.position;
        p.generate_moves(s);
        let mut rng = rand::thread_rng();
        let mut between = Range::new(0, 10);
        let mut best_score = -20000;
        while let Some(m) = s.pop() {
            if p.do_move(m) {
                let score = -p.evaluate_quiescence(-20000, 20000).0 + between.sample(&mut rng);
                if score >= best_score {
                    best_score = score;
                    // self.replies.push(EngineReply::Info(vec![("info".to_string(),
                    //                                           format!("{} -> {}",
                    //                                                   m.notation(),
                    //                                                   score))]));
                    self.best_move = m.notation();
                }
                p.undo_move();
            }
        }

        // Start a new search
        self.commands
            .send(search::Command::Search(search::Parameters {
                id: 0,
                position: p.clone(),
                depth: 5,
                lower_bound: -20000,
                upper_bound: 20000,
            }))
            .unwrap();
        self.ponder = ponder;
        self.infinite = infinite;
        self.is_thinking = true;
    }

    fn ponder_hit(&mut self) {
        if self.is_thinking && self.ponder {
            self.ponder = false;
        }
    }

    fn stop(&mut self) {
        if self.is_thinking {
            self.commands.send(search::Command::Stop).unwrap();
            self.ponder = false;
            self.infinite = false;
        }
    }

    fn is_thinking(&self) -> bool {
        self.is_thinking
    }

    fn get_reply(&mut self) -> Option<EngineReply> {
        if !self.ponder {
            if let Ok(_) = self.results.try_recv() {
                self.try_to_get_best_move_from_tt();
                self.replies.push(EngineReply::BestMove {
                    best_move: self.best_move.clone(),
                    ponder_move: None,
                });
                self.best_move = "0000".to_string();
                self.is_thinking = false;
            }
        }
        self.replies.pop()
    }
}

impl Engine {
    fn try_to_get_best_move_from_tt(&mut self) {
        if let Some(entry) = self.tt.probe(self.position.hash()) {
            let m = entry.move16();
            if m != 0 {
                let move_type = ((m & 0b1100000000000000) >> 14) as MoveType;
                let orig_square = ((m & 0b0011111100000000) >> 8) as Square;
                let dest_square = ((m & 0b0000000011111100) >> 2) as Square;
                let promoted_piece = match m & 0b11 {
                    0 => "q",
                    1 => "r",
                    2 => "b",
                    3 => "n",
                    _ => panic!("invalid promoted piece code"),
                };
                self.best_move = format!("{}{}{}",
                                         notation(orig_square),
                                         notation(dest_square),
                                         if move_type == MOVE_PROMOTION {
                                             promoted_piece
                                         } else {
                                             ""
                                         })
            }
        }
    }
}

/// Implements `UciEngineFactory` trait.
pub struct EngineFactory;

impl UciEngineFactory<Engine> for EngineFactory {
    fn name(&self) -> String {
        format!("Alcibiades {}", VERSION)
    }

    fn author(&self) -> String {
        "Evgeni Pandurski".to_string()
    }

    fn options(&self) -> Vec<(OptionName, OptionDescription)> {
        vec![
            // TODO: Calculate a sane limit for the hash size.
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 2048, default: 16 }),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
            ("MultiPV".to_string(), OptionDescription::Spin { min: 1, max: 32, default: 1 }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}
<commit_msg>more work on engine<commit_after>pub mod search;

use std::thread;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
use basetypes::*;
use uci::{UciEngine, UciEngineFactory, EngineReply, OptionName, OptionDescription};
use position::Position;
use chess_move::*;
use tt::TranspositionTable;
use rand;
use rand::distributions::{Sample, Range};

const VERSION: &'static str = "0.1";

/// Implements `UciEngine` trait.
pub struct Engine {
    position: Position,
    pondering_is_allowed: bool,
    multi_pv: usize,
    replies: Vec<EngineReply>,
    is_thinking: bool,
    infinite: bool,
    ponder: bool,
    best_move: String,
    move_stack: Vec<Move>,
    tt: Arc<TranspositionTable>,
    commands: Sender<search::Command>,
    reports: Receiver<search::Progress>,
    results: Receiver<search::Done>,
    searcher: thread::JoinHandle<()>,
}

impl Engine {
    /// Creates a new instance.
    ///
    /// `tt_size_mb` is the preferred size of the transposition
    /// table in Mbytes.
    pub fn new(tt_size_mb: usize) -> Engine {
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        let (results_tx, results_rx) = channel();
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt = Arc::new(tt);
        Engine {
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            pondering_is_allowed: false,
            multi_pv: 1,
            is_thinking: false,
            replies: vec![],
            infinite: false,
            ponder: false,
            best_move: "0000".to_string(),
            move_stack: Vec::new(),
            tt: tt.clone(),
            commands: commands_tx,
            reports: reports_rx,
            results: results_rx,
            searcher: thread::spawn(move || {
                search::run(&tt, commands_rx, reports_tx, results_tx);
            }),
        }
    }
}

impl UciEngine for Engine {
    #[allow(unused_variables)]
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            // We do not support re-sizing of the transposition table
            // once the engine had started, so we do nothing.
            "Hash" => (),

            // Tells the engine that it will be allowed to ponder.
            // This option is needed because the engine might change
            // its time management algorithm when pondering is
            // allowed.
            "Ponder" => {
                self.pondering_is_allowed = value == "true";
            }

            // Tells the engine to output multiple best lines.
            // the default value is `1`.)
            "MultiPV" => {
                self.multi_pv = value.parse::<usize>().ok().unwrap_or(1);
            }

            // An invalid option.
            _ => {
                self.replies.push(EngineReply::Info(vec![("string".to_string(),
                                                          format!("Invalid option \"{}\"", name))]));
            }
        }
    }

    fn new_game(&mut self) {
        // Clearing the transposition table would not change anything,
        // so we do nothing.
    }

    #[allow(unused_variables)]
    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if self.is_thinking {
            return;
        }
        if let Ok(p) = Position::from_history(fen, moves) {
            self.position = p;
        }
    }

    #[allow(unused_variables)]
    fn go(&mut self,
          searchmoves: Option<Vec<String>>,
          ponder: bool,
          wtime: Option<u64>,
          btime: Option<u64>,
          winc: Option<u64>,
          binc: Option<u64>,
          movestogo: Option<u64>,
          depth: Option<u64>,
          nodes: Option<u64>,
          mate: Option<u64>,
          movetime: Option<u64>,
          infinite: bool) {
        if self.is_thinking {
            return;
        }

        // Get some legal move
        let s = &mut self.move_stack;
        let p = &mut self.position;
        p.generate_moves(s);
        let mut rng = rand::thread_rng();
        let mut between = Range::new(0, 10);
        let mut best_score = -20000;
        while let Some(m) = s.pop() {
            if p.do_move(m) {
                let score = -p.evaluate_quiescence(-20000, 20000).0 + between.sample(&mut rng);
                if score >= best_score {
                    best_score = score;
                    // self.replies.push(EngineReply::Info(vec![("info".to_string(),
                    //                                           format!("{} -> {}",
                    //                                                   m.notation(),
                    //                                                   score))]));
                    self.best_move = m.notation();
                }
                p.undo_move();
            }
        }

        // Start a new search
        self.commands
            .send(search::Command::Search(search::Parameters {
                id: 0,
                position: p.clone(),
                depth: 5,
                lower_bound: -20000,
                upper_bound: 20000,
            }))
            .unwrap();
        self.ponder = ponder;
        self.infinite = infinite;
        self.is_thinking = true;
    }

    fn ponder_hit(&mut self) {
        if self.is_thinking && self.ponder {
            self.ponder = false;
        }
    }

    fn stop(&mut self) {
        if self.is_thinking {
            self.commands.send(search::Command::Stop).unwrap();
            self.ponder = false;
            self.infinite = false;
        }
    }

    fn is_thinking(&self) -> bool {
        self.is_thinking
    }

    fn get_reply(&mut self) -> Option<EngineReply> {
        if !self.ponder {
            if let Ok(_) = self.results.try_recv() {
                self.try_to_get_best_move_from_tt();
                self.replies.push(EngineReply::BestMove {
                    best_move: self.best_move.clone(),
                    ponder_move: None,
                });
                self.best_move = "0000".to_string();
                self.is_thinking = false;
            }
        }
        self.replies.pop()
    }
}

impl Engine {
    fn try_to_get_best_move_from_tt(&mut self) {
        if let Some(entry) = self.tt.probe(self.position.hash()) {
            let m = entry.move16();
            if m != 0 {
                let move_type = ((m & 0b1100000000000000) >> 14) as MoveType;
                let orig_square = ((m & 0b0011111100000000) >> 8) as Square;
                let dest_square = ((m & 0b0000000011111100) >> 2) as Square;
                let promoted_piece = match m & 0b11 {
                    0 => "q",
                    1 => "r",
                    2 => "b",
                    3 => "n",
                    _ => panic!("invalid promoted piece code"),
                };
                self.best_move = format!("{}{}{}",
                                         notation(orig_square),
                                         notation(dest_square),
                                         if move_type == MOVE_PROMOTION {
                                             promoted_piece
                                         } else {
                                             ""
                                         })
            }
        }
    }
}

/// Implements `UciEngineFactory` trait.
pub struct EngineFactory;

impl UciEngineFactory<Engine> for EngineFactory {
    fn name(&self) -> String {
        format!("Alcibiades {}", VERSION)
    }

    fn author(&self) -> String {
        "Evgeni Pandurski".to_string()
    }

    fn options(&self) -> Vec<(OptionName, OptionDescription)> {
        vec![
            // TODO: Calculate a sane limit for the hash size.
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 2048, default: 16 }),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
            ("MultiPV".to_string(), OptionDescription::Spin { min: 1, max: 32, default: 1 }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}
<|endoftext|>