<commit_before>#![feature(proc_macro)]
#[macro_use]
extern crate serde_derive;

extern crate futures;
extern crate tokio_core;

use std::thread;
use std::sync::mpsc::channel;
//use futures::stream::Stream;
//use tokio_core::reactor::Core;
//use tokio_core::net::TcpListener;
use std::net::TcpStream;
use std::net::TcpListener;
use std::collections::HashMap;
use std::io::Write;
use std::time::Duration;
use std::option::Option;
//use std::io::Read;
use std::io;
use std::io::prelude::*;
use std::io::BufReader;
//use futures::Future;
use std::sync::{Arc, Mutex};
use std::io::{Error, ErrorKind};

mod userchannel;
mod messaging;

fn main() {
    let address = "127.0.0.1:8080";
    let listener = TcpListener::bind(&address).unwrap();
    let mut spawned = HashMap::new();
    let mut cnt = 0;
    let option: Option<Duration> = Some(Duration::from_millis(50));

    let addr = listener.local_addr().unwrap();
    println!("Listening for connections on {}", addr);
    let (tx, rx) = channel();

    for socket in listener.incoming() {
        /* 1. Read incoming msg channel
         * 2. If it does not exist, spawn new thread
         * 3. Send socket to thread
         */
        let mut result_str = String::new();
        let mut socket = socket.unwrap();
        let mut stream = BufReader::new(socket.try_clone().unwrap());
        let len = stream.read_line(&mut result_str);
        println!("Got string {}", result_str);
        socket.set_read_timeout(option);
        if !spawned.contains_key(&result_str) {
            let tx = tx.clone();
            let result_str = result_str.clone();
            thread::spawn(move|| {
                let (thr_tx, thr_rx) = channel();
                println!("Spawned: Sending socket");
                let mut users = HashMap::new();
                tx.send(thr_tx.clone()).unwrap();
                println!("Got thr sock");
                loop {
                    let mut removals = Vec::new();
                    match thr_rx.try_recv() {
                        Ok(val) => {
                            let thr: TcpStream = val;
                            cnt += 1;
                            users.insert(cnt, thr);
                        },
                        Err(_) => {}
                    }
                    let mut bufstr = String::new();
                    for (user, thr_socket) in &users {
                        let mut stream = BufReader::new(thr_socket.try_clone().unwrap());
                        match stream.read_line(&mut bufstr) {
                            Ok(len) => {
                                if(len > 0) {
                                    for (another_user, mut thr_sock) in &users {
                                        if user != another_user {
                                            let mut msg = another_user.to_string();
                                            msg.push_str(":");
                                            msg.push_str(&bufstr);
                                            thr_sock.write(&msg.into_bytes()).unwrap();
                                        }
                                    }
                                }
                                else {
                                    /* Socket closed, drop user */
                                    removals.push(user.clone());
                                }
                            },
                            Err(_) => {
                                //println!("{:?}", err.kind());
                            }
                        }
                    }
                    for removal in &removals {
                        users.remove(removal);
                    }
                    if cnt > 0 && users.is_empty() {
                        break;
                    }
                }
            });
            let thr_feed = rx.recv().unwrap();
            println!("Got thr_feed");
            spawned.insert(result_str.clone(), thr_feed);
        }
        let thr_socket = spawned.get_mut(&result_str).unwrap();
        thr_socket.send(socket).unwrap();
    }
}
<commit_msg>User id fix.<commit_after>#![feature(proc_macro)]
#[macro_use]
extern crate serde_derive;

extern crate futures;
extern crate tokio_core;

use std::thread;
use std::sync::mpsc::channel;
//use futures::stream::Stream;
//use tokio_core::reactor::Core;
//use tokio_core::net::TcpListener;
use std::net::TcpStream;
use std::net::TcpListener;
use std::collections::HashMap;
use std::io::Write;
use std::time::Duration;
use std::option::Option;
//use std::io::Read;
use std::io;
use std::io::prelude::*;
use std::io::BufReader;
//use futures::Future;
use std::sync::{Arc, Mutex};
use std::io::{Error, ErrorKind};

mod userchannel;
mod messaging;

fn main() {
    let address = "127.0.0.1:8080";
    let listener = TcpListener::bind(&address).unwrap();
    let mut spawned = HashMap::new();
    let mut cnt = 0;
    let option: Option<Duration> = Some(Duration::from_millis(50));

    let addr = listener.local_addr().unwrap();
    println!("Listening for connections on {}", addr);
    let (tx, rx) = channel();

    for socket in listener.incoming() {
        /* 1. Read incoming msg channel
         * 2. If it does not exist, spawn new thread
         * 3. Send socket to thread
         */
        let mut result_str = String::new();
        let mut socket = socket.unwrap();
        let mut stream = BufReader::new(socket.try_clone().unwrap());
        let len = stream.read_line(&mut result_str);
        println!("Got string {}", result_str);
        socket.set_read_timeout(option);
        if !spawned.contains_key(&result_str) {
            let tx = tx.clone();
            let result_str = result_str.clone();
            thread::spawn(move|| {
                let (thr_tx, thr_rx) = channel();
                println!("Spawned: Sending socket");
                let mut users = HashMap::new();
                tx.send(thr_tx.clone()).unwrap();
                println!("Got thr sock");
                loop {
                    let mut removals = Vec::new();
                    match thr_rx.try_recv() {
                        Ok(val) => {
                            let thr: TcpStream = val;
                            cnt += 1;
                            users.insert(cnt, thr);
                        },
                        Err(_) => {}
                    }
                    let mut bufstr = String::new();
                    for (user, thr_socket) in &users {
                        let mut stream = BufReader::new(thr_socket.try_clone().unwrap());
                        match stream.read_line(&mut bufstr) {
                            Ok(len) => {
                                if(len > 0) {
                                    for (another_user, mut thr_sock) in &users {
                                        if user != another_user {
                                            let mut msg = user.to_string();
                                            msg.push_str(":");
                                            msg.push_str(&bufstr);
                                            thr_sock.write(&msg.into_bytes()).unwrap();
                                        }
                                    }
                                }
                                else {
                                    /* Socket closed, drop user */
                                    removals.push(user.clone());
                                }
                            },
                            Err(_) => {
                                //println!("{:?}", err.kind());
                            }
                        }
                    }
                    for removal in &removals {
                        users.remove(removal);
                    }
                    if cnt > 0 && users.is_empty() {
                        break;
                    }
                }
            });
            let thr_feed = rx.recv().unwrap();
            println!("Got thr_feed");
            spawned.insert(result_str.clone(), thr_feed);
        }
        let thr_socket = spawned.get_mut(&result_str).unwrap();
        thr_socket.send(socket).unwrap();
    }
}
<|endoftext|>


<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/

//! Loader for a single run, with one or more event files.

use std::collections::{BTreeMap, HashMap, HashSet};
use std::fs::File;
use std::io::BufReader;
use std::path::PathBuf;
use std::sync::RwLock;

use crate::commit;
use crate::data_compat::{EventValue, SummaryValue};
use crate::event_file::EventFileReader;
use crate::proto::tensorboard as pb;
use crate::reservoir::StageReservoir;
use crate::types::{Step, Tag, WallTime};

/// A loader to accumulate reservoir-sampled events in a single TensorBoard run.
///
/// For now, a run loader always reads from [`File`]s on disk. In the future, this may be
/// parameterized over a filesystem interface.
#[derive(Debug)]
pub struct RunLoader {
    /// The earliest event `wall_time` seen in any event file in this run.
    ///
    /// This is `None` if and only if no events have been seen. Its value may decrease as new
    /// events are read, but in practice this is expected to be the wall time of the first
    /// `file_version` event in the first event file.
    start_time: Option<WallTime>,

    /// The event files in this run.
    ///
    /// Event files are sorted and read lexicographically by name, which is designed to coincide
    /// with actual start time. See [`EventFile::Dead`] for conditions under which an event file
    /// may be dead. Once an event file is added to this map, it may become dead, but it will not
    /// be removed entirely. This way, we know not to just re-open it again at the next load cycle.
    files: BTreeMap<PathBuf, EventFile<BufReader<File>>>,

    /// Reservoir-sampled data and metadata for each time series.
    time_series: HashMap<Tag, TimeSeries>,
}

#[derive(Debug)]
enum EventFile<R> {
    /// An event file that may still have more valid data.
    Active(EventFileReader<R>),
    /// An event file that can no longer be read.
    ///
    /// This can be because of a non-recoverable read error (e.g., a bad length checksum), due to
    /// the last-read record being very old (note: not yet implemented), or due to the file being
    /// deleted.
    Dead,
}

#[derive(Debug)]
struct TimeSeries {
    data_class: pb::DataClass,
    metadata: Box<pb::SummaryMetadata>,
    rsv: StageReservoir<StageValue>,
}

/// A value staged in the reservoir.
///
/// This is kept as close as possible to the on-disk event representation, since every record in
/// the stream is converted into this format.
#[derive(Debug)]
struct StageValue {
    wall_time: WallTime,
    payload: EventValue,
}

impl TimeSeries {
    fn new(metadata: Box<pb::SummaryMetadata>) -> Self {
        let data_class =
            pb::DataClass::from_i32(metadata.data_class).unwrap_or(pb::DataClass::Unknown);
        let capacity = match data_class {
            pb::DataClass::Scalar => 1000,
            pb::DataClass::Tensor => 100,
            pb::DataClass::BlobSequence => 10,
            _ => 0,
        };
        Self {
            data_class,
            metadata,
            rsv: StageReservoir::new(capacity),
        }
    }

    /// Writes all staged data for this time series into the commit.
    fn commit(&mut self, tag: &Tag, run: &mut commit::RunData) {
        use pb::DataClass;
        match self.data_class {
            DataClass::Scalar => self.commit_to(tag, &mut run.scalars, EventValue::into_scalar),
            DataClass::Tensor => todo!("tensor time series not yet supported"),
            DataClass::BlobSequence => todo!("blob sequence time series not yet supported"),
            _ => (),
        };
    }

    /// Helper for `commit`: writes staged data for this time series into storage for a statically
    /// known data class.
    fn commit_to<V, F: FnMut(EventValue) -> Result<V, commit::DataLoss>>(
        &mut self,
        tag: &Tag,
        store: &mut commit::TagStore<V>,
        mut enrich: F,
    ) {
        let commit_ts = store
            .entry(tag.clone())
            .or_insert_with(|| commit::TimeSeries::new(self.metadata.clone()));
        self.rsv
            .commit_map(&mut commit_ts.basin, |StageValue { wall_time, payload }| {
                (wall_time, enrich(payload))
            });
    }
}

impl RunLoader {
    pub fn new() -> Self {
        Self {
            start_time: None,
            files: BTreeMap::new(),
            time_series: HashMap::new(),
        }
    }

    /// Loads new data given the current set of event files.
    ///
    /// The provided filenames should correspond to the entire set of event files currently part of
    /// this run.
    ///
    /// The given commit must have an entry for this run (the entry may be empty).
    ///
    /// # Panics
    ///
    /// If we need to access `run_data` but the lock is poisoned.
    pub fn reload(&mut self, filenames: Vec<PathBuf>, run_data: &RwLock<commit::RunData>) {
        self.update_file_set(filenames);
        self.reload_files();
        self.commit_all(run_data);
    }

    /// Updates the active key set of `self.files` to match the given filenames.
    ///
    /// After this function returns, `self.files` may still have keys not in `filenames`, but they
    /// will all map to [`EventFile::Dead`].
    fn update_file_set(&mut self, filenames: Vec<PathBuf>) {
        // Remove any discarded files.
        let new_file_set: HashSet<&PathBuf> = filenames.iter().collect();
        for (k, v) in self.files.iter_mut() {
            if !new_file_set.contains(k) {
                *v = EventFile::Dead;
            }
        }

        // Open readers for any new files.
        for filename in filenames {
            use std::collections::btree_map::Entry;
            match self.files.entry(filename) {
                Entry::Occupied(_) => {}
                Entry::Vacant(v) => {
                    let event_file = match File::open(v.key()) {
                        Ok(file) => {
                            let reader = EventFileReader::new(BufReader::new(file));
                            EventFile::Active(reader)
                        }
                        // TODO(@wchargin): Improve error handling?
                        Err(e) => {
                            eprintln!("failed to open event file {:?}: {:?}", v.key(), e);
                            EventFile::Dead
                        }
                    };
                    v.insert(event_file);
                }
            };
        }
    }

    /// Reads data from all active event files.
    fn reload_files(&mut self) {
        for (filename, ef) in self.files.iter_mut() {
            let reader = match ef {
                EventFile::Dead => continue,
                EventFile::Active(reader) => reader,
            };

            loop {
                use crate::event_file::ReadEventError::ReadRecordError;
                use crate::tf_record::ReadRecordError::Truncated;
                let event = match reader.read_event() {
                    Ok(event) => event,
                    Err(ReadRecordError(Truncated)) => break,
                    Err(e) => {
                        // TODO(@wchargin): Improve error handling?
                        eprintln!("read error in {}: {:?}", filename.display(), e);
                        *ef = EventFile::Dead;
                        break;
                    }
                };
                read_event(&mut self.time_series, &mut self.start_time, event);
            }
        }
    }

    fn commit_all(&mut self, run_data: &RwLock<commit::RunData>) {
        let mut run = run_data.write().expect("acquiring run data lock");
        for (tag, ts) in &mut self.time_series {
            ts.commit(tag, &mut *run);
        }
    }
}

/// Reads a single event into the structures of a run loader.
///
/// This is a standalone function because it's called from `reload_files` in a context that already
/// has an exclusive reference into `self.files`, and so can't call methods on `&mut self`.
fn read_event(
    time_series: &mut HashMap<Tag, TimeSeries>,
    start_time: &mut Option<WallTime>,
    e: pb::Event,
) {
    let step = Step(e.step);
    let wall_time = match WallTime::new(e.wall_time) {
        None => {
            // TODO(@wchargin): Improve error handling.
            eprintln!(
                "dropping event at step {} with invalid wall time {}",
                e.step, e.wall_time
            );
            return;
        }
        Some(wt) => wt,
    };
    if start_time.map(|start| wall_time < start).unwrap_or(true) {
        *start_time = Some(wall_time);
    }
    match e.what {
        Some(pb::event::What::GraphDef(_)) => {
            // TODO(@wchargin): Handle run graphs.
            eprintln!("graph_def events not yet handled");
        }
        Some(pb::event::What::Summary(sum)) => {
            for mut summary_pb_value in sum.value {
                let summary_value = match summary_pb_value.value {
                    None => continue,
                    Some(v) => SummaryValue(Box::new(v)),
                };

                use std::collections::hash_map::Entry;
                let ts = match time_series.entry(Tag(summary_pb_value.tag)) {
                    Entry::Occupied(o) => o.into_mut(),
                    Entry::Vacant(v) => {
                        let metadata =
                            summary_value.initial_metadata(summary_pb_value.metadata.take());
                        v.insert(TimeSeries::new(metadata))
                    }
                };
                let sv = StageValue {
                    wall_time,
                    payload: EventValue::Summary(summary_value),
                };
                ts.rsv.offer(step, sv);
            }
        }
        _ => {}
    }
}

impl Default for RunLoader {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::fs::File;
    use std::io::BufWriter;

    use crate::commit::Commit;
    use crate::types::Run;
    use crate::writer::SummaryWriteExt;

    #[test]
    fn test() -> Result<(), Box<dyn std::error::Error>> {
        let logdir = tempfile::tempdir()?;
        let f1_name = logdir.path().join("tfevents.123");
        let f2_name = logdir.path().join("tfevents.456");
        let mut f1 = BufWriter::new(File::create(&f1_name)?);
        let mut f2 = BufWriter::new(File::create(&f2_name)?);

        // Write file versions.
        for (f, wall_time) in &mut [(&mut f1, 1234.0), (&mut f2, 2345.0)] {
            let file_version = pb::Event {
                wall_time: *wall_time,
                what: Some(pb::event::What::FileVersion("brain.Event:2".to_string())),
                ..Default::default()
            };
            f.write_event(&file_version)?;
        }

        // Write some data points across both files.
        let run = Run("train".to_string());
        let tag = Tag("accuracy".to_string());
        f1.write_scalar(&tag, Step(0), WallTime::new(1235.0).unwrap(), 0.25)?;
        f1.write_scalar(&tag, Step(1), WallTime::new(1236.0).unwrap(), 0.50)?;
        f1.write_scalar(&tag, Step(2), WallTime::new(1237.0).unwrap(), 0.75)?;
        f1.write_scalar(&tag, Step(3), WallTime::new(1238.0).unwrap(), 1.00)?;
        // preempt!
        f2.write_scalar(&tag, Step(2), WallTime::new(2346.0).unwrap(), 0.70)?;
        f2.write_scalar(&tag, Step(3), WallTime::new(2347.0).unwrap(), 0.85)?;
        f2.write_scalar(&tag, Step(4), WallTime::new(2348.0).unwrap(), 0.90)?;
        // flush, so that the data's there when we read it
        f1.into_inner()?.sync_all()?;
        f2.into_inner()?.sync_all()?;

        let mut loader = RunLoader::new();
        let commit = Commit::new();
        commit
            .runs
            .write()
            .expect("write-locking runs map")
            .insert(run.clone(), Default::default());
        loader.reload(vec![f1_name, f2_name], &commit.runs.read().unwrap()[&run]);

        // Start time should be that of the file version event, even though that didn't correspond
        // to any time series.
        assert_eq!(loader.start_time, Some(WallTime::new(1234.0).unwrap()));

        let runs = commit.runs.read().expect("read-locking runs map");
        let run_data: &commit::RunData = &*runs
            .get(&run)
            .expect("looking up data for run")
            .read()
            .expect("read-locking run data map");

        assert_eq!(run_data.scalars.keys().collect::<Vec<_>>(), vec![&tag]);
        let ts = run_data.scalars.get(&tag).unwrap();
        assert_eq!(
            *ts.metadata,
            pb::SummaryMetadata {
                plugin_data: Some(pb::summary_metadata::PluginData {
                    plugin_name: crate::data_compat::SCALARS_PLUGIN_NAME.to_string(),
                    ..Default::default()
                }),
                data_class: pb::DataClass::Scalar.into(),
                ..Default::default()
            }
        );

        // Points should be as expected (no downsampling at these sizes).
        let scalar = |f: f32| commit::ScalarValue(f64::from(f));
        assert_eq!(
            ts.valid_values().collect::<Vec<_>>(),
            vec![
                (Step(0), WallTime::new(1235.0).unwrap(), &scalar(0.25)),
                (Step(1), WallTime::new(1236.0).unwrap(), &scalar(0.50)),
                (Step(2), WallTime::new(2346.0).unwrap(), &scalar(0.70)),
                (Step(3), WallTime::new(2347.0).unwrap(), &scalar(0.85)),
                (Step(4), WallTime::new(2348.0).unwrap(), &scalar(0.90)),
            ]
        );

        Ok(())
    }
}
<commit_msg>rust: rename `run::TimeSeries` to `StageTimeSeries` (#4384)<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/

//! Loader for a single run, with one or more event files.

use std::collections::{BTreeMap, HashMap, HashSet};
use std::fs::File;
use std::io::BufReader;
use std::path::PathBuf;
use std::sync::RwLock;

use crate::commit;
use crate::data_compat::{EventValue, SummaryValue};
use crate::event_file::EventFileReader;
use crate::proto::tensorboard as pb;
use crate::reservoir::StageReservoir;
use crate::types::{Step, Tag, WallTime};

/// A loader to accumulate reservoir-sampled events in a single TensorBoard run.
///
/// For now, a run loader always reads from [`File`]s on disk. In the future, this may be
/// parameterized over a filesystem interface.
#[derive(Debug)]
pub struct RunLoader {
    /// The earliest event `wall_time` seen in any event file in this run.
    ///
    /// This is `None` if and only if no events have been seen. Its value may decrease as new
    /// events are read, but in practice this is expected to be the wall time of the first
    /// `file_version` event in the first event file.
    start_time: Option<WallTime>,

    /// The event files in this run.
    ///
    /// Event files are sorted and read lexicographically by name, which is designed to coincide
    /// with actual start time. See [`EventFile::Dead`] for conditions under which an event file
    /// may be dead. Once an event file is added to this map, it may become dead, but it will not
    /// be removed entirely. This way, we know not to just re-open it again at the next load cycle.
    files: BTreeMap<PathBuf, EventFile<BufReader<File>>>,

    /// Reservoir-sampled data and metadata for each time series.
    time_series: HashMap<Tag, StageTimeSeries>,
}

#[derive(Debug)]
enum EventFile<R> {
    /// An event file that may still have more valid data.
    Active(EventFileReader<R>),
    /// An event file that can no longer be read.
    ///
    /// This can be because of a non-recoverable read error (e.g., a bad length checksum), due to
    /// the last-read record being very old (note: not yet implemented), or due to the file being
    /// deleted.
    Dead,
}

#[derive(Debug)]
struct StageTimeSeries {
    data_class: pb::DataClass,
    metadata: Box<pb::SummaryMetadata>,
    rsv: StageReservoir<StageValue>,
}

/// A value staged in the reservoir.
///
/// This is kept as close as possible to the on-disk event representation, since every record in
/// the stream is converted into this format.
#[derive(Debug)]
struct StageValue {
    wall_time: WallTime,
    payload: EventValue,
}

impl StageTimeSeries {
    fn new(metadata: Box<pb::SummaryMetadata>) -> Self {
        let data_class =
            pb::DataClass::from_i32(metadata.data_class).unwrap_or(pb::DataClass::Unknown);
        let capacity = match data_class {
            pb::DataClass::Scalar => 1000,
            pb::DataClass::Tensor => 100,
            pb::DataClass::BlobSequence => 10,
            _ => 0,
        };
        Self {
            data_class,
            metadata,
            rsv: StageReservoir::new(capacity),
        }
    }

    /// Writes all staged data for this time series into the commit.
    fn commit(&mut self, tag: &Tag, run: &mut commit::RunData) {
        use pb::DataClass;
        match self.data_class {
            DataClass::Scalar => self.commit_to(tag, &mut run.scalars, EventValue::into_scalar),
            DataClass::Tensor => todo!("tensor time series not yet supported"),
            DataClass::BlobSequence => todo!("blob sequence time series not yet supported"),
            _ => (),
        };
    }

    /// Helper for `commit`: writes staged data for this time series into storage for a statically
    /// known data class.
    fn commit_to<V, F: FnMut(EventValue) -> Result<V, commit::DataLoss>>(
        &mut self,
        tag: &Tag,
        store: &mut commit::TagStore<V>,
        mut enrich: F,
    ) {
        let commit_ts = store
            .entry(tag.clone())
            .or_insert_with(|| commit::TimeSeries::new(self.metadata.clone()));
        self.rsv
            .commit_map(&mut commit_ts.basin, |StageValue { wall_time, payload }| {
                (wall_time, enrich(payload))
            });
    }
}

impl RunLoader {
    pub fn new() -> Self {
        Self {
            start_time: None,
            files: BTreeMap::new(),
            time_series: HashMap::new(),
        }
    }

    /// Loads new data given the current set of event files.
    ///
    /// The provided filenames should correspond to the entire set of event files currently part of
    /// this run.
    ///
    /// The given commit must have an entry for this run (the entry may be empty).
    ///
    /// # Panics
    ///
    /// If we need to access `run_data` but the lock is poisoned.
    pub fn reload(&mut self, filenames: Vec<PathBuf>, run_data: &RwLock<commit::RunData>) {
        self.update_file_set(filenames);
        self.reload_files();
        self.commit_all(run_data);
    }

    /// Updates the active key set of `self.files` to match the given filenames.
    ///
    /// After this function returns, `self.files` may still have keys not in `filenames`, but they
    /// will all map to [`EventFile::Dead`].
    fn update_file_set(&mut self, filenames: Vec<PathBuf>) {
        // Remove any discarded files.
        let new_file_set: HashSet<&PathBuf> = filenames.iter().collect();
        for (k, v) in self.files.iter_mut() {
            if !new_file_set.contains(k) {
                *v = EventFile::Dead;
            }
        }

        // Open readers for any new files.
        for filename in filenames {
            use std::collections::btree_map::Entry;
            match self.files.entry(filename) {
                Entry::Occupied(_) => {}
                Entry::Vacant(v) => {
                    let event_file = match File::open(v.key()) {
                        Ok(file) => {
                            let reader = EventFileReader::new(BufReader::new(file));
                            EventFile::Active(reader)
                        }
                        // TODO(@wchargin): Improve error handling?
                        Err(e) => {
                            eprintln!("failed to open event file {:?}: {:?}", v.key(), e);
                            EventFile::Dead
                        }
                    };
                    v.insert(event_file);
                }
            };
        }
    }

    /// Reads data from all active event files.
    fn reload_files(&mut self) {
        for (filename, ef) in self.files.iter_mut() {
            let reader = match ef {
                EventFile::Dead => continue,
                EventFile::Active(reader) => reader,
            };

            loop {
                use crate::event_file::ReadEventError::ReadRecordError;
                use crate::tf_record::ReadRecordError::Truncated;
                let event = match reader.read_event() {
                    Ok(event) => event,
                    Err(ReadRecordError(Truncated)) => break,
                    Err(e) => {
                        // TODO(@wchargin): Improve error handling?
                        eprintln!("read error in {}: {:?}", filename.display(), e);
                        *ef = EventFile::Dead;
                        break;
                    }
                };
                read_event(&mut self.time_series, &mut self.start_time, event);
            }
        }
    }

    fn commit_all(&mut self, run_data: &RwLock<commit::RunData>) {
        let mut run = run_data.write().expect("acquiring run data lock");
        for (tag, ts) in &mut self.time_series {
            ts.commit(tag, &mut *run);
        }
    }
}

/// Reads a single event into the structures of a run loader.
///
/// This is a standalone function because it's called from `reload_files` in a context that already
/// has an exclusive reference into `self.files`, and so can't call methods on `&mut self`.
fn read_event(
    time_series: &mut HashMap<Tag, StageTimeSeries>,
    start_time: &mut Option<WallTime>,
    e: pb::Event,
) {
    let step = Step(e.step);
    let wall_time = match WallTime::new(e.wall_time) {
        None => {
            // TODO(@wchargin): Improve error handling.
            eprintln!(
                "dropping event at step {} with invalid wall time {}",
                e.step, e.wall_time
            );
            return;
        }
        Some(wt) => wt,
    };
    if start_time.map(|start| wall_time < start).unwrap_or(true) {
        *start_time = Some(wall_time);
    }
    match e.what {
        Some(pb::event::What::GraphDef(_)) => {
            // TODO(@wchargin): Handle run graphs.
            eprintln!("graph_def events not yet handled");
        }
        Some(pb::event::What::Summary(sum)) => {
            for mut summary_pb_value in sum.value {
                let summary_value = match summary_pb_value.value {
                    None => continue,
                    Some(v) => SummaryValue(Box::new(v)),
                };

                use std::collections::hash_map::Entry;
                let ts = match time_series.entry(Tag(summary_pb_value.tag)) {
                    Entry::Occupied(o) => o.into_mut(),
                    Entry::Vacant(v) => {
                        let metadata =
                            summary_value.initial_metadata(summary_pb_value.metadata.take());
                        v.insert(StageTimeSeries::new(metadata))
                    }
                };
                let sv = StageValue {
                    wall_time,
                    payload: EventValue::Summary(summary_value),
                };
                ts.rsv.offer(step, sv);
            }
        }
        _ => {}
    }
}

impl Default for RunLoader {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::fs::File;
    use std::io::BufWriter;

    use crate::commit::Commit;
    use crate::types::Run;
    use crate::writer::SummaryWriteExt;

    #[test]
    fn test() -> Result<(), Box<dyn std::error::Error>> {
        let logdir = tempfile::tempdir()?;
        let f1_name = logdir.path().join("tfevents.123");
        let f2_name = logdir.path().join("tfevents.456");
        let mut f1 = BufWriter::new(File::create(&f1_name)?);
        let mut f2 = BufWriter::new(File::create(&f2_name)?);

        // Write file versions.
        for (f, wall_time) in &mut [(&mut f1, 1234.0), (&mut f2, 2345.0)] {
            let file_version = pb::Event {
                wall_time: *wall_time,
                what: Some(pb::event::What::FileVersion("brain.Event:2".to_string())),
                ..Default::default()
            };
            f.write_event(&file_version)?;
        }

        // Write some data points across both files.
        let run = Run("train".to_string());
        let tag = Tag("accuracy".to_string());
        f1.write_scalar(&tag, Step(0), WallTime::new(1235.0).unwrap(), 0.25)?;
        f1.write_scalar(&tag, Step(1), WallTime::new(1236.0).unwrap(), 0.50)?;
        f1.write_scalar(&tag, Step(2), WallTime::new(1237.0).unwrap(), 0.75)?;
        f1.write_scalar(&tag, Step(3), WallTime::new(1238.0).unwrap(), 1.00)?;
        // preempt!
        f2.write_scalar(&tag, Step(2), WallTime::new(2346.0).unwrap(), 0.70)?;
        f2.write_scalar(&tag, Step(3), WallTime::new(2347.0).unwrap(), 0.85)?;
        f2.write_scalar(&tag, Step(4), WallTime::new(2348.0).unwrap(), 0.90)?;
        // flush, so that the data's there when we read it
        f1.into_inner()?.sync_all()?;
        f2.into_inner()?.sync_all()?;

        let mut loader = RunLoader::new();
        let commit = Commit::new();
        commit
            .runs
            .write()
            .expect("write-locking runs map")
            .insert(run.clone(), Default::default());
        loader.reload(vec![f1_name, f2_name], &commit.runs.read().unwrap()[&run]);

        // Start time should be that of the file version event, even though that didn't correspond
        // to any time series.
        assert_eq!(loader.start_time, Some(WallTime::new(1234.0).unwrap()));

        let runs = commit.runs.read().expect("read-locking runs map");
        let run_data: &commit::RunData = &*runs
            .get(&run)
            .expect("looking up data for run")
            .read()
            .expect("read-locking run data map");

        assert_eq!(run_data.scalars.keys().collect::<Vec<_>>(), vec![&tag]);
        let ts = run_data.scalars.get(&tag).unwrap();
        assert_eq!(
            *ts.metadata,
            pb::SummaryMetadata {
                plugin_data: Some(pb::summary_metadata::PluginData {
                    plugin_name: crate::data_compat::SCALARS_PLUGIN_NAME.to_string(),
                    ..Default::default()
                }),
                data_class: pb::DataClass::Scalar.into(),
                ..Default::default()
            }
        );

        // Points should be as expected (no downsampling at these sizes).
        let scalar = |f: f32| commit::ScalarValue(f64::from(f));
        assert_eq!(
            ts.valid_values().collect::<Vec<_>>(),
            vec![
                (Step(0), WallTime::new(1235.0).unwrap(), &scalar(0.25)),
                (Step(1), WallTime::new(1236.0).unwrap(), &scalar(0.50)),
                (Step(2), WallTime::new(2346.0).unwrap(), &scalar(0.70)),
                (Step(3), WallTime::new(2347.0).unwrap(), &scalar(0.85)),
                (Step(4), WallTime::new(2348.0).unwrap(), &scalar(0.90)),
            ]
        );

        Ok(())
    }
}
<|endoftext|>


<commit_before>use spin::Mutex;

pub struct Allocator {
    heap_start: usize,
    heap_size: usize
}

impl Allocator {
    pub const fn new(heap_start: usize, heap_size: usize) -> Allocator {
        Allocator {
            heap_start: heap_start,
            heap_size: heap_size
        }
    }

    pub fn alloc(&mut self, size: usize, align: usize) -> Option<*mut u8> {
        // TODO
        unimplemented!();
    }

    pub fn dealloc(&mut self, ptr: *mut u8, size: usize, align: usize) {
        // TODO
        unimplemented!();
    }
}<commit_msg>Implemeted the basic allocation logic in jem_alloc.<commit_after>use spin::Mutex;
use core::ptr;

const HEADER_SIZE: usize = 8;
const MIN_BLOCK_SIZE: usize = 16;

pub struct Allocator {
    heap_start: usize,
    heap_size: usize,
    position: usize
}

impl Allocator {
    pub const fn new(heap_start: usize, heap_size: usize) -> Allocator {
        Allocator {
            heap_start: heap_start,
            heap_size: heap_size,
            position: heap_start
        }
    }

    pub fn alloc(&mut self, size: usize, align: usize) -> Option<*mut u8> {
        // Get the size to allocate (the block size).
        let bsize = get_req_size(size);

        // Find the next fitting block and allocate if we can. Otherwise OOM.
        // TODO - Maybe on failure try some kind of compression / de-fragmentation
        if self.next_fit(bsize) {
            self.allocate_position(bsize);
            Some(self.position as *mut u8)
        } else {
            None
        }
    }

    pub fn dealloc(&mut self, ptr: *mut u8, size: usize, align: usize) {
        // TODO
        unimplemented!();
    }

    fn next_fit(&mut self, size: usize) -> bool {
        // Starting at self.position find the next hole that fits. If we pass where we started then
        // we are out of memory and must fail.
        let start = self.position;

        loop {
            // If the next block is free and big enough return here. Because blocks must be at least
            // a multiple of the MIN_BLOCK_SIZE we can be sure the rest of the free space will be
            // at least this big.
            let current_header = unsafe { *(self.position as *const usize) };
            let bsize = block_size(current_header);

            if block_is_free(current_header) {
                // If the next block size is zero it has never been allocated yet. Check we have
                // enough heap left then allocate.
                if bsize == 0 && self.position + size < self.heap_size {
                    return true;
                }

                // Previously freed block that is big enough. Success.
                if bsize > size {
                    return true;
                }
            }

            // Onto the next block...
            let mut new_position = (bsize / MIN_BLOCK_SIZE) + self.position;

            // If we have reached the end, time to go back to the start.
            if new_position > self.heap_size {
                new_position = self.heap_start;
            }

            // Are we passing the initial start? If so then we have failed in our quest to allocate
            // memory. Return false.
            if self.position < start && new_position >= start {
                return false;
            }

            // Update the position and try again...
            self.position = new_position;
        }
    }

    fn allocate_position(&self, size: usize) {
        let header = alloc_header(size);
        let position = self.position as *mut usize;
        unsafe {
            *position = header;
        }
    }

}

const ALLOCATED_MASK: usize = 0x8000_0000_0000_0000;

#[inline]
fn block_is_free(header: usize) -> bool {
    // Free if the first bit is not set
    (header & ALLOCATED_MASK) < 0
}

#[inline]
fn block_size(header: usize) -> usize {
    // Free if the first bit is not set
    (header ^ ALLOCATED_MASK) as usize
}

#[inline]
fn alloc_header(size: usize) -> usize {
    size | ALLOCATED_MASK
}

fn get_req_size(size: usize) -> usize {
    // Block size must be at-least the minimum block size or the smallest multiple.
    if size < MIN_BLOCK_SIZE {
        MIN_BLOCK_SIZE
    } else {
        ((size / MIN_BLOCK_SIZE) + 1) * MIN_BLOCK_SIZE
    }
}<|endoftext|>


<commit_before>//! The compression algorithm.
//!
//! We make use of hash tables to find duplicates. This gives a reasonable compression ratio with a
//! high performance. It has fixed memory usage, which contrary to other approachs, makes it less
//! memory hungry.

use byteorder::{NativeEndian, ByteOrder};

/// Duplication dictionary size.
///
/// Every four bytes is assigned an entry. When this number is lower, fewer entries exists, and
/// thus collisions are more likely, hurting the compression ratio.
const DICTIONARY_SIZE: usize = 4096;

/// A LZ4 block.
///
/// This defines a single compression "unit", consisting of two parts, a number of raw literals,
/// and possibly a pointer to the already encoded buffer from which to copy.
struct Block {
    /// The length (in bytes) of the literals section.
    lit_len: usize,
    /// The duplicates section if any.
    ///
    /// Only the last block in a stream can lack of the duplicates section.
    dup: Option<Duplicate>,
}

/// A consecutive sequence of bytes found in already encoded part of the input.
#[derive(Copy, Clone)]
struct Duplicate {
    /// The number of bytes before our cursor, where the duplicate starts.
    offset: u16,
    /// The length beyond the four first bytes.
    ///
    /// Adding four to this number yields the actual length.
    extra_bytes: usize,
}

/// An LZ4 encoder.
struct Encoder<'a> {
    /// The raw uncompressed input.
    input: &'a [u8],
    /// The compressed output.
    output: &'a mut Vec<u8>,
    /// The number of bytes from the input that are encoded.
    cur: usize,
    /// The dictionary of previously encoded sequences.
    ///
    /// This is used to find duplicates in the stream so they are not written multiple times.
    ///
    /// Every four bytes are hashed, and in the resulting slot their position in the input buffer
    /// is placed. This way we can easily look up a candidate to back references.
    dict: [usize; DICTIONARY_SIZE],
}

impl<'a> Encoder<'a> {
    /// Go forward by some number of bytes.
    ///
    /// This will update the cursor and dictionary to reflect the now processed bytes.
    ///
    /// This returns `false` if all the input bytes are processed.
    fn go_forward(&mut self, steps: usize) -> bool {
        // Go over all the bytes we are skipping and update the cursor and dictionary.
        for _ in 0..steps {
            // Increment the cursor.
            self.cur += 1;
            // Insert the cursor position into the dictionary.
            self.insert_cursor();
        }

        // Return `true` if there's more to read.
        self.cur < self.input.len()
    }

    /// Insert the batch under the cursor into the dictionary.
    fn insert_cursor(&mut self) {
        // Make sure that there is at least one batch remaining.
        if self.remaining_batch() {
            // Insert the cursor into the table.
            self.dict[self.get_cur_hash()] = self.cur;
        }
    }

    /// Check if there are any remaining batches.
    fn remaining_batch(&self) -> bool {
        self.cur + 4 < self.input.len()
    }

    /// Get the hash of the current four bytes below the cursor.
    ///
    /// This is guaranteed to be below `DICTIONARY_SIZE`.
    fn get_cur_hash(&self) -> usize {
        // Use LCG hashing to generate a relatively good hash of the four bytes batch at the
        // cursor.
        self.get_batch_at_cursor().wrapping_mul(134775813).wrapping_add(1) as usize % DICTIONARY_SIZE
    }

    /// Read a 4-byte "batch" from some position.
    ///
    /// This will read a native-endian 4-byte integer from some position.
    fn get_batch(&self, n: usize) -> u32 {
        debug_assert!(self.remaining_batch(), "Reading a partial batch.");

        NativeEndian::read_u32(&self.input[n..])
    }

    /// Read the batch at the cursor.
    fn get_batch_at_cursor(&self) -> u32 {
        self.get_batch(self.cur)
    }

    /// Find a duplicate of the current batch.
    ///
    /// If any duplicate is found, a tuple `(position, size - 4)` is returned.
    fn find_duplicate(&self) -> Option<Duplicate> {
        // If there is no remaining batch, we return none.
        if !self.remaining_batch() {
            println!("No remaining batches.");
            return None;
        }

        // Find a candidate in the dictionary by hashing the current four bytes.
        let candidate = self.dict[self.get_cur_hash()];

        // We need to ensure two things:
        // - We should not return a position which is merely a hash collision, so we need to check
        //   that the candidate actually matches what we search for.
        // - We can address up to 16-bit offset, hence we are only able to address the candidate if
        //   its offset is less than or equals to 0xFFFF.
        if self.get_batch(candidate) == self.get_batch_at_cursor() && self.cur - candidate <= 0xFFFF {
            // Calculate the "extension bytes", i.e. the duplicate bytes beyond the batch.
            let ext = self.input[self.cur + 4..].iter().zip(&self.input[candidate + 4..]).count();

            Some(Duplicate {
                offset: (self.cur - candidate) as u16,
                extra_bytes: ext,
            })
        } else { None }
    }

    /// Write an integer to the output in LSIC format.
    fn write_integer(&mut self, mut n: usize) {
        // Write the 0xFF bytes as long as the integer is higher than said value.
        while n > 0xFF {
            n -= 0xFF;
            self.output.push(0xFF);
        }

        // Write the remaining byte.
        self.output.push(n as u8);
    }

    /// Read the block of the top of the stream.
    fn pop_block(&mut self) -> Block {
        // The length of the literals section.
        let mut lit = 0;

        loop {
            // Search for a duplicate.
            if let Some(dup) = self.find_duplicate() {
                // We found a duplicate, so the literals section is over...

                // Move forward. Note that `ext` is actually the steps minus 4, because of the
                // minimum matchlenght, so we need to add 4.
                self.go_forward(dup.extra_bytes + 4);

                return Block {
                    lit_len: lit,
                    dup: Some(dup),
                };
            }

            // Try to move forward.
            if !self.go_forward(1) {
                // We reached the end of the stream, and no duplicates section follows.
                return Block {
                    lit_len: lit,
                    dup: None,
                };
            }

            // No duplicates found yet, so extend the literals section.
            lit += 1;
        }
    }

    /// Complete the encoding into `self.output`.
    fn complete(&mut self) {
        // Construct one block at a time.
        loop {
            // The start of the literals section.
            let start = self.cur;

            // Read the next block into two sections, the literals and the duplicates.
            let block = self.pop_block();

            // Generate the higher half of the token.
            let mut token = if block.lit_len < 0xF {
                // Since we can fit the literals length into it, there is no need for saturation.
                (block.lit_len as u8) << 4
            } else {
                // We were unable to fit the literals into it, so we saturate to 0xF. We will later
                // write the extensional value through LSIC encoding.
                0xF0
            };

            // Generate the lower half of the token, the duplicates length.
            let dup_extra_len = block.dup.map_or(0, |x| x.extra_bytes);
            token |= if dup_extra_len < 0xF {
                // We could fit it in.
                dup_extra_len as u8
            } else {
                // We were unable to fit it in, so we default to 0xF, which will later be extended
                // by LSIC encoding.
                0xF
            };

            // Push the token to the output stream.
            self.output.push(token);

            // If we were unable to fit the literals length into the token, write the extensional
            // part through LSIC.
            if block.lit_len >= 0xF {
                self.write_integer(block.lit_len - 0xF);
            }

            // Now, write the actual literals.
            self.output.extend_from_slice(&self.input[start..start + block.lit_len]);

            if let Some(Duplicate { offset, .. }) = block.dup {
                // Wait! There's more. Now, we encode the duplicates section.

                // Push the offset in little endian.
                self.output.push(offset as u8);
                self.output.push((offset >> 8) as u8);

                // If we were unable to fit the duplicates length into the token, write the
                // extensional part through LSIC.
                if dup_extra_len >= 0xF {
                    self.write_integer(dup_extra_len - 0xF);
                }
            } else {
                break;
            }
        }
    }
}

/// Compress all bytes of `input`.
pub fn compress(input: &[u8]) -> Vec<u8> {
    // In most cases, the compression won't expand the size, so we set the input size as capacity.
    let mut vec = Vec::with_capacity(input.len());

    Encoder {
        input: input,
        output: &mut vec,
        cur: 0,
        dict: [0; DICTIONARY_SIZE],
    }.complete();

    vec
}
<commit_msg>Remove log point that slipped though debugging.<commit_after>//! The compression algorithm.
//!
//! We make use of hash tables to find duplicates. This gives a reasonable compression ratio with a
//! high performance. It has fixed memory usage, which contrary to other approachs, makes it less
//! memory hungry.

use byteorder::{NativeEndian, ByteOrder};

/// Duplication dictionary size.
///
/// Every four bytes is assigned an entry. When this number is lower, fewer entries exists, and
/// thus collisions are more likely, hurting the compression ratio.
const DICTIONARY_SIZE: usize = 4096;

/// A LZ4 block.
///
/// This defines a single compression "unit", consisting of two parts, a number of raw literals,
/// and possibly a pointer to the already encoded buffer from which to copy.
struct Block {
    /// The length (in bytes) of the literals section.
    lit_len: usize,
    /// The duplicates section if any.
    ///
    /// Only the last block in a stream can lack of the duplicates section.
    dup: Option<Duplicate>,
}

/// A consecutive sequence of bytes found in already encoded part of the input.
#[derive(Copy, Clone)]
struct Duplicate {
    /// The number of bytes before our cursor, where the duplicate starts.
    offset: u16,
    /// The length beyond the four first bytes.
    ///
    /// Adding four to this number yields the actual length.
    extra_bytes: usize,
}

/// An LZ4 encoder.
struct Encoder<'a> {
    /// The raw uncompressed input.
    input: &'a [u8],
    /// The compressed output.
    output: &'a mut Vec<u8>,
    /// The number of bytes from the input that are encoded.
    cur: usize,
    /// The dictionary of previously encoded sequences.
    ///
    /// This is used to find duplicates in the stream so they are not written multiple times.
    ///
    /// Every four bytes are hashed, and in the resulting slot their position in the input buffer
    /// is placed. This way we can easily look up a candidate to back references.
    dict: [usize; DICTIONARY_SIZE],
}

impl<'a> Encoder<'a> {
    /// Go forward by some number of bytes.
    ///
    /// This will update the cursor and dictionary to reflect the now processed bytes.
    ///
    /// This returns `false` if all the input bytes are processed.
    fn go_forward(&mut self, steps: usize) -> bool {
        // Go over all the bytes we are skipping and update the cursor and dictionary.
        for _ in 0..steps {
            // Increment the cursor.
            self.cur += 1;
            // Insert the cursor position into the dictionary.
            self.insert_cursor();
        }

        // Return `true` if there's more to read.
        self.cur < self.input.len()
    }

    /// Insert the batch under the cursor into the dictionary.
    fn insert_cursor(&mut self) {
        // Make sure that there is at least one batch remaining.
        if self.remaining_batch() {
            // Insert the cursor into the table.
            self.dict[self.get_cur_hash()] = self.cur;
        }
    }

    /// Check if there are any remaining batches.
    fn remaining_batch(&self) -> bool {
        self.cur + 4 < self.input.len()
    }

    /// Get the hash of the current four bytes below the cursor.
    ///
    /// This is guaranteed to be below `DICTIONARY_SIZE`.
    fn get_cur_hash(&self) -> usize {
        // Use LCG hashing to generate a relatively good hash of the four bytes batch at the
        // cursor.
        self.get_batch_at_cursor().wrapping_mul(134775813).wrapping_add(1) as usize % DICTIONARY_SIZE
    }

    /// Read a 4-byte "batch" from some position.
    ///
    /// This will read a native-endian 4-byte integer from some position.
    fn get_batch(&self, n: usize) -> u32 {
        debug_assert!(self.remaining_batch(), "Reading a partial batch.");

        NativeEndian::read_u32(&self.input[n..])
    }

    /// Read the batch at the cursor.
    fn get_batch_at_cursor(&self) -> u32 {
        self.get_batch(self.cur)
    }

    /// Find a duplicate of the current batch.
    ///
    /// If any duplicate is found, a tuple `(position, size - 4)` is returned.
    fn find_duplicate(&self) -> Option<Duplicate> {
        // If there is no remaining batch, we return none.
        if !self.remaining_batch() {
            return None;
        }

        // Find a candidate in the dictionary by hashing the current four bytes.
        let candidate = self.dict[self.get_cur_hash()];

        // We need to ensure two things:
        // - We should not return a position which is merely a hash collision, so we need to check
        //   that the candidate actually matches what we search for.
        // - We can address up to 16-bit offset, hence we are only able to address the candidate if
        //   its offset is less than or equals to 0xFFFF.
        if self.get_batch(candidate) == self.get_batch_at_cursor() && self.cur - candidate <= 0xFFFF {
            // Calculate the "extension bytes", i.e. the duplicate bytes beyond the batch.
            let ext = self.input[self.cur + 4..].iter().zip(&self.input[candidate + 4..]).count();

            Some(Duplicate {
                offset: (self.cur - candidate) as u16,
                extra_bytes: ext,
            })
        } else { None }
    }

    /// Write an integer to the output in LSIC format.
    fn write_integer(&mut self, mut n: usize) {
        // Write the 0xFF bytes as long as the integer is higher than said value.
        while n > 0xFF {
            n -= 0xFF;
            self.output.push(0xFF);
        }

        // Write the remaining byte.
        self.output.push(n as u8);
    }

    /// Read the block of the top of the stream.
    fn pop_block(&mut self) -> Block {
        // The length of the literals section.
        let mut lit = 0;

        loop {
            // Search for a duplicate.
            if let Some(dup) = self.find_duplicate() {
                // We found a duplicate, so the literals section is over...

                // Move forward. Note that `ext` is actually the steps minus 4, because of the
                // minimum matchlenght, so we need to add 4.
                self.go_forward(dup.extra_bytes + 4);

                return Block {
                    lit_len: lit,
                    dup: Some(dup),
                };
            }

            // Try to move forward.
            if !self.go_forward(1) {
                // We reached the end of the stream, and no duplicates section follows.
                return Block {
                    lit_len: lit,
                    dup: None,
                };
            }

            // No duplicates found yet, so extend the literals section.
            lit += 1;
        }
    }

    /// Complete the encoding into `self.output`.
    fn complete(&mut self) {
        // Construct one block at a time.
        loop {
            // The start of the literals section.
            let start = self.cur;

            // Read the next block into two sections, the literals and the duplicates.
            let block = self.pop_block();

            // Generate the higher half of the token.
            let mut token = if block.lit_len < 0xF {
                // Since we can fit the literals length into it, there is no need for saturation.
                (block.lit_len as u8) << 4
            } else {
                // We were unable to fit the literals into it, so we saturate to 0xF. We will later
                // write the extensional value through LSIC encoding.
                0xF0
            };

            // Generate the lower half of the token, the duplicates length.
            let dup_extra_len = block.dup.map_or(0, |x| x.extra_bytes);
            token |= if dup_extra_len < 0xF {
                // We could fit it in.
                dup_extra_len as u8
            } else {
                // We were unable to fit it in, so we default to 0xF, which will later be extended
                // by LSIC encoding.
                0xF
            };

            // Push the token to the output stream.
            self.output.push(token);

            // If we were unable to fit the literals length into the token, write the extensional
            // part through LSIC.
            if block.lit_len >= 0xF {
                self.write_integer(block.lit_len - 0xF);
            }

            // Now, write the actual literals.
            self.output.extend_from_slice(&self.input[start..start + block.lit_len]);

            if let Some(Duplicate { offset, .. }) = block.dup {
                // Wait! There's more. Now, we encode the duplicates section.

                // Push the offset in little endian.
                self.output.push(offset as u8);
                self.output.push((offset >> 8) as u8);

                // If we were unable to fit the duplicates length into the token, write the
                // extensional part through LSIC.
                if dup_extra_len >= 0xF {
                    self.write_integer(dup_extra_len - 0xF);
                }
            } else {
                break;
            }
        }
    }
}

/// Compress all bytes of `input`.
pub fn compress(input: &[u8]) -> Vec<u8> {
    // In most cases, the compression won't expand the size, so we set the input size as capacity.
    let mut vec = Vec::with_capacity(input.len());

    Encoder {
        input: input,
        output: &mut vec,
        cur: 0,
        dict: [0; DICTIONARY_SIZE],
    }.complete();

    vec
}
<|endoftext|>


<commit_before>//! Memory layout of matrices
//!
//! Different from ndarray format which consists of shape and strides,
//! matrix format in LAPACK consists of row or column size and leading dimension.
//!
//! ndarray format and stride
//! --------------------------
//!
//! Let us consider 3-dimensional array for explaining ndarray structure.
//! The address of `(x,y,z)`-element in ndarray satisfies following relation:
//!
//! ```text
//! shape = [Nx, Ny, Nz]
//!     where Nx > 0, Ny > 0, Nz > 0
//! stride = [Sx, Sy, Sz]
//!
//! &data[(x, y, z)] = &data[(0, 0, 0)] + Sx*x + Sy*y + Sz*z
//!     for x < Nx, y < Ny, z < Nz
//! ```
//!
//! The array is called
//!
//! - C-continuous if `[Sx, Sy, Sz] = [Nz*Ny, Nz, 1]`
//! - F(Fortran)-continuous if `[Sx, Sy, Sz] = [1, Nx, Nx*Ny]`
//!
//! Strides of ndarray `[Sx, Sy, Sz]` take arbitrary value,
//! e.g. it can be non-ordered `Sy > Sx > Sz`, or can be negative `Sx < 0`.
//! If the minimum of `[Sx, Sy, Sz]` equals to `1`,
//! the value of elements fills `data` memory region and called "continuous".
//! Non-continuous ndarray is useful to get sub-array without copying data.
//!
//! Matrix layout for LAPACK
//! -------------------------
//!
//! LAPACK interface focuses on the linear algebra operations for F-continuous 2-dimensional array.
//! Under this restriction, stride becomes far simpler; we only have to consider the case `[1, S]`
//! This `S` for a matrix `A` is called "leading dimension of the array A" in LAPACK document, and denoted by `lda`.
//!

use cauchy::Scalar;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MatrixLayout {
    C { row: i32, lda: i32 },
    F { col: i32, lda: i32 },
}

impl MatrixLayout {
    pub fn size(&self) -> (i32, i32) {
        match *self {
            MatrixLayout::C { row, lda } => (row, lda),
            MatrixLayout::F { col, lda } => (lda, col),
        }
    }

    pub fn resized(&self, row: i32, col: i32) -> MatrixLayout {
        match *self {
            MatrixLayout::C { .. } => MatrixLayout::C { row, lda: col },
            MatrixLayout::F { .. } => MatrixLayout::F { col, lda: row },
        }
    }

    pub fn lda(&self) -> i32 {
        std::cmp::max(
            1,
            match *self {
                MatrixLayout::C { lda, .. } | MatrixLayout::F { lda, .. } => lda,
            },
        )
    }

    pub fn len(&self) -> i32 {
        match *self {
            MatrixLayout::C { row, .. } => row,
            MatrixLayout::F { col, .. } => col,
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn lapacke_layout(&self) -> lapacke::Layout {
        match *self {
            MatrixLayout::C { .. } => lapacke::Layout::RowMajor,
            MatrixLayout::F { .. } => lapacke::Layout::ColumnMajor,
        }
    }

    pub fn same_order(&self, other: &MatrixLayout) -> bool {
        self.lapacke_layout() == other.lapacke_layout()
    }

    pub fn toggle_order(&self) -> Self {
        match *self {
            MatrixLayout::C { row, lda } => MatrixLayout::F { lda: row, col: lda },
            MatrixLayout::F { col, lda } => MatrixLayout::C { row: lda, lda: col },
        }
    }
}

/// In-place transpose of a square matrix by keeping F/C layout
///
/// Transpose for C-continuous array
///
/// ```rust
/// # use lax::layout::*;
/// let layout = MatrixLayout::C { row: 2, lda: 2 };
/// let mut a = vec![1., 2., 3., 4.];
/// square_transpose(layout, &mut a);
/// assert_eq!(a, &[1., 3., 2., 4.]);
/// ```
///
/// Transpose for F-continuous array
///
/// ```rust
/// # use lax::layout::*;
/// let layout = MatrixLayout::F { col: 2, lda: 2 };
/// let mut a = vec![1., 3., 2., 4.];
/// square_transpose(layout, &mut a);
/// assert_eq!(a, &[1., 2., 3., 4.]);
/// ```
///
/// Panics
/// ------
/// - If size of `a` and `layout` size mismatch
///
pub fn square_transpose<T: Scalar>(layout: MatrixLayout, a: &mut [T]) {
    let (m, n) = layout.size();
    let n = n as usize;
    let m = m as usize;
    assert_eq!(a.len(), n * m);
    for i in 0..m {
        for j in (i + 1)..n {
            let a_ij = a[i * n + j];
            let a_ji = a[j * m + i];
            a[i * n + j] = a_ji.conj();
            a[j * m + i] = a_ij.conj();
        }
    }
}
<commit_msg>Out-place transpose<commit_after>//! Memory layout of matrices
//!
//! Different from ndarray format which consists of shape and strides,
//! matrix format in LAPACK consists of row or column size and leading dimension.
//!
//! ndarray format and stride
//! --------------------------
//!
//! Let us consider 3-dimensional array for explaining ndarray structure.
//! The address of `(x,y,z)`-element in ndarray satisfies following relation:
//!
//! ```text
//! shape = [Nx, Ny, Nz]
//!     where Nx > 0, Ny > 0, Nz > 0
//! stride = [Sx, Sy, Sz]
//!
//! &data[(x, y, z)] = &data[(0, 0, 0)] + Sx*x + Sy*y + Sz*z
//!     for x < Nx, y < Ny, z < Nz
//! ```
//!
//! The array is called
//!
//! - C-continuous if `[Sx, Sy, Sz] = [Nz*Ny, Nz, 1]`
//! - F(Fortran)-continuous if `[Sx, Sy, Sz] = [1, Nx, Nx*Ny]`
//!
//! Strides of ndarray `[Sx, Sy, Sz]` take arbitrary value,
//! e.g. it can be non-ordered `Sy > Sx > Sz`, or can be negative `Sx < 0`.
//! If the minimum of `[Sx, Sy, Sz]` equals to `1`,
//! the value of elements fills `data` memory region and called "continuous".
//! Non-continuous ndarray is useful to get sub-array without copying data.
//!
//! Matrix layout for LAPACK
//! -------------------------
//!
//! LAPACK interface focuses on the linear algebra operations for F-continuous 2-dimensional array.
//! Under this restriction, stride becomes far simpler; we only have to consider the case `[1, S]`
//! This `S` for a matrix `A` is called "leading dimension of the array A" in LAPACK document, and denoted by `lda`.
//!

use cauchy::Scalar;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MatrixLayout {
    C { row: i32, lda: i32 },
    F { col: i32, lda: i32 },
}

impl MatrixLayout {
    pub fn size(&self) -> (i32, i32) {
        match *self {
            MatrixLayout::C { row, lda } => (row, lda),
            MatrixLayout::F { col, lda } => (lda, col),
        }
    }

    pub fn resized(&self, row: i32, col: i32) -> MatrixLayout {
        match *self {
            MatrixLayout::C { .. } => MatrixLayout::C { row, lda: col },
            MatrixLayout::F { .. } => MatrixLayout::F { col, lda: row },
        }
    }

    pub fn lda(&self) -> i32 {
        std::cmp::max(
            1,
            match *self {
                MatrixLayout::C { lda, .. } | MatrixLayout::F { lda, .. } => lda,
            },
        )
    }

    pub fn len(&self) -> i32 {
        match *self {
            MatrixLayout::C { row, .. } => row,
            MatrixLayout::F { col, .. } => col,
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn lapacke_layout(&self) -> lapacke::Layout {
        match *self {
            MatrixLayout::C { .. } => lapacke::Layout::RowMajor,
            MatrixLayout::F { .. } => lapacke::Layout::ColumnMajor,
        }
    }

    pub fn same_order(&self, other: &MatrixLayout) -> bool {
        self.lapacke_layout() == other.lapacke_layout()
    }

    pub fn toggle_order(&self) -> Self {
        match *self {
            MatrixLayout::C { row, lda } => MatrixLayout::F { lda: row, col: lda },
            MatrixLayout::F { col, lda } => MatrixLayout::C { row: lda, lda: col },
        }
    }

    /// Transpose without changing memory representation
    ///
    /// C-contigious row=2, lda=3
    ///
    /// ```text
    /// [[1, 2, 3]
    ///  [4, 5, 6]]
    /// ```
    ///
    /// and F-contigious col=2, lda=3
    ///
    /// ```text
    /// [[1, 4]
    ///  [2, 5]
    ///  [3, 6]]
    /// ```
    ///
    /// have same memory representation `[1, 2, 3, 4, 5, 6]`, and this toggles them.
    ///
    /// ```
    /// # use lax::layout::*;
    /// let layout = MatrixLayout::C { row: 2, lda: 3 };
    /// assert_eq!(layout.t(), MatrixLayout::F { col: 2, lda: 3 });
    /// ```
    pub fn t(&self) -> Self {
        match *self {
            MatrixLayout::C { row, lda } => MatrixLayout::F { col: row, lda },
            MatrixLayout::F { col, lda } => MatrixLayout::C { row: col, lda },
        }
    }
}

/// In-place transpose of a square matrix by keeping F/C layout
///
/// Transpose for C-continuous array
///
/// ```rust
/// # use lax::layout::*;
/// let layout = MatrixLayout::C { row: 2, lda: 2 };
/// let mut a = vec![1., 2., 3., 4.];
/// square_transpose(layout, &mut a);
/// assert_eq!(a, &[1., 3., 2., 4.]);
/// ```
///
/// Transpose for F-continuous array
///
/// ```rust
/// # use lax::layout::*;
/// let layout = MatrixLayout::F { col: 2, lda: 2 };
/// let mut a = vec![1., 3., 2., 4.];
/// square_transpose(layout, &mut a);
/// assert_eq!(a, &[1., 2., 3., 4.]);
/// ```
///
/// Panics
/// ------
/// - If size of `a` and `layout` size mismatch
///
pub fn square_transpose<T: Scalar>(layout: MatrixLayout, a: &mut [T]) {
    let (m, n) = layout.size();
    let n = n as usize;
    let m = m as usize;
    assert_eq!(a.len(), n * m);
    for i in 0..m {
        for j in (i + 1)..n {
            let a_ij = a[i * n + j];
            let a_ji = a[j * m + i];
            a[i * n + j] = a_ji.conj();
            a[j * m + i] = a_ij.conj();
        }
    }
}

/// Out-place transpose for general matrix
///
/// Inplace transpose of non-square matrices is hard.
/// See also: https://en.wikipedia.org/wiki/In-place_matrix_transposition
///
/// ```rust
/// # use lax::layout::*;
/// let layout = MatrixLayout::C { row: 2, lda: 3 };
/// let a = vec![1., 2., 3., 4., 5., 6.];
/// let mut b = vec![0.0; a.len()];
/// let l = transpose(layout, &a, &mut b);
/// assert_eq!(l, MatrixLayout::F { col: 3, lda: 2 });
/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
/// ```
///
/// ```rust
/// # use lax::layout::*;
/// let layout = MatrixLayout::F { col: 2, lda: 3 };
/// let a = vec![1., 2., 3., 4., 5., 6.];
/// let mut b = vec![0.0; a.len()];
/// let l = transpose(layout, &a, &mut b);
/// assert_eq!(l, MatrixLayout::C { row: 3, lda: 2 });
/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
/// ```
///
/// Panics
/// ------
/// - If size of `a` and `layout` size mismatch
///
pub fn transpose<T: Scalar>(layout: MatrixLayout, from: &[T], to: &mut [T]) -> MatrixLayout {
    let (m, n) = layout.size();
    let transposed = layout.resized(n, m).t();
    let m = m as usize;
    let n = n as usize;
    assert_eq!(from.len(), m * n);
    assert_eq!(to.len(), m * n);

    match layout {
        MatrixLayout::C { .. } => {
            for i in 0..m {
                for j in 0..n {
                    to[j * m + i] = from[i * n + j];
                }
            }
        }
        MatrixLayout::F { .. } => {
            for i in 0..m {
                for j in 0..n {
                    to[i * n + j] = from[j * m + i];
                }
            }
        }
    }
    transposed
}
<|endoftext|>


<commit_before>mod statistics;
mod planner;

use kite::query::Query;
use kite::collectors::{Collector, DocumentMatch};
use byteorder::{ByteOrder, BigEndian};
use rocksdb;

use key_builder::KeyBuilder;
use document_index::DocRef;
use super::RocksDBIndexReader;
use doc_id_set::DocIdSet;
use search::statistics::StatisticsReader;
use search::planner::{SearchPlan, plan_query};
use search::planner::boolean_query::BooleanQueryOp;
use search::planner::score_function::{CombinatorScorer, ScoreFunctionOp};

impl<'a> RocksDBIndexReader<'a> {
    fn search_segment_boolean_phase(&self, boolean_query: &Vec<BooleanQueryOp>, is_negated: bool, segment: u32) -> Result<DocIdSet, rocksdb::Error> {
        // Execute boolean query
        let mut stack = Vec::new();
        for op in boolean_query.iter() {
            match *op {
                BooleanQueryOp::PushEmpty => {
                    stack.push(DocIdSet::new_filled(0));
                }
                BooleanQueryOp::PushFull => {
                    stack.push(DocIdSet::new_filled(65536));
                }
                BooleanQueryOp::PushTermDirectory(field_ref, term_ref) => {
                    let kb = KeyBuilder::segment_dir_list(segment, field_ref.ord(), term_ref.ord());
                    match try!(self.snapshot.get(&kb.key())) {
                        Some(docid_set) => {
                            let data = DocIdSet::FromRDB(docid_set);
                            stack.push(data);
                        }
                        None => stack.push(DocIdSet::new_filled(0)),
                    }
                }
                BooleanQueryOp::PushDeletionList => {
                    let kb = KeyBuilder::segment_del_list(segment);
                    match try!(self.snapshot.get(&kb.key())) {
                        Some(deletion_list) => {
                            let data = DocIdSet::FromRDB(deletion_list);
                            stack.push(data);
                        }
                        None => stack.push(DocIdSet::new_filled(0)),
                    }
                }
                BooleanQueryOp::And => {
                    let b = stack.pop().expect("boolean query executor: stack underflow");
                    let a = stack.pop().expect("boolean query executor: stack underflow");
                    stack.push(a.intersection(&b));
                }
                BooleanQueryOp::Or => {
                    let b = stack.pop().expect("boolean query executor: stack underflow");
                    let a = stack.pop().expect("boolean query executor: stack underflow");
                    stack.push(a.union(&b));
                }
                BooleanQueryOp::AndNot => {
                    let b = stack.pop().expect("boolean query executor: stack underflow");
                    let a = stack.pop().expect("boolean query executor: stack underflow");
                    stack.push(a.exclusion(&b));
                }
            }
        }

        if !stack.len() == 1 {
            // This shouldn't be possible unless there's a bug in the planner
            panic!("boolean query executor: stack size too big ({})", stack.len());
        }
        let mut matches = stack.pop().unwrap();

        // Invert the list if the query is negated
        if is_negated {
            let kb = KeyBuilder::segment_stat(segment, b"total_docs");
            let total_docs = match try!(self.snapshot.get(&kb.key())) {
                Some(total_docs) => BigEndian::read_i64(&total_docs) as u32,
                None => 0,
            };

            let all_docs = DocIdSet::new_filled(total_docs);
            matches = all_docs.exclusion(&matches);
        }

        Ok(matches)
    }

    fn score_doc(&self, doc_id: u16, score_function: &Vec<ScoreFunctionOp>, segment: u32, mut stats: &mut StatisticsReader) -> Result<f64, rocksdb::Error> {
        // Execute score function
        let mut stack = Vec::new();
        for op in score_function.iter() {
            match *op {
                ScoreFunctionOp::Literal(val) => stack.push(val),
                ScoreFunctionOp::TermScorer(field_ref, term_ref, ref scorer) => {
                    let kb = KeyBuilder::segment_dir_list(segment, field_ref.ord(), term_ref.ord());
                    match try!(self.snapshot.get(&kb.key())) {
                        Some(data) => {
                            let docid_set = DocIdSet::FromRDB(data);

                            if docid_set.contains_doc(doc_id) {
                                // Read field length
                                // TODO: we only need this for BM25
                                let kb = KeyBuilder::stored_field_value(segment, doc_id, field_ref.ord(), b"len");
                                let field_length = match try!(self.snapshot.get(&kb.key())) {
                                    Some(value) => {
                                        let length_sqrt = (value[0] as f64) / 3.0 + 1.0;
                                        length_sqrt * length_sqrt
                                    }
                                    None => 1.0,
                                };

                                // Read term frequency
                                let mut value_type = vec![b't', b'f'];
                                value_type.extend(term_ref.ord().to_string().as_bytes());
                                let kb = KeyBuilder::stored_field_value(segment, doc_id, field_ref.ord(), &value_type);
                                let term_frequency = match try!(self.snapshot.get(&kb.key())) {
                                    Some(value) => BigEndian::read_i64(&value),
                                    None => 1,
                                };

                                let score = scorer.similarity_model.score(term_frequency as u32, field_length, try!(stats.total_tokens(field_ref)) as u64, try!(stats.total_docs(field_ref)) as u64, try!(stats.term_document_frequency(field_ref, term_ref)) as u64);
                                stack.push(score * scorer.boost);
                            } else {
                                stack.push(0.0f64);
                            }
                        }
                        None => stack.push(0.0f64),
                    }
                }
                ScoreFunctionOp::CombinatorScorer(num_vals, ref scorer) => {
                    let score = match *scorer {
                        CombinatorScorer::Avg => {
                            let mut total_score = 0.0f64;

                            for _ in 0..num_vals {
                                total_score += stack.pop().expect("document scorer: stack underflow");
                            }

                            total_score / num_vals as f64
                        }
                        CombinatorScorer::Max => {
                            let mut max_score = 0.0f64;

                            for _ in 0..num_vals {
                                let score = stack.pop().expect("document scorer: stack underflow");
                                if score > max_score {
                                    max_score = score
                                }
                            }

                            max_score
                        }
                    };

                    stack.push(score);
                }
            }
        }

        if !stack.len() == 1 {
            // This shouldn't be possible unless there's a bug in the planner
            panic!("document scorer: stack size too big ({})", stack.len());
        }

        Ok(stack.pop().expect("document scorer: stack underflow"))
    }

    fn search_segment<C: Collector>(&self, collector: &mut C, plan: &SearchPlan, segment: u32, mut stats: &mut StatisticsReader) -> Result<(), rocksdb::Error> {
        let matches = try!(self.search_segment_boolean_phase(&plan.boolean_query, plan.boolean_query_is_negated, segment));

        // Score documents and pass to collector
        for doc in matches.iter() {
            let score = try!(self.score_doc(doc, &plan.score_function, segment, &mut stats));

            let doc_ref = DocRef::from_segment_ord(segment, doc);
            let doc_match = DocumentMatch::new_scored(doc_ref.as_u64(), score);
            collector.collect(doc_match);
        }

        Ok(())
    }

    pub fn search<C: Collector>(&self, collector: &mut C, query: &Query) -> Result<(), rocksdb::Error> {
        // Plan query
        let plan = plan_query(&self, query, collector.needs_score());

        // Initialise statistics reader
        let mut stats = StatisticsReader::new(&self);

        // Run query on each segment
        for segment in self.store.segments.iter_active(&self.snapshot) {
            try!(self.search_segment(collector, &plan, segment, &mut stats));
        }

        Ok(())
    }
}
<commit_msg>Use Segment stuct in search_segment_boolean_phase<commit_after>mod statistics;
mod planner;

use kite::query::Query;
use kite::collectors::{Collector, DocumentMatch};
use byteorder::{ByteOrder, BigEndian};
use rocksdb;

use key_builder::KeyBuilder;
use document_index::DocRef;
use super::RocksDBIndexReader;
use doc_id_set::DocIdSet;
use segment::Segment;
use search::statistics::StatisticsReader;
use search::planner::{SearchPlan, plan_query};
use search::planner::boolean_query::BooleanQueryOp;
use search::planner::score_function::{CombinatorScorer, ScoreFunctionOp};

impl<'a> RocksDBIndexReader<'a> {
    fn search_segment_boolean_phase(&self, boolean_query: &Vec<BooleanQueryOp>, is_negated: bool, segment: &Segment) -> Result<DocIdSet, rocksdb::Error> {
        // Execute boolean query
        let mut stack = Vec::new();
        for op in boolean_query.iter() {
            match *op {
                BooleanQueryOp::PushEmpty => {
                    stack.push(DocIdSet::new_filled(0));
                }
                BooleanQueryOp::PushFull => {
                    stack.push(DocIdSet::new_filled(65536));
                }
                BooleanQueryOp::PushTermDirectory(field_ref, term_ref) => {
                    match try!(segment.load_term_directory(field_ref, term_ref)) {
                        Some(doc_id_set) => stack.push(doc_id_set),
                        None => stack.push(DocIdSet::new_filled(0)),
                    }
                }
                BooleanQueryOp::PushDeletionList => {
                     match try!(segment.load_deletion_list()) {
                        Some(doc_id_set) => stack.push(doc_id_set),
                        None => stack.push(DocIdSet::new_filled(0)),
                    }
                }
                BooleanQueryOp::And => {
                    let b = stack.pop().expect("boolean query executor: stack underflow");
                    let a = stack.pop().expect("boolean query executor: stack underflow");
                    stack.push(a.intersection(&b));
                }
                BooleanQueryOp::Or => {
                    let b = stack.pop().expect("boolean query executor: stack underflow");
                    let a = stack.pop().expect("boolean query executor: stack underflow");
                    stack.push(a.union(&b));
                }
                BooleanQueryOp::AndNot => {
                    let b = stack.pop().expect("boolean query executor: stack underflow");
                    let a = stack.pop().expect("boolean query executor: stack underflow");
                    stack.push(a.exclusion(&b));
                }
            }
        }

        if !stack.len() == 1 {
            // This shouldn't be possible unless there's a bug in the planner
            panic!("boolean query executor: stack size too big ({})", stack.len());
        }
        let mut matches = stack.pop().unwrap();

        // Invert the list if the query is negated
        if is_negated {
            let total_docs = try!(segment.load_statistic(b"total_docs")).unwrap_or(0);
            let all_docs = DocIdSet::new_filled(total_docs as u32);
            matches = all_docs.exclusion(&matches);
        }

        Ok(matches)
    }

    fn score_doc(&self, doc_id: u16, score_function: &Vec<ScoreFunctionOp>, segment: u32, mut stats: &mut StatisticsReader) -> Result<f64, rocksdb::Error> {
        // Execute score function
        let mut stack = Vec::new();
        for op in score_function.iter() {
            match *op {
                ScoreFunctionOp::Literal(val) => stack.push(val),
                ScoreFunctionOp::TermScorer(field_ref, term_ref, ref scorer) => {
                    let kb = KeyBuilder::segment_dir_list(segment, field_ref.ord(), term_ref.ord());
                    match try!(self.snapshot.get(&kb.key())) {
                        Some(data) => {
                            let docid_set = DocIdSet::FromRDB(data);

                            if docid_set.contains_doc(doc_id) {
                                // Read field length
                                // TODO: we only need this for BM25
                                let kb = KeyBuilder::stored_field_value(segment, doc_id, field_ref.ord(), b"len");
                                let field_length = match try!(self.snapshot.get(&kb.key())) {
                                    Some(value) => {
                                        let length_sqrt = (value[0] as f64) / 3.0 + 1.0;
                                        length_sqrt * length_sqrt
                                    }
                                    None => 1.0,
                                };

                                // Read term frequency
                                let mut value_type = vec![b't', b'f'];
                                value_type.extend(term_ref.ord().to_string().as_bytes());
                                let kb = KeyBuilder::stored_field_value(segment, doc_id, field_ref.ord(), &value_type);
                                let term_frequency = match try!(self.snapshot.get(&kb.key())) {
                                    Some(value) => BigEndian::read_i64(&value),
                                    None => 1,
                                };

                                let score = scorer.similarity_model.score(term_frequency as u32, field_length, try!(stats.total_tokens(field_ref)) as u64, try!(stats.total_docs(field_ref)) as u64, try!(stats.term_document_frequency(field_ref, term_ref)) as u64);
                                stack.push(score * scorer.boost);
                            } else {
                                stack.push(0.0f64);
                            }
                        }
                        None => stack.push(0.0f64),
                    }
                }
                ScoreFunctionOp::CombinatorScorer(num_vals, ref scorer) => {
                    let score = match *scorer {
                        CombinatorScorer::Avg => {
                            let mut total_score = 0.0f64;

                            for _ in 0..num_vals {
                                total_score += stack.pop().expect("document scorer: stack underflow");
                            }

                            total_score / num_vals as f64
                        }
                        CombinatorScorer::Max => {
                            let mut max_score = 0.0f64;

                            for _ in 0..num_vals {
                                let score = stack.pop().expect("document scorer: stack underflow");
                                if score > max_score {
                                    max_score = score
                                }
                            }

                            max_score
                        }
                    };

                    stack.push(score);
                }
            }
        }

        if !stack.len() == 1 {
            // This shouldn't be possible unless there's a bug in the planner
            panic!("document scorer: stack size too big ({})", stack.len());
        }

        Ok(stack.pop().expect("document scorer: stack underflow"))
    }

    fn search_segment<C: Collector>(&self, collector: &mut C, plan: &SearchPlan, segment: u32, mut stats: &mut StatisticsReader) -> Result<(), rocksdb::Error> {
        let matches = try!(self.search_segment_boolean_phase(&plan.boolean_query, plan.boolean_query_is_negated, &Segment::new(self, segment)));

        // Score documents and pass to collector
        for doc in matches.iter() {
            let score = try!(self.score_doc(doc, &plan.score_function, segment, &mut stats));

            let doc_ref = DocRef::from_segment_ord(segment, doc);
            let doc_match = DocumentMatch::new_scored(doc_ref.as_u64(), score);
            collector.collect(doc_match);
        }

        Ok(())
    }

    pub fn search<C: Collector>(&self, collector: &mut C, query: &Query) -> Result<(), rocksdb::Error> {
        // Plan query
        let plan = plan_query(&self, query, collector.needs_score());

        // Initialise statistics reader
        let mut stats = StatisticsReader::new(&self);

        // Run query on each segment
        for segment in self.store.segments.iter_active(&self.snapshot) {
            try!(self.search_segment(collector, &plan, segment, &mut stats));
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>//! The "main crate" of the Rust compiler. This crate contains common
//! type definitions that are used by the other crates in the rustc
//! "family". Some prominent examples (note that each of these modules
//! has their own README with further details).
//!
//! - **HIR.** The "high-level (H) intermediate representation (IR)" is
//!   defined in the `hir` module.
//! - **MIR.** The "mid-level (M) intermediate representation (IR)" is
//!   defined in the `mir` module. This module contains only the
//!   *definition* of the MIR; the passes that transform and operate
//!   on MIR are found in `librustc_mir` crate.
//! - **Types.** The internal representation of types used in rustc is
//!   defined in the `ty` module. This includes the **type context**
//!   (or `tcx`), which is the central context during most of
//!   compilation, containing the interners and other things.
//! - **Traits.** Trait resolution is implemented in the `traits` module.
//! - **Type inference.** The type inference code can be found in the `infer` module;
//!   this code handles low-level equality and subtyping operations. The
//!   type check pass in the compiler is found in the `librustc_typeck` crate.
//!
//! For more information about how rustc works, see the [rustc guide].
//!
//! [rustc guide]: https://rust-lang.github.io/rustc-guide/
//!
//! # Note
//!
//! This API is completely unstable and subject to change.

#![doc(html_root_url = "https://doc.rust-lang.org/nightly/")]

#![deny(rust_2018_idioms)]
#![allow(explicit_outlives_requirements)]

#![feature(box_patterns)]
#![feature(box_syntax)]
#![feature(core_intrinsics)]
#![feature(drain_filter)]
#![cfg_attr(windows, feature(libc))]
#![feature(never_type)]
#![feature(exhaustive_patterns)]
#![feature(extern_types)]
#![feature(nll)]
#![feature(non_exhaustive)]
#![feature(proc_macro_internals)]
#![feature(optin_builtin_traits)]
#![feature(refcell_replace_swap)]
#![feature(rustc_diagnostic_macros)]
#![feature(rustc_attrs)]
#![feature(slice_patterns)]
#![feature(specialization)]
#![feature(unboxed_closures)]
#![feature(thread_local)]
#![feature(trace_macros)]
#![feature(trusted_len)]
#![feature(vec_remove_item)]
#![feature(step_trait)]
#![feature(stmt_expr_attributes)]
#![feature(integer_atomics)]
#![feature(test)]
#![feature(in_band_lifetimes)]
#![feature(crate_visibility_modifier)]

#![recursion_limit="512"]

#![warn(elided_lifetimes_in_paths)]

#[macro_use] extern crate bitflags;
extern crate getopts;
#[macro_use] extern crate lazy_static;
#[macro_use] extern crate scoped_tls;
#[cfg(windows)]
extern crate libc;
#[macro_use] extern crate rustc_data_structures;

#[macro_use] extern crate log;
#[macro_use] extern crate syntax;

// FIXME: This import is used by deriving `RustcDecodable` and `RustcEncodable`. Removing this
// results in a bunch of "failed to resolve" errors. Hopefully, the compiler moves to serde or
// something, and we can get rid of this.
#[allow(rust_2018_idioms)]
extern crate serialize as rustc_serialize;

#[macro_use] extern crate smallvec;

// Note that librustc doesn't actually depend on these crates, see the note in
// `Cargo.toml` for this crate about why these are here.
#[allow(unused_extern_crates)]
extern crate flate2;
#[allow(unused_extern_crates)]
extern crate test;

#[macro_use]
mod macros;

// N.B., this module needs to be declared first so diagnostics are
// registered before they are used.
pub mod diagnostics;

pub mod cfg;
pub mod dep_graph;
pub mod hir;
pub mod ich;
pub mod infer;
pub mod lint;

pub mod middle {
    pub mod allocator;
    pub mod borrowck;
    pub mod expr_use_visitor;
    pub mod cstore;
    pub mod dead;
    pub mod dependency_format;
    pub mod entry;
    pub mod exported_symbols;
    pub mod free_region;
    pub mod intrinsicck;
    pub mod lib_features;
    pub mod lang_items;
    pub mod liveness;
    pub mod mem_categorization;
    pub mod privacy;
    pub mod reachable;
    pub mod region;
    pub mod recursion_limit;
    pub mod resolve_lifetime;
    pub mod stability;
    pub mod weak_lang_items;
}

pub mod mir;
pub mod session;
pub mod traits;
pub mod ty;

pub mod util {
    pub mod captures;
    pub mod common;
    pub mod ppaux;
    pub mod nodemap;
    pub mod time_graph;
    pub mod profiling;
    pub mod bug;
}

// A private module so that macro-expanded idents like
// `::rustc::lint::Lint` will also work in `rustc` itself.
//
// `libstd` uses the same trick.
#[doc(hidden)]
mod rustc {
    pub use crate::lint;
}

// FIXME(#27438): right now the unit tests of librustc don't refer to any actual
//                functions generated in librustc_data_structures (all
//                references are through generic functions), but statics are
//                referenced from time to time. Due to this bug we won't
//                actually correctly link in the statics unless we also
//                reference a function, so be sure to reference a dummy
//                function.
#[test]
fn noop() {
    rustc_data_structures::__noop_fix_for_27438();
}

// Build the diagnostics array at the end so that the metadata includes error use sites.
__build_diagnostic_array! { librustc, DIAGNOSTICS }
<commit_msg>librustc: deny(elided_lifetimes_in_paths)<commit_after>//! The "main crate" of the Rust compiler. This crate contains common
//! type definitions that are used by the other crates in the rustc
//! "family". Some prominent examples (note that each of these modules
//! has their own README with further details).
//!
//! - **HIR.** The "high-level (H) intermediate representation (IR)" is
//!   defined in the `hir` module.
//! - **MIR.** The "mid-level (M) intermediate representation (IR)" is
//!   defined in the `mir` module. This module contains only the
//!   *definition* of the MIR; the passes that transform and operate
//!   on MIR are found in `librustc_mir` crate.
//! - **Types.** The internal representation of types used in rustc is
//!   defined in the `ty` module. This includes the **type context**
//!   (or `tcx`), which is the central context during most of
//!   compilation, containing the interners and other things.
//! - **Traits.** Trait resolution is implemented in the `traits` module.
//! - **Type inference.** The type inference code can be found in the `infer` module;
//!   this code handles low-level equality and subtyping operations. The
//!   type check pass in the compiler is found in the `librustc_typeck` crate.
//!
//! For more information about how rustc works, see the [rustc guide].
//!
//! [rustc guide]: https://rust-lang.github.io/rustc-guide/
//!
//! # Note
//!
//! This API is completely unstable and subject to change.

#![doc(html_root_url = "https://doc.rust-lang.org/nightly/")]

#![deny(rust_2018_idioms)]
#![allow(explicit_outlives_requirements)]

#![feature(box_patterns)]
#![feature(box_syntax)]
#![feature(core_intrinsics)]
#![feature(drain_filter)]
#![cfg_attr(windows, feature(libc))]
#![feature(never_type)]
#![feature(exhaustive_patterns)]
#![feature(extern_types)]
#![feature(nll)]
#![feature(non_exhaustive)]
#![feature(proc_macro_internals)]
#![feature(optin_builtin_traits)]
#![feature(refcell_replace_swap)]
#![feature(rustc_diagnostic_macros)]
#![feature(rustc_attrs)]
#![feature(slice_patterns)]
#![feature(specialization)]
#![feature(unboxed_closures)]
#![feature(thread_local)]
#![feature(trace_macros)]
#![feature(trusted_len)]
#![feature(vec_remove_item)]
#![feature(step_trait)]
#![feature(stmt_expr_attributes)]
#![feature(integer_atomics)]
#![feature(test)]
#![feature(in_band_lifetimes)]
#![feature(crate_visibility_modifier)]

#![recursion_limit="512"]

#[macro_use] extern crate bitflags;
extern crate getopts;
#[macro_use] extern crate lazy_static;
#[macro_use] extern crate scoped_tls;
#[cfg(windows)]
extern crate libc;
#[macro_use] extern crate rustc_data_structures;

#[macro_use] extern crate log;
#[macro_use] extern crate syntax;

// FIXME: This import is used by deriving `RustcDecodable` and `RustcEncodable`. Removing this
// results in a bunch of "failed to resolve" errors. Hopefully, the compiler moves to serde or
// something, and we can get rid of this.
#[allow(rust_2018_idioms)]
extern crate serialize as rustc_serialize;

#[macro_use] extern crate smallvec;

// Note that librustc doesn't actually depend on these crates, see the note in
// `Cargo.toml` for this crate about why these are here.
#[allow(unused_extern_crates)]
extern crate flate2;
#[allow(unused_extern_crates)]
extern crate test;

#[macro_use]
mod macros;

// N.B., this module needs to be declared first so diagnostics are
// registered before they are used.
pub mod diagnostics;

pub mod cfg;
pub mod dep_graph;
pub mod hir;
pub mod ich;
pub mod infer;
pub mod lint;

pub mod middle {
    pub mod allocator;
    pub mod borrowck;
    pub mod expr_use_visitor;
    pub mod cstore;
    pub mod dead;
    pub mod dependency_format;
    pub mod entry;
    pub mod exported_symbols;
    pub mod free_region;
    pub mod intrinsicck;
    pub mod lib_features;
    pub mod lang_items;
    pub mod liveness;
    pub mod mem_categorization;
    pub mod privacy;
    pub mod reachable;
    pub mod region;
    pub mod recursion_limit;
    pub mod resolve_lifetime;
    pub mod stability;
    pub mod weak_lang_items;
}

pub mod mir;
pub mod session;
pub mod traits;
pub mod ty;

pub mod util {
    pub mod captures;
    pub mod common;
    pub mod ppaux;
    pub mod nodemap;
    pub mod time_graph;
    pub mod profiling;
    pub mod bug;
}

// A private module so that macro-expanded idents like
// `::rustc::lint::Lint` will also work in `rustc` itself.
//
// `libstd` uses the same trick.
#[doc(hidden)]
mod rustc {
    pub use crate::lint;
}

// FIXME(#27438): right now the unit tests of librustc don't refer to any actual
//                functions generated in librustc_data_structures (all
//                references are through generic functions), but statics are
//                referenced from time to time. Due to this bug we won't
//                actually correctly link in the statics unless we also
//                reference a function, so be sure to reference a dummy
//                function.
#[test]
fn noop() {
    rustc_data_structures::__noop_fix_for_27438();
}

// Build the diagnostics array at the end so that the metadata includes error use sites.
__build_diagnostic_array! { librustc, DIAGNOSTICS }
<|endoftext|>