<commit_before>use std::os::unix::io::RawFd;
use std::sync::RwLock;
use std::ptr;

use nix::sys::mman;

pub struct Pool {
    map: RwLock<MemMap>
}

pub enum ResizeError {
    InvalidSize,
    MremapFailed
}

impl Pool {
    pub fn new(fd: RawFd, size: usize) -> Result<Pool,()> {
        let memmap = MemMap::new(fd, size)?;
        Ok(Pool {
            map: RwLock::new(memmap)
        })
    }

    pub fn resize(&self, newsize: i32) -> Result<(),ResizeError> {
        let mut guard = self.map.write().unwrap();
        if newsize <= 0 || guard.size() > (newsize as usize) {
            return Err(ResizeError::InvalidSize)
        }
        guard.remap(newsize as usize).map_err(|()| ResizeError::MremapFailed)
    }

    pub fn with_data_slice<F: FnOnce(&[u8])>(&self, f: F) -> Result<(),()> {
        // TODO: handle SIGBUS
        let guard = self.map.read().unwrap();

        let slice = guard.get_slice();
        f(slice);

        Ok(())
    }
}

struct MemMap {
    ptr: *mut u8,
    fd: RawFd,
    size: usize
}

impl MemMap {
    fn new(fd: RawFd, size: usize) -> Result<MemMap,()> {
        Ok(MemMap {
            ptr: unsafe { map(fd, size) }?,
            fd: fd,
            size: size
        })
    }

    fn remap(&mut self, newsize: usize) -> Result<(),()> {
        if self.ptr.is_null() {
            return Err(())
        }
        // memunmap cannot fail, as we are unmapping a pre-existing map
        let _ = unsafe { unmap(self.ptr, self.size) };
        // remap the fd with the new size
        match unsafe { map(self.fd, newsize) } {
            Ok(ptr) => {
                // update the parameters
                self.ptr = ptr;
                self.size = newsize;
                Ok(())
            },
            Err(()) => {
                // set ourselves in an empty state
                self.ptr = ptr::null_mut();
                self.size = 0;
                self.fd = -1;
                Err(())
            }
        }
    }

    fn size(&self) -> usize {
        self.size
    }

    fn get_slice(&self) -> &[u8] {
        // if we are in the 'invalid state', self.size == 0 and we return &[]
        // which is perfectly safe even if self.ptr is null
        unsafe { ::std::slice::from_raw_parts(self.ptr, self.size) }
    }
}

impl Drop for MemMap {
    fn drop(&mut self) {
        if !self.ptr.is_null() {
            let _ = unsafe { unmap(self.ptr, self.size) };
        }
    }
}

// mman::mmap should really be unsafe... why isn't it?
unsafe fn map(fd: RawFd, size: usize) -> Result<*mut u8, ()> {
    let ret = mman::mmap(
        ptr::null_mut(),
        size,
        mman::PROT_READ,
        mman::MAP_SHARED,
        fd,
        0
    );
    ret.map(|p| p as *mut u8).map_err(|_| ())
}

// mman::munmap should really be unsafe... why isn't it?
unsafe fn unmap(ptr: *mut u8, size: usize) -> Result<(),()> {
    let ret = mman::munmap(ptr as *mut _, size);
    ret.map_err(|_| ())
}
<commit_msg>shm: SIGBUS handling machinery<commit_after>use std::cell::Cell;
use std::os::unix::io::RawFd;
use std::sync::{RwLock, Once, ONCE_INIT};
use std::ptr;

use nix::{c_int, c_void, libc};
use nix::sys::mman;
use nix::sys::signal::{self, SigAction, Signal, SigHandler};

thread_local!(static SIGBUS_GUARD: Cell<(*const MemMap, bool)> = Cell::new((ptr::null_mut(), false)));

static SIGBUS_INIT: Once = ONCE_INIT;
static mut OLD_SIGBUS_HANDLER: *mut SigAction = 0 as *mut SigAction;

pub struct Pool {
    map: RwLock<MemMap>
}

pub enum ResizeError {
    InvalidSize,
    MremapFailed
}

impl Pool {
    pub fn new(fd: RawFd, size: usize) -> Result<Pool,()> {
        let memmap = MemMap::new(fd, size)?;
        Ok(Pool {
            map: RwLock::new(memmap)
        })
    }

    pub fn resize(&self, newsize: i32) -> Result<(),ResizeError> {
        let mut guard = self.map.write().unwrap();
        if newsize <= 0 || guard.size() > (newsize as usize) {
            return Err(ResizeError::InvalidSize)
        }
        guard.remap(newsize as usize).map_err(|()| ResizeError::MremapFailed)
    }

    pub fn with_data_slice<F: FnOnce(&[u8])>(&self, f: F) -> Result<(),()> {
        // Place the sigbus handler
        SIGBUS_INIT.call_once(|| {
            unsafe { place_sigbus_handler(); }
        });

        let pool_guard = self.map.read().unwrap();

        // Prepare the access
        SIGBUS_GUARD.with(|guard| {
            let (p,_) = guard.get();
            if !p.is_null() {
                // Recursive call of this method is not supported
                panic!("Recursive access to a SHM pool content is not supported.");
            }
            guard.set((&*pool_guard as *const MemMap, false))
        });

        let slice = pool_guard.get_slice();
        f(slice);

        // Cleanup Post-access
        SIGBUS_GUARD.with(|guard| {
            let (_, triggered) = guard.get();
            guard.set((ptr::null_mut(), false));
            if triggered { Err(()) } else { Ok(()) }
        })
    }
}

struct MemMap {
    ptr: *mut u8,
    fd: RawFd,
    size: usize
}

impl MemMap {
    fn new(fd: RawFd, size: usize) -> Result<MemMap,()> {
        Ok(MemMap {
            ptr: unsafe { map(fd, size) }?,
            fd: fd,
            size: size
        })
    }

    fn remap(&mut self, newsize: usize) -> Result<(),()> {
        if self.ptr.is_null() {
            return Err(())
        }
        // memunmap cannot fail, as we are unmapping a pre-existing map
        let _ = unsafe { unmap(self.ptr, self.size) };
        // remap the fd with the new size
        match unsafe { map(self.fd, newsize) } {
            Ok(ptr) => {
                // update the parameters
                self.ptr = ptr;
                self.size = newsize;
                Ok(())
            },
            Err(()) => {
                // set ourselves in an empty state
                self.ptr = ptr::null_mut();
                self.size = 0;
                self.fd = -1;
                Err(())
            }
        }
    }

    fn size(&self) -> usize {
        self.size
    }

    fn get_slice(&self) -> &[u8] {
        // if we are in the 'invalid state', self.size == 0 and we return &[]
        // which is perfectly safe even if self.ptr is null
        unsafe { ::std::slice::from_raw_parts(self.ptr, self.size) }
    }

    fn contains(&self, ptr: *mut u8) -> bool {
        ptr >= self.ptr && ptr < unsafe { self.ptr.offset(self.size as isize) }
    }

    fn nullify(&self) -> Result<(),()> {
        unsafe { nullify_map(self.ptr, self.size) }
    }
}

impl Drop for MemMap {
    fn drop(&mut self) {
        if !self.ptr.is_null() {
            let _ = unsafe { unmap(self.ptr, self.size) };
        }
    }
}

// mman::mmap should really be unsafe... why isn't it?
unsafe fn map(fd: RawFd, size: usize) -> Result<*mut u8, ()> {
    let ret = mman::mmap(
        ptr::null_mut(),
        size,
        mman::PROT_READ,
        mman::MAP_SHARED,
        fd,
        0
    );
    ret.map(|p| p as *mut u8).map_err(|_| ())
}

// mman::munmap should really be unsafe... why isn't it?
unsafe fn unmap(ptr: *mut u8, size: usize) -> Result<(),()> {
    let ret = mman::munmap(ptr as *mut _, size);
    ret.map_err(|_| ())
}

unsafe fn nullify_map(ptr: *mut u8, size: usize) -> Result<(), ()> {
    let ret = mman::mmap(
        ptr as *mut _,
        size,
        mman::PROT_READ,
        mman::MAP_ANONYMOUS | mman::MAP_PRIVATE | mman::MAP_FIXED,
        -1,
        0
    );
    ret.map(|_| ()).map_err(|_| ())
}

unsafe fn place_sigbus_handler() {
    // create our sigbus handler
    let action = SigAction::new(
        SigHandler::SigAction(sigbus_handler),
        signal::SA_NODEFER,
        signal::SigSet::empty()
    );
    match signal::sigaction(Signal::SIGBUS, &action) {
        Ok(old_signal) => {
            OLD_SIGBUS_HANDLER = Box::into_raw(Box::new(old_signal));
        },
        Err(e) => {
            panic!("sigaction failed sor SIGBUS handler: {:?}", e)
        }
    }
}

unsafe fn reraise_sigbus() {
    // reset the old sigaction
    let _ = signal::sigaction(Signal::SIGBUS, &*OLD_SIGBUS_HANDLER);
    let _ = signal::raise(Signal::SIGBUS);
}

extern "C" fn sigbus_handler(_signum: c_int, info: *mut libc::siginfo_t, _context: *mut c_void) {
    let faulty_ptr = unsafe { siginfo_si_addr(info) } as *mut u8;
    SIGBUS_GUARD.with(|guard| {
        let (memmap, _) = guard.get();
        match unsafe { memmap.as_ref() }.map(|m| (m, m.contains(faulty_ptr))) {
            Some((m, true)) => {
                // we are in a faulty memory pool !
                // remember that it was faulty
                guard.set((memmap, true));
                // nullify the pool
                if m.nullify().is_err() {
                    // something terrible occured !
                    unsafe { reraise_sigbus() }
                }
            },
            _ => {
                // something else occured, let's die honorably
                unsafe { reraise_sigbus() }
            }
        }
    });
}

// This was shamelessly stolen from rustc's source
// so I expect it to work whevener rust works
// I guess it's good enough?

#[cfg(any(target_os = "linux", target_os = "android"))]
unsafe fn siginfo_si_addr(info: *mut libc::siginfo_t) -> *mut c_void {
	#[repr(C)]
	struct siginfo_t {
		a: [libc::c_int; 3], // si_signo, si_errno, si_code
		si_addr: *mut libc::c_void,
	}

	(*(info as *const siginfo_t)).si_addr
}

#[cfg(not(any(target_os = "linux", target_os = "android")))]
unsafe fn siginfo_si_addr(info: *mut libc::siginfo_t) -> *mut c_void {
	(*info).si_addr
}
<|endoftext|>


<commit_before>use std::ops::{Add, Sub, Div};
use std::collections::HashMap;
use std::iter::FromIterator;
use std::iter::repeat;
use std::fmt::{Debug, Formatter};
use std::fmt;
use std::cmp::Ordering::Equal;
use std::ops::Drop;
use std::cmp;
use std::path::Path;
use std::io::{BufWriter, Write, BufReader, Read, BufRead};
use std::io;
use std::fs::File;
use std::mem;

#[derive(Clone)]
pub struct WordVec {
    pub word: String,
    vec: Vec<f32>,
}

impl PartialEq for WordVec {
    fn eq(&self, other: &WordVec) -> bool {
        self.vec == other.vec
    }
}

impl Debug for WordVec {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        try!(self.word.fmt(fmt));
        try!(fmt.write_str(": "));
        try!(fmt.write_str(", "));
        self.vec.iter().take(5).collect::<Vec<_>>().fmt(fmt)
    }
}

impl WordVec {
    pub fn distance(&self, other: &WordVec) -> f32 {
        self.vec
            .iter()
            .zip(other.vec.iter())
            .map(|(x, y)| (x - y) * (x - y))
            .fold(0.0, |x, y| x + y)
            .sqrt()
    }

    fn new(word: String, num_words: usize) -> WordVec {
        WordVec {
            word: word,
            vec: repeat(0.0).take(num_words).collect(),
        }
    }

    fn cooccurences(&self) -> f32 {
        self.vec.iter().fold(0.0, |x, y| {
            assert!(!y.is_nan());
            x + *y
        })
    }

    fn normalize(&mut self) {
        for f in &mut self.vec {
            *f = (*f + 1.0).log2();
        }
    }

    fn inc(&mut self, i: usize, dist: usize) {
        assert!(dist != 0);
        self.vec[i] += 1.0 / dist as f32;
    }
}

impl Div<i32> for WordVec {
    type Output = WordVec;

    fn div(mut self, i: i32) -> WordVec {
        for f in &mut self.vec {
            *f /= i as f32;
        }

        self
    }
}

impl<'a> Add<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn add(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} + {}", &self.word, &other.word),
            vec: self.vec,
        };

        for i in 0..newvec.vec.len() {
            newvec.vec[i] += other.vec[i];
        }

        newvec
    }
}

impl<'a> Sub<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn sub(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} - {}", &self.word, &other.word),
            vec: self.vec,
        };

        for i in 0..(newvec.vec.len()) {
            newvec.vec[i] -= other.vec[i];
        }

        newvec
    }
}

#[derive(Clone, PartialEq, Debug)]
pub struct LanguageModel {
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
}

pub struct LanguageModelBuilder {
    window_radius: usize,
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
    sentence: Vec<Option<usize>>,
}

pub struct WordAcceptor<'a> {
    builder: &'a mut LanguageModelBuilder,
}

impl LanguageModelBuilder {
    pub fn new(window_radius: usize, words: Vec<String>) -> LanguageModelBuilder {

        let word_vecs = words.iter()
                             .map(|s| WordVec::new(s.clone(), words.len()))
                             .collect();

        let words = HashMap::from_iter(words.into_iter()
                                            .enumerate()
                                            .map(|(a, b)| (b, a)));

        LanguageModelBuilder {
            window_radius: window_radius,
            words: words,
            word_vecs: word_vecs,
            sentence: vec![],
        }
    }

    pub fn build(mut self) -> LanguageModel {
        for vec in self.word_vecs.iter_mut() {
            vec.normalize();
        }

        LanguageModel {
            words: self.words,
            word_vecs: self.word_vecs,
        }
    }

    pub fn new_sentence<'a>(&'a mut self) -> WordAcceptor<'a> {
        WordAcceptor { builder: self }
    }

    pub fn save(&self, path: &Path) -> io::Result<()> {
        let mut file = BufWriter::new(File::create(path).unwrap());

        try!(write_raw(self.word_vecs.len(), &mut file));
        try!(file.write(&[b'\n']));

        for vec in &self.word_vecs {
            try!(file.write(vec.word.as_bytes()));
            try!(file.write(&[b':']));
            for f in &vec.vec {
                try!(write_raw(*f, &mut file));
            }
            try!(file.write(&[b'\n']));
        }
        Ok(())
    }

    pub fn load(path: &Path) -> io::Result<LanguageModelBuilder> {
        let mut file = BufReader::new(File::open(path).unwrap());
        let mut word_vecs = Vec::new();

        let size: u64 = unsafe { try!(read_raw::<u64, _>(&mut file)) };
        try!(read_byte(b'\n', &mut file));

        for _ in 0..size {
            let mut word: Vec<u8> = Vec::new();
            file.read_until(b':', &mut word).unwrap();
            assert_eq!(Some(b':'), word.pop());
            let word = String::from_utf8(word).unwrap();

            let mut vec: Vec<f32> = repeat(0f32).take(size as usize).collect();
            for f in &mut vec {
                unsafe {
                    *f = try!(read_raw::<f32, _>(&mut file));
                }
            }

            try!(read_byte(b'\n', &mut file));

            word_vecs.push(WordVec {
                word: word,
                vec: vec,
            });
        }

        let words: HashMap<_, _> = word_vecs.iter()
                                            .map(|v| v.word.clone())
                                            .enumerate()
                                            .map(|(i, w)| (w, i))
                                            .collect();

        Ok(LanguageModelBuilder {
            window_radius: 0,
            words: words,
            word_vecs: word_vecs,
            sentence: Vec::new(),
        })
    }
}

impl<'a> WordAcceptor<'a> {
    pub fn add_word(&mut self, word: &str) {
        let idx_opt = self.builder.words.get(word).map(|w| *w);

        self.builder.sentence.push(idx_opt);
    }
}

impl<'a> Drop for WordAcceptor<'a> {
    fn drop(&mut self) {
        let LanguageModelBuilder {
            ref mut sentence,
            ref mut word_vecs,
            window_radius,
            ..
        } = *self.builder;

        for (i, &from_word) in sentence.iter().enumerate() {
            let start = cmp::max(0, i as isize - window_radius as isize) as usize;
            let end = cmp::min(sentence.len(), i + window_radius + 1);
            for j in (start..i).chain(i + 1..end) {
                let dist = cmp::max(i, j) - cmp::min(i, j);

                if let (Some(f), Some(t)) = (from_word, sentence[j]) {
                    word_vecs[f].inc(t, dist);
                }
            }
        }

        sentence.clear();
    }
}

impl LanguageModel {
    pub fn get(&self, word: &str) -> Option<&WordVec> {
        self.words.get(word).map(|i| &self.word_vecs[*i])
    }

    pub fn nearest_words(&self, word: &WordVec) -> Vec<&WordVec> {
        let mut vec_refs = self.word_vecs
                               .iter()
                               .filter_map(|w| {
                                   let dist = w.distance(word);
                                   if w.word != word.word {
                                       Some((dist, w))
                                   } else {
                                       None
                                   }
                               })
                               .collect::<Vec<_>>();
        vec_refs.sort_by(|a, b| {
            a.0.partial_cmp(&b.0).unwrap_or_else(|| {
                println!("bad vector: {:?}, {:?}", a, b);
                Equal
            })
        });
        vec_refs.into_iter().map(|x| x.1).collect()
    }
}

fn read_byte<R: Read>(b: u8, read: &mut R) -> io::Result<()> {
    let buf = &mut [0];
    while try!(read.read(buf)) == 0 {}
    assert_eq!(&[b], buf);
    Ok(())
}

unsafe fn read_raw<T: Copy, R: Read>(reader: &mut BufReader<R>) -> io::Result<T> {
    let mut buffer = [0u8; 64];
    let t_size = mem::size_of::<T>();
    assert!(t_size <= buffer.len());

    let mut remainder = t_size;
    while remainder > 0 {
        let bytes_read = try!(reader.read(&mut buffer[(t_size - remainder)..t_size]));
        remainder -= bytes_read;
    }

    let bptr: *mut T = mem::transmute(buffer.as_ptr());
    Ok(*bptr)
}

fn write_raw<T: Copy, F: Write>(t: T, writer: &mut F) -> io::Result<usize> {
    let buffer = [0u8; 64];
    let t_size = mem::size_of::<T>();
    assert!(t_size <= buffer.len());
    unsafe {
        let bptr: *mut T = mem::transmute(buffer.as_ptr());
        *bptr = t;
    }

    let mut start = 0;
    while start != t_size {
        start += try!(writer.write(&buffer[start..t_size]));
    }
    Ok((start))
}

#[cfg(test)]
mod test {
    use super::LanguageModelBuilder;
    use std::path::Path;

    fn get_builder() -> LanguageModelBuilder {
        let words = "foo bar baz blort".split(" ").map(|w| w.to_string()).collect::<Vec<_>>();
        let mut builder = LanguageModelBuilder::new(1, words);

        let input = "x foo bar baz x x x x x x x blort".split(" ");

        {
            let mut acc = builder.new_sentence();
            for word in input {
                acc.add_word(word);
            }
        }

        builder
    }

    #[test]
    fn test_accept_sentences() {
        let model = get_builder().build();
        let foo = model.get("foo").unwrap();
        let baz = model.get("baz").unwrap();
        let bar = model.get("bar").unwrap();
        let blort = model.get("blort").unwrap();

        assert!(foo.distance(baz) < foo.distance(blort));
        assert!(foo.distance(bar) == bar.distance(baz));
    }

    #[test]
    fn test_serialization() {
        let builder = get_builder();
        let path = Path::new("/tmp/model.data");
        builder.save(&path).unwrap();

        let loaded_model = LanguageModelBuilder::load(path).unwrap().build();
        assert_eq!(builder.build(), loaded_model);
    }
}
<commit_msg>Allow this method to be dead since I'l playing around.<commit_after>use std::ops::{Add, Sub, Div};
use std::collections::HashMap;
use std::iter::FromIterator;
use std::iter::repeat;
use std::fmt::{Debug, Formatter};
use std::fmt;
use std::cmp::Ordering::Equal;
use std::ops::Drop;
use std::cmp;
use std::path::Path;
use std::io::{BufWriter, Write, BufReader, Read, BufRead};
use std::io;
use std::fs::File;
use std::mem;

#[derive(Clone)]
pub struct WordVec {
    pub word: String,
    vec: Vec<f32>,
}

impl PartialEq for WordVec {
    fn eq(&self, other: &WordVec) -> bool {
        self.vec == other.vec
    }
}

impl Debug for WordVec {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        try!(self.word.fmt(fmt));
        try!(fmt.write_str(": "));
        try!(fmt.write_str(", "));
        self.vec.iter().take(5).collect::<Vec<_>>().fmt(fmt)
    }
}

impl WordVec {
    pub fn distance(&self, other: &WordVec) -> f32 {
        self.vec
            .iter()
            .zip(other.vec.iter())
            .map(|(x, y)| (x - y) * (x - y))
            .fold(0.0, |x, y| x + y)
            .sqrt()
    }

    fn new(word: String, num_words: usize) -> WordVec {
        WordVec {
            word: word,
            vec: repeat(0.0).take(num_words).collect(),
        }
    }

    #[allow(dead_code)]
    fn cooccurences(&self) -> f32 {
        self.vec.iter().fold(0.0, |x, y| {
            assert!(!y.is_nan());
            x + *y
        })
    }

    fn normalize(&mut self) {
        for f in &mut self.vec {
            *f = (*f + 1.0).log2();
        }
    }

    fn inc(&mut self, i: usize, dist: usize) {
        assert!(dist != 0);
        self.vec[i] += 1.0 / dist as f32;
    }
}

impl Div<i32> for WordVec {
    type Output = WordVec;

    fn div(mut self, i: i32) -> WordVec {
        for f in &mut self.vec {
            *f /= i as f32;
        }

        self
    }
}

impl<'a> Add<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn add(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} + {}", &self.word, &other.word),
            vec: self.vec,
        };

        for i in 0..newvec.vec.len() {
            newvec.vec[i] += other.vec[i];
        }

        newvec
    }
}

impl<'a> Sub<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn sub(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} - {}", &self.word, &other.word),
            vec: self.vec,
        };

        for i in 0..(newvec.vec.len()) {
            newvec.vec[i] -= other.vec[i];
        }

        newvec
    }
}

#[derive(Clone, PartialEq, Debug)]
pub struct LanguageModel {
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
}

pub struct LanguageModelBuilder {
    window_radius: usize,
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
    sentence: Vec<Option<usize>>,
}

pub struct WordAcceptor<'a> {
    builder: &'a mut LanguageModelBuilder,
}

impl LanguageModelBuilder {
    pub fn new(window_radius: usize, words: Vec<String>) -> LanguageModelBuilder {

        let word_vecs = words.iter()
                             .map(|s| WordVec::new(s.clone(), words.len()))
                             .collect();

        let words = HashMap::from_iter(words.into_iter()
                                            .enumerate()
                                            .map(|(a, b)| (b, a)));

        LanguageModelBuilder {
            window_radius: window_radius,
            words: words,
            word_vecs: word_vecs,
            sentence: vec![],
        }
    }

    pub fn build(mut self) -> LanguageModel {
        for vec in self.word_vecs.iter_mut() {
            vec.normalize();
        }

        LanguageModel {
            words: self.words,
            word_vecs: self.word_vecs,
        }
    }

    pub fn new_sentence<'a>(&'a mut self) -> WordAcceptor<'a> {
        WordAcceptor { builder: self }
    }

    pub fn save(&self, path: &Path) -> io::Result<()> {
        let mut file = BufWriter::new(File::create(path).unwrap());

        try!(write_raw(self.word_vecs.len(), &mut file));
        try!(file.write(&[b'\n']));

        for vec in &self.word_vecs {
            try!(file.write(vec.word.as_bytes()));
            try!(file.write(&[b':']));
            for f in &vec.vec {
                try!(write_raw(*f, &mut file));
            }
            try!(file.write(&[b'\n']));
        }
        Ok(())
    }

    pub fn load(path: &Path) -> io::Result<LanguageModelBuilder> {
        let mut file = BufReader::new(File::open(path).unwrap());
        let mut word_vecs = Vec::new();

        let size: u64 = unsafe { try!(read_raw::<u64, _>(&mut file)) };
        try!(read_byte(b'\n', &mut file));

        for _ in 0..size {
            let mut word: Vec<u8> = Vec::new();
            file.read_until(b':', &mut word).unwrap();
            assert_eq!(Some(b':'), word.pop());
            let word = String::from_utf8(word).unwrap();

            let mut vec: Vec<f32> = repeat(0f32).take(size as usize).collect();
            for f in &mut vec {
                unsafe {
                    *f = try!(read_raw::<f32, _>(&mut file));
                }
            }

            try!(read_byte(b'\n', &mut file));

            word_vecs.push(WordVec {
                word: word,
                vec: vec,
            });
        }

        let words: HashMap<_, _> = word_vecs.iter()
                                            .map(|v| v.word.clone())
                                            .enumerate()
                                            .map(|(i, w)| (w, i))
                                            .collect();

        Ok(LanguageModelBuilder {
            window_radius: 0,
            words: words,
            word_vecs: word_vecs,
            sentence: Vec::new(),
        })
    }
}

impl<'a> WordAcceptor<'a> {
    pub fn add_word(&mut self, word: &str) {
        let idx_opt = self.builder.words.get(word).map(|w| *w);

        self.builder.sentence.push(idx_opt);
    }
}

impl<'a> Drop for WordAcceptor<'a> {
    fn drop(&mut self) {
        let LanguageModelBuilder {
            ref mut sentence,
            ref mut word_vecs,
            window_radius,
            ..
        } = *self.builder;

        for (i, &from_word) in sentence.iter().enumerate() {
            let start = cmp::max(0, i as isize - window_radius as isize) as usize;
            let end = cmp::min(sentence.len(), i + window_radius + 1);
            for j in (start..i).chain(i + 1..end) {
                let dist = cmp::max(i, j) - cmp::min(i, j);

                if let (Some(f), Some(t)) = (from_word, sentence[j]) {
                    word_vecs[f].inc(t, dist);
                }
            }
        }

        sentence.clear();
    }
}

impl LanguageModel {
    pub fn get(&self, word: &str) -> Option<&WordVec> {
        self.words.get(word).map(|i| &self.word_vecs[*i])
    }

    pub fn nearest_words(&self, word: &WordVec) -> Vec<&WordVec> {
        let mut vec_refs = self.word_vecs
                               .iter()
                               .filter_map(|w| {
                                   let dist = w.distance(word);
                                   if w.word != word.word {
                                       Some((dist, w))
                                   } else {
                                       None
                                   }
                               })
                               .collect::<Vec<_>>();
        vec_refs.sort_by(|a, b| {
            a.0.partial_cmp(&b.0).unwrap_or_else(|| {
                println!("bad vector: {:?}, {:?}", a, b);
                Equal
            })
        });
        vec_refs.into_iter().map(|x| x.1).collect()
    }
}

fn read_byte<R: Read>(b: u8, read: &mut R) -> io::Result<()> {
    let buf = &mut [0];
    while try!(read.read(buf)) == 0 {}
    assert_eq!(&[b], buf);
    Ok(())
}

unsafe fn read_raw<T: Copy, R: Read>(reader: &mut BufReader<R>) -> io::Result<T> {
    let mut buffer = [0u8; 64];
    let t_size = mem::size_of::<T>();
    assert!(t_size <= buffer.len());

    let mut remainder = t_size;
    while remainder > 0 {
        let bytes_read = try!(reader.read(&mut buffer[(t_size - remainder)..t_size]));
        remainder -= bytes_read;
    }

    let bptr: *mut T = mem::transmute(buffer.as_ptr());
    Ok(*bptr)
}

fn write_raw<T: Copy, F: Write>(t: T, writer: &mut F) -> io::Result<usize> {
    let buffer = [0u8; 64];
    let t_size = mem::size_of::<T>();
    assert!(t_size <= buffer.len());
    unsafe {
        let bptr: *mut T = mem::transmute(buffer.as_ptr());
        *bptr = t;
    }

    let mut start = 0;
    while start != t_size {
        start += try!(writer.write(&buffer[start..t_size]));
    }
    Ok((start))
}

#[cfg(test)]
mod test {
    use super::LanguageModelBuilder;
    use std::path::Path;

    fn get_builder() -> LanguageModelBuilder {
        let words = "foo bar baz blort".split(" ").map(|w| w.to_string()).collect::<Vec<_>>();
        let mut builder = LanguageModelBuilder::new(1, words);

        let input = "x foo bar baz x x x x x x x blort".split(" ");

        {
            let mut acc = builder.new_sentence();
            for word in input {
                acc.add_word(word);
            }
        }

        builder
    }

    #[test]
    fn test_accept_sentences() {
        let model = get_builder().build();
        let foo = model.get("foo").unwrap();
        let baz = model.get("baz").unwrap();
        let bar = model.get("bar").unwrap();
        let blort = model.get("blort").unwrap();

        assert!(foo.distance(baz) < foo.distance(blort));
        assert!(foo.distance(bar) == bar.distance(baz));
    }

    #[test]
    fn test_serialization() {
        let builder = get_builder();
        let path = Path::new("/tmp/model.data");
        builder.save(&path).unwrap();

        let loaded_model = LanguageModelBuilder::load(path).unwrap().build();
        assert_eq!(builder.build(), loaded_model);
    }
}
<|endoftext|>


<commit_before>#[derive(Copy, Clone, Debug)]
pub enum TockValue<E> {
    Expected(E),
    Unexpected(isize),
}

pub type TockResult<T, E> = Result<T, TockValue<E>>;

pub trait TockResultExt<T, E>: Sized {
    fn as_expected(self) -> Result<T, E>;
}

impl<T, E> TockResultExt<T, E> for TockResult<T, E> {
    fn as_expected(self) -> Result<T, E> {
        match self {
            Ok(ok) => Ok(ok),
            Err(TockValue::Expected(err)) => Err(err),
            Err(TockValue::Unexpected(_)) => panic!("Unexpected error"),
        }
    }
}

pub const SUCCESS: isize = 0;
pub const FAIL: isize = -1;
pub const EBUSY: isize = -2;
pub const EALREADY: isize = -3;
pub const EINVAL: isize = -6;
pub const ESIZE: isize = -7;
pub const ENOMEM: isize = -9;
<commit_msg>Replace TockValue's auto-derived Debug implementation with a hand-written implementation optimized for size.<commit_after>#[derive(Copy, Clone)]
pub enum TockValue<E> {
    Expected(E),
    Unexpected(isize),
}

// Size-optimized implementation of Debug.
impl<E: core::fmt::Debug> core::fmt::Debug for TockValue<E> {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> Result<(), core::fmt::Error> {
        match self {
            // Printing out the value of E would cause TockResult::unwrap() to
            // use a &dyn core::fmt::Debug, which defeats LLVM's
            // devirtualization and prevents LTO from removing unused Debug
            // implementations. Unfortunately, that generates too much code
            // bloat (several kB), so we cannot display the value contained in
            // this TockValue.
            TockValue::Expected(_) => f.write_str("Expected(...)"),

            TockValue::Unexpected(n) => {
                f.write_str("Unexpected(")?;
                n.fmt(f)?;
                f.write_str(")")
            }
        }
    }
}

pub type TockResult<T, E> = Result<T, TockValue<E>>;

pub trait TockResultExt<T, E>: Sized {
    fn as_expected(self) -> Result<T, E>;
}

impl<T, E> TockResultExt<T, E> for TockResult<T, E> {
    fn as_expected(self) -> Result<T, E> {
        match self {
            Ok(ok) => Ok(ok),
            Err(TockValue::Expected(err)) => Err(err),
            Err(TockValue::Unexpected(_)) => panic!("Unexpected error"),
        }
    }
}

pub const SUCCESS: isize = 0;
pub const FAIL: isize = -1;
pub const EBUSY: isize = -2;
pub const EALREADY: isize = -3;
pub const EINVAL: isize = -6;
pub const ESIZE: isize = -7;
pub const ENOMEM: isize = -9;
<|endoftext|>


<commit_before>use std::{ mem, slice, string };
use handle::Owned;
use raw;
use raw::{ IV, UV, NV };
use raw::{ SVt_PVAV, SVt_PVHV, SVt_PVCV, SVt_PVGV };
use array::{ AV };
use hash::{ HV };
use convert::{ IntoSV, FromSV };

/// Perl scalar object.
pub struct SV(Owned<raw::SV>);

impl SV {
    method! {
        /// Return true if SV is a real scalar value.
        simple fn is_scalar() -> bool = sv_type() < SVt_PVAV
    }
    method! {
        /// Return true if SV contains array.
        simple fn is_array() -> bool = sv_type() == SVt_PVAV
    }
    method! {
        /// Return true if SV contains hash.
        simple fn is_hash() -> bool = sv_type() == SVt_PVHV
    }
    method! {
        /// Return true if SV contains subroutine.
        simple fn is_code() -> bool = sv_type() == SVt_PVCV
    }
    method! {
        /// Return true if SV contains glob.
        simple fn is_glob() -> bool = sv_type() == SVt_PVGV
    }
    method! {
        /// Return true if SV is defined.
        ///
        /// [`SvOK`](http://perldoc.perl.org/perlapi.html#SvOK).
        simple fn ok() -> bool = sv_ok() != 0
    }
    method! {
        /// Return true if SV contains a signed integer.
        ///
        /// Perl macro:[`SvIOK`](http://perldoc.perl.org/perlapi.html#SvIOK).
        simple fn iv_ok() -> bool = sv_iok() != 0
    }
    method! {
        /// Coerce the given SV to an integer and return it.
        ///
        /// Perl macro: [`SvIV`](http://perldoc.perl.org/perlapi.html#SvIV).
        simple fn iv() -> IV = sv_iv()
    }
    method! {
        /// Return true if SV contains a unsigned integer.
        ///
        /// Perl macro:[`SvUOK`](http://perldoc.perl.org/perlapi.html#SvUOK).
        simple fn uv_ok() -> bool = sv_uok() != 0
    }
    method! {
        /// Coerce the given SV to an unsigned integer and return it.
        ///
        /// Perl macro: [`SvUV`](http://perldoc.perl.org/perlapi.html#SvUV).
        simple fn uv() -> UV = sv_uv()
    }
    method! {
        /// Return true if SV contains a floating point value.
        ///
        /// Perl macro:[`SvNOK`](http://perldoc.perl.org/perlapi.html#SvNOK).
        simple fn nv_ok() -> bool = sv_nok() != 0
    }
    method! {
        /// Coerce the given SV to a floating point value and return it.
        ///
        /// Perl macro: [`SvNV`](http://perldoc.perl.org/perlapi.html#SvNV).
        simple fn nv() -> NV = sv_nv()
    }
    method! {
        /// Return true if SV contains a string.
        ///
        /// Perl macro:[`SvPOK`](http://perldoc.perl.org/perlapi.html#SvPOK).
        simple fn pv_ok() -> bool = sv_pok() != 0
    }
    method! {
        /// Return UTF8 flag on the SV.
        ///
        /// You should use this after a call to `pv()` or `str()`, in case any call to string
        /// overloading updates the internal flag.
        ///
        /// Perl macro: [`SvUTF8`](http://perldoc.perl.org/perlapi.html#SvUTF8).
        simple fn utf8() -> bool = sv_utf8() != 0
    }

    /// Return a copy of string in the SV as a vector of bytes.
    ///
    /// Perl macro: [`SvPV`](http://perldoc.perl.org/perlapi.html#SvPV).
    #[inline]
    pub fn pv(&self) -> Vec<u8> {
        unsafe {
            let mut len = 0;
            let ptr = self.pthx().sv_pv(self.as_ptr(), &mut len);
            slice::from_raw_parts(ptr as *const u8, len as usize).to_owned()
        }
    }

    /// Return a copy of string in the SV.
    ///
    /// Perl macro: [`SvPV`](http://perldoc.perl.org/perlapi.html#SvPV).
    pub fn str(&self) -> Result<String, string::FromUtf8Error> {
        String::from_utf8(self.pv())
    }

    method! {
        /// Return true if SV contains a Perl reference.
        ///
        /// Perl macro:[`SvROK`](http://perldoc.perl.org/perlapi.html#SvROK).
        simple fn rv_ok() -> bool = sv_rok() != 0
    }

    #[inline]
    unsafe fn deref_raw(&self) -> *mut raw::SV {
        self.pthx().sv_rv(self.as_ptr())
    }

    /// Dereference RV.
    ///
    /// Return `None` if `self` is not a valid Perl reference.
    #[inline]
    pub fn deref(&self) -> Option<SV> {
        if self.rv_ok() {
            Some(unsafe { SV::from_raw_borrowed(self.pthx(), self.deref_raw()) })
        } else {
            None
        }
    }

    /// Dereference RV into AV.
    ///
    /// Return `None` if `self` is not an array reference.
    #[inline]
    pub fn deref_av(&self) -> Option<AV> {
        self.deref().and_then(|sv| sv.into_av())
    }

    /// Dereference RV into HV.
    ///
    /// Return `None` if `self` is not a hash reference.
    #[inline]
    pub fn deref_hv(&self) -> Option<HV> {
        self.deref().and_then(|sv| sv.into_hv())
    }

    /// Cast SV into AV.
    #[inline]
    pub fn into_av(self) -> Option<AV> {
        if self.is_array() {
            Some(unsafe { AV::from_raw_owned(self.pthx(), self.into_raw() as *mut _) })
        } else {
            None
        }
    }

    /// Cast SV into HV.
    #[inline]
    pub fn into_hv(self) -> Option<HV> {
        if self.is_hash() {
            Some(unsafe { HV::from_raw_owned(self.pthx(), self.into_raw() as *mut _) })
        } else {
            None
        }
    }

    /// Consume SV and convert into raw pointer.
    ///
    /// Does not decrement reference count. Returned pointer must be correctly disposed of to avoid
    /// memory leaks.
    #[inline]
    pub fn into_raw(self) -> *mut raw::SV {
        let raw = self.0.as_ptr();
        mem::forget(self);
        raw
    }

    /// Construct new instance from a raw SV pointer without incrementing reference counter.
    ///
    /// Owned SV pointers are returned by assorted
    /// [`newSV`](http://perldoc.perl.org/perlapi.html#newSV) functions.
    #[inline]
    pub unsafe fn from_raw_owned(pthx: raw::Interpreter, raw: *mut raw::SV) -> SV {
        SV(Owned::from_raw_owned(pthx, raw))
    }

    /// Construct new instance from a raw SV pointer and increment reference counter.
    ///
    /// Borrowed SV pointers exist on stack and are returned by functions like
    /// [`av_fetch`](http://perldoc.perl.org/perlapi.html#av_fetch) or
    /// [`hv_fetch`](http://perldoc.perl.org/perlapi.html#av_fetch).
    #[inline]
    pub unsafe fn from_raw_borrowed(pthx: raw::Interpreter, raw: *mut raw::SV) -> SV {
        SV(Owned::from_raw_borrowed(pthx, raw))
    }

    fn pthx(&self) -> raw::Interpreter { self.0.pthx() }

    fn as_ptr(&self) -> *mut raw::SV { self.0.as_ptr() }
}

impl FromSV for IV {
    #[inline]
    unsafe fn from_sv(pthx: raw::Interpreter, raw: *mut raw::SV) -> IV {
        pthx.sv_iv(raw)
    }
}

impl FromSV for UV {
    #[inline]
    unsafe fn from_sv(pthx: raw::Interpreter, raw: *mut raw::SV) -> UV {
        pthx.sv_uv(raw)
    }
}

impl FromSV for NV {
    #[inline]
    unsafe fn from_sv(pthx: raw::Interpreter, raw: *mut raw::SV) -> NV {
        pthx.sv_nv(raw)
    }
}

impl FromSV for SV {
    #[inline]
    unsafe fn from_sv(pthx: raw::Interpreter, raw: *mut raw::SV) -> SV {
        SV::from_raw_borrowed(pthx, raw)
    }
}

impl IntoSV for IV {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe { SV::from_raw_owned(pthx, pthx.new_sv_iv(self)) }
    }
}

impl IntoSV for UV {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe { SV::from_raw_owned(pthx, pthx.new_sv_uv(self)) }
    }
}

impl IntoSV for NV {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe { SV::from_raw_owned(pthx, pthx.new_sv_nv(self)) }
    }
}

impl IntoSV for bool {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe {
            let raw = if self { pthx.sv_yes() } else { pthx.sv_no() };
            SV::from_raw_owned(pthx, raw)
        }
    }
}

impl<'a> IntoSV for &'a str {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe {
            let svp = pthx.new_sv_pvn(self.as_ptr() as *const i8,
                                      self.len() as raw::STRLEN,
                                      raw::SVf_UTF8 as raw::U32);
            SV::from_raw_owned(pthx, svp)
        }
    }
}

impl IntoSV for SV {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        assert!(self.pthx() == pthx);
        self
    }
}

impl IntoSV for () {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe {
            let svp = pthx.new_sv(0);
            SV::from_raw_owned(pthx, svp)
        }
    }
}
<commit_msg>Fix perlapi link.<commit_after>use std::{ mem, slice, string };
use handle::Owned;
use raw;
use raw::{ IV, UV, NV };
use raw::{ SVt_PVAV, SVt_PVHV, SVt_PVCV, SVt_PVGV };
use array::{ AV };
use hash::{ HV };
use convert::{ IntoSV, FromSV };

/// Perl scalar object.
pub struct SV(Owned<raw::SV>);

impl SV {
    method! {
        /// Return true if SV is a real scalar value.
        simple fn is_scalar() -> bool = sv_type() < SVt_PVAV
    }
    method! {
        /// Return true if SV contains array.
        simple fn is_array() -> bool = sv_type() == SVt_PVAV
    }
    method! {
        /// Return true if SV contains hash.
        simple fn is_hash() -> bool = sv_type() == SVt_PVHV
    }
    method! {
        /// Return true if SV contains subroutine.
        simple fn is_code() -> bool = sv_type() == SVt_PVCV
    }
    method! {
        /// Return true if SV contains glob.
        simple fn is_glob() -> bool = sv_type() == SVt_PVGV
    }
    method! {
        /// Return true if SV is defined.
        ///
        /// [`SvOK`](http://perldoc.perl.org/perlapi.html#SvOK).
        simple fn ok() -> bool = sv_ok() != 0
    }
    method! {
        /// Return true if SV contains a signed integer.
        ///
        /// Perl macro:[`SvIOK`](http://perldoc.perl.org/perlapi.html#SvIOK).
        simple fn iv_ok() -> bool = sv_iok() != 0
    }
    method! {
        /// Coerce the given SV to an integer and return it.
        ///
        /// Perl macro: [`SvIV`](http://perldoc.perl.org/perlapi.html#SvIV).
        simple fn iv() -> IV = sv_iv()
    }
    method! {
        /// Return true if SV contains a unsigned integer.
        ///
        /// Perl macro:[`SvUOK`](http://perldoc.perl.org/perlapi.html#SvUOK).
        simple fn uv_ok() -> bool = sv_uok() != 0
    }
    method! {
        /// Coerce the given SV to an unsigned integer and return it.
        ///
        /// Perl macro: [`SvUV`](http://perldoc.perl.org/perlapi.html#SvUV).
        simple fn uv() -> UV = sv_uv()
    }
    method! {
        /// Return true if SV contains a floating point value.
        ///
        /// Perl macro:[`SvNOK`](http://perldoc.perl.org/perlapi.html#SvNOK).
        simple fn nv_ok() -> bool = sv_nok() != 0
    }
    method! {
        /// Coerce the given SV to a floating point value and return it.
        ///
        /// Perl macro: [`SvNV`](http://perldoc.perl.org/perlapi.html#SvNV).
        simple fn nv() -> NV = sv_nv()
    }
    method! {
        /// Return true if SV contains a string.
        ///
        /// Perl macro:[`SvPOK`](http://perldoc.perl.org/perlapi.html#SvPOK).
        simple fn pv_ok() -> bool = sv_pok() != 0
    }
    method! {
        /// Return UTF8 flag on the SV.
        ///
        /// You should use this after a call to `pv()` or `str()`, in case any call to string
        /// overloading updates the internal flag.
        ///
        /// Perl macro: [`SvUTF8`](http://perldoc.perl.org/perlapi.html#SvUTF8).
        simple fn utf8() -> bool = sv_utf8() != 0
    }

    /// Return a copy of string in the SV as a vector of bytes.
    ///
    /// Perl macro: [`SvPV`](http://perldoc.perl.org/perlapi.html#SvPV).
    #[inline]
    pub fn pv(&self) -> Vec<u8> {
        unsafe {
            let mut len = 0;
            let ptr = self.pthx().sv_pv(self.as_ptr(), &mut len);
            slice::from_raw_parts(ptr as *const u8, len as usize).to_owned()
        }
    }

    /// Return a copy of string in the SV.
    ///
    /// Perl macro: [`SvPV`](http://perldoc.perl.org/perlapi.html#SvPV).
    pub fn str(&self) -> Result<String, string::FromUtf8Error> {
        String::from_utf8(self.pv())
    }

    method! {
        /// Return true if SV contains a Perl reference.
        ///
        /// Perl macro:[`SvROK`](http://perldoc.perl.org/perlapi.html#SvROK).
        simple fn rv_ok() -> bool = sv_rok() != 0
    }

    #[inline]
    unsafe fn deref_raw(&self) -> *mut raw::SV {
        self.pthx().sv_rv(self.as_ptr())
    }

    /// Dereference RV.
    ///
    /// Return `None` if `self` is not a valid Perl reference.
    #[inline]
    pub fn deref(&self) -> Option<SV> {
        if self.rv_ok() {
            Some(unsafe { SV::from_raw_borrowed(self.pthx(), self.deref_raw()) })
        } else {
            None
        }
    }

    /// Dereference RV into AV.
    ///
    /// Return `None` if `self` is not an array reference.
    #[inline]
    pub fn deref_av(&self) -> Option<AV> {
        self.deref().and_then(|sv| sv.into_av())
    }

    /// Dereference RV into HV.
    ///
    /// Return `None` if `self` is not a hash reference.
    #[inline]
    pub fn deref_hv(&self) -> Option<HV> {
        self.deref().and_then(|sv| sv.into_hv())
    }

    /// Cast SV into AV.
    #[inline]
    pub fn into_av(self) -> Option<AV> {
        if self.is_array() {
            Some(unsafe { AV::from_raw_owned(self.pthx(), self.into_raw() as *mut _) })
        } else {
            None
        }
    }

    /// Cast SV into HV.
    #[inline]
    pub fn into_hv(self) -> Option<HV> {
        if self.is_hash() {
            Some(unsafe { HV::from_raw_owned(self.pthx(), self.into_raw() as *mut _) })
        } else {
            None
        }
    }

    /// Consume SV and convert into raw pointer.
    ///
    /// Does not decrement reference count. Returned pointer must be correctly disposed of to avoid
    /// memory leaks.
    #[inline]
    pub fn into_raw(self) -> *mut raw::SV {
        let raw = self.0.as_ptr();
        mem::forget(self);
        raw
    }

    /// Construct new instance from a raw SV pointer without incrementing reference counter.
    ///
    /// Owned SV pointers are returned by assorted
    /// [`newSV`](http://perldoc.perl.org/perlapi.html#newSV) functions.
    #[inline]
    pub unsafe fn from_raw_owned(pthx: raw::Interpreter, raw: *mut raw::SV) -> SV {
        SV(Owned::from_raw_owned(pthx, raw))
    }

    /// Construct new instance from a raw SV pointer and increment reference counter.
    ///
    /// Borrowed SV pointers exist on stack and are returned by functions like
    /// [`av_fetch`](http://perldoc.perl.org/perlapi.html#av_fetch) or
    /// [`hv_fetch`](http://perldoc.perl.org/perlapi.html#hv_fetch).
    #[inline]
    pub unsafe fn from_raw_borrowed(pthx: raw::Interpreter, raw: *mut raw::SV) -> SV {
        SV(Owned::from_raw_borrowed(pthx, raw))
    }

    fn pthx(&self) -> raw::Interpreter { self.0.pthx() }

    fn as_ptr(&self) -> *mut raw::SV { self.0.as_ptr() }
}

impl FromSV for IV {
    #[inline]
    unsafe fn from_sv(pthx: raw::Interpreter, raw: *mut raw::SV) -> IV {
        pthx.sv_iv(raw)
    }
}

impl FromSV for UV {
    #[inline]
    unsafe fn from_sv(pthx: raw::Interpreter, raw: *mut raw::SV) -> UV {
        pthx.sv_uv(raw)
    }
}

impl FromSV for NV {
    #[inline]
    unsafe fn from_sv(pthx: raw::Interpreter, raw: *mut raw::SV) -> NV {
        pthx.sv_nv(raw)
    }
}

impl FromSV for SV {
    #[inline]
    unsafe fn from_sv(pthx: raw::Interpreter, raw: *mut raw::SV) -> SV {
        SV::from_raw_borrowed(pthx, raw)
    }
}

impl IntoSV for IV {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe { SV::from_raw_owned(pthx, pthx.new_sv_iv(self)) }
    }
}

impl IntoSV for UV {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe { SV::from_raw_owned(pthx, pthx.new_sv_uv(self)) }
    }
}

impl IntoSV for NV {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe { SV::from_raw_owned(pthx, pthx.new_sv_nv(self)) }
    }
}

impl IntoSV for bool {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe {
            let raw = if self { pthx.sv_yes() } else { pthx.sv_no() };
            SV::from_raw_owned(pthx, raw)
        }
    }
}

impl<'a> IntoSV for &'a str {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe {
            let svp = pthx.new_sv_pvn(self.as_ptr() as *const i8,
                                      self.len() as raw::STRLEN,
                                      raw::SVf_UTF8 as raw::U32);
            SV::from_raw_owned(pthx, svp)
        }
    }
}

impl IntoSV for SV {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        assert!(self.pthx() == pthx);
        self
    }
}

impl IntoSV for () {
    #[inline]
    fn into_sv(self, pthx: raw::Interpreter) -> SV {
        unsafe {
            let svp = pthx.new_sv(0);
            SV::from_raw_owned(pthx, svp)
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashMap, HashSet};
use std::collections::hash_map::Entry::*;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::iter;
use std::str::CharIndices;
use std::ops;
use itertools::{GroupBy, Itertools};
use coalesce::Coalesce;

/// A Document contains an id and content.
/// Hashing and equality are based only on the id field.
#[derive(Clone, Eq, Debug, RustcEncodable, RustcDecodable)]
pub struct Document {
    id: String,
    content: String,
}

impl Document {
    /// Construct a new Document from an id and content.
    /// Both two arguments can be anything that can be turned into a String.
    pub fn new<S, T>(id: S, content: T) -> Document
        where S: Into<String>,
              T: Into<String>
    {
        Document { id: id.into(), content: content.into(), }
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn content(&self) -> &str {
        &self.content
    }
}

impl Hash for Document {
    // Documents are unique only upon their id
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.id.hash(state);
    }
}

impl PartialEq<Document> for Document {
    // Documents are unique only upon their id
    fn eq(&self, other: &Document) -> bool {
        self.id == other.id
    }
}

/// A SearchResult is the representation of a Document returned for a specific set of search
/// terms. It is unique upon the document and the vec of highlight indices. It also contains a
/// score and a String of the document's content post-highlighting.
#[derive(Clone, Debug, RustcEncodable)]
pub struct SearchResult {
    doc: Arc<Document>,
    highlights: Vec<(usize, usize)>,
    score: f32,
}

impl PartialEq for SearchResult {
    // SearchResult is unique upon its document and highlights
    fn eq(&self, other: &SearchResult) -> bool {
        self.doc == other.doc && self.highlights == other.highlights
    }
}

impl Eq for SearchResult {}

impl Hash for SearchResult {
    // SearchResult is unique upon its document and highlights
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.doc.hash(state);
        self.highlights.hash(state);
    }
}

impl SearchResult {
    fn new(doc: Arc<Document>, mut highlights: Vec<(usize, usize)>) -> SearchResult {
        SearchResult {
            score: highlights.iter()
                .map(|&(begin, end)| end - begin)
                .sum::<usize>() as f32 / (doc.content.len() as f32).sqrt(),
            doc: doc,
            highlights: { highlights.sort(); highlights },
        }
    }

    /// Returns the document
    pub fn doc(&self) -> &Arc<Document> {
        &self.doc
    }

    /// Returns the highlighted indices.
    ///
    /// Each `(usize, usize)` indicates the start and end of a term in the document's content
    /// that should be highlighted.
    pub fn highlights(&self) -> &Vec<(usize, usize)> {
        &self.highlights
    }

    /// Returns the search result's score.
    ///
    /// Score is computed by the product of the summed length of the matching terms and the inverse
    /// square root of the length of the document. Taking the square root of the document's length
    /// helps to combat bias toward short content.
    pub fn score(&self) -> f32 {
        self.score
    }

    pub fn highlight(&self, before: &str, after: &str) -> String {
        let mut begin_idx = 0;
        let mut parts = String::new();
        for &(begin, end) in &self.highlights {
            parts.push_str(&self.doc.content[begin_idx..begin]);
            parts.push_str(before);
            parts.push_str(&self.doc.content[begin..end]);
            parts.push_str(after);
            begin_idx = end;
        }
        parts.push_str(&self.doc.content[begin_idx..]);
        parts
    }
}

/// A basic implementation of an `Index`, the inverted index is a data structure that maps
/// from words to sets of Documents.
#[derive(Debug)]
pub struct InvertedIndex {
    index: BTreeMap<String, BTreeMap<String, Vec<(usize, usize)>>>,
    docs: BTreeMap<String, Arc<Document>>,
}

impl InvertedIndex {
    pub fn new() -> InvertedIndex {
        InvertedIndex {
            index: BTreeMap::new(),
            docs: BTreeMap::new(),
        }
    }

    /// A basic implementation of index, splits the document's content into whitespace-separated
    /// words, and inserts each word-document pair into the map.
    pub fn index(&mut self, doc: Document) {
        let doc = Arc::new(doc);
        let analyzed  = analyze_doc(doc.content());
        let previous_version = self.docs.insert(doc.id.clone(), doc.clone());
        if let Some(previous_version) = previous_version {
            let previous_analyzed = analyze_doc(previous_version.content());
            for (ngram, _) in previous_analyzed {
                let is_empty = {
                    let docs_for_ngram = self.index.get_mut(&ngram).unwrap();
                    docs_for_ngram.remove(&doc.id);
                    docs_for_ngram.is_empty()
                };
                if is_empty {
                    self.index.remove(&ngram);
                }
            }
        }

        for (ngram, highlighted) in analyzed {
            let mut highlights = self.index.entry(ngram)
                .or_insert_with(BTreeMap::new)
                .entry(doc.id.clone())
                .or_insert_with(Vec::new);
            let coalesce_idx = highlights.binary_search(&highlighted).err().unwrap();
            highlights.coalesce(coalesce_idx, highlighted);
        }
    }

    /// A basic search implementation that splits the query's content into whitespace-separated
    /// words, looks up the set of Documents for each word, and then concatenates the sets.
    pub fn search(&self, query: &str) -> Vec<SearchResult> {
        let unique_terms: HashSet<_> = query.split_whitespace().map(str::to_lowercase).collect();
        let map: HashMap<String, Vec<(usize, usize)>> = unique_terms.into_iter()
            .flat_map(|word| self.index.get(&word))
            .flat_map(BTreeMap::iter)
            .fold(HashMap::new(), |mut map, (doc_id, highlights)| {
                match map.entry(doc_id.clone()) {
                    Vacant(entry) => { entry.insert(highlights.clone()); },
                    Occupied(mut entry) => {
                        let entry: &mut Vec<(usize, usize)> = entry.get_mut();
                        for highlight in highlights {
                            let coalesce_idx = entry.binary_search(highlight).err().unwrap();
                            entry.coalesce(coalesce_idx, highlight.clone());
                        }
                    }
                }
                map
            });
        let mut results: Vec<_> = map.into_iter()
            .map(|(doc_id, index_map)| SearchResult::new(self.docs[&doc_id].clone(), index_map.into_iter().collect()))
            .collect();
        results.sort_by(|result1, result2| result2.score.partial_cmp(&result1.score).unwrap());
        results
    }
}

fn analyze_doc(doc: &str)
-> iter::FlatMap<
    iter::Filter<
        GroupBy<bool, CharIndices, fn(&(usize, char)) -> bool>,
        fn(&(bool, Vec<(usize, char)>)) -> bool>,
    iter::Map<ops::Range<usize>, Ngrams>,
    fn((bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams>>
{
    doc.char_indices()
        .group_by(is_whitespace as fn(&(usize, char)) -> bool)
        .filter(not_whitespace as fn(&(bool, Vec<(usize, char)>)) -> bool)
        .flat_map(ngrams)
}

fn ngrams((_, chars): (bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams> {
    (1..chars.len() + 1).map(Ngrams::new(chars))
}

fn not_whitespace(&(is_whitespace, _): &(bool, Vec<(usize, char)>)) -> bool {
    !is_whitespace
}

fn is_whitespace(&(_, c): &(usize, char)) -> bool {
    c.is_whitespace()
}

struct Ngrams {
    chars: Vec<(usize, char)>
}

impl Ngrams {
    fn new(chars: Vec<(usize, char)>) -> Ngrams {
        Ngrams { chars: chars }
    }
}

impl Fn<(usize,)> for Ngrams {
    extern "rust-call" fn call(&self, (to,): (usize,)) -> (String, (usize, usize)) {
        let word = self.chars[..to].iter().flat_map(|&(_, c)| c.to_lowercase()).collect();
        let start = self.chars[0].0;
        let (last_idx, last_char) = self.chars[to - 1];
        let finish = last_idx + last_char.len_utf8();
        (word, (start, finish))
    }
}

impl FnMut<(usize,)> for Ngrams {
    extern "rust-call" fn call_mut(&mut self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

impl FnOnce<(usize,)> for Ngrams {
    type Output = (String, (usize, usize));
    extern "rust-call" fn call_once(self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

#[cfg(test)]
mod bench {
    use super::*;
    use std::sync::Arc;
    use std::collections::HashSet;

    #[test]
    fn test_search1() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("1", "learn to program in rust today");
        index.index(doc1.clone());
        let doc2 = Document::new("2", "what did you today do");
        index.index(doc2.clone());
        let search_results = index.search("to");
        let expected = [
            SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
            SearchResult::new(Arc::new(doc2), vec![(13, 15)])
        ];
        assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
        assert_eq!("learn to program in rust \
                   today",
                   expected[0].highlight("", ""));
    }

    #[test]
    fn test_ngrams() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("1", "learn to program in rust today");
        let doc2 = Document::new("2", "what did you today do");
        index.index(doc1.clone());
        index.index(doc2.clone());
        let search_results = index.search("to");
        let expected = [
            SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
            SearchResult::new(Arc::new(doc2), vec![(13, 15)]),
        ];
        assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
        assert_eq!("learn to program in rust \
                   today",
                   expected[0].highlight("", ""));

    }

    #[test]
    fn test_highlight1() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("2", "what to do today");
        let doc2 = Document::new("3", "hey today");
        index.index(doc1.clone());
        index.index(doc2.clone());
        let search_results = index.search("to");
        let expected = [
            SearchResult::new(Arc::new(doc1), vec![(5, 7), (11, 13)]),
            SearchResult::new(Arc::new(doc2), vec![(4, 6)]),
        ];
        assert_eq!(search_results.into_iter().collect::<HashSet<_>>(),
                   expected.iter().cloned().collect());
        assert_eq!("what to do today",
                   expected[0].highlight("", ""));
    }

    #[test]
    fn test_highlight2() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("2", "Won\u{2019}t this split the ecosystem? Will everyone use?");
        index.index(doc1.clone());
        let expected =  "Won\u{2019}t this split the *e*cosystem? Will *e*veryone use?";
        let search_results = index.search("e");
        assert_eq!(1, search_results.len());
        assert_eq!(search_results[0].highlight("*", "*"), expected);
    }

    #[test]
    fn test_unicode() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "嗨, 您好");
        index.index(doc.clone());
        let to_search = "您";
        let search_results = index.search(to_search);
        let &SearchResult { ref doc, ref highlights, .. } = search_results.iter().next().unwrap();
        let (begin, end) = highlights[0];
        assert_eq!(&doc.content()[begin..end], to_search);
    }

    #[test]
    fn test_update_doc() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "abc åäö");
        index.index(doc);
        let doc = Document::new("0", "different");
        index.index(doc);
        let search_results = index.search("å");
        assert!(search_results.is_empty());
        assert_eq!(index.docs.len(), 1);
    }

    #[test]
    fn test_ranking() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let doc2 = Document::new("1", "beast");
        index.index(doc2.clone());
        let search_results = index.search("be");
        assert_eq!(index.docs.len(), 2);
        assert_eq!(search_results.into_iter().map(|r| r.doc).collect::<HashSet<_>>(),
                [Arc::new(doc), Arc::new(doc2)].iter().cloned().collect());
    }

    #[test]
    fn test_duplicate_term() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let search_results = index.search("be be");
        assert_eq!(search_results.len(), 1);
    }

    #[test]
    fn test_duplicate_term2() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let search_results = index.search("be b");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
    }

    #[test]
    fn test_lowercase_search() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "BeAt");
        index.index(doc.clone());
        let search_results = index.search("bE");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
    }

    #[test]
    fn test_lowercase_index() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "BeAt");
        index.index(doc.clone());
        let search_results = index.search("be");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
    }
}
<commit_msg>Remove some traits from SearchResult that didn't make sense<commit_after>use std::collections::{BTreeMap, HashMap, HashSet};
use std::collections::hash_map::Entry::*;
use std::hash::{Hash, Hasher};
use std::iter;
use std::str::CharIndices;
use std::ops;
use itertools::{GroupBy, Itertools};
use coalesce::Coalesce;

/// A Document contains an id and content.
/// Hashing and equality are based only on the id field.
#[derive(Clone, Debug, Eq, Ord, RustcEncodable, RustcDecodable)]
pub struct Document {
    id: String,
    content: String,
}

impl Document {
    /// Construct a new Document from an id and content.
    /// Both two arguments can be anything that can be turned into a String.
    pub fn new<S, T>(id: S, content: T) -> Document
        where S: Into<String>,
              T: Into<String>
    {
        Document { id: id.into(), content: content.into(), }
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn content(&self) -> &str {
        &self.content
    }
}

impl Hash for Document {
    // Documents are unique only upon their id
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.id.hash(state);
    }
}

impl PartialEq for Document {
    // Documents are unique only upon their id
    fn eq(&self, other: &Document) -> bool {
        self.id == other.id
    }
}

impl PartialOrd for Document {
    fn partial_cmp(&self, other: &Document) -> Option<::std::cmp::Ordering> {
        self.id.partial_cmp(&other.id)
    }
}

/// A SearchResult is the representation of a Document returned for a specific set of search
/// terms. It is unique upon the document and the vec of highlight indices. It also contains a
/// score and a String of the document's content post-highlighting.
#[derive(Clone, Debug, RustcEncodable)]
pub struct SearchResult<'a> {
    doc: &'a Document,
    highlights: Vec<(usize, usize)>,
    score: f32,
}

impl<'a> SearchResult<'a> {
    fn new(doc: &'a Document, mut highlights: Vec<(usize, usize)>) -> SearchResult<'a> {
        SearchResult {
            score: highlights.iter()
                .map(|&(begin, end)| end - begin)
                .sum::<usize>() as f32 / (doc.content.len() as f32).sqrt(),
            doc: doc,
            highlights: { highlights.sort(); highlights },
        }
    }

    /// Returns the document
    pub fn doc(&self) -> &Document {
        &self.doc
    }

    /// Returns the highlighted indices.
    ///
    /// Each `(usize, usize)` indicates the start and end of a term in the document's content
    /// that should be highlighted.
    pub fn highlights(&self) -> &Vec<(usize, usize)> {
        &self.highlights
    }

    /// Returns the search result's score.
    ///
    /// Score is computed by the product of the summed length of the matching terms and the inverse
    /// square root of the length of the document. Taking the square root of the document's length
    /// helps to combat bias toward short content.
    pub fn score(&self) -> f32 {
        self.score
    }

    pub fn highlight(&self, before: &str, after: &str) -> String {
        let mut begin_idx = 0;
        let mut parts = String::new();
        for &(begin, end) in &self.highlights {
            parts.push_str(&self.doc.content[begin_idx..begin]);
            parts.push_str(before);
            parts.push_str(&self.doc.content[begin..end]);
            parts.push_str(after);
            begin_idx = end;
        }
        parts.push_str(&self.doc.content[begin_idx..]);
        parts
    }
}

/// A basic implementation of an `Index`, the inverted index is a data structure that maps
/// from words to sets of Documents.
#[derive(Debug)]
pub struct InvertedIndex {
    index: BTreeMap<String, BTreeMap<String, Vec<(usize, usize)>>>,
    docs: BTreeMap<String, Document>,
}

impl InvertedIndex {
    pub fn new() -> InvertedIndex {
        InvertedIndex {
            index: BTreeMap::new(),
            docs: BTreeMap::new(),
        }
    }

    /// A basic implementation of index, splits the document's content into whitespace-separated
    /// words, and inserts each word-document pair into the map.
    pub fn index(&mut self, doc: Document) {
        let analyzed  = analyze_doc(doc.content());
        let previous_version = self.docs.insert(doc.id.clone(), doc.clone());
        if let Some(previous_version) = previous_version {
            let previous_analyzed = analyze_doc(previous_version.content());
            for (ngram, _) in previous_analyzed {
                let is_empty = {
                    let docs_for_ngram = self.index.get_mut(&ngram).unwrap();
                    docs_for_ngram.remove(&doc.id);
                    docs_for_ngram.is_empty()
                };
                if is_empty {
                    self.index.remove(&ngram);
                }
            }
        }

        for (ngram, highlighted) in analyzed {
            let mut highlights = self.index.entry(ngram)
                .or_insert_with(BTreeMap::new)
                .entry(doc.id.clone())
                .or_insert_with(Vec::new);
            let coalesce_idx = highlights.binary_search(&highlighted).err().unwrap();
            highlights.coalesce(coalesce_idx, highlighted);
        }
    }

    /// A basic search implementation that splits the query's content into whitespace-separated
    /// words, looks up the set of Documents for each word, and then concatenates the sets.
    pub fn search(&self, query: &str) -> Vec<SearchResult> {
        let unique_terms: HashSet<_> = query.split_whitespace().map(str::to_lowercase).collect();
        let map: HashMap<String, Vec<(usize, usize)>> = unique_terms.into_iter()
            .flat_map(|word| self.index.get(&word))
            .flat_map(BTreeMap::iter)
            .fold(HashMap::new(), |mut map, (doc_id, highlights)| {
                match map.entry(doc_id.clone()) {
                    Vacant(entry) => { entry.insert(highlights.clone()); },
                    Occupied(mut entry) => {
                        let entry: &mut Vec<(usize, usize)> = entry.get_mut();
                        for highlight in highlights {
                            let coalesce_idx = entry.binary_search(highlight).err().unwrap();
                            entry.coalesce(coalesce_idx, highlight.clone());
                        }
                    }
                }
                map
            });
        let mut results: Vec<_> = map.into_iter()
            .map(|(doc_id, index_map)| SearchResult::new(&self.docs[&doc_id], index_map.into_iter().collect()))
            .collect();
        results.sort_by(|result1, result2| result2.score.partial_cmp(&result1.score).unwrap());
        results
    }
}

fn analyze_doc(doc: &str)
-> iter::FlatMap<
    iter::Filter<
        GroupBy<bool, CharIndices, fn(&(usize, char)) -> bool>,
        fn(&(bool, Vec<(usize, char)>)) -> bool>,
    iter::Map<ops::Range<usize>, Ngrams>,
    fn((bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams>>
{
    doc.char_indices()
        .group_by(is_whitespace as fn(&(usize, char)) -> bool)
        .filter(not_whitespace as fn(&(bool, Vec<(usize, char)>)) -> bool)
        .flat_map(ngrams)
}

fn ngrams((_, chars): (bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams> {
    (1..chars.len() + 1).map(Ngrams::new(chars))
}

fn not_whitespace(&(is_whitespace, _): &(bool, Vec<(usize, char)>)) -> bool {
    !is_whitespace
}

fn is_whitespace(&(_, c): &(usize, char)) -> bool {
    c.is_whitespace()
}

struct Ngrams {
    chars: Vec<(usize, char)>
}

impl Ngrams {
    fn new(chars: Vec<(usize, char)>) -> Ngrams {
        Ngrams { chars: chars }
    }
}

impl Fn<(usize,)> for Ngrams {
    extern "rust-call" fn call(&self, (to,): (usize,)) -> (String, (usize, usize)) {
        let word = self.chars[..to].iter().flat_map(|&(_, c)| c.to_lowercase()).collect();
        let start = self.chars[0].0;
        let (last_idx, last_char) = self.chars[to - 1];
        let finish = last_idx + last_char.len_utf8();
        (word, (start, finish))
    }
}

impl FnMut<(usize,)> for Ngrams {
    extern "rust-call" fn call_mut(&mut self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

impl FnOnce<(usize,)> for Ngrams {
    type Output = (String, (usize, usize));
    extern "rust-call" fn call_once(self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::collections::BTreeMap;

    #[test]
    fn test_ngrams() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("1", "learn to program in rust today");
        let doc2 = Document::new("2", "what did you today do");
        index.index(doc1.clone());
        index.index(doc2.clone());
        let search_results = index.search("to");
        let expected: BTreeMap<_, _> = [
            (doc1.clone(), vec![(6, 8), (25, 27)]),
            (doc2, vec![(13, 15)]),
        ].iter().cloned().collect();
        assert_eq!(search_results.len(), expected.len());
        for search_result in &search_results {
            assert_eq!(&search_result.highlights, &expected[search_result.doc])
        }
        assert_eq!("learn to program in rust \
                   today",
                   search_results.iter()
                        .find(|search_result| search_result.doc == &doc1)
                        .unwrap()
                        .highlight("", ""));

    }

    #[test]
    fn test_highlight() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("2", "Won\u{2019}t this split the ecosystem? Will everyone use?");
        index.index(doc1.clone());
        let expected =  "Won\u{2019}t this split the *e*cosystem? Will *e*veryone use?";
        let search_results = index.search("e");
        assert_eq!(1, search_results.len());
        assert_eq!(search_results[0].highlight("*", "*"), expected);
    }

    #[test]
    fn test_unicode() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "嗨, 您好");
        index.index(doc.clone());
        let to_search = "您";
        let search_results = index.search(to_search);
        let &SearchResult { ref doc, ref highlights, .. } = search_results.iter().next().unwrap();
        let (begin, end) = highlights[0];
        assert_eq!(&doc.content()[begin..end], to_search);
    }

    #[test]
    fn test_update_doc() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "abc åäö");
        index.index(doc);
        let doc = Document::new("0", "different");
        index.index(doc);
        let search_results = index.search("å");
        assert!(search_results.is_empty());
        assert_eq!(index.docs.len(), 1);
    }

    #[test]
    fn test_ranking() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let doc2 = Document::new("1", "beast");
        index.index(doc2);
        let search_results = index.search("be");
        assert_eq!(index.docs.len(), 2);
        // "beat" should be first, since it's a closer match
        assert_eq!(search_results[0].doc, &doc);
    }

    #[test]
    fn test_duplicate_term() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let search_results = index.search("be be");
        assert_eq!(search_results.len(), 1);
    }

    #[test]
    fn test_duplicate_term2() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let search_results = index.search("be b");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0].highlights, vec![(0, 2)]);
    }

    #[test]
    fn test_lowercase_search() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "BeAt");
        index.index(doc.clone());
        let search_results = index.search("bE");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0].highlights, vec![(0, 2)]);
    }

    #[test]
    fn test_lowercase_index() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "BeAt");
        index.index(doc.clone());
        let search_results = index.search("be");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0].highlights, vec![(0, 2)]);
    }
}
<|endoftext|>