<commit_before>#![feature(integer_atomics)]

extern crate abra;
extern crate rocksdb;
extern crate rustc_serialize;
#[macro_use]
extern crate maplit;
extern crate byteorder;

use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicU32, Ordering};
use std::collections::BTreeMap;

use rocksdb::{DB, WriteBatch, Writable, Options, MergeOperands};
use abra::{Term, Document};
use abra::schema::{Schema, FieldType, FieldRef, AddFieldError};
use rustc_serialize::{json, Encodable};
use byteorder::{ByteOrder, BigEndian};

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct TermRef(u32);

impl TermRef {
    pub fn ord(&self) -> u32 {
        self.0
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct DocRef(u32, u16);

impl DocRef {
    pub fn chunk(&self) -> u32 {
        self.0
    }

    pub fn ord(&self) -> u16 {
        self.1
    }
}

fn merge_keys(key: &[u8], existing_val: Option<&[u8]>, operands: &mut MergeOperands) -> Vec<u8> {
    match key[0] {
        b'd' | b'x' => {
            // Sequence of two byte document ids
            // d = directory
            // x = deletion list

            // Allocate vec for new Value
            let new_size = match existing_val {
                Some(existing_val) => existing_val.len(),
                None => 0,
            } + operands.size_hint().0 * 2;

            let mut new_val = Vec::with_capacity(new_size);

            // Push existing value
            existing_val.map(|v| {
                for b in v {
                    new_val.push(*b);
                }
            });

            // Append new entries
            for op in operands {
                for b in op {
                    new_val.push(*b);
                }
            }

            new_val
        }
        _ => {
            // Unrecognised key, fallback to emulating a put operation (by taking the last value)
            operands.last().unwrap().iter().cloned().collect()
        }
    }
}

pub struct RocksDBIndexStore {
    schema: Arc<Schema>,
    db: DB,
    next_term_ref: AtomicU32,
    term_dictionary: RwLock<BTreeMap<Vec<u8>, TermRef>>,
    next_chunk: AtomicU32,
    doc_key_mapping: RwLock<BTreeMap<Vec<u8>, DocRef>>,
    deleted_docs: RwLock<Vec<DocRef>>,
}

impl RocksDBIndexStore {
    pub fn create(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        opts.create_if_missing(true);
        let db = try!(DB::open(&opts, path));

        // Schema
        let schema = Schema::new();
        db.put(b".schema", json::encode(&schema).unwrap().as_bytes());

        // Next term ref
        db.put(b".next_term_ref", b"1");

        // Next chunk
        db.put(b".next_chunk", b"1");

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(1),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(1),
            doc_key_mapping: RwLock::new(BTreeMap::new()),
            deleted_docs: RwLock::new(Vec::new()),
        })
    }

    pub fn open(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        let db = try!(DB::open(&opts, path));

        let schema = match db.get(b".schema") {
            Ok(Some(schema)) => {
                let schema = schema.to_utf8().unwrap().to_string();
                json::decode(&schema).unwrap()
            }
            Ok(None) => Schema::new(),  // TODO: error
            Err(_) => Schema::new(),  // TODO: error
        };

        let next_term_ref = match db.get(b".next_term_ref") {
            Ok(Some(next_term_ref)) => {
                next_term_ref.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        let next_chunk = match db.get(b".next_chunk") {
            Ok(Some(next_chunk)) => {
                next_chunk.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(next_term_ref),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(next_chunk),
            doc_key_mapping: RwLock::new(BTreeMap::new()),
            deleted_docs: RwLock::new(Vec::new()),
        })
    }

    pub fn add_field(&mut self, name: String, field_type: FieldType) -> Result<FieldRef, AddFieldError> {
        let mut schema_copy = (*self.schema).clone();
        let field_ref = try!(schema_copy.add_field(name, field_type));
        self.schema = Arc::new(schema_copy);

        self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());

        Ok(field_ref)
    }

    pub fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        let mut schema_copy = (*self.schema).clone();
        let field_removed = schema_copy.remove_field(field_ref);

        if field_removed {
            self.schema = Arc::new(schema_copy);
            self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());
        }

        field_removed
    }

    fn get_or_create_term(&mut self, term: &Term) -> TermRef {
        let term_bytes = term.to_bytes();

        if let Some(term_ref) = self.term_dictionary.read().unwrap().get(&term_bytes) {
            return *term_ref;
        }

        // Term doesn't exist in the term dictionary

        // Increment next_term_ref
        let next_term_ref = self.next_term_ref.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_term_ref", (next_term_ref + 1).to_string().as_bytes());

        // Create term ref
        let term_ref = TermRef(next_term_ref);

        // Get exclusive lock to term dictionary
        let mut term_dictionary = self.term_dictionary.write().unwrap();

        // It's possible that another thread has written the term to the dictionary
        // since we checked earlier. If this is the case, We should forget about
        // writing our TermRef and use the one that has been inserted already.
        if let Some(term_ref) = term_dictionary.get(&term_bytes) {
            return *term_ref;
        }

        // Write it to the on-disk term dictionary
        let mut key = Vec::with_capacity(1 + term_bytes.len());
        key.push(b't');
        for byte in term_bytes.iter() {
            key.push(*byte);
        }
        self.db.put(&key, next_term_ref.to_string().as_bytes());

        // Write it to the term dictionary
        term_dictionary.insert(term_bytes, term_ref);

        term_ref
    }

    pub fn insert_or_update_document(&mut self, doc: Document) {
        // Allocate a new chunk for the document

        // Chunk merges are very slow so we should avoid doing them at runtime
        // which is why each new document is created in a fresh chunk.
        // Later on, a background process will come and merge any small chunks
        // together.

        // For best performance, documents should be inserted in batches.

        // Increment next_chunk
        let next_chunk = self.next_chunk.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_chunk", (next_chunk + 1).to_string().as_bytes());

        // Create doc ref
        let doc_ref = DocRef(next_chunk, 0);

        // Start write batch
        let mut write_batch = WriteBatch::default();

        // Insert contents
        for (field_ref, tokens) in doc.fields.iter() {
            for token in tokens.iter() {
                let term_ref = self.get_or_create_term(&token.term);

                let mut key = Vec::with_capacity(20);
                key.push(b'd');
                for byte in field_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in term_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in doc_ref.chunk().to_string().as_bytes() {
                    key.push(*byte);
                }

                let mut doc_id_bytes = [0; 2];
                BigEndian::write_u16(&mut doc_id_bytes, doc_ref.ord());
                write_batch.merge(&key, &doc_id_bytes);
            }
        }

        // Write document data
        self.db.write(write_batch);

        // Update doc_key_mapping
        let mut write_batch = WriteBatch::default();
        let previous_doc_ref = self.doc_key_mapping.write().unwrap().insert(doc.key.as_bytes().iter().cloned().collect(), doc_ref);

        let mut key = Vec::with_capacity(1 + doc.key.len());
        key.push(b'k');
        for byte in doc.key.as_bytes() {
            key.push(*byte);
        }
        let mut doc_ref_bytes = [0; 6];
        BigEndian::write_u32(&mut doc_ref_bytes, doc_ref.chunk());
        BigEndian::write_u16(&mut doc_ref_bytes[4..], doc_ref.ord());
        write_batch.put(&key, &doc_ref_bytes);

        // If there was a document there previously, mark it as deleted
        if let Some(previous_doc_ref) = previous_doc_ref {
            self.deleted_docs.write().unwrap().push(previous_doc_ref);

            let mut key = Vec::with_capacity(5);
            key.push(b'x');
            for byte in previous_doc_ref.chunk().to_string().as_bytes() {
                key.push(*byte);
            }

            let mut previous_doc_id_bytes = [0; 2];
            BigEndian::write_u16(&mut previous_doc_id_bytes, previous_doc_ref.ord());
            write_batch.merge(&key, &previous_doc_id_bytes);
        }

        // Write document data
        self.db.write(write_batch);
    }
}

#[cfg(test)]
mod tests {
    use rocksdb::{DB, Options, IteratorMode};
    use abra::{Term, Token, Document};
    use abra::schema::{Schema, FieldType, FieldRef};

    use super::RocksDBIndexStore;

    fn clean_test_indices() {
        use std::fs::remove_dir_all;

        remove_dir_all("test_indices");
    }

    #[test]
    fn test_create() {
        clean_test_indices();

        let store = RocksDBIndexStore::create("test_indices/test_create");
        assert!(store.is_ok());
    }

    #[test]
    fn test_open() {
        clean_test_indices();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_err());

        // Create DB
        let mut opts = Options::default();
        opts.create_if_missing(true);
        DB::open(&opts, "test_indices/test_open").unwrap();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_ok());
    }

    fn make_test_store(path: &str) -> RocksDBIndexStore {
        let mut store = RocksDBIndexStore::create(path).unwrap();
        let mut title_field = store.add_field("title".to_string(), FieldType::Text).unwrap();
        let mut body_field = store.add_field("body".to_string(), FieldType::Text).unwrap();

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("hello".to_string()), position: 1 },
                    Token { term: Term::String("world".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store.insert_or_update_document(Document {
            key: "another_test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("howdy".to_string()), position: 1 },
                    Token { term: Term::String("partner".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store
    }

    pub fn print_keys(db: &DB) {
        fn bytes_to_string(bytes: &Box<[u8]>) -> String {
            use std::char;

            let mut string = String::new();

            for byte in bytes.iter() {
                if *byte < 128 {
                    // ASCII character
                    string.push(char::from_u32(*byte as u32).unwrap());
                } else {
                    string.push('?');
                }
            }

            string
        }

        for (key, value) in db.iterator(IteratorMode::Start) {
            println!("{} = {:?}", bytes_to_string(&key), value);
        }
    }

    #[test]
    fn test() {
        clean_test_indices();

        let store = make_test_store("test_indices/test");

        print_keys(&store.db);

    }
}
<commit_msg>Added statistics<commit_after>#![feature(integer_atomics)]

extern crate abra;
extern crate rocksdb;
extern crate rustc_serialize;
#[macro_use]
extern crate maplit;
extern crate byteorder;

use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicU32, Ordering};
use std::collections::BTreeMap;

use rocksdb::{DB, WriteBatch, Writable, Options, MergeOperands};
use abra::{Term, Document};
use abra::schema::{Schema, FieldType, FieldRef, AddFieldError};
use rustc_serialize::{json, Encodable};
use byteorder::{ByteOrder, BigEndian};

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct TermRef(u32);

impl TermRef {
    pub fn ord(&self) -> u32 {
        self.0
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct DocRef(u32, u16);

impl DocRef {
    pub fn chunk(&self) -> u32 {
        self.0
    }

    pub fn ord(&self) -> u16 {
        self.1
    }
}

fn merge_keys(key: &[u8], existing_val: Option<&[u8]>, operands: &mut MergeOperands) -> Vec<u8> {
    match key[0] {
        b'd' | b'x' => {
            // Sequence of two byte document ids
            // d = directory
            // x = deletion list

            // Allocate vec for new Value
            let new_size = match existing_val {
                Some(existing_val) => existing_val.len(),
                None => 0,
            } + operands.size_hint().0 * 2;

            let mut new_val = Vec::with_capacity(new_size);

            // Push existing value
            existing_val.map(|v| {
                for b in v {
                    new_val.push(*b);
                }
            });

            // Append new entries
            for op in operands {
                for b in op {
                    new_val.push(*b);
                }
            }

            new_val
        }
        b's' => {
            // Statistic
            // An i64 number that can be incremented or decremented
            let mut value = match existing_val {
                Some(existing_val) => BigEndian::read_i64(existing_val),
                None => 0
            };

            for op in operands {
                value += BigEndian::read_i64(op);
            }

            let mut buf = [0; 8];
            BigEndian::write_i64(&mut buf, value);
            buf.iter().cloned().collect()
        }
        _ => {
            // Unrecognised key, fallback to emulating a put operation (by taking the last value)
            operands.last().unwrap().iter().cloned().collect()
        }
    }
}

pub struct RocksDBIndexStore {
    schema: Arc<Schema>,
    db: DB,
    next_term_ref: AtomicU32,
    term_dictionary: RwLock<BTreeMap<Vec<u8>, TermRef>>,
    next_chunk: AtomicU32,
    doc_key_mapping: RwLock<BTreeMap<Vec<u8>, DocRef>>,
    deleted_docs: RwLock<Vec<DocRef>>,
}

impl RocksDBIndexStore {
    pub fn create(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        opts.create_if_missing(true);
        let db = try!(DB::open(&opts, path));

        // Schema
        let schema = Schema::new();
        db.put(b".schema", json::encode(&schema).unwrap().as_bytes());

        // Next term ref
        db.put(b".next_term_ref", b"1");

        // Next chunk
        db.put(b".next_chunk", b"1");

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(1),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(1),
            doc_key_mapping: RwLock::new(BTreeMap::new()),
            deleted_docs: RwLock::new(Vec::new()),
        })
    }

    pub fn open(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        let db = try!(DB::open(&opts, path));

        let schema = match db.get(b".schema") {
            Ok(Some(schema)) => {
                let schema = schema.to_utf8().unwrap().to_string();
                json::decode(&schema).unwrap()
            }
            Ok(None) => Schema::new(),  // TODO: error
            Err(_) => Schema::new(),  // TODO: error
        };

        let next_term_ref = match db.get(b".next_term_ref") {
            Ok(Some(next_term_ref)) => {
                next_term_ref.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        let next_chunk = match db.get(b".next_chunk") {
            Ok(Some(next_chunk)) => {
                next_chunk.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(next_term_ref),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(next_chunk),
            doc_key_mapping: RwLock::new(BTreeMap::new()),
            deleted_docs: RwLock::new(Vec::new()),
        })
    }

    pub fn add_field(&mut self, name: String, field_type: FieldType) -> Result<FieldRef, AddFieldError> {
        let mut schema_copy = (*self.schema).clone();
        let field_ref = try!(schema_copy.add_field(name, field_type));
        self.schema = Arc::new(schema_copy);

        self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());

        Ok(field_ref)
    }

    pub fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        let mut schema_copy = (*self.schema).clone();
        let field_removed = schema_copy.remove_field(field_ref);

        if field_removed {
            self.schema = Arc::new(schema_copy);
            self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());
        }

        field_removed
    }

    fn get_or_create_term(&mut self, term: &Term) -> TermRef {
        let term_bytes = term.to_bytes();

        if let Some(term_ref) = self.term_dictionary.read().unwrap().get(&term_bytes) {
            return *term_ref;
        }

        // Term doesn't exist in the term dictionary

        // Increment next_term_ref
        let next_term_ref = self.next_term_ref.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_term_ref", (next_term_ref + 1).to_string().as_bytes());

        // Create term ref
        let term_ref = TermRef(next_term_ref);

        // Get exclusive lock to term dictionary
        let mut term_dictionary = self.term_dictionary.write().unwrap();

        // It's possible that another thread has written the term to the dictionary
        // since we checked earlier. If this is the case, We should forget about
        // writing our TermRef and use the one that has been inserted already.
        if let Some(term_ref) = term_dictionary.get(&term_bytes) {
            return *term_ref;
        }

        // Write it to the on-disk term dictionary
        let mut key = Vec::with_capacity(1 + term_bytes.len());
        key.push(b't');
        for byte in term_bytes.iter() {
            key.push(*byte);
        }
        self.db.put(&key, next_term_ref.to_string().as_bytes());

        // Write it to the term dictionary
        term_dictionary.insert(term_bytes, term_ref);

        term_ref
    }

    pub fn insert_or_update_document(&mut self, doc: Document) {
        // Allocate a new chunk for the document

        // Chunk merges are very slow so we should avoid doing them at runtime
        // which is why each new document is created in a fresh chunk.
        // Later on, a background process will come and merge any small chunks
        // together.

        // For best performance, documents should be inserted in batches.

        // Increment next_chunk
        let next_chunk = self.next_chunk.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_chunk", (next_chunk + 1).to_string().as_bytes());

        // Create doc ref
        let doc_ref = DocRef(next_chunk, 0);

        // Start write batch
        let mut write_batch = WriteBatch::default();

        // Insert contents
        let mut token_count: i64 = 0;
        for (field_ref, tokens) in doc.fields.iter() {
            for token in tokens.iter() {
                token_count += 1;
                let term_ref = self.get_or_create_term(&token.term);

                let mut key = Vec::with_capacity(20);
                key.push(b'd');
                for byte in field_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in term_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in doc_ref.chunk().to_string().as_bytes() {
                    key.push(*byte);
                }

                let mut doc_id_bytes = [0; 2];
                BigEndian::write_u16(&mut doc_id_bytes, doc_ref.ord());
                write_batch.merge(&key, &doc_id_bytes);
            }
        }

        // Increment total docs
        let mut key = Vec::with_capacity(5);
        for byte in b"stotal_docs".iter() {
            key.push(*byte);
        }
        key.push(b'/');
        for byte in doc_ref.chunk().to_string().as_bytes() {
            key.push(*byte);
        }
        let mut inc_bytes = [0; 8];
        BigEndian::write_i64(&mut inc_bytes, 1);
        write_batch.merge(&key, &inc_bytes);

        // Increment total tokens
        let mut key = Vec::with_capacity(5);
        for byte in b"stotal_tokens".iter() {
            key.push(*byte);
        }
        key.push(b'/');
        for byte in doc_ref.chunk().to_string().as_bytes() {
            key.push(*byte);
        }
        let mut inc_bytes = [0; 8];
        BigEndian::write_i64(&mut inc_bytes, token_count);
        write_batch.merge(&key, &inc_bytes);

        // Write document data
        self.db.write(write_batch);

        // Update doc_key_mapping
        let mut write_batch = WriteBatch::default();
        let previous_doc_ref = self.doc_key_mapping.write().unwrap().insert(doc.key.as_bytes().iter().cloned().collect(), doc_ref);

        let mut key = Vec::with_capacity(1 + doc.key.len());
        key.push(b'k');
        for byte in doc.key.as_bytes() {
            key.push(*byte);
        }
        let mut doc_ref_bytes = [0; 6];
        BigEndian::write_u32(&mut doc_ref_bytes, doc_ref.chunk());
        BigEndian::write_u16(&mut doc_ref_bytes[4..], doc_ref.ord());
        write_batch.put(&key, &doc_ref_bytes);

        // If there was a document there previously, mark it as deleted
        if let Some(previous_doc_ref) = previous_doc_ref {
            self.deleted_docs.write().unwrap().push(previous_doc_ref);

            let mut key = Vec::with_capacity(5);
            key.push(b'x');
            for byte in previous_doc_ref.chunk().to_string().as_bytes() {
                key.push(*byte);
            }

            let mut previous_doc_id_bytes = [0; 2];
            BigEndian::write_u16(&mut previous_doc_id_bytes, previous_doc_ref.ord());
            write_batch.merge(&key, &previous_doc_id_bytes);
        }

        // Write document data
        self.db.write(write_batch);
    }
}

#[cfg(test)]
mod tests {
    use rocksdb::{DB, Options, IteratorMode};
    use abra::{Term, Token, Document};
    use abra::schema::{Schema, FieldType, FieldRef};

    use super::RocksDBIndexStore;

    fn clean_test_indices() {
        use std::fs::remove_dir_all;

        remove_dir_all("test_indices");
    }

    #[test]
    fn test_create() {
        clean_test_indices();

        let store = RocksDBIndexStore::create("test_indices/test_create");
        assert!(store.is_ok());
    }

    #[test]
    fn test_open() {
        clean_test_indices();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_err());

        // Create DB
        let mut opts = Options::default();
        opts.create_if_missing(true);
        DB::open(&opts, "test_indices/test_open").unwrap();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_ok());
    }

    fn make_test_store(path: &str) -> RocksDBIndexStore {
        let mut store = RocksDBIndexStore::create(path).unwrap();
        let mut title_field = store.add_field("title".to_string(), FieldType::Text).unwrap();
        let mut body_field = store.add_field("body".to_string(), FieldType::Text).unwrap();

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("hello".to_string()), position: 1 },
                    Token { term: Term::String("world".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store.insert_or_update_document(Document {
            key: "another_test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("howdy".to_string()), position: 1 },
                    Token { term: Term::String("partner".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store
    }

    pub fn print_keys(db: &DB) {
        fn bytes_to_string(bytes: &Box<[u8]>) -> String {
            use std::char;

            let mut string = String::new();

            for byte in bytes.iter() {
                if *byte < 128 {
                    // ASCII character
                    string.push(char::from_u32(*byte as u32).unwrap());
                } else {
                    string.push('?');
                }
            }

            string
        }

        for (key, value) in db.iterator(IteratorMode::Start) {
            println!("{} = {:?}", bytes_to_string(&key), value);
        }
    }

    #[test]
    fn test() {
        clean_test_indices();

        let store = make_test_store("test_indices/test");

        print_keys(&store.db);

    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

#[derive(Debug, Serialize, Deserialize)]
pub enum EncryptionMethod {
    ChaCha20Poly1305IETF {
        // **ChaCha20-Poly1305-IETF** cypher in blocks per chunk_size bytes
        salt: Vec<u8>,
        // pwhash_argon2i13::Salt as bytes. Random salt used for deriving of key from passphrase
        nonce: Vec<u8>,
        // chacha20poly1305_ietf::Nonce as bytes. Random start nonce. We increment nonce for each chunk to be sure in export file consistency
        chunk_size: usize,
        // size of encrypted chunk
    },
    ChaCha20Poly1305IETFWithSimplify {
        // **ChaCha20-Poly1305-IETF** cypher in blocks per chunk_size bytes
        salt: Vec<u8>,
        // pwhash_argon2i13::Salt as bytes. Random salt used for deriving of key from passphrase
        nonce: Vec<u8>,
        // chacha20poly1305_ietf::Nonce as bytes. Random start nonce. We increment nonce for each chunk to be sure in export file consistency
        chunk_size: usize
    },
}

impl EncryptionMethod {
    pub fn simplify_security(&self) -> bool {
        match self {
            EncryptionMethod::ChaCha20Poly1305IETF { .. } => false,
            EncryptionMethod::ChaCha20Poly1305IETFWithSimplify { .. } => true,
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Header {
    pub encryption_method: EncryptionMethod,
    // Method of encryption for encrypted stream
    pub time: u64,
    // Export time in seconds from UNIX Epoch
    pub version: u32
}

// Note that we use externally tagged enum serialization and header will be represented as:
//
// {
//   "encryption_method": {
//     "ChaCha20Poly1305IETF": {
//       "salt": ..,
//       "nonce": ..,
//       "chunk_size": ..,
//     },
//   },
//   "time": ..,
//   "version": ..,
// }

#[derive(Debug, Serialize, Deserialize)]
pub struct Record {
    // Wallet record type
    #[serde(rename = "type")]
    pub type_: String,
    // Wallet record id
    pub id: String,
    // Wallet record value
    pub value: String,
    // Wallet record tags
    pub tags: HashMap<String, String>,
}<commit_msg>Corrected comments<commit_after>use std::collections::HashMap;

#[derive(Debug, Serialize, Deserialize)]
pub enum EncryptionMethod {
    // **ChaCha20-Poly1305-IETF** cypher in blocks per chunk_size bytes
    ChaCha20Poly1305IETF {
        // pwhash_argon2i13::Salt as bytes. Random salt used for deriving of key from passphrase
        salt: Vec<u8>,
        // chacha20poly1305_ietf::Nonce as bytes. Random start nonce. We increment nonce for each chunk to be sure in export file consistency
        nonce: Vec<u8>,
        // size of encrypted chunk
        chunk_size: usize,
    },
    // **ChaCha20-Poly1305-IETF with simplified key derivation** cypher in blocks per chunk_size bytes
    ChaCha20Poly1305IETFWithSimplify {
        // pwhash_argon2i13::Salt as bytes. Random salt used for deriving of key from passphrase
        salt: Vec<u8>,
        // chacha20poly1305_ietf::Nonce as bytes. Random start nonce. We increment nonce for each chunk to be sure in export file consistency
        nonce: Vec<u8>,
        // size of encrypted chunk
        chunk_size: usize,
    },
}

impl EncryptionMethod {
    pub fn simplify_security(&self) -> bool {
        match self {
            EncryptionMethod::ChaCha20Poly1305IETF { .. } => false,
            EncryptionMethod::ChaCha20Poly1305IETFWithSimplify { .. } => true,
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Header {
    // Method of encryption for encrypted stream
    pub encryption_method: EncryptionMethod,
    // Export time in seconds from UNIX Epoch
    pub time: u64,
    // Version of header
    pub version: u32
}

// Note that we use externally tagged enum serialization and header will be represented as:
//
// {
//   "encryption_method": {
//     "ChaCha20Poly1305IETF": {
//       "salt": ..,
//       "nonce": ..,
//       "chunk_size": ..,
//     },
//   },
//   "time": ..,
//   "version": ..,
// }

#[derive(Debug, Serialize, Deserialize)]
pub struct Record {
    // Wallet record type
    #[serde(rename = "type")]
    pub type_: String,
    // Wallet record id
    pub id: String,
    // Wallet record value
    pub value: String,
    // Wallet record tags
    pub tags: HashMap<String, String>,
}<|endoftext|>


<commit_before>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::SocketAddr;
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;
use extra::priority_queue;
use extra::md4;
use extra::sort;
use std::path;

static PORT:    int = 4414;
static IP: &'static str = "127.0.0.1";

static MAX_CACHE_SIZE_BYTES: u64 = 50000000;
static CACHE_MANAGER_A_RATE: u64 = 2000;
static CACHE_MANAGER_B_RATE: u64 = 6000;

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath,
    topPriority: int,
    fileSize: u64
}

impl Ord for sched_msg {
    fn lt(&self, other: &sched_msg) -> bool {

        let mut retVal:bool = false;

        //If self has a lower geographic priority, then it is less than
        if (self.topPriority < other.topPriority) {
            retVal = true;
        }
        else if (self.topPriority == other.topPriority) {

            //If the two have equal priorities, then we will sort by file size
            //Smaller files get a higher priority
            if (self.fileSize > other.fileSize) {
                retVal = true;
            }
        }

        retVal
    }
}

struct cache_item {
    name: ~str,
    in_use_flag: bool,
    ssi_flag: bool,
    hash: ~str,
    data: ~[u8],
    count: uint,
    size: u64
}

impl cache_item {
    fn le(&self, other: &cache_item) -> bool {
        let mut retVal:bool = false;

        if(self.count <= other.count) {
            retVal = true;
        }

        retVal
    }
}

fn le(this: &cache_item, other: &cache_item) -> bool {
        let mut retVal:bool = false;

        if(this.count <= other.count) {
            retVal = true;
        }

        retVal
    }

fn main() {
    let req_vec: ~priority_queue::PriorityQueue<sched_msg> = ~priority_queue::PriorityQueue::new();
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    let safe_count: ~[uint] = ~[0];
    let shared_visit = arc::RWArc::new(safe_count);

    let cache_list: ~[cache_item] = ~[];
    let shared_cache_list = arc::RWArc::new(cache_list);

    let cache_manager_a = shared_cache_list.clone();
    let cache_manager_b = shared_cache_list.clone();

    //CACHE UPDATE MANAGER (Manager B)
    //This will handle making sure that items in the cache are up-to-date in case they are changed
    do spawn {

        loop {

            do cache_manager_b.write |vec| {

                for i in range(0, (*vec).len()) {

                    //If it is in the cache and in use, we will check to see if the file has been updated
                    if((*vec)[i].in_use_flag) {

                        let curr_path = ~path::Path((*vec)[i].name);
                        if os::path_exists(curr_path) {

                            match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {

                                        let temp_md4 = md4::md4_str(file_data.to_owned());

                                        if(temp_md4 != (*vec)[i].hash)
                                        {
                                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                                Some(s) => s,
                                                None => fail!("Could not access file stats for cache")
                                            };

                                            println(fmt!("===== UPDATING FILE: %?", (*vec)[i].name));

                                            (*vec)[i].data = file_data;
                                            (*vec)[i].size = fileInfo.size;
                                            (*vec)[i].hash = temp_md4;
                                        }
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                        }
                    }
                }
            }

            timer::sleep(CACHE_MANAGER_B_RATE);
        }
    }

    //MAIN CACHE MANAGER (Manager A)
    do spawn {

        loop {
            do cache_manager_a.write |vec| {

                //Quick sort sorts in-place, so we don't need to worry about memory overhead
                //Just time overhead
                sort::quick_sort((*vec), le);

                let mut cache_remaining = MAX_CACHE_SIZE_BYTES;

                for i in range(0, (*vec).len()) {

                    if((*vec)[i].size <= cache_remaining && !(*vec)[i].in_use_flag) {

                        let curr_path = ~path::Path((*vec)[i].name);

                        if os::path_exists(curr_path) {

                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats for cache")
                            };

                            if(fileInfo.size <= cache_remaining) {

                                match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {

                                        //NEED TO ADD CHECK IF IT HAS SSI
                                        //PERHAPS EXCLUDE ALL SSI FROM CACHE?
                                        //ADD TO IF: size < remaining, !in use, !ssi?

                                        (*vec)[i].data = file_data.to_owned();
                                        (*vec)[i].size = fileInfo.size;
                                        (*vec)[i].hash = md4::md4_str(file_data);
                                        (*vec)[i].in_use_flag = true;

                                        cache_remaining = cache_remaining - fileInfo.size;
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                            }
                        }
                    }
                    else if((*vec)[i].size <= cache_remaining && (*vec)[i].in_use_flag) {
                        cache_remaining = cache_remaining - (*vec)[i].size;
                    }
                    else if((*vec)[i].size > cache_remaining) {
                        (*vec)[i].in_use_flag = false;
                        (*vec)[i].data = ~[];
                    }

                }
            }

            timer::sleep(CACHE_MANAGER_A_RATE);
        }
    }

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle

                //Check if file is in cache
                //Will do so using RWArc
                let mut serve_from_cache: bool = false;

                do shared_cache_list.write |vec| {

                    let mut found: bool = false;

                    for i in range(0, (*vec).len()) {

                        if( (*vec)[i].name == tf.filepath.to_str() && (*vec)[i].in_use_flag) {
                            serve_from_cache = true;
                            found = true;

                            println(fmt!("===== SERVING FROM CACHE: %?", tf.filepath.to_str()));

                            tf.stream.write((*vec)[i].data);

                            (*vec)[i].count += 1;
                        }
                        else if( (*vec)[i].name == tf.filepath.to_str() && !(*vec)[i].in_use_flag) {
                            (*vec)[i].count += 1;
                            found = true;
                        }
                    }

                    //If it isn't found, we will create a blank entry and enter basic info, and add the
                    //data and md4 later
                    if(!found) {

                        println(fmt!("===== ADDING ITEM %?", tf.filepath.to_str()));

                        let new_cache_item: cache_item = cache_item{name: tf.filepath.to_str(), in_use_flag: false,
                            ssi_flag: false, hash: ~"", data: ~[], count: 1, size: tf.fileSize};

                        (*vec).push(new_cache_item);

                    }
                }

                if(!serve_from_cache) {

                    match io::read_whole_file(tf.filepath) { // killed if file size is larger than memory size.
                        Ok(file_data) => {
                            //println(fmt!("begin serving file [%?]", tf.filepath));

                            tf.stream.write(file_data);
                            //println(fmt!("finish file [%?]", tf.filepath));

                            println(fmt!("===== SERVING FROM DISK: %?", tf.filepath.to_str()));

                        }
                        Err(err) => {
                            println("ERROR IN SEND");
                            println(err);
                        }
                    }
                }
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {

                    //Since we are using a priority queue, we will use pop()
                    let tf = (*vec).pop();
                    //println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.
                }
            }
        }
    }

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    let mut acceptor = socket.listen().unwrap();

    loop {
        let stream = acceptor.accept();
        let stream = Cell::new(stream);

        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let child_add_vec = add_vec.clone();
        let child_arc = shared_visit.clone();

        do spawn {

            let mut actual_count:uint = 0;

            //We will use the cloned RWARC
            do child_arc.write |count_vec| {

                let prev_count = (*count_vec).pop();
                let new_count: uint = prev_count + 1;

                actual_count = new_count;

                (*count_vec).push(new_count);
            }

            let stream = stream.take();

            match stream {
                Some(s) => {

                    let mut stream = s;

                    let mut buf = [0, ..500];
                    stream.read(buf);
                    let request_str = str::from_utf8(buf);

                    let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
                    if req_group.len() > 2 {
                        let path = req_group[1];
                        //println(fmt!("Request for path: \n%?", path));

                        let file_path = ~os::getcwd().push(path.replace("/../", ""));
                        if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                            //println(fmt!("Request received:\n%s", request_str));
                            let response: ~str = fmt!(
                                "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                                 <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                                 \r\n", actual_count);

                            stream.write(response.as_bytes());
                        }
                        else {
                            // Requests scheduling

                            let mut streamPriority: int = 0;

                            //Retrieving the requesting IP address
                            let ipStr: ~str = match (stream).peer_name() {
                                Some(pr) => pr.ip.to_str(),
                                None => ~"0.0.0.0"
                            };

                            //Split the IP address by '.' so that we can compare
                            let ipSplit: ~[~str] = ipStr.split_iter('.').filter(|&x| x != "")
                                 .map(|x| x.to_owned()).collect();

                            //Assign priority based on geography or if localhost
                            if ( (ipSplit[0] == ~"127" && ipSplit[1] == ~"0") || (ipSplit[0] == ~"128" && ipSplit[1] == ~"143")
                                || (ipSplit[0] == ~"137" && ipSplit[1] == ~"54") ) {

                                streamPriority = 1;
                            }

                            //Get filetype in order to determine HTTP header type
                            let fileType = match file_path.filetype() {
                                Some(s) => s,
                                None => &""
                            };

                            //In order to optimize for the benchmark, we will send the HTTP header quickly before adding to the queue
                            let httpHeader: ~str = match fileType {
                                ".html" | ".htm" | ".php" => ~"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n",
                                _ => ~"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n"
                            };

                            stream.write(httpHeader.as_bytes());
                            stream.flush();

                            //Retrieve file info for additional latency fixes
                            let fileInfo = match std::rt::io::file::stat(file_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats")
                            };

                            let msg: sched_msg = sched_msg{stream: Some(stream), filepath: file_path.clone(), topPriority: streamPriority, fileSize: fileInfo.size};
                            let (sm_port, sm_chan) = std::comm::stream();
                            sm_chan.send(msg);

                            do child_add_vec.write |vec| {
                                let msg = sm_port.recv();
                                (*vec).push(msg); // enqueue new request.
                                //println("add to queue");
                            }
                            child_chan.send(""); //notify the new arriving request.
                            //println(fmt!("get file request: %?", file_path));
                        }
                    }
                    println!("connection terminates")

                },
                None => ()
            }
        }
    }
}<commit_msg>Removed some legacy functions, grouped variable declarations<commit_after>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::SocketAddr;
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;
use extra::priority_queue;
use extra::md4;
use extra::sort;
use std::path;

static PORT:    int = 4414;
static IP: &'static str = "127.0.0.1";

static MAX_CACHE_SIZE_BYTES: u64 = 50000000;
static CACHE_MANAGER_A_RATE: u64 = 2000;
static CACHE_MANAGER_B_RATE: u64 = 6000;

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath,
    topPriority: int,
    fileSize: u64
}

impl Ord for sched_msg {
    fn lt(&self, other: &sched_msg) -> bool {

        let mut retVal:bool = false;

        //If self has a lower geographic priority, then it is less than
        if (self.topPriority < other.topPriority) {
            retVal = true;
        }
        else if (self.topPriority == other.topPriority) {

            //If the two have equal priorities, then we will sort by file size
            //Smaller files get a higher priority
            if (self.fileSize > other.fileSize) {
                retVal = true;
            }
        }

        retVal
    }
}

struct cache_item {
    name: ~str,
    in_use_flag: bool,
    ssi_flag: bool,
    hash: ~str,
    data: ~[u8],
    count: uint,
    size: u64
}

fn le(this: &cache_item, other: &cache_item) -> bool {
        let mut retVal:bool = false;

        if(this.count <= other.count) {
            retVal = true;
        }

        retVal
    }

fn main() {
    let req_vec: ~priority_queue::PriorityQueue<sched_msg> = ~priority_queue::PriorityQueue::new();
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    //Variables pertaining to safe visitor counting
    let safe_count: ~[uint] = ~[0];
    let shared_visit = arc::RWArc::new(safe_count);

    //Variables pertaining to cache and cache manager
    let cache_list: ~[cache_item] = ~[];
    let shared_cache_list = arc::RWArc::new(cache_list);
    let cache_manager_a = shared_cache_list.clone();
    let cache_manager_b = shared_cache_list.clone();

    //CACHE UPDATE MANAGER (Manager B)
    //This will handle making sure that items in the cache are up-to-date in case they are changed
    do spawn {

        loop {

            do cache_manager_b.write |vec| {

                for i in range(0, (*vec).len()) {

                    //If it is in the cache and in use, we will check to see if the file has been updated
                    if((*vec)[i].in_use_flag) {

                        let curr_path = ~path::Path((*vec)[i].name);
                        if os::path_exists(curr_path) {

                            match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {

                                        let temp_md4 = md4::md4_str(file_data.to_owned());

                                        if(temp_md4 != (*vec)[i].hash)
                                        {
                                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                                Some(s) => s,
                                                None => fail!("Could not access file stats for cache")
                                            };

                                            println(fmt!("===== UPDATING FILE: %?", (*vec)[i].name));

                                            (*vec)[i].data = file_data;
                                            (*vec)[i].size = fileInfo.size;
                                            (*vec)[i].hash = temp_md4;
                                        }
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                        }
                    }
                }
            }

            timer::sleep(CACHE_MANAGER_B_RATE);
        }
    }

    //MAIN CACHE MANAGER (Manager A)
    do spawn {

        loop {
            do cache_manager_a.write |vec| {

                //Quick sort sorts in-place, so we don't need to worry about memory overhead
                //Just time overhead
                sort::quick_sort((*vec), le);

                let mut cache_remaining = MAX_CACHE_SIZE_BYTES;

                for i in range(0, (*vec).len()) {

                    if((*vec)[i].size <= cache_remaining && !(*vec)[i].in_use_flag) {

                        let curr_path = ~path::Path((*vec)[i].name);

                        if os::path_exists(curr_path) {

                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats for cache")
                            };

                            if(fileInfo.size <= cache_remaining) {

                                match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {

                                        //NEED TO ADD CHECK IF IT HAS SSI
                                        //PERHAPS EXCLUDE ALL SSI FROM CACHE?
                                        //ADD TO IF: size < remaining, !in use, !ssi?

                                        (*vec)[i].data = file_data.to_owned();
                                        (*vec)[i].size = fileInfo.size;
                                        (*vec)[i].hash = md4::md4_str(file_data);
                                        (*vec)[i].in_use_flag = true;

                                        cache_remaining = cache_remaining - fileInfo.size;
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                            }
                        }
                    }
                    else if((*vec)[i].size <= cache_remaining && (*vec)[i].in_use_flag) {
                        cache_remaining = cache_remaining - (*vec)[i].size;
                    }
                    else if((*vec)[i].size > cache_remaining) {
                        (*vec)[i].in_use_flag = false;
                        (*vec)[i].data = ~[];
                    }

                }
            }

            timer::sleep(CACHE_MANAGER_A_RATE);
        }
    }

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle

                //Check if file is in cache
                //Will do so using RWArc
                let mut serve_from_cache: bool = false;

                do shared_cache_list.write |vec| {

                    let mut found: bool = false;

                    for i in range(0, (*vec).len()) {

                        if( (*vec)[i].name == tf.filepath.to_str() && (*vec)[i].in_use_flag) {
                            serve_from_cache = true;
                            found = true;

                            println(fmt!("===== SERVING FROM CACHE: %?", tf.filepath.to_str()));

                            tf.stream.write((*vec)[i].data);

                            (*vec)[i].count += 1;
                        }
                        else if( (*vec)[i].name == tf.filepath.to_str() && !(*vec)[i].in_use_flag) {
                            (*vec)[i].count += 1;
                            found = true;
                        }
                    }

                    //If it isn't found, we will create a blank entry and enter basic info, and add the
                    //data and md4 later
                    if(!found) {

                        println(fmt!("===== ADDING ITEM %?", tf.filepath.to_str()));

                        let new_cache_item: cache_item = cache_item{name: tf.filepath.to_str(), in_use_flag: false,
                            ssi_flag: false, hash: ~"", data: ~[], count: 1, size: tf.fileSize};

                        (*vec).push(new_cache_item);

                    }
                }

                if(!serve_from_cache) {

                    match io::read_whole_file(tf.filepath) { // killed if file size is larger than memory size.
                        Ok(file_data) => {
                            //println(fmt!("begin serving file [%?]", tf.filepath));

                            tf.stream.write(file_data);
                            //println(fmt!("finish file [%?]", tf.filepath));

                            println(fmt!("===== SERVING FROM DISK: %?", tf.filepath.to_str()));

                        }
                        Err(err) => {
                            println("ERROR IN SEND");
                            println(err);
                        }
                    }
                }
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {

                    //Since we are using a priority queue, we will use pop()
                    let tf = (*vec).pop();
                    //println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.
                }
            }
        }
    }

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    let mut acceptor = socket.listen().unwrap();

    loop {
        let stream = acceptor.accept();
        let stream = Cell::new(stream);

        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let child_add_vec = add_vec.clone();
        let child_arc = shared_visit.clone();

        do spawn {

            let mut actual_count:uint = 0;

            //We will use the cloned RWARC
            do child_arc.write |count_vec| {

                let prev_count = (*count_vec).pop();
                let new_count: uint = prev_count + 1;

                actual_count = new_count;

                (*count_vec).push(new_count);
            }

            let stream = stream.take();

            match stream {
                Some(s) => {

                    let mut stream = s;

                    let mut buf = [0, ..500];
                    stream.read(buf);
                    let request_str = str::from_utf8(buf);

                    let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
                    if req_group.len() > 2 {
                        let path = req_group[1];
                        //println(fmt!("Request for path: \n%?", path));

                        let file_path = ~os::getcwd().push(path.replace("/../", ""));
                        if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                            //println(fmt!("Request received:\n%s", request_str));
                            let response: ~str = fmt!(
                                "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                                 <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                                 \r\n", actual_count);

                            stream.write(response.as_bytes());
                        }
                        else {
                            // Requests scheduling

                            let mut streamPriority: int = 0;

                            //Retrieving the requesting IP address
                            let ipStr: ~str = match (stream).peer_name() {
                                Some(pr) => pr.ip.to_str(),
                                None => ~"0.0.0.0"
                            };

                            //Split the IP address by '.' so that we can compare
                            let ipSplit: ~[~str] = ipStr.split_iter('.').filter(|&x| x != "")
                                 .map(|x| x.to_owned()).collect();

                            //Assign priority based on geography or if localhost
                            if ( (ipSplit[0] == ~"127" && ipSplit[1] == ~"0") || (ipSplit[0] == ~"128" && ipSplit[1] == ~"143")
                                || (ipSplit[0] == ~"137" && ipSplit[1] == ~"54") ) {

                                streamPriority = 1;
                            }

                            //Get filetype in order to determine HTTP header type
                            let fileType = match file_path.filetype() {
                                Some(s) => s,
                                None => &""
                            };

                            //In order to optimize for the benchmark, we will send the HTTP header quickly before adding to the queue
                            let httpHeader: ~str = match fileType {
                                ".html" | ".htm" | ".php" => ~"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n",
                                _ => ~"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n"
                            };

                            stream.write(httpHeader.as_bytes());
                            stream.flush();

                            //Retrieve file info for additional latency fixes
                            let fileInfo = match std::rt::io::file::stat(file_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats")
                            };

                            let msg: sched_msg = sched_msg{stream: Some(stream), filepath: file_path.clone(), topPriority: streamPriority, fileSize: fileInfo.size};
                            let (sm_port, sm_chan) = std::comm::stream();
                            sm_chan.send(msg);

                            do child_add_vec.write |vec| {
                                let msg = sm_port.recv();
                                (*vec).push(msg); // enqueue new request.
                                //println("add to queue");
                            }
                            child_chan.send(""); //notify the new arriving request.
                            //println(fmt!("get file request: %?", file_path));
                        }
                    }
                    println!("connection terminates")

                },
                None => ()
            }
        }
    }
}<|endoftext|>


<commit_before>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.1

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::{SocketAddr, Ipv4Addr};
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;

static PORT:    int = 4414;
static IPV4_LOOPBACK: &'static str = "127.0.0.1";

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath
}

fn main() {
    let req_vec: ~[sched_msg] = ~[];
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    // add file requests into queue.
    do spawn {
        while(true) {
            do add_vec.write |vec| {
                let tf:sched_msg = port.recv();
                (*vec).push(tf);
                println("add to queue");
            }
        }
    }

    // take file requests from queue, and send a response.
    do spawn {
        while(true) {
            do take_vec.write |vec| {
                let mut tf = (*vec).pop();

                match io::read_whole_file(tf.filepath) {
                    Ok(file_data) => {
                        tf.stream.write(file_data);
                    }
                    Err(err) => {
                        println(err);
                    }
                }
            }
        }
    }

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: Ipv4Addr(127,0,0,1), port: PORT as u16});

    println(fmt!("Listening on tcp port %d ...", PORT));
    let mut acceptor = socket.listen().unwrap();

    let mycount: uint = 0;
    let shared_mycount = arc::RWArc::new(mycount);

    // we can limit the incoming connection count.
    //for stream in acceptor.incoming().take(10 as uint) {

    for stream in acceptor.incoming() {

        let stream = Cell::new(stream);

        let update_count = shared_mycount.clone();

        // Start a new task to handle the connection
        let child_chan = chan.clone();
        do spawn{

            let mut vcount = 0;
            do update_count.write |c| {
                *c += 1;
                println(fmt!("incremented count, %u", *c));
                vcount = *c;
            }

            let mut stream = stream.take();

            let mut buf = [0, ..500];
            stream.read(buf);
            let request_str = str::from_utf8(buf);

            let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
            if req_group.len() > 2 {
                let path = req_group[1];
                println(fmt!("Request for path: \n%?", path));

                let file_path = ~os::getcwd().push(path.replace("/../", ""));
                if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                    println(fmt!("Request received:\n%s", request_str));
                    let response: ~str = fmt!(
                        "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                         <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                         \r\n", vcount );

                    stream.write(response.as_bytes());
                }
                else {
                    // may do scheduling here
                    let msg: sched_msg = sched_msg{stream: stream, filepath: file_path.clone()};
                    child_chan.send(msg);

                    println(fmt!("get file request: %?", file_path));
                }
            }
            println!("connection terminates")
        }
    }
}
<commit_msg>got peer_name working to see IP address<commit_after>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.1

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::{SocketAddr, Ipv4Addr};
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;

static PORT:    int = 4414;
static IPV4_LOOPBACK: &'static str = "127.0.0.1";

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath
}

fn main() {
    let req_vec: ~[sched_msg] = ~[];
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    // add file requests into queue.
    do spawn {
        while(true) {
            do add_vec.write |vec| {
                let tf:sched_msg = port.recv();
                (*vec).push(tf);
                println("add to queue");
            }
        }
    }

    // take file requests from queue, and send a response.
    do spawn {
        while(true) {
            do take_vec.write |vec| {
                let mut tf = (*vec).pop();

                match io::read_whole_file(tf.filepath) {
                    Ok(file_data) => {
                        tf.stream.write(file_data);
                    }
                    Err(err) => {
                        println(err);
                    }
                }
            }
        }
    }

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: Ipv4Addr(127,0,0,1), port: PORT as u16});

    println(fmt!("Listening on tcp port %d ...", PORT));
    let mut acceptor = socket.listen().unwrap();

    let mycount: uint = 0;
    let shared_mycount = arc::RWArc::new(mycount);

    // we can limit the incoming connection count.
    //for stream in acceptor.incoming().take(10 as uint) {

    for stream in acceptor.incoming() {

        let stream = Cell::new(stream);

        let update_count = shared_mycount.clone();

        // Start a new task to handle the connection
        let child_chan = chan.clone();
        do spawn{

            let mut vcount = 0;
            do update_count.write |c| {
                *c += 1;
                println(fmt!("incremented count, %u", *c));
                vcount = *c;
            }

// let v = ~[1]; let x = Some(v); match x { Some(ref y) => {} None => {}} x

            let mut stream = stream.take();

            stream.and_then_mut_ref(|x| {
                    let pn = x.peer_name();
                    println(fmt!("Peer address: %s", pn.to_str()));
                    Some(x)
                });

            let mut buf = [0, ..500];
            stream.read(buf);
            let request_str = str::from_utf8(buf);

            let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
            if req_group.len() > 2 {
                let path = req_group[1];
                println(fmt!("Request for path: \n%?", path));

                let file_path = ~os::getcwd().push(path.replace("/../", ""));
                if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                    println(fmt!("Request received:\n%s", request_str));
                    let response: ~str = fmt!(
                        "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                         <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                         \r\n", vcount );

                    stream.write(response.as_bytes());
                }
                else {
                    // may do scheduling here
                    let msg: sched_msg = sched_msg{stream: stream, filepath: file_path.clone()};
                    child_chan.send(msg);

                    println(fmt!("get file request: %?", file_path));
                }
            }
            println!("connection terminates")
        }
    }
}
<|endoftext|>


<commit_before>//! From the NLL RFC: "The deep [aka 'supporting'] prefixes for an
//! place are formed by stripping away fields and derefs, except that
//! we stop when we reach the deref of a shared reference. [...] "
//!
//! "Shallow prefixes are found by stripping away fields, but stop at
//! any dereference. So: writing a path like `a` is illegal if `a.b`
//! is borrowed. But: writing `a` is legal if `*a` is borrowed,
//! whether or not `a` is a shared or mutable reference. [...] "

use super::MirBorrowckCtxt;

use rustc::hir;
use rustc::ty::{self, TyCtxt};
use rustc::mir::{Body, Place, PlaceBase, PlaceRef, ProjectionElem};

pub trait IsPrefixOf<'cx, 'tcx> {
    fn is_prefix_of(&self, other: PlaceRef<'cx, 'tcx>) -> bool;
}

impl<'cx, 'tcx> IsPrefixOf<'cx, 'tcx> for PlaceRef<'cx, 'tcx> {
    fn is_prefix_of(&self, other: PlaceRef<'cx, 'tcx>) -> bool {
        let mut cursor = other.projection;
        loop {
            if self.projection == cursor {
                return self.base == other.base;
            }

            match cursor {
                None => return false,
                Some(proj) => cursor = &proj.base,
            }
        }
    }
}

pub(super) struct Prefixes<'cx, 'tcx> {
    body: &'cx Body<'tcx>,
    tcx: TyCtxt<'tcx>,
    kind: PrefixSet,
    next: Option<(PlaceRef<'cx, 'tcx>)>,
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
#[allow(dead_code)]
pub(super) enum PrefixSet {
    /// Doesn't stop until it returns the base case (a Local or
    /// Static prefix).
    All,
    /// Stops at any dereference.
    Shallow,
    /// Stops at the deref of a shared reference.
    Supporting,
}

impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
    /// Returns an iterator over the prefixes of `place`
    /// (inclusive) from longest to smallest, potentially
    /// terminating the iteration early based on `kind`.
    pub(super) fn prefixes(
        &self,
        place_ref: PlaceRef<'cx, 'tcx>,
        kind: PrefixSet,
    ) -> Prefixes<'cx, 'tcx> {
        Prefixes {
            next: Some(place_ref),
            kind,
            body: self.body,
            tcx: self.infcx.tcx,
        }
    }
}

impl<'cx, 'tcx> Iterator for Prefixes<'cx, 'tcx> {
    type Item = PlaceRef<'cx, 'tcx>;
    fn next(&mut self) -> Option<Self::Item> {
        let mut cursor = self.next?;

        // Post-processing `place`: Enqueue any remaining
        // work. Also, `place` may not be a prefix itself, but
        // may hold one further down (e.g., we never return
        // downcasts here, but may return a base of a downcast).

        'cursor: loop {
            let proj = match &cursor {
                PlaceRef {
                    base: PlaceBase::Local(_),
                    projection: None,
                }
                | // search yielded this leaf
                PlaceRef {
                    base: PlaceBase::Static(_),
                    projection: None,
                } => {
                    self.next = None;
                    return Some(cursor);
                }
                PlaceRef {
                    base: _,
                    projection: Some(proj),
                } => proj,
            };

            match proj.elem {
                ProjectionElem::Field(_ /*field*/, _ /*ty*/) => {
                    // FIXME: add union handling
                    self.next = Some(PlaceRef {
                        base: cursor.base,
                        projection: &proj.base,
                    });
                    return Some(cursor);
                }
                ProjectionElem::Downcast(..) |
                ProjectionElem::Subslice { .. } |
                ProjectionElem::ConstantIndex { .. } |
                ProjectionElem::Index(_) => {
                    cursor = PlaceRef {
                        base: cursor.base,
                        projection: &proj.base,
                    };
                    continue 'cursor;
                }
                ProjectionElem::Deref => {
                    // (handled below)
                }
            }

            assert_eq!(proj.elem, ProjectionElem::Deref);

            match self.kind {
                PrefixSet::Shallow => {
                    // shallow prefixes are found by stripping away
                    // fields, but stop at *any* dereference.
                    // So we can just stop the traversal now.
                    self.next = None;
                    return Some(cursor);
                }
                PrefixSet::All => {
                    // all prefixes: just blindly enqueue the base
                    // of the projection
                    self.next = Some(PlaceRef {
                        base: cursor.base,
                        projection: &proj.base,
                    });
                    return Some(cursor);
                }
                PrefixSet::Supporting => {
                    // fall through!
                }
            }

            assert_eq!(self.kind, PrefixSet::Supporting);
            // supporting prefixes: strip away fields and
            // derefs, except we stop at the deref of a shared
            // reference.

            let ty = Place::ty_from(cursor.base, &proj.base, self.body, self.tcx).ty;
            match ty.sty {
                ty::RawPtr(_) |
                ty::Ref(
                    _, /*rgn*/
                    _, /*ty*/
                    hir::MutImmutable
                    ) => {
                    // don't continue traversing over derefs of raw pointers or shared borrows.
                    self.next = None;
                    return Some(cursor);
                }

                ty::Ref(
                    _, /*rgn*/
                    _, /*ty*/
                    hir::MutMutable,
                    ) => {
                    self.next = Some(PlaceRef {
                        base: cursor.base,
                        projection: &proj.base,
                    });
                    return Some(cursor);
                }

                ty::Adt(..) if ty.is_box() => {
                    self.next = Some(PlaceRef {
                        base: cursor.base,
                        projection: &proj.base,
                    });
                    return Some(cursor);
                }

                _ => panic!("unknown type fed to Projection Deref."),
            }
        }
    }
}
<commit_msg>Add a period at the end of the sentence<commit_after>//! From the NLL RFC: "The deep [aka 'supporting'] prefixes for an
//! place are formed by stripping away fields and derefs, except that
//! we stop when we reach the deref of a shared reference. [...] "
//!
//! "Shallow prefixes are found by stripping away fields, but stop at
//! any dereference. So: writing a path like `a` is illegal if `a.b`
//! is borrowed. But: writing `a` is legal if `*a` is borrowed,
//! whether or not `a` is a shared or mutable reference. [...] "

use super::MirBorrowckCtxt;

use rustc::hir;
use rustc::ty::{self, TyCtxt};
use rustc::mir::{Body, Place, PlaceBase, PlaceRef, ProjectionElem};

pub trait IsPrefixOf<'cx, 'tcx> {
    fn is_prefix_of(&self, other: PlaceRef<'cx, 'tcx>) -> bool;
}

impl<'cx, 'tcx> IsPrefixOf<'cx, 'tcx> for PlaceRef<'cx, 'tcx> {
    fn is_prefix_of(&self, other: PlaceRef<'cx, 'tcx>) -> bool {
        let mut cursor = other.projection;
        loop {
            if self.projection == cursor {
                return self.base == other.base;
            }

            match cursor {
                None => return false,
                Some(proj) => cursor = &proj.base,
            }
        }
    }
}

pub(super) struct Prefixes<'cx, 'tcx> {
    body: &'cx Body<'tcx>,
    tcx: TyCtxt<'tcx>,
    kind: PrefixSet,
    next: Option<(PlaceRef<'cx, 'tcx>)>,
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
#[allow(dead_code)]
pub(super) enum PrefixSet {
    /// Doesn't stop until it returns the base case (a Local or
    /// Static prefix).
    All,
    /// Stops at any dereference.
    Shallow,
    /// Stops at the deref of a shared reference.
    Supporting,
}

impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
    /// Returns an iterator over the prefixes of `place`
    /// (inclusive) from longest to smallest, potentially
    /// terminating the iteration early based on `kind`.
    pub(super) fn prefixes(
        &self,
        place_ref: PlaceRef<'cx, 'tcx>,
        kind: PrefixSet,
    ) -> Prefixes<'cx, 'tcx> {
        Prefixes {
            next: Some(place_ref),
            kind,
            body: self.body,
            tcx: self.infcx.tcx,
        }
    }
}

impl<'cx, 'tcx> Iterator for Prefixes<'cx, 'tcx> {
    type Item = PlaceRef<'cx, 'tcx>;
    fn next(&mut self) -> Option<Self::Item> {
        let mut cursor = self.next?;

        // Post-processing `place`: Enqueue any remaining
        // work. Also, `place` may not be a prefix itself, but
        // may hold one further down (e.g., we never return
        // downcasts here, but may return a base of a downcast).

        'cursor: loop {
            let proj = match &cursor {
                PlaceRef {
                    base: PlaceBase::Local(_),
                    projection: None,
                }
                | // search yielded this leaf
                PlaceRef {
                    base: PlaceBase::Static(_),
                    projection: None,
                } => {
                    self.next = None;
                    return Some(cursor);
                }
                PlaceRef {
                    base: _,
                    projection: Some(proj),
                } => proj,
            };

            match proj.elem {
                ProjectionElem::Field(_ /*field*/, _ /*ty*/) => {
                    // FIXME: add union handling
                    self.next = Some(PlaceRef {
                        base: cursor.base,
                        projection: &proj.base,
                    });
                    return Some(cursor);
                }
                ProjectionElem::Downcast(..) |
                ProjectionElem::Subslice { .. } |
                ProjectionElem::ConstantIndex { .. } |
                ProjectionElem::Index(_) => {
                    cursor = PlaceRef {
                        base: cursor.base,
                        projection: &proj.base,
                    };
                    continue 'cursor;
                }
                ProjectionElem::Deref => {
                    // (handled below)
                }
            }

            assert_eq!(proj.elem, ProjectionElem::Deref);

            match self.kind {
                PrefixSet::Shallow => {
                    // shallow prefixes are found by stripping away
                    // fields, but stop at *any* dereference.
                    // So we can just stop the traversal now.
                    self.next = None;
                    return Some(cursor);
                }
                PrefixSet::All => {
                    // all prefixes: just blindly enqueue the base
                    // of the projection.
                    self.next = Some(PlaceRef {
                        base: cursor.base,
                        projection: &proj.base,
                    });
                    return Some(cursor);
                }
                PrefixSet::Supporting => {
                    // fall through!
                }
            }

            assert_eq!(self.kind, PrefixSet::Supporting);
            // supporting prefixes: strip away fields and
            // derefs, except we stop at the deref of a shared
            // reference.

            let ty = Place::ty_from(cursor.base, &proj.base, self.body, self.tcx).ty;
            match ty.sty {
                ty::RawPtr(_) |
                ty::Ref(
                    _, /*rgn*/
                    _, /*ty*/
                    hir::MutImmutable
                    ) => {
                    // don't continue traversing over derefs of raw pointers or shared borrows.
                    self.next = None;
                    return Some(cursor);
                }

                ty::Ref(
                    _, /*rgn*/
                    _, /*ty*/
                    hir::MutMutable,
                    ) => {
                    self.next = Some(PlaceRef {
                        base: cursor.base,
                        projection: &proj.base,
                    });
                    return Some(cursor);
                }

                ty::Adt(..) if ty.is_box() => {
                    self.next = Some(PlaceRef {
                        base: cursor.base,
                        projection: &proj.base,
                    });
                    return Some(cursor);
                }

                _ => panic!("unknown type fed to Projection Deref."),
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::cell::{Cell, RefCell};
use std::error::Error;
use std::fmt;
use std::future::Future;
use std::marker::PhantomData;

thread_local!(static ENTERED: Cell<bool> = Cell::new(false));

/// Represents an executor context.
///
/// For more details, see [`enter` documentation](fn.enter.html)
pub struct Enter {
    _p: PhantomData<RefCell<()>>,
}

/// An error returned by `enter` if an execution scope has already been
/// entered.
pub struct EnterError {
    _a: (),
}

impl fmt::Debug for EnterError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("EnterError")
            .field("reason", &format!("{}", self))
            .finish()
    }
}

impl fmt::Display for EnterError {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            fmt,
            "attempted to run an executor while another executor is already running"
        )
    }
}

impl Error for EnterError {}

/// Marks the current thread as being within the dynamic extent of an
/// executor.
///
/// Executor implementations should call this function before blocking the
/// thread. If `None` is returned, the executor should fail by panicking or
/// taking some other action without blocking the current thread. This prevents
/// deadlocks due to multiple executors competing for the same thread.
///
/// # Error
///
/// Returns an error if the current thread is already marked
pub fn enter() -> Result<Enter, EnterError> {
    ENTERED.with(|c| {
        if c.get() {
            Err(EnterError { _a: () })
        } else {
            c.set(true);

            Ok(Enter { _p: PhantomData })
        }
    })
}

impl Enter {
    /// Blocks the thread on the specified future, returning the value with
    /// which that future completes.
    pub fn block_on<F: Future>(&mut self, mut f: F) -> F::Output {
        use crate::park::{Park, ParkThread};
        use std::pin::Pin;
        use std::task::Context;
        use std::task::Poll::Ready;

        let park = ParkThread::new();
        let waker = park.unpark().into_waker();
        let mut cx = Context::from_waker(&waker);

        // `block_on` takes ownership of `f`. Once it is pinned here, the original `f` binding can
        // no longer be accessed, making the pinning safe.
        let mut f = unsafe { Pin::new_unchecked(&mut f) };

        loop {
            if let Ready(v) = f.as_mut().poll(&mut cx) {
                return v;
            }
        }
    }
}

impl fmt::Debug for Enter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Enter").finish()
    }
}

impl Drop for Enter {
    fn drop(&mut self) {
        ENTERED.with(|c| {
            assert!(c.get());
            c.set(false);
        });
    }
}
<commit_msg>executor: block thread when needed in block fn (#1303)<commit_after>use std::cell::{Cell, RefCell};
use std::error::Error;
use std::fmt;
use std::future::Future;
use std::marker::PhantomData;

thread_local!(static ENTERED: Cell<bool> = Cell::new(false));

/// Represents an executor context.
///
/// For more details, see [`enter` documentation](fn.enter.html)
pub struct Enter {
    _p: PhantomData<RefCell<()>>,
}

/// An error returned by `enter` if an execution scope has already been
/// entered.
pub struct EnterError {
    _a: (),
}

impl fmt::Debug for EnterError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("EnterError")
            .field("reason", &format!("{}", self))
            .finish()
    }
}

impl fmt::Display for EnterError {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            fmt,
            "attempted to run an executor while another executor is already running"
        )
    }
}

impl Error for EnterError {}

/// Marks the current thread as being within the dynamic extent of an
/// executor.
///
/// Executor implementations should call this function before blocking the
/// thread. If `None` is returned, the executor should fail by panicking or
/// taking some other action without blocking the current thread. This prevents
/// deadlocks due to multiple executors competing for the same thread.
///
/// # Error
///
/// Returns an error if the current thread is already marked
pub fn enter() -> Result<Enter, EnterError> {
    ENTERED.with(|c| {
        if c.get() {
            Err(EnterError { _a: () })
        } else {
            c.set(true);

            Ok(Enter { _p: PhantomData })
        }
    })
}

impl Enter {
    /// Blocks the thread on the specified future, returning the value with
    /// which that future completes.
    pub fn block_on<F: Future>(&mut self, mut f: F) -> F::Output {
        use crate::park::{Park, ParkThread};
        use std::pin::Pin;
        use std::task::Context;
        use std::task::Poll::Ready;

        let mut park = ParkThread::new();
        let waker = park.unpark().into_waker();
        let mut cx = Context::from_waker(&waker);

        // `block_on` takes ownership of `f`. Once it is pinned here, the original `f` binding can
        // no longer be accessed, making the pinning safe.
        let mut f = unsafe { Pin::new_unchecked(&mut f) };

        loop {
            if let Ready(v) = f.as_mut().poll(&mut cx) {
                return v;
            }
            park.park().unwrap();
        }
    }
}

impl fmt::Debug for Enter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Enter").finish()
    }
}

impl Drop for Enter {
    fn drop(&mut self) {
        ENTERED.with(|c| {
            assert!(c.get());
            c.set(false);
        });
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// ignore-stage1
// ignore-cross-compile
#![feature(quote, rustc_private)]

extern crate syntax;
extern crate rustc_data_structures;

use syntax::ext::base::{ExtCtxt, DummyResolver};
use syntax::ext::expand::ExpansionConfig;
use syntax::parse::ParseSess;
use syntax::codemap::{FilePathMapping, dummy_spanned};
use syntax::print::pprust::expr_to_string;
use syntax::ast::{Expr, ExprKind, LitKind, StrStyle, RangeLimits};
use syntax::symbol::Symbol;
use syntax::ptr::P;

use rustc_data_structures::sync::Lrc;

fn main() {
    syntax::with_globals(|| run());
}

fn run() {
    let parse_sess = ParseSess::new(FilePathMapping::empty());
    let exp_cfg = ExpansionConfig::default("issue_35829".to_owned());
    let mut resolver = DummyResolver;
    let cx = ExtCtxt::new(&parse_sess, exp_cfg, &mut resolver);

    // check byte string
    let byte_string = quote_expr!(&cx, b"one");
    let byte_string_lit_kind = LitKind::ByteStr(Lrc::new(b"one".to_vec()));
    assert_eq!(byte_string.node, ExprKind::Lit(P(dummy_spanned(byte_string_lit_kind))));

    // check raw byte string
    let raw_byte_string = quote_expr!(&cx, br###"#"two"#"###);
    let raw_byte_string_lit_kind = LitKind::ByteStr(Lrc::new(b"#\"two\"#".to_vec()));
    assert_eq!(raw_byte_string.node, ExprKind::Lit(P(dummy_spanned(raw_byte_string_lit_kind))));

    // check dotdoteq
    let closed_range = quote_expr!(&cx, 0 ..= 1);
    assert_eq!(closed_range.node, ExprKind::Range(
        Some(quote_expr!(&cx, 0)),
        Some(quote_expr!(&cx, 1)),
        RangeLimits::Closed
    ));

    // test case from 35829
    let expr_35829 = quote_expr!(&cx, std::io::stdout().write(b"one"));
    assert_eq!(expr_to_string(&expr_35829), r#"std::io::stdout().write(b"one")"#);
}
<commit_msg>Remove unused imports<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// ignore-stage1
// ignore-cross-compile
#![feature(quote, rustc_private)]

extern crate syntax;
extern crate rustc_data_structures;

use syntax::ext::base::{ExtCtxt, DummyResolver};
use syntax::ext::expand::ExpansionConfig;
use syntax::parse::ParseSess;
use syntax::codemap::{FilePathMapping, dummy_spanned};
use syntax::print::pprust::expr_to_string;
use syntax::ast::{ExprKind, LitKind, RangeLimits};
use syntax::ptr::P;

use rustc_data_structures::sync::Lrc;

fn main() {
    syntax::with_globals(|| run());
}

fn run() {
    let parse_sess = ParseSess::new(FilePathMapping::empty());
    let exp_cfg = ExpansionConfig::default("issue_35829".to_owned());
    let mut resolver = DummyResolver;
    let cx = ExtCtxt::new(&parse_sess, exp_cfg, &mut resolver);

    // check byte string
    let byte_string = quote_expr!(&cx, b"one");
    let byte_string_lit_kind = LitKind::ByteStr(Lrc::new(b"one".to_vec()));
    assert_eq!(byte_string.node, ExprKind::Lit(P(dummy_spanned(byte_string_lit_kind))));

    // check raw byte string
    let raw_byte_string = quote_expr!(&cx, br###"#"two"#"###);
    let raw_byte_string_lit_kind = LitKind::ByteStr(Lrc::new(b"#\"two\"#".to_vec()));
    assert_eq!(raw_byte_string.node, ExprKind::Lit(P(dummy_spanned(raw_byte_string_lit_kind))));

    // check dotdoteq
    let closed_range = quote_expr!(&cx, 0 ..= 1);
    assert_eq!(closed_range.node, ExprKind::Range(
        Some(quote_expr!(&cx, 0)),
        Some(quote_expr!(&cx, 1)),
        RangeLimits::Closed
    ));

    // test case from 35829
    let expr_35829 = quote_expr!(&cx, std::io::stdout().write(b"one"));
    assert_eq!(expr_to_string(&expr_35829), r#"std::io::stdout().write(b"one")"#);
}
<|endoftext|>