<commit_before>use std::io::prelude::*;
use rand::OsRng;
use num::traits::FromPrimitive;

use tls_result::{TlsResult, TlsError, TlsErrorKind};
use tls_result::TlsErrorKind::{UnexpectedMessage, RecordOverflow, BadRecordMac, AlertReceived};
use alert::{self, Alert};
use handshake::{Handshake, HandshakeBuffer};
use util::u64_be_array;
use util::{ReadExt, WriteExt};
use cipher::{Encryptor, Decryptor};
use tls_item::TlsItem;

use self::ContentType::{ChangeCipherSpecTy, AlertTy, HandshakeTy, ApplicationDataTy};
use self::Message::{HandshakeMessage, ChangeCipherSpecMessage, AlertMessage,
                    ApplicationDataMessage};

pub static TLS_VERSION: (u8, u8) = (3, 3);

enum_from_primitive! {
    #[repr(u8)]
    #[derive(Copy, Clone, PartialEq, Debug)]
    pub enum ContentType {
        ChangeCipherSpecTy = 20,
        AlertTy = 21,
        HandshakeTy = 22,
        ApplicationDataTy = 23,
        // HeartBeat = 24, RFC 6520 extension :-)
    }
}

/// maximum length of Record (excluding content_type, version, length fields)
pub const RECORD_MAX_LEN: usize = 1 << 14;

/// maximum length of EncryptedRecord (excluding content_type, version, length fields)
pub const ENC_RECORD_MAX_LEN: usize = (1 << 14) + 2048;

/// corresponds to `TLSPlaintext` in Section 6.2.1.
#[derive(Debug)]
pub struct Record {
    pub content_type: ContentType,
    pub ver_major: u8,
    pub ver_minor: u8,
    // fragment length < 2^14
    pub fragment: Vec<u8>,
}

impl Record {
    pub fn new(content_type: ContentType,
               ver_major: u8,
               ver_minor: u8,
               fragment: Vec<u8>) -> Record {
        let len = fragment.len();
        if len > RECORD_MAX_LEN {
            panic!("record too long: {} > 2^14", len);
        }

        Record {
            content_type: content_type,
            ver_major: ver_major,
            ver_minor: ver_minor,
            fragment: fragment,
        }
    }
}

pub struct RecordWriter<W: Write> {
    writer: W,
    // if encryptor is None, handshake is not done yet.
    encryptor: Option<Box<Encryptor + Send + 'static>>,
    write_count: u64,
}

impl<W: Write> RecordWriter<W> {
    pub fn new(writer: W) -> RecordWriter<W> {
        RecordWriter {
            writer: writer,
            encryptor: None,
            write_count: 0,
        }
    }

    #[inline]
    pub fn get_mut(&mut self) -> &mut W {
        &mut self.writer
    }

    pub fn set_encryptor(&mut self, encryptor: Box<Encryptor + Send + 'static>) {
        self.encryptor = Some(encryptor);
        self.write_count = 0;
    }

    pub fn write_record(&mut self, record: Record) -> TlsResult<()> {
        let encrypted_fragment = match self.encryptor {
            None => record.fragment,
            Some(ref mut encryptor) => {
                let seq_num = u64_be_array(self.write_count);

                let mut ad = Vec::new();
                ad.extend(&seq_num);
                ad.push(record.content_type as u8);
                ad.push(record.ver_major);
                ad.push(record.ver_minor);
                let frag_len = record.fragment.len() as u16;
                ad.push((frag_len >> 8) as u8);
                ad.push(frag_len as u8);

                let encrypted_fragment = encryptor.encrypt(&seq_num,
                                                           &record.fragment,
                                                           &ad);
                encrypted_fragment
            }
        };

        let fragment_len = encrypted_fragment.len();
        if fragment_len > ENC_RECORD_MAX_LEN {
            panic!("record too long: {} > 2^14 + 2048", fragment_len);
        }

        try!(self.writer.write_u8(record.content_type as u8));
        try!(self.writer.write_u8(record.ver_major));
        try!(self.writer.write_u8(record.ver_minor));
        try!(self.writer.write_be_u16(fragment_len as u16));
        try!(self.writer.write_all(&encrypted_fragment));

        self.write_count += 1;

        Ok(())
    }

    pub fn write_data(&mut self, ty: ContentType, data: &[u8]) -> TlsResult<()> {
        let (major, minor) = TLS_VERSION;
        // TODO: configurable maxlen
        for fragment in data.chunks(RECORD_MAX_LEN) {
            let fragment = fragment.to_vec();
            let record = Record::new(ty, major, minor, fragment);
            try!(self.write_record(record));
        }

        Ok(())
    }

    pub fn write_handshake(&mut self, handshake: &Handshake) -> TlsResult<()> {
        let mut data = Vec::new();
        try!(handshake.tls_write(&mut data));
        self.write_data(HandshakeTy, &data)
    }

    pub fn write_alert(&mut self, alert: &Alert) -> TlsResult<()> {
        let mut data = Vec::new();
        try!(alert.tls_write(&mut data));
        self.write_data(AlertTy, &data)
    }

    pub fn write_change_cipher_spec(&mut self) -> TlsResult<()> {
        self.write_data(ChangeCipherSpecTy, &[1u8])
    }

    pub fn write_application_data(&mut self, data: &[u8]) -> TlsResult<()> {
        if self.encryptor.is_none() {
            panic!("attempted to write ApplicationData before handshake");
        }
        self.write_data(ApplicationDataTy, data)
    }
}

pub enum Message {
    HandshakeMessage(Handshake),
    ChangeCipherSpecMessage,
    AlertMessage(Alert),
    ApplicationDataMessage(Vec<u8>),
}

pub struct RecordReader<R: ReadExt> {
    reader: R,
    // if decryptor is none, handshake is not done yet.
    decryptor: Option<Box<Decryptor + Send + 'static>>,
    read_count: u64,
    handshake_buffer: HandshakeBuffer,
}

impl<R: ReadExt> RecordReader<R> {
    pub fn new(reader: R) -> RecordReader<R> {
        RecordReader {
            reader: reader,
            decryptor: None,
            read_count: 0,
            handshake_buffer: HandshakeBuffer::new(),
        }
    }

    #[inline]
    pub fn get_mut(&mut self) -> &mut R {
        &mut self.reader
    }

    pub fn set_decryptor(&mut self, decryptor: Box<Decryptor + Send + 'static>) {
        self.decryptor = Some(decryptor);
        self.read_count = 0;
    }

    fn read_record(&mut self) -> TlsResult<Record> {
        let content_type = {
            let ty = try!(self.reader.read_u8());
            let ct: Option<ContentType> = FromPrimitive::from_u8(ty);
            match ct {
                Some(ty) => ty,
                None => return tls_err!(UnexpectedMessage, "unexpected ContentType: {}", ty),
            }
        };

        let major = try!(self.reader.read_u8());
        let minor = try!(self.reader.read_u8());

        let len = {
            let len = try!(self.reader.read_be_u16()) as usize;
            if len > ENC_RECORD_MAX_LEN {
                return tls_err!(RecordOverflow, "TLSEncryptedText too long: {}", len);
            }
            len
        };

        let fragment = try!(self.reader.read_exact(len as usize));

        let record = match self.decryptor {
            None => {
                if fragment.len() > RECORD_MAX_LEN {
                    return tls_err!(RecordOverflow,
                                    "decrypted record too long: {}",
                                    fragment.len());
                }
                Record::new(content_type, major, minor, fragment)
            }
            Some(ref mut decryptor) => {
                let seq_num = u64_be_array(self.read_count);

                let mut ad = Vec::new();
                ad.extend(&seq_num);
                ad.push(content_type as u8); // TLSCompressed.type
                ad.push(major);
                ad.push(minor);

                let mac_len = decryptor.mac_len();
                let total_len = fragment.len();
                if total_len < mac_len {
                    return tls_err!(BadRecordMac, "encrypted message too short: {}", total_len);
                }
                let frag_len = (total_len - mac_len) as u16;
                ad.push((frag_len >> 8) as u8);
                ad.push(frag_len as u8);

                // TODO: "seq_num as nonce" is chacha20poly1305-specific
                let data = try!(decryptor.decrypt(&seq_num, &fragment, &ad));
                if data.len() > RECORD_MAX_LEN {
                    // decryption routine went wrong.
                    return panic!("decrypted record too long: {}", data.len());
                }

                Record::new(content_type, major, minor, data)
            }
        };

        self.read_count += 1;

        Ok(record)
    }

    /// read records until a "complete" message is found, and return the message.
    /// if invalid ChangeCipherSpec/Alert/Handshake message is found, return Err.
    /// (application record is always considered "complete" and "valid"
    /// since it is opaque to TLS layer.)
    pub fn read_message(&mut self) -> TlsResult<Message> {
        match try!(self.handshake_buffer.get_message()) {
            Some(handshake_msg) => return Ok(HandshakeMessage(handshake_msg)),
            None => {}
        }

        // ok, no message found. read it from network!
        loop {
            // TODO: what if handshake record is present in buffer then
            // other record comes? is it legal?

            let record = try!(self.read_record());
            match record.content_type {
                ChangeCipherSpecTy => {
                    if record.fragment.len() != 1 || record.fragment[0] != 1 {
                        return tls_err!(UnexpectedMessage, "invalid ChangeCipherSpec arrived");
                    }
                    return Ok(ChangeCipherSpecMessage);
                }
                AlertTy => {
                    let len = record.fragment.len();
                    if len == 0 {
                        return tls_err!(UnexpectedMessage, "zero-length Alert record arrived");
                    } else if len < 2 {
                        // alert packet can be broken into several records,
                        // buf it is rarely used and may cause alert attack
                        // if carelessly implemented:
                        // http://www.mitls.org/wsgi/alert-attack
                        // we just don't accept such record for simplicity.
                        // If alert messages are long, use the first two bytes.
                        return tls_err!(UnexpectedMessage, "awkward Alert record arrived");
                    }
                    let level = FromPrimitive::from_u8(record.fragment[0]);
                    let desc = FromPrimitive::from_u8(record.fragment[1]);
                    match (level, desc) {
                        (Some(level), Some(desc)) => {
                            return Ok(AlertMessage(try!(Alert::new(level, desc))));
                        }
                        _ => return tls_err!(UnexpectedMessage,
                                             "unknown alert: {:?}",
                                             record.fragment),
                    }
                }
                HandshakeTy => {
                    if record.fragment.len() == 0 {
                        return tls_err!(UnexpectedMessage, "zero-length Handshake arrived");
                    }
                    self.handshake_buffer.add_record(&record.fragment);

                    match try!(self.handshake_buffer.get_message()) {
                        Some(handshake_msg) => return Ok(HandshakeMessage(handshake_msg)),
                        _ => {}
                    }
                }
                ApplicationDataTy => {
                    return Ok(ApplicationDataMessage(record.fragment));
                }
            }
        }
    }

    pub fn read_application_data(&mut self) -> TlsResult<Vec<u8>> {
        if self.decryptor.is_none() {
            panic!("ApplicationData called before handshake");
        }
        loop {
            let msg = try!(self.read_message());
            match msg {
                ApplicationDataMessage(msg) => return Ok(msg),
                // TODO: handle other cases
                AlertMessage(..) => unimplemented!(),
                ChangeCipherSpecMessage(..) => unimplemented!(), // this should not come here
                HandshakeMessage(..) => unimplemented!(), // TODO: re-handshake
            }
        }
    }

    pub fn read_handshake(&mut self) -> TlsResult<Handshake> {
        match try!(self.read_message()) {
            HandshakeMessage(handshake) => Ok(handshake),
            AlertMessage(alert) => tls_err!(AlertReceived, "alert: {:?}", alert.description),
            _ => tls_err!(UnexpectedMessage, "expected Handshake"),
        }
    }

    pub fn read_change_cipher_spec(&mut self) -> TlsResult<()> {
        match try!(self.read_message()) {
            ChangeCipherSpecMessage => Ok(()),
            _ => tls_err!(UnexpectedMessage, "expected ChangeCipherSpec"),
        }
    }
}

pub struct Tls<R: Read, W: Write> {
    pub writer: RecordWriter<W>,
    pub reader: RecordReader<R>,
    pub rng: OsRng,
}

impl<R: Read, W: Write> Tls<R, W> {
    pub fn new(reader: R, writer: W, rng: OsRng) -> Tls<R, W> {
        let writer = RecordWriter::new(writer);
        let reader = RecordReader::new(reader);
        Tls {
            writer: writer,
            reader: reader,
            rng: rng,
        }
    }

    pub fn close(&mut self) -> TlsResult<()> {
        let alert_data = alert::Alert {
            level: alert::AlertLevel::fatal,
            description: alert::AlertDescription::close_notify,
        };
        try!(self.writer.write_alert(&alert_data));
        Ok(())
    }

    // send fatal alert and return error
    // (it may be different to `err`, because writing alert can fail)
    pub fn send_tls_alert(&mut self, err: TlsError) -> TlsError {
        match err.kind {
            TlsErrorKind::IoFailure => return err,
            _ => {
                let alert = alert::Alert::from_tls_err(&err);
                let result = self.writer.write_alert(&alert);
                match result {
                    Ok(()) => return err,
                    Err(err) => return err,
                }
            }
        }
    }
}
<commit_msg>tls: Add doccomments and assertions<commit_after>use std::io::prelude::*;
use rand::OsRng;
use num::traits::FromPrimitive;

use tls_result::{TlsResult, TlsError, TlsErrorKind};
use tls_result::TlsErrorKind::{UnexpectedMessage, RecordOverflow, BadRecordMac, AlertReceived};
use alert::{self, Alert};
use handshake::{Handshake, HandshakeBuffer};
use util::u64_be_array;
use util::{ReadExt, WriteExt};
use cipher::{Encryptor, Decryptor};
use tls_item::TlsItem;

use self::ContentType::{ChangeCipherSpecTy, AlertTy, HandshakeTy, ApplicationDataTy};
use self::Message::{HandshakeMessage, ChangeCipherSpecMessage, AlertMessage,
                    ApplicationDataMessage};

pub static TLS_VERSION: (u8, u8) = (3, 3);

enum_from_primitive! {
    #[repr(u8)]
    #[derive(Copy, Clone, PartialEq, Debug)]
    pub enum ContentType {
        ChangeCipherSpecTy = 20,
        AlertTy = 21,
        HandshakeTy = 22,
        ApplicationDataTy = 23,
        // HeartBeat = 24, RFC 6520 extension :-)
    }
}

/// maximum length of Record (excluding content_type, version, length fields)
pub const RECORD_MAX_LEN: usize = 1 << 14;

/// maximum length of EncryptedRecord (excluding content_type, version, length fields)
pub const ENC_RECORD_MAX_LEN: usize = (1 << 14) + 2048;

/// corresponds to `TLSPlaintext` in Section 6.2.1.
#[derive(Debug)]
pub struct Record {
    pub content_type: ContentType,
    pub ver_major: u8,
    pub ver_minor: u8,
    // fragment length < 2^14
    pub fragment: Vec<u8>,
}

impl Record {
    pub fn new(content_type: ContentType,
               ver_major: u8,
               ver_minor: u8,
               fragment: Vec<u8>) -> Record {
        let len = fragment.len();
        if len > RECORD_MAX_LEN {
            panic!("record too long: {} > 2^14", len);
        }

        Record {
            content_type: content_type,
            ver_major: ver_major,
            ver_minor: ver_minor,
            fragment: fragment,
        }
    }
}

/// Writes `Record` or higher-layer message to a writable object.
/// Record is internally encrypted before written.
pub struct RecordWriter<W: Write> {
    writer: W,
    // if encryptor is None, handshake is not done yet.
    encryptor: Option<Box<Encryptor + Send + 'static>>,
    write_count: u64,
}

impl<W: Write> RecordWriter<W> {
    /// Create new `RecordWriter` with null encryption.
    /// Invoke `set_encryptor` to set encryptor.
    pub fn new(writer: W) -> RecordWriter<W> {
        RecordWriter {
            writer: writer,
            encryptor: None,
            write_count: 0,
        }
    }

    #[inline]
    pub fn get_mut(&mut self) -> &mut W {
        &mut self.writer
    }

    /// Set encryptor and reset count.
    /// This must be called only once.
    pub fn set_encryptor(&mut self, encryptor: Box<Encryptor + Send + 'static>) {
        assert!(self.encryptor.is_none());
        self.encryptor = Some(encryptor);
        self.write_count = 0;
    }

    pub fn write_record(&mut self, record: Record) -> TlsResult<()> {
        let encrypted_fragment = match self.encryptor {
            None => record.fragment,
            Some(ref mut encryptor) => {
                let seq_num = u64_be_array(self.write_count);

                let mut ad = Vec::new();
                ad.extend(&seq_num);
                ad.push(record.content_type as u8);
                ad.push(record.ver_major);
                ad.push(record.ver_minor);
                let frag_len = record.fragment.len() as u16;
                ad.push((frag_len >> 8) as u8);
                ad.push(frag_len as u8);

                let encrypted_fragment = encryptor.encrypt(&seq_num,
                                                           &record.fragment,
                                                           &ad);
                encrypted_fragment
            }
        };

        let fragment_len = encrypted_fragment.len();
        if fragment_len > ENC_RECORD_MAX_LEN {
            panic!("record too long: {} > 2^14 + 2048", fragment_len);
        }

        try!(self.writer.write_u8(record.content_type as u8));
        try!(self.writer.write_u8(record.ver_major));
        try!(self.writer.write_u8(record.ver_minor));
        try!(self.writer.write_be_u16(fragment_len as u16));
        try!(self.writer.write_all(&encrypted_fragment));

        self.write_count += 1;

        Ok(())
    }

    pub fn write_data(&mut self, ty: ContentType, data: &[u8]) -> TlsResult<()> {
        let (major, minor) = TLS_VERSION;
        // TODO: configurable maxlen
        for fragment in data.chunks(RECORD_MAX_LEN) {
            let fragment = fragment.to_vec();
            let record = Record::new(ty, major, minor, fragment);
            try!(self.write_record(record));
        }

        Ok(())
    }

    pub fn write_handshake(&mut self, handshake: &Handshake) -> TlsResult<()> {
        let mut data = Vec::new();
        try!(handshake.tls_write(&mut data));
        self.write_data(HandshakeTy, &data)
    }

    pub fn write_alert(&mut self, alert: &Alert) -> TlsResult<()> {
        let mut data = Vec::new();
        try!(alert.tls_write(&mut data));
        self.write_data(AlertTy, &data)
    }

    pub fn write_change_cipher_spec(&mut self) -> TlsResult<()> {
        self.write_data(ChangeCipherSpecTy, &[1u8])
    }

    pub fn write_application_data(&mut self, data: &[u8]) -> TlsResult<()> {
        if self.encryptor.is_none() {
            panic!("attempted to write ApplicationData before handshake");
        }
        self.write_data(ApplicationDataTy, data)
    }
}

/// Return type of `RecordReader.read_record()`.
pub enum Message {
    HandshakeMessage(Handshake),
    ChangeCipherSpecMessage,
    AlertMessage(Alert),
    ApplicationDataMessage(Vec<u8>),
}

pub struct RecordReader<R: ReadExt> {
    reader: R,
    // if decryptor is none, handshake is not done yet.
    decryptor: Option<Box<Decryptor + Send + 'static>>,
    read_count: u64,
    handshake_buffer: HandshakeBuffer,
}

/// Reads `Record` or `Message` from a readable object.
/// Record is internally decrypted after read.
impl<R: ReadExt> RecordReader<R> {
    pub fn new(reader: R) -> RecordReader<R> {
        RecordReader {
            reader: reader,
            decryptor: None,
            read_count: 0,
            handshake_buffer: HandshakeBuffer::new(),
        }
    }

    #[inline]
    pub fn get_mut(&mut self) -> &mut R {
        &mut self.reader
    }

    /// Set decryptor and reset count.
    /// This must be called only once.
    pub fn set_decryptor(&mut self, decryptor: Box<Decryptor + Send + 'static>) {
        assert!(self.decryptor.is_none());
        self.decryptor = Some(decryptor);
        self.read_count = 0;
    }

    /// Read a record from readable stream.
    ///
    /// Any record with unknown content type is treated as an error.
    fn read_record(&mut self) -> TlsResult<Record> {
        let content_type = {
            let ty = try!(self.reader.read_u8());
            let ct: Option<ContentType> = FromPrimitive::from_u8(ty);
            match ct {
                Some(ty) => ty,
                None => return tls_err!(UnexpectedMessage, "unexpected ContentType: {}", ty),
            }
        };

        let major = try!(self.reader.read_u8());
        let minor = try!(self.reader.read_u8());

        let len = {
            let len = try!(self.reader.read_be_u16()) as usize;
            if len > ENC_RECORD_MAX_LEN {
                return tls_err!(RecordOverflow, "TLSEncryptedText too long: {}", len);
            }
            len
        };

        let fragment = try!(self.reader.read_exact(len as usize));

        let record = match self.decryptor {
            None => {
                if fragment.len() > RECORD_MAX_LEN {
                    return tls_err!(RecordOverflow,
                                    "decrypted record too long: {}",
                                    fragment.len());
                }
                Record::new(content_type, major, minor, fragment)
            }
            Some(ref mut decryptor) => {
                let seq_num = u64_be_array(self.read_count);

                let mut ad = Vec::new();
                ad.extend(&seq_num);
                ad.push(content_type as u8); // TLSCompressed.type
                ad.push(major);
                ad.push(minor);

                let mac_len = decryptor.mac_len();
                let total_len = fragment.len();
                if total_len < mac_len {
                    return tls_err!(BadRecordMac, "encrypted message too short: {}", total_len);
                }
                let frag_len = (total_len - mac_len) as u16;
                ad.push((frag_len >> 8) as u8);
                ad.push(frag_len as u8);

                // TODO: "seq_num as nonce" is chacha20poly1305-specific
                let data = try!(decryptor.decrypt(&seq_num, &fragment, &ad));
                if data.len() > RECORD_MAX_LEN {
                    // decryption routine went wrong.
                    return panic!("decrypted record too long: {}", data.len());
                }

                Record::new(content_type, major, minor, data)
            }
        };

        self.read_count += 1;

        Ok(record)
    }

    /// Read records until a "complete" message is found, then return the message.
    ///
    /// if invalid ChangeCipherSpec/Alert/Handshake message is found, return Err.
    /// (application record is always considered "complete" and "valid"
    /// since it is opaque to TLS layer.)
    ///
    /// Note: In theory, `Alert` message can be broken into several records.
    /// It is not useful in practice and requires more complex routines.
    /// (Incorrect handling leads to [Alert attack](http://www.mitls.org/wsgi/alert-attack).)
    ///
    /// We treat partial alert message as an error and returns `UnexpectedMessage`.
    pub fn read_message(&mut self) -> TlsResult<Message> {
        match try!(self.handshake_buffer.get_message()) {
            Some(handshake_msg) => return Ok(HandshakeMessage(handshake_msg)),
            None => {}
        }

        // ok, no message found. read it from network!
        loop {
            // TODO: what if handshake record is present in buffer then
            // other record comes? is it legal?

            let record = try!(self.read_record());
            match record.content_type {
                ChangeCipherSpecTy => {
                    if record.fragment.len() != 1 || record.fragment[0] != 1 {
                        return tls_err!(UnexpectedMessage, "invalid ChangeCipherSpec arrived");
                    }
                    return Ok(ChangeCipherSpecMessage);
                }
                AlertTy => {
                    let len = record.fragment.len();
                    if len == 0 {
                        return tls_err!(UnexpectedMessage, "zero-length Alert record arrived");
                    } else if len < 2 {
                        // alert attack
                        return tls_err!(UnexpectedMessage, "awkward Alert record arrived");
                    }
                    let level = FromPrimitive::from_u8(record.fragment[0]);
                    let desc = FromPrimitive::from_u8(record.fragment[1]);
                    match (level, desc) {
                        (Some(level), Some(desc)) => {
                            return Ok(AlertMessage(try!(Alert::new(level, desc))));
                        }
                        _ => return tls_err!(UnexpectedMessage,
                                             "unknown alert: {:?}",
                                             record.fragment),
                    }
                }
                HandshakeTy => {
                    if record.fragment.len() == 0 {
                        return tls_err!(UnexpectedMessage, "zero-length Handshake arrived");
                    }
                    self.handshake_buffer.add_record(&record.fragment);

                    match try!(self.handshake_buffer.get_message()) {
                        Some(handshake_msg) => return Ok(HandshakeMessage(handshake_msg)),
                        _ => {}
                    }
                }
                ApplicationDataTy => {
                    return Ok(ApplicationDataMessage(record.fragment));
                }
            }
        }
    }

    pub fn read_application_data(&mut self) -> TlsResult<Vec<u8>> {
        if self.decryptor.is_none() {
            panic!("ApplicationData called before handshake");
        }
        loop {
            let msg = try!(self.read_message());
            match msg {
                ApplicationDataMessage(msg) => return Ok(msg),
                // TODO: handle other cases
                AlertMessage(..) => unimplemented!(),
                ChangeCipherSpecMessage(..) => unimplemented!(), // this should not come here
                HandshakeMessage(..) => unimplemented!(), // TODO: re-handshake
            }
        }
    }

    pub fn read_handshake(&mut self) -> TlsResult<Handshake> {
        match try!(self.read_message()) {
            HandshakeMessage(handshake) => Ok(handshake),
            AlertMessage(alert) => tls_err!(AlertReceived, "alert: {:?}", alert.description),
            _ => tls_err!(UnexpectedMessage, "expected Handshake"),
        }
    }

    pub fn read_change_cipher_spec(&mut self) -> TlsResult<()> {
        match try!(self.read_message()) {
            ChangeCipherSpecMessage => Ok(()),
            _ => tls_err!(UnexpectedMessage, "expected ChangeCipherSpec"),
        }
    }
}

pub struct Tls<R: Read, W: Write> {
    pub writer: RecordWriter<W>,
    pub reader: RecordReader<R>,
    pub rng: OsRng,
}

impl<R: Read, W: Write> Tls<R, W> {
    pub fn new(reader: R, writer: W, rng: OsRng) -> Tls<R, W> {
        let writer = RecordWriter::new(writer);
        let reader = RecordReader::new(reader);
        Tls {
            writer: writer,
            reader: reader,
            rng: rng,
        }
    }

    pub fn close(&mut self) -> TlsResult<()> {
        let alert_data = alert::Alert {
            level: alert::AlertLevel::fatal,
            description: alert::AlertDescription::close_notify,
        };
        try!(self.writer.write_alert(&alert_data));
        Ok(())
    }

    // send fatal alert and return error
    // (it may be different to `err`, because writing alert can fail)
    pub fn send_tls_alert(&mut self, err: TlsError) -> TlsError {
        match err.kind {
            TlsErrorKind::IoFailure => return err,
            _ => {
                let alert = alert::Alert::from_tls_err(&err);
                let result = self.writer.write_alert(&alert);
                match result {
                    Ok(()) => return err,
                    Err(err) => return err,
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>//
// tui.rs
// Copyright (C) 2017 Szymon Urbaś <szymon.urbas@aol.com>
// Distributed under terms of the BSD (2-clause) license.
//
// Created on: 16 Mar 2017 22:14:17 +0100 (CET)
//

use cursive::Cursive;
use cursive::view::*;
use cursive::views::*;

use process::*;

pub struct Tui {
  renderer: Cursive,
}

pub fn new() -> Tui {
  let mut renderer = Cursive::new();

  let process_list = ListView::new().with_id("process_list").min_width(32);
  let results = TextView::new("").scroll_strategy(ScrollStrategy::StickToBottom)
    .with_id("results").full_height().min_width(48);
  let header = TextView::new("").with_id("header");
  let debug = TextView::new("").scroll_strategy(ScrollStrategy::StickToBottom).with_id("debug").full_height();

  let mut layout = LinearLayout::horizontal();
  let mut right_pane =
    LinearLayout::vertical()
      .child(Dialog::around(header).title("Header"))
      .child(Dialog::around(results).title("Results"))
  ;

  if cfg!(feature = "debug-info"){
    right_pane.add_child(Dialog::around(debug).title("Debug info"));
  }

  layout.add_child(Dialog::around(process_list).title("Process list"));
  layout.add_child(right_pane);

  renderer.set_fps(60);
  renderer.add_layer(layout);
  renderer.add_global_callback('q', |tui| tui.quit());

  Tui {
    renderer: renderer,
  }
}

impl Tui {
  pub fn update(&mut self) {
    self.renderer.step();
  }

  pub fn draw_process_list(&mut self, process_list: &[Process]) {
    let mut process_list_view = self.renderer.find_id::<ListView>("process_list").unwrap();

    process_list_view.clear();

    for p in process_list.iter() {
      let progress_bar_value = ((p.execution_time as f64 / p.burst_time as f64) as f64 * 100f64) as usize;

      process_list_view.add_child(p.name.clone().as_str(),
        ProgressBar::new().with_value(Counter::new(progress_bar_value)));
    }
  }

  pub fn set_header(&mut self, text: String) {
    let mut header = self.renderer.find_id::<TextView>("header").unwrap();

    header.set_content(text.clone());
  }

  pub fn debug(&mut self, text: String) {
    if cfg!(feature = "debug-info") {
      let mut debug = self.renderer.find_id::<TextView>("debug").unwrap();

      debug.append_content(&text);
      debug.append_content("\n");
    }
  }

  pub fn add_result(&mut self, text: String) {
    let mut debug = self.renderer.find_id::<TextView>("results").unwrap();

    debug.append_content(&text);
    debug.append_content("\n");
  }
}

/*
 * vi: ts=2 sw=2 expandtab
 */

<commit_msg>Slight rework of setting up the design<commit_after>//
// tui.rs
// Copyright (C) 2017 Szymon Urbaś <szymon.urbas@aol.com>
// Distributed under terms of the BSD (2-clause) license.
//
// Created on: 16 Mar 2017 22:14:17 +0100 (CET)
//

use cursive::Cursive;
use cursive::view::*;
use cursive::views::*;

use process::*;

pub struct Tui {
  renderer: Cursive,
}

pub fn new() -> Tui {
  let mut renderer = Cursive::new();

  let process_list = ListView::new().with_id("process_list").min_width(32);
  let results = TextView::new("").scroll_strategy(ScrollStrategy::StickToBottom)
    .with_id("results").full_height().min_width(48);
  let header = TextView::new("").with_id("header");
  let debug = TextView::new("").scroll_strategy(ScrollStrategy::StickToBottom).with_id("debug").full_height();
  let scenarios = ListView::new().with_id("scenarios").max_height(4);

  renderer.set_fps(60);

  if cfg!(feature = "debug-info") {
    renderer.add_layer(debug);
  }

  renderer.add_layer(
    LinearLayout::vertical().child(
        LinearLayout::horizontal()
          .child(Dialog::around(process_list).title("Process list"))
          .child(
            LinearLayout::vertical()
              .child(Dialog::around(header).title("Header"))
              .child(Dialog::around(results).title("Results"))
          )
      ).child(Dialog::around(scenarios).title("Scenarios")));

  renderer.add_global_callback('q', |tui| tui.quit());

  Tui {
    renderer: renderer,
  }
}

impl Tui {
  pub fn update(&mut self) {
    self.renderer.step();
  }

  pub fn draw_process_list(&mut self, process_list: &[Process]) {
    let mut process_list_view = self.renderer.find_id::<ListView>("process_list").unwrap();

    process_list_view.clear();

    for p in process_list.iter() {
      let progress_bar_value = ((p.execution_time as f64 / p.burst_time as f64) as f64 * 100f64) as usize;

      process_list_view.add_child(p.name.clone().as_str(),
        ProgressBar::new().with_value(Counter::new(progress_bar_value)));
    }
  }

  pub fn set_header(&mut self, text: String) {
    let mut header = self.renderer.find_id::<TextView>("header").unwrap();

    header.set_content(text.clone());
  }

  pub fn debug(&mut self, text: String) {
    if cfg!(feature = "debug-info") {
      let mut debug = self.renderer.find_id::<TextView>("debug").unwrap();

      debug.append_content(&text);
      debug.append_content("\n");
    }
  }

  pub fn add_result(&mut self, text: String) {
    let mut debug = self.renderer.find_id::<TextView>("results").unwrap();

    debug.append_content(&text);
    debug.append_content("\n");
  }
}

/*
 * vi: ts=2 sw=2 expandtab
 */

<|endoftext|>


<commit_before>
use iron::{Iron, IronResult, Listening, status};
use iron::error::HttpResult;
use iron::response::Response;
use iron::request::Request;
use iron::middleware::Handler;
use iron::mime::Mime;

use sysinfo::{System, SystemExt};
use std::sync::{Arc, Mutex};
use SysinfoExt;
use serde_json;

const INDEX_HTML: &'static str = include_str!("index.html");

struct SysinfoIronHandler(Arc<Mutex<System>>);

impl Handler for SysinfoIronHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        match req.url.path().last() {
            Some(path) => {
                if *path == "" {
                    Ok(Response::with((status::Ok,
                                       "text/html".parse::<Mime>().unwrap(),
                                       INDEX_HTML)))
                } else {
                    let mut system = self.0.lock().unwrap();
                    system.refresh_all();
                    let sysinfo = SysinfoExt::new(&system);
                    Ok(Response::with((status::Ok,
                                       "application/json".parse::<Mime>().unwrap(),
                                       serde_json::to_string(&sysinfo).unwrap_or(String::new()))))
                }
            },
            None => Ok(Response::with((status::NotFound, "Not found")))
        }
    }
}

pub fn start_web_server(sock_addr: Option<String>) -> HttpResult<Listening> {
    let system = Arc::new(Mutex::new(System::new()));
    let mut iron = Iron::new(SysinfoIronHandler(system));
    iron.threads = 4;
    iron.http(sock_addr.unwrap_or("localhost:3000".to_owned()))
}

#[test]
fn test_sysinfo_json_str() {
    assert!(!sysinfo_json_str().is_empty());
}
<commit_msg>Remove obsolete test<commit_after>
use iron::{Iron, IronResult, Listening, status};
use iron::error::HttpResult;
use iron::response::Response;
use iron::request::Request;
use iron::middleware::Handler;
use iron::mime::Mime;

use sysinfo::{System, SystemExt};
use std::sync::{Arc, Mutex};
use SysinfoExt;
use serde_json;

const INDEX_HTML: &'static str = include_str!("index.html");

struct SysinfoIronHandler(Arc<Mutex<System>>);

impl Handler for SysinfoIronHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        match req.url.path().last() {
            Some(path) => {
                if *path == "" {
                    Ok(Response::with((status::Ok,
                                       "text/html".parse::<Mime>().unwrap(),
                                       INDEX_HTML)))
                } else {
                    let mut system = self.0.lock().unwrap();
                    system.refresh_all();
                    let sysinfo = SysinfoExt::new(&system);
                    Ok(Response::with((status::Ok,
                                       "application/json".parse::<Mime>().unwrap(),
                                       serde_json::to_string(&sysinfo).unwrap_or(String::new()))))
                }
            },
            None => Ok(Response::with((status::NotFound, "Not found")))
        }
    }
}

pub fn start_web_server(sock_addr: Option<String>) -> HttpResult<Listening> {
    let system = Arc::new(Mutex::new(System::new()));
    let mut iron = Iron::new(SysinfoIronHandler(system));
    iron.threads = 4;
    iron.http(sock_addr.unwrap_or("localhost:3000".to_owned()))
}
<|endoftext|>


<commit_before>// This file is part of the uutils coreutils package.
//
// (c) Derek Chiang <derekchiang93@gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

use clap::{crate_version, Arg, Command};
use std::path::Path;
use uucore::display::print_verbatim;
use uucore::error::{UResult, UUsageError};
use uucore::{format_usage, InvalidEncodingHandling};

static ABOUT: &str = "strip last component from file name";
const USAGE: &str = "{} [OPTION] NAME...";

mod options {
    pub const ZERO: &str = "zero";
    pub const DIR: &str = "dir";
}

fn get_long_usage() -> String {
    String::from(
        "Output each NAME with its last non-slash component and trailing slashes
        removed; if NAME contains no /'s, output '.' (meaning the current directory).",
    )
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let after_help = get_long_usage();

    let matches = uu_app().after_help(&after_help[..]).get_matches_from(args);

    let separator = if matches.is_present(options::ZERO) {
        "\0"
    } else {
        "\n"
    };

    let dirnames: Vec<String> = matches
        .values_of(options::DIR)
        .unwrap_or_default()
        .map(str::to_owned)
        .collect();

    if !dirnames.is_empty() {
        for path in &dirnames {
            let p = Path::new(path);
            match p.parent() {
                Some(d) => {
                    if d.components().next() == None {
                        print!(".");
                    } else {
                        print_verbatim(d).unwrap();
                    }
                }
                None => {
                    if p.is_absolute() || path == "/" {
                        print!("/");
                    } else {
                        print!(".");
                    }
                }
            }
            print!("{}", separator);
        }
    } else {
        return Err(UUsageError::new(1, "missing operand"));
    }

    Ok(())
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .about(ABOUT)
        .version(crate_version!())
        .override_usage(format_usage(USAGE))
        .infer_long_args(true)
        .arg(
            Arg::new(options::ZERO)
                .long(options::ZERO)
                .short('z')
                .help("separate output with NUL rather than newline"),
        )
        .arg(
            Arg::new(options::DIR)
                .hide(true)
                .multiple_occurrences(true)
                .value_hint(clap::ValueHint::AnyPath),
        )
}
<commit_msg>dirname: add missing "\n\" to usage message<commit_after>// This file is part of the uutils coreutils package.
//
// (c) Derek Chiang <derekchiang93@gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

use clap::{crate_version, Arg, Command};
use std::path::Path;
use uucore::display::print_verbatim;
use uucore::error::{UResult, UUsageError};
use uucore::{format_usage, InvalidEncodingHandling};

static ABOUT: &str = "strip last component from file name";
const USAGE: &str = "{} [OPTION] NAME...";

mod options {
    pub const ZERO: &str = "zero";
    pub const DIR: &str = "dir";
}

fn get_long_usage() -> String {
    String::from(
        "Output each NAME with its last non-slash component and trailing slashes \n\
        removed; if NAME contains no /'s, output '.' (meaning the current directory).",
    )
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let after_help = get_long_usage();

    let matches = uu_app().after_help(&after_help[..]).get_matches_from(args);

    let separator = if matches.is_present(options::ZERO) {
        "\0"
    } else {
        "\n"
    };

    let dirnames: Vec<String> = matches
        .values_of(options::DIR)
        .unwrap_or_default()
        .map(str::to_owned)
        .collect();

    if !dirnames.is_empty() {
        for path in &dirnames {
            let p = Path::new(path);
            match p.parent() {
                Some(d) => {
                    if d.components().next() == None {
                        print!(".");
                    } else {
                        print_verbatim(d).unwrap();
                    }
                }
                None => {
                    if p.is_absolute() || path == "/" {
                        print!("/");
                    } else {
                        print!(".");
                    }
                }
            }
            print!("{}", separator);
        }
    } else {
        return Err(UUsageError::new(1, "missing operand"));
    }

    Ok(())
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .about(ABOUT)
        .version(crate_version!())
        .override_usage(format_usage(USAGE))
        .infer_long_args(true)
        .arg(
            Arg::new(options::ZERO)
                .long(options::ZERO)
                .short('z')
                .help("separate output with NUL rather than newline"),
        )
        .arg(
            Arg::new(options::DIR)
                .hide(true)
                .multiple_occurrences(true)
                .value_hint(clap::ValueHint::AnyPath),
        )
}
<|endoftext|>


<commit_before>/*!
This module defines the various error types that can be produced by a failed conversion.
*/

use std::any::Any;
use std::error::Error;
use std::fmt::{self, Debug, Display};
use misc::{Saturated, InvalidSentinel, SignedInfinity};

/// A general error enumeration that subsumes all of the below.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum GeneralError {
    /// Input underflowed the target type.
    Underflow,

    /// Input overflowed the target type.
    Overflow,

    /// Input was not representable in the target type.
    Unrepresentable,
}

impl Display for GeneralError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.description())
    }
}

impl Error for GeneralError {
    fn description(&self) -> &str {
        use self::GeneralError::*;
        match *self {
            Underflow => "conversion resulted in underflow",
            Overflow => "conversion resulted in overflow",
            Unrepresentable => "could not convert unrepresentable value",
        }
    }
}

impl From<NoError> for GeneralError {
    fn from(_: NoError) -> GeneralError {
        panic!("cannot convert NoError into GeneralError");
    }
}

impl<T> From<Unrepresentable<T>> for GeneralError {
    fn from(_: Unrepresentable<T>) -> GeneralError {
        GeneralError::Unrepresentable
    }
}

impl From<Underflow> for GeneralError {
    fn from(_: Underflow) -> GeneralError {
        GeneralError::Underflow
    }
}

impl From<Overflow> for GeneralError {
    fn from(_: Overflow) -> GeneralError {
        GeneralError::Overflow
    }
}

impl From<RangeError> for GeneralError {
    fn from(e: RangeError) -> GeneralError {
        use self::RangeError as R;
        use self::GeneralError as G;
        match e {
            R::Underflow => G::Underflow,
            R::Overflow => G::Overflow,
        }
    }
}

impl From<FloatError> for GeneralError {
    fn from(e: FloatError) -> GeneralError {
        use self::FloatError as F;
        use self::GeneralError as G;
        match e {
            F::Underflow => G::Underflow,
            F::Overflow => G::Overflow,
            F::NotANumber => G::Unrepresentable,
        }
    }
}

/// Indicates that it is not possible for the conversion to fail.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum NoError {}

impl Display for NoError {
    fn fmt(&self, _: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        unreachable!()
    }
}

impl Error for NoError {
    fn description(&self) -> &str {
        unreachable!()
    }
}

/// Indicates that the conversion failed because the value was not representable.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Unrepresentable<T>(pub T);

impl<T> Display for Unrepresentable<T>
where T: Display {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "could not convert unrepresentable value: {}", self.0)
    }
}

impl<T> Error for Unrepresentable<T>
where T: Debug + Display + Any {
    fn description(&self) -> &str {
        "could not convert unrepresentable value"
    }
}

/// Indicates that the conversion failed due to an underflow.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Underflow;

impl Display for Underflow {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.description())
    }
}

impl Error for Underflow {
    fn description(&self) -> &str {
        "conversion resulted in underflow"
    }
}

impl From<NoError> for Underflow {
    fn from(_: NoError) -> Underflow {
        panic!("cannot convert NoError into Underflow");
    }
}

/// Indicates that the conversion failed due to an overflow.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Overflow;

impl Display for Overflow {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.description())
    }
}

impl Error for Overflow {
    fn description(&self) -> &str {
        "conversion resulted in overflow"
    }
}

impl From<NoError> for Overflow {
    fn from(_: NoError) -> Overflow {
        panic!("cannot convert NoError into Overflow");
    }
}

/**
Indicates that a conversion from a floating point type failed.
*/
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum FloatError {
    /// Input underflowed the target type.
    Underflow,

    /// Input overflowed the target type.
    Overflow,

    /// Input was not-a-number, which the target type could not represent.
    NotANumber,
}

impl Display for FloatError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.description())
    }
}

impl Error for FloatError {
    fn description(&self) -> &str {
        use self::FloatError::*;
        match *self {
            Underflow => "conversion resulted in underflow",
            Overflow => "conversion resulted in overflow",
            NotANumber => "conversion target does not support not-a-number",
        }
    }
}

impl From<NoError> for FloatError {
    fn from(_: NoError) -> FloatError {
        panic!("cannot convert NoError into FloatError");
    }
}

impl From<Underflow> for FloatError {
    fn from(_: Underflow) -> FloatError {
        FloatError::Underflow
    }
}

impl From<Overflow> for FloatError {
    fn from(_: Overflow) -> FloatError {
        FloatError::Overflow
    }
}

impl From<RangeError> for FloatError {
    fn from(e: RangeError) -> FloatError {
        use self::RangeError as R;
        use self::FloatError as F;
        match e {
            R::Underflow => F::Underflow,
            R::Overflow => F::Overflow,
        }
    }
}

/**
Indicates that a conversion failed due to a range error.
*/
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum RangeError {
    /// Input underflowed the target type.
    Underflow,

    /// Input overflowed the target type.
    Overflow,
}

impl Display for RangeError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.description())
    }
}

impl Error for RangeError {
    fn description(&self) -> &str {
        use self::RangeError::*;
        match *self {
            Underflow => "conversion resulted in underflow",
            Overflow => "conversion resulted in overflow",
        }
    }
}

impl From<NoError> for RangeError {
    fn from(_: NoError) -> RangeError {
        panic!("cannot convert NoError into RangeError");
    }
}

impl From<Underflow> for RangeError {
    fn from(_: Underflow) -> RangeError {
        RangeError::Underflow
    }
}

impl From<Overflow> for RangeError {
    fn from(_: Overflow) -> RangeError {
        RangeError::Overflow
    }
}

/**
Safely unwrap a `Result` that cannot contain an error.
*/
pub trait UnwrapOk<T> {
    /**
    Unwraps a `Result` without possibility of failing.

    Technically, this is not necessary; it's provided simply to make user code a little clearer.
    */
    fn unwrap_ok(self) -> T;
}

impl<T> UnwrapOk<T> for Result<T, NoError> {
    fn unwrap_ok(self) -> T {
        match self {
            Ok(v) => v,
            Err(..) => loop {},
        }
    }
}

/**
Unwrap a conversion by saturating to infinity.
*/
pub trait UnwrapOrInf {
    /// The result of unwrapping.
    type Output;

    /**
    Either unwraps the successfully converted value, or saturates to infinity in the "direction" of overflow/underflow.
    */
    fn unwrap_or_inf(self) -> Self::Output;
}

/**
Unwrap a conversion by replacing a failure with an invalid sentinel value.
*/
pub trait UnwrapOrInvalid {
    /// The result of unwrapping.
    type Output;

    /**
    Either unwraps the successfully converted value, or returns the output type's invalid sentinel value.
    */
    fn unwrap_or_invalid(self) -> Self::Output;
}

/**
Unwrap a conversion by saturating.
*/
pub trait UnwrapOrSaturate {
    /// The result of unwrapping.
    type Output;

    /**
    Either unwraps the successfully converted value, or saturates in the "direction" of overflow/underflow.
    */
    fn unwrap_or_saturate(self) -> Self::Output;
}

impl<T, E> UnwrapOrInf for Result<T, E>
where T: SignedInfinity, E: Into<RangeError> {
    type Output = T;
    fn unwrap_or_inf(self) -> T {
        use self::RangeError::*;
        match self.map_err(Into::into) {
            Ok(v) => v,
            Err(Underflow) => T::neg_infinity(),
            Err(Overflow) => T::pos_infinity(),
        }
    }
}

impl<T, E> UnwrapOrInvalid for Result<T, E>
where T: InvalidSentinel {
    type Output = T;
    fn unwrap_or_invalid(self) -> T {
        match self {
            Ok(v) => v,
            Err(..) => T::invalid_sentinel(),
        }
    }
}

impl<T, E> UnwrapOrSaturate for Result<T, E>
where T: Saturated, E: Into<RangeError> {
    type Output = T;
    fn unwrap_or_saturate(self) -> T {
        use self::RangeError::*;
        match self.map_err(Into::into) {
            Ok(v) => v,
            Err(Underflow) => T::saturated_min(),
            Err(Overflow) => T::saturated_max(),
        }
    }
}
<commit_msg>Fixed slightly dodgy GeneralError description.<commit_after>/*!
This module defines the various error types that can be produced by a failed conversion.
*/

use std::any::Any;
use std::error::Error;
use std::fmt::{self, Debug, Display};
use misc::{Saturated, InvalidSentinel, SignedInfinity};

/**
A general error enumeration that subsumes all other conversion errors.

This exists primarily as a "catch-all" for reliably unifying various different kinds of conversion errors.
*/
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum GeneralError {
    /// Input underflowed the target type.
    Underflow,

    /// Input overflowed the target type.
    Overflow,

    /// Input was not representable in the target type.
    Unrepresentable,
}

impl Display for GeneralError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.description())
    }
}

impl Error for GeneralError {
    fn description(&self) -> &str {
        use self::GeneralError::*;
        match *self {
            Underflow => "conversion resulted in underflow",
            Overflow => "conversion resulted in overflow",
            Unrepresentable => "could not convert unrepresentable value",
        }
    }
}

impl From<NoError> for GeneralError {
    fn from(_: NoError) -> GeneralError {
        panic!("cannot convert NoError into GeneralError");
    }
}

impl<T> From<Unrepresentable<T>> for GeneralError {
    fn from(_: Unrepresentable<T>) -> GeneralError {
        GeneralError::Unrepresentable
    }
}

impl From<Underflow> for GeneralError {
    fn from(_: Underflow) -> GeneralError {
        GeneralError::Underflow
    }
}

impl From<Overflow> for GeneralError {
    fn from(_: Overflow) -> GeneralError {
        GeneralError::Overflow
    }
}

impl From<RangeError> for GeneralError {
    fn from(e: RangeError) -> GeneralError {
        use self::RangeError as R;
        use self::GeneralError as G;
        match e {
            R::Underflow => G::Underflow,
            R::Overflow => G::Overflow,
        }
    }
}

impl From<FloatError> for GeneralError {
    fn from(e: FloatError) -> GeneralError {
        use self::FloatError as F;
        use self::GeneralError as G;
        match e {
            F::Underflow => G::Underflow,
            F::Overflow => G::Overflow,
            F::NotANumber => G::Unrepresentable,
        }
    }
}

/// Indicates that it is not possible for the conversion to fail.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum NoError {}

impl Display for NoError {
    fn fmt(&self, _: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        unreachable!()
    }
}

impl Error for NoError {
    fn description(&self) -> &str {
        unreachable!()
    }
}

/// Indicates that the conversion failed because the value was not representable.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Unrepresentable<T>(pub T);

impl<T> Display for Unrepresentable<T>
where T: Display {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "could not convert unrepresentable value: {}", self.0)
    }
}

impl<T> Error for Unrepresentable<T>
where T: Debug + Display + Any {
    fn description(&self) -> &str {
        "could not convert unrepresentable value"
    }
}

/// Indicates that the conversion failed due to an underflow.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Underflow;

impl Display for Underflow {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.description())
    }
}

impl Error for Underflow {
    fn description(&self) -> &str {
        "conversion resulted in underflow"
    }
}

impl From<NoError> for Underflow {
    fn from(_: NoError) -> Underflow {
        panic!("cannot convert NoError into Underflow");
    }
}

/// Indicates that the conversion failed due to an overflow.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Overflow;

impl Display for Overflow {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.description())
    }
}

impl Error for Overflow {
    fn description(&self) -> &str {
        "conversion resulted in overflow"
    }
}

impl From<NoError> for Overflow {
    fn from(_: NoError) -> Overflow {
        panic!("cannot convert NoError into Overflow");
    }
}

/**
Indicates that a conversion from a floating point type failed.
*/
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum FloatError {
    /// Input underflowed the target type.
    Underflow,

    /// Input overflowed the target type.
    Overflow,

    /// Input was not-a-number, which the target type could not represent.
    NotANumber,
}

impl Display for FloatError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.description())
    }
}

impl Error for FloatError {
    fn description(&self) -> &str {
        use self::FloatError::*;
        match *self {
            Underflow => "conversion resulted in underflow",
            Overflow => "conversion resulted in overflow",
            NotANumber => "conversion target does not support not-a-number",
        }
    }
}

impl From<NoError> for FloatError {
    fn from(_: NoError) -> FloatError {
        panic!("cannot convert NoError into FloatError");
    }
}

impl From<Underflow> for FloatError {
    fn from(_: Underflow) -> FloatError {
        FloatError::Underflow
    }
}

impl From<Overflow> for FloatError {
    fn from(_: Overflow) -> FloatError {
        FloatError::Overflow
    }
}

impl From<RangeError> for FloatError {
    fn from(e: RangeError) -> FloatError {
        use self::RangeError as R;
        use self::FloatError as F;
        match e {
            R::Underflow => F::Underflow,
            R::Overflow => F::Overflow,
        }
    }
}

/**
Indicates that a conversion failed due to a range error.
*/
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum RangeError {
    /// Input underflowed the target type.
    Underflow,

    /// Input overflowed the target type.
    Overflow,
}

impl Display for RangeError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.description())
    }
}

impl Error for RangeError {
    fn description(&self) -> &str {
        use self::RangeError::*;
        match *self {
            Underflow => "conversion resulted in underflow",
            Overflow => "conversion resulted in overflow",
        }
    }
}

impl From<NoError> for RangeError {
    fn from(_: NoError) -> RangeError {
        panic!("cannot convert NoError into RangeError");
    }
}

impl From<Underflow> for RangeError {
    fn from(_: Underflow) -> RangeError {
        RangeError::Underflow
    }
}

impl From<Overflow> for RangeError {
    fn from(_: Overflow) -> RangeError {
        RangeError::Overflow
    }
}

/**
Safely unwrap a `Result` that cannot contain an error.
*/
pub trait UnwrapOk<T> {
    /**
    Unwraps a `Result` without possibility of failing.

    Technically, this is not necessary; it's provided simply to make user code a little clearer.
    */
    fn unwrap_ok(self) -> T;
}

impl<T> UnwrapOk<T> for Result<T, NoError> {
    fn unwrap_ok(self) -> T {
        match self {
            Ok(v) => v,
            Err(..) => loop {},
        }
    }
}

/**
Unwrap a conversion by saturating to infinity.
*/
pub trait UnwrapOrInf {
    /// The result of unwrapping.
    type Output;

    /**
    Either unwraps the successfully converted value, or saturates to infinity in the "direction" of overflow/underflow.
    */
    fn unwrap_or_inf(self) -> Self::Output;
}

/**
Unwrap a conversion by replacing a failure with an invalid sentinel value.
*/
pub trait UnwrapOrInvalid {
    /// The result of unwrapping.
    type Output;

    /**
    Either unwraps the successfully converted value, or returns the output type's invalid sentinel value.
    */
    fn unwrap_or_invalid(self) -> Self::Output;
}

/**
Unwrap a conversion by saturating.
*/
pub trait UnwrapOrSaturate {
    /// The result of unwrapping.
    type Output;

    /**
    Either unwraps the successfully converted value, or saturates in the "direction" of overflow/underflow.
    */
    fn unwrap_or_saturate(self) -> Self::Output;
}

impl<T, E> UnwrapOrInf for Result<T, E>
where T: SignedInfinity, E: Into<RangeError> {
    type Output = T;
    fn unwrap_or_inf(self) -> T {
        use self::RangeError::*;
        match self.map_err(Into::into) {
            Ok(v) => v,
            Err(Underflow) => T::neg_infinity(),
            Err(Overflow) => T::pos_infinity(),
        }
    }
}

impl<T, E> UnwrapOrInvalid for Result<T, E>
where T: InvalidSentinel {
    type Output = T;
    fn unwrap_or_invalid(self) -> T {
        match self {
            Ok(v) => v,
            Err(..) => T::invalid_sentinel(),
        }
    }
}

impl<T, E> UnwrapOrSaturate for Result<T, E>
where T: Saturated, E: Into<RangeError> {
    type Output = T;
    fn unwrap_or_saturate(self) -> T {
        use self::RangeError::*;
        match self.map_err(Into::into) {
            Ok(v) => v,
            Err(Underflow) => T::saturated_min(),
            Err(Overflow) => T::saturated_max(),
        }
    }
}
<|endoftext|>


<commit_before>/// Exports a given list of functions to a Erlang module.
///
/// This should be called exactly once in every NIF library. It will wrap and export the given rust
/// functions into the Erlang module.
///
/// The first argument is a string specifying what Erlang/Elixir module you want the function
/// exported into. In Erlang this will simply be the atom you named your module. In Elixir, all
/// modules are prefixed with `Elixir.<module path>`
///
/// The second argument is a list of 3-tuples. Each tuple contains information on a single exported
/// NIF function. The first tuple item is the name you want to export the function into, the second
/// is the arity (number of arguments) of the exported function. The third argument is a
/// indentifier of a rust function. This is where your actual NIF will be implemented.
///
/// The third argument is an `Option<fn(env: &NifEnv, load_info: NifTerm) -> bool>`. If this is
/// `Some`, the function will execute when the NIF is first loaded by the BEAM.
#[macro_export]
macro_rules! rustler_export_nifs {
    ($name:expr, [$( ($nif_name:expr, $nif_arity:expr, $nif_fun:path) ),*], $on_load:expr) => (
        rustler_export_nifs!($name, [$( ($nif_name, $nif_arity, $nif_fun, rustler::wrapper::nif_interface::ErlNifTaskFlags::ERL_NIF_NORMAL_JOB) ),*], $on_load);
    );
    ($name:expr, [$( ($nif_name:expr, $nif_arity:expr, $nif_fun:path, $nif_flag:expr) ),*], $on_load:expr) => (
        static mut NIF_ENTRY: Option<rustler::wrapper::nif_interface::DEF_NIF_ENTRY> = None;

        #[no_mangle]
        pub extern "C" fn nif_init() -> *const rustler::wrapper::nif_interface::DEF_NIF_ENTRY {
            // TODO: If a NIF ever gets unloaded, there will be a memory leak! Fix this!
            // TODO: If an unwrap ever happens, we will unwind right into C! Fix this!

            extern "C" fn nif_load(
                env: rustler::wrapper::nif_interface::NIF_ENV,
                _priv_data: *mut *mut rustler::codegen_runtime::c_void,
                load_info: rustler::wrapper::nif_interface::NIF_TERM)
                -> rustler::codegen_runtime::c_int {
                    rustler::codegen_runtime::handle_nif_init_call($on_load, env, load_info)
                }

            let fun_entries = [
                $(
                    rustler::wrapper::nif_interface::DEF_NIF_FUNC {
                        name: ::std::ffi::CString::new($nif_name).unwrap().into_raw() as *const u8,
                        arity: $nif_arity,
                        function: {
                            extern "C" fn nif_func(
                                env: rustler::wrapper::nif_interface::NIF_ENV,
                                argc: rustler::codegen_runtime::c_int,
                                argv: *const rustler::wrapper::nif_interface::NIF_TERM)
                                -> rustler::wrapper::nif_interface::NIF_TERM {
                                    rustler::codegen_runtime::handle_nif_call($nif_fun, $nif_arity, env, argc, argv)
                                }
                            nif_func
                        },
                        flags: $nif_flag as u32,
                    }
                 ),*
            ];
            let fun_entries_len = fun_entries.len();
            let fun_entries_ptr = Box::into_raw(Box::new(fun_entries));

            let entry = rustler::wrapper::nif_interface::DEF_NIF_ENTRY {
                major: rustler::wrapper::nif_interface::NIF_MAJOR_VERSION,
                minor: rustler::wrapper::nif_interface::NIF_MINOR_VERSION,
                name: ::std::ffi::CString::new($name).unwrap().into_raw() as *const u8,
                num_of_funcs: fun_entries_len as rustler::codegen_runtime::c_int,
                funcs: fun_entries_ptr as *const rustler::wrapper::nif_interface::DEF_NIF_FUNC,
                load: Some(nif_load),
                reload: None,
                upgrade: None,
                unload: None,
                vm_variant: b"beam.vanilla\x00" as *const u8,
                options: 0,
            };
            unsafe { NIF_ENTRY = Some(entry) };

            unsafe { NIF_ENTRY.as_ref().unwrap() }
        }
    );
}
<commit_msg>fix scheduling flag in rustler_export_nifs! macro<commit_after>/// Exports a given list of functions to a Erlang module.
///
/// This should be called exactly once in every NIF library. It will wrap and export the given rust
/// functions into the Erlang module.
///
/// The first argument is a string specifying what Erlang/Elixir module you want the function
/// exported into. In Erlang this will simply be the atom you named your module. In Elixir, all
/// modules are prefixed with `Elixir.<module path>`
///
/// The second argument is a list of 3-tuples. Each tuple contains information on a single exported
/// NIF function. The first tuple item is the name you want to export the function into, the second
/// is the arity (number of arguments) of the exported function. The third argument is a
/// indentifier of a rust function. This is where your actual NIF will be implemented.
///
/// The third argument is an `Option<fn(env: &NifEnv, load_info: NifTerm) -> bool>`. If this is
/// `Some`, the function will execute when the NIF is first loaded by the BEAM.
#[macro_export]
macro_rules! rustler_export_nifs {
    ($name:expr, [$( $exported_nif:tt ),*], $on_load:expr) => {
        static mut NIF_ENTRY: Option<rustler::wrapper::nif_interface::DEF_NIF_ENTRY> = None;

        #[no_mangle]
        pub extern "C" fn nif_init() -> *const rustler::wrapper::nif_interface::DEF_NIF_ENTRY {
            // TODO: If a NIF ever gets unloaded, there will be a memory leak! Fix this!
            // TODO: If an unwrap ever happens, we will unwind right into C! Fix this!

            extern "C" fn nif_load(
                env: rustler::wrapper::nif_interface::NIF_ENV,
                _priv_data: *mut *mut rustler::codegen_runtime::c_void,
                load_info: rustler::wrapper::nif_interface::NIF_TERM)
                -> rustler::codegen_runtime::c_int {
                    rustler::codegen_runtime::handle_nif_init_call($on_load, env, load_info)
                }

            let fun_entries = [
                $(rustler_export_nifs!(internal, $exported_nif)),*
            ];
            let fun_entries_len = fun_entries.len();
            let fun_entries_ptr = Box::into_raw(Box::new(fun_entries));

            let entry = rustler::wrapper::nif_interface::DEF_NIF_ENTRY {
                major: rustler::wrapper::nif_interface::NIF_MAJOR_VERSION,
                minor: rustler::wrapper::nif_interface::NIF_MINOR_VERSION,
                name: ::std::ffi::CString::new($name).unwrap().into_raw() as *const u8,
                num_of_funcs: fun_entries_len as rustler::codegen_runtime::c_int,
                funcs: fun_entries_ptr as *const rustler::wrapper::nif_interface::DEF_NIF_FUNC,
                load: Some(nif_load),
                reload: None,
                upgrade: None,
                unload: None,
                vm_variant: b"beam.vanilla\x00" as *const u8,
                options: 0,
            };
            unsafe { NIF_ENTRY = Some(entry) };

            unsafe { NIF_ENTRY.as_ref().unwrap() }
        }
    };

    (internal, ($nif_name:expr, $nif_arity:expr, $nif_fun:path)) => {
        rustler_export_nifs!(internal, ($nif_name, $nif_arity, $nif_fun, rustler::wrapper::nif_interface::ErlNifTaskFlags::ERL_NIF_NORMAL_JOB))
    };
    (internal, ($nif_name:expr, $nif_arity:expr, $nif_fun:path, $nif_flag:expr)) => {
        rustler::wrapper::nif_interface::DEF_NIF_FUNC {
            name: ::std::ffi::CString::new($nif_name).unwrap().into_raw() as *const u8,
            arity: $nif_arity,
            function: {
                extern "C" fn nif_func(
                    env: rustler::wrapper::nif_interface::NIF_ENV,
                    argc: rustler::codegen_runtime::c_int,
                    argv: *const rustler::wrapper::nif_interface::NIF_TERM)
                    -> rustler::wrapper::nif_interface::NIF_TERM {
                    rustler::codegen_runtime::handle_nif_call($nif_fun, $nif_arity, env, argc, argv)
                }
                nif_func
            },
            flags: $nif_flag as u32,
        }
    };
}
<|endoftext|>


<commit_before>use std;
use concurrent_help::to_concurrent_on_section;
use num_cpus;
// macro_rules! if_any_divides {
//      ( $( $div:tt ),* | $n:ident $if_b:block ) => { if $(($n % $div == 0))|* $if_b};
// }
// macro_rules! if_sm_prime_divides{
//     ( $n:ident $if_b:block ) => {if_any_divides!(
//     2, 3, 5, 7, 11, 13,
//     17, 19, 23, 29, 31,
//     37, 41, 43, 47, 53,
//     59, 61, 67, 71, 73,
//     79, 83, 89, 97, 101,
//     103, 107, 109, 113, 127,
//     131, 137, 139, 149, 151,
//     157, 163, 167, 173, 179,
//     181, 191, 193, 197, 199 | $n $if_b)};
// }
macro_rules! set_true_if_in_range{
    ( $( $i:tt ),* => $filter:ident + $offset:expr, $min:expr, $max:expr) => {
        $(if ($min <= $i) & ($max > $i) {$filter[$i-$offset] = true;})*
    };
    ( $( $i:tt ),* => $filter:ident, $min:expr, $max:expr) => {
        $(if ($min <= $i) & ($max > $i) {$filter[$i] = true;})*
    };
}
// macro_rules! set_sm_primes_true {
//     ($filter:expr, $min:expr, $max:expr) => {
//         set_true_if_in_range!(
//         2, 3, 5, 7, 11, 13,
//         17, 19, 23, 29, 31,
//         37, 41, 43, 47, 53,
//         59, 61, 67, 71, 73,
//         79, 83, 89, 97, 101,
//         103, 107, 109, 113, 127,
//         131, 137, 139, 149, 151,
//         157, 163, 167, 173, 179,
//         181, 191, 193, 197, 199 => $filter + $min, $min, $max)
//     };
// }
pub fn prime_filter_concurrently(max_num: usize, threads: usize) -> Vec<bool>{
    prime_filter_section_concurrently(0, max_num, threads)
}
pub fn prime_filter_section_concurrently(min_num:usize, max_num: usize, threads: usize) -> Vec<bool>{
    to_concurrent_on_section(prime_filter_section_sequentially, min_num, max_num, threads)
}

fn int_sqrt(n:usize) -> usize{
    if n < (1 << 53) {
        return (n as f64).sqrt() as usize;
    }
    match n {
        0 => 0,
        1 ... 3 => 1,
        4 ... 8 => 2,
        9 ... 15 => 3,
        k => {
        let mut x = k;
            loop{
                x = match (x + n/x) >> 1 {
                    new_x if new_x == x => break,
                    new_x if new_x*new_x == n + 1 => {x = new_x - 1; break},
                    new_x => new_x,
                };
            }
            x
        },
    }

}

fn ceil_sqrt(n:usize) -> usize{
    match int_sqrt(n){
        sqrt if sqrt*sqrt == n => sqrt,
        sqrt => sqrt+1,
    }
}
pub fn prime_filter(max_num: usize) -> Vec<bool>{
    prime_filter_concurrently(max_num, num_cpus::get())
}

pub fn prime_filter_section(min_num:usize, max_num: usize) -> Vec<bool>{
    prime_filter_section_concurrently(min_num, max_num, num_cpus::get())
}

pub fn prime_filter_sequentially(max_num: usize) -> Vec<bool>{
    if max_num<100{
        slow_prime_filter(max_num)
    }else{
        prime_filter_section(0, max_num)
    }
}
pub fn prime_filter_section_sequentially(min_num:usize, max_num: usize) -> Vec<bool>{
    //Sieve of Atkin
    assert!(min_num<max_num);
    let mut prime_filter = vec![false; max_num-min_num];

    set_true_if_in_range!(2, 3, 5 => prime_filter + min_num, min_num, max_num);
    //Macro equivilent:
    // if (min_num <= 2) & (max_num > 2) {prime_filter[2-min_num] = true;}
    // if (min_num <= 3) & (max_num > 3) {prime_filter[3-min_num] = true;}
    // if (min_num <= 5) & (max_num > 5) {prime_filter[5-min_num] = true;}

    let (mut y_sq, mut to_next_y_sq) = (1, 3);

    while y_sq<max_num {
        if y_sq%2 == 1 {
            //CASE 1
            //n_1 = 4x^2 + y^2 === 1 (mod 4)
            let (mut n_1, mut to_next_n_1) = match y_sq < min_num {
                false => (y_sq+4, 12),
                _ => {
                    let min_num_x = (ceil_sqrt(min_num - y_sq) +1)/2;
                    (4*min_num_x*min_num_x + y_sq, 8*min_num_x + 4)
                },
            };
            loop{

                match n_1{
                    n if n >= max_num => break,
                    n => {match n%60{
                        1 | 13 | 17 | 29 | 37 | 41 | 49 | 53 => prime_filter[n-min_num] ^= true,
                        _ => (),
                    };},
                };

                n_1 += to_next_n_1;
                to_next_n_1 += 8;
            };
        };
        if y_sq%3 == 1 {
            //CASE 2
            //n_2 = 3x^2 + y^2 === 1 (mod 6)
            let (mut n_2, mut to_next_n_2) = match y_sq < min_num {
                false => (y_sq+3, 9),
                _ => {
                    let min_num_x = (ceil_sqrt((min_num - y_sq)*3)+2)/3;
                    (3*min_num_x*min_num_x + y_sq, 6*min_num_x + 3)
                },
            };
            loop {
                match n_2{
                    n if n >= max_num => break,
                    n => {match n%60{
                            7 | 19 | 31 | 43 => prime_filter[n-min_num] ^= true,
                            _ => (),
                        };}
                };
                n_2 += to_next_n_2;
                to_next_n_2 += 6;
            };
            //CASE 3
            //n_3 = 3x^2 - y^2 === 11 (mod 12)
            //Initially, we set x = y+1 -> n_3 = 3(y+1)^2 - y^2 = 2*y^2 + 6*y + 3
            //Amd then hop x by 2 each iteration.
            let (mut n_3, mut to_next_n_3) = match (y_sq << 1) < min_num {
                false => (2*y_sq+3*to_next_y_sq, 6*to_next_y_sq+18),
                _ => {
                    let min_num_x = match (ceil_sqrt((min_num + y_sq)*3) +2)/3{
                        mx if (mx+y_sq)%2 == 0 => mx + 1,
                        mx => mx,
                    };
                    (3*min_num_x*min_num_x - y_sq, 12*min_num_x + 12)
                },
            };
            loop {
                match n_3{
                    n if n >= max_num => break,
                    n => {match n%60{
                            11 | 23 | 47 | 59 => prime_filter[n-min_num] ^= true,
                            _ => (),
                    };},
                };
                n_3 += to_next_n_3;
                to_next_n_3 += 24;
            };
        };
        while{ //Do-while
            y_sq += to_next_y_sq;
            to_next_y_sq += 2;
            y_sq%6 == 0
        } {};
    };

    //Elimin_numate non-squarefree numbers
    let mut n_sq = 49; // 7^2
    let mut next_n_sq = 32; //9^2 - 7^2, skip even numbers.
    while n_sq < max_num {
        let mut non_sq_free = n_sq;
        while non_sq_free < max_num {
            if non_sq_free >= min_num {
                prime_filter[non_sq_free - min_num] = false;
            }
            while{ //Do-while
                non_sq_free += n_sq + n_sq;
                (non_sq_free%3==0) | (non_sq_free%5==0)
            } {};
        };
        while{ //Do-while
            n_sq += next_n_sq;
            next_n_sq += 8;
            (n_sq%3==0) | (n_sq%5 == 0)
        } {};
    }
    prime_filter
}

#[cfg(test)]
pub fn old_prime_filter(max_num: usize) -> std::vec::Vec<bool>{
    slow_prime_filter(max_num)
}

#[test]
fn private_filter_test(){
    assert_eq!(5, ceil_sqrt(24));
    assert_eq!(2, int_sqrt(4));
    assert_eq!(4, int_sqrt(24));
    assert_eq!(10, int_sqrt(101));
    assert_eq!(1, int_sqrt(1));
    assert_eq!(10, int_sqrt(100));
    assert_eq!(3, int_sqrt(13));
}

fn slow_prime_filter(max_num: usize) -> std::vec::Vec<bool>{
    if max_num < 5 {
         let mut ret = vec![false, false, true, true];
         ret.truncate(max_num);
         return ret
     }
    let mut prime_filter = vec![true; max_num];
    prime_filter[0] = false;
    prime_filter[1] = false;
    let mut cur_num = 2;
    'outer: loop{
        for i in (cur_num+1)..max_num{
            if 0 == i%cur_num { prime_filter[i] = false; }
        }
        cur_num += 1;
        while !prime_filter[cur_num]{
            if cur_num*cur_num > max_num {
                break 'outer
            }
            cur_num += 1;
        }
    };
    prime_filter
}
<commit_msg>Code speedup...<commit_after>use std;
use concurrent_help::to_concurrent_on_section;
use num_cpus;
// macro_rules! if_any_divides {
//      ( $( $div:tt ),* | $n:ident $if_b:block ) => { if $(($n % $div == 0))|* $if_b};
// }
// macro_rules! if_sm_prime_divides{
//     ( $n:ident $if_b:block ) => {if_any_divides!(
//     2, 3, 5, 7, 11, 13,
//     17, 19, 23, 29, 31,
//     37, 41, 43, 47, 53,
//     59, 61, 67, 71, 73,
//     79, 83, 89, 97, 101,
//     103, 107, 109, 113, 127,
//     131, 137, 139, 149, 151,
//     157, 163, 167, 173, 179,
//     181, 191, 193, 197, 199 | $n $if_b)};
// }
macro_rules! set_true_if_in_range{
    ( $( $i:tt ),* => $filter:ident + $offset:expr, $min:expr, $max:expr) => {
        $(if ($min <= $i) & ($max > $i) {$filter[$i-$offset] = true;})*
    };
    ( $( $i:tt ),* => $filter:ident, $min:expr, $max:expr) => {
        $(if ($min <= $i) & ($max > $i) {$filter[$i] = true;})*
    };
}
// macro_rules! set_sm_primes_true {
//     ($filter:expr, $min:expr, $max:expr) => {
//         set_true_if_in_range!(
//         2, 3, 5, 7, 11, 13,
//         17, 19, 23, 29, 31,
//         37, 41, 43, 47, 53,
//         59, 61, 67, 71, 73,
//         79, 83, 89, 97, 101,
//         103, 107, 109, 113, 127,
//         131, 137, 139, 149, 151,
//         157, 163, 167, 173, 179,
//         181, 191, 193, 197, 199 => $filter + $min, $min, $max)
//     };
// }
pub fn prime_filter_concurrently(max_num: usize, threads: usize) -> Vec<bool>{
    prime_filter_section_concurrently(0, max_num, threads)
}
pub fn prime_filter_section_concurrently(min_num:usize, max_num: usize, threads: usize) -> Vec<bool>{
    to_concurrent_on_section(prime_filter_section_sequentially, min_num, max_num, threads)
}

fn int_sqrt(n:usize) -> usize{
    if n < (1 << 53) {
        return (n as f64).sqrt() as usize;
    }
    match n {
        0 => 0,
        1 ... 3 => 1,
        4 ... 8 => 2,
        9 ... 15 => 3,
        k => {
        let mut x = k;
            loop{
                x = match (x + n/x) >> 1 {
                    new_x if new_x == x => break,
                    new_x if new_x*new_x == n + 1 => {x = new_x - 1; break},
                    new_x => new_x,
                };
            }
            x
        },
    }

}

fn ceil_sqrt(n:usize) -> usize{
    if n == 0{
        return 0;
    }
    match int_sqrt(n){
        sqrt if sqrt*sqrt == n => sqrt,
        sqrt => sqrt+1,
    }
}
pub fn prime_filter(max_num: usize) -> Vec<bool>{
    prime_filter_concurrently(max_num, num_cpus::get())
}

pub fn prime_filter_section(min_num:usize, max_num: usize) -> Vec<bool>{
    prime_filter_section_concurrently(min_num, max_num, num_cpus::get())
}

pub fn prime_filter_sequentially(max_num: usize) -> Vec<bool>{
    if max_num<100{
        slow_prime_filter(max_num)
    }else{
        prime_filter_section(0, max_num)
    }
}
pub fn prime_filter_section_sequentially(min_num:usize, max_num: usize) -> Vec<bool>{
    //Sieve of Atkin
    assert!(min_num<max_num);
    let mut prime_filter = vec![false; max_num-min_num];

    set_true_if_in_range!(2, 3, 5 => prime_filter + min_num, min_num, max_num);
    //Macro equivilent:
    // if (min_num <= 2) & (max_num > 2) {prime_filter[2-min_num] = true;}
    // if (min_num <= 3) & (max_num > 3) {prime_filter[3-min_num] = true;}
    // if (min_num <= 5) & (max_num > 5) {prime_filter[5-min_num] = true;}

    let (mut y_sq, mut to_next_y_sq) = (1, 3);

    while y_sq<max_num {
        if y_sq%2 == 1 {
            //CASE 1
            //n_1 = 4x^2 + y^2 === 1 (mod 4)
            let (mut n_1, mut to_next_n_1) = match y_sq < min_num {
                false => (y_sq+4, 12),
                _ => {
                    let min_num_x = (ceil_sqrt(min_num - y_sq) +1)/2;
                    (4*min_num_x*min_num_x + y_sq, 8*min_num_x + 4)
                },
            };
            loop{

                match n_1{
                    n if n >= max_num => break,
                    n => {match n%60{
                        1 | 13 | 17 | 29 | 37 | 41 | 49 | 53 => prime_filter[n-min_num] ^= true,
                        _ => (),
                    };},
                };

                n_1 += to_next_n_1;
                to_next_n_1 += 8;
            };
        };
        if y_sq%3 == 1 {
            //CASE 2
            //n_2 = 3x^2 + y^2 === 1 (mod 6)
            let (mut n_2, mut to_next_n_2) = match y_sq < min_num {
                false => (y_sq+3, 9),
                _ => {
                    let min_num_x = (ceil_sqrt((min_num - y_sq)*3)+2)/3;
                    (3*min_num_x*min_num_x + y_sq, 6*min_num_x + 3)
                },
            };
            loop {
                match n_2{
                    n if n >= max_num => break,
                    n => {match n%60{
                            7 | 19 | 31 | 43 => prime_filter[n-min_num] ^= true,
                            _ => (),
                        };}
                };
                n_2 += to_next_n_2;
                to_next_n_2 += 6;
            };
            //CASE 3
            //n_3 = 3x^2 - y^2 === 11 (mod 12)
            //Initially, we set x = y+1 -> n_3 = 3(y+1)^2 - y^2 = 2*y^2 + 6*y + 3
            //Amd then hop x by 2 each iteration.
            let (mut n_3, mut to_next_n_3) = match (y_sq << 1) < min_num {
                false => (2*y_sq+3*to_next_y_sq, 6*to_next_y_sq+18),
                _ => {
                    let min_num_x = match (ceil_sqrt((min_num + y_sq)*3) +2)/3{
                        mx if (mx+y_sq)%2 == 0 => mx + 1,
                        mx => mx,
                    };
                    (3*min_num_x*min_num_x - y_sq, 12*min_num_x + 12)
                },
            };
            loop {
                match n_3{
                    n if n >= max_num => break,
                    n => {match n%60{
                            11 | 23 | 47 | 59 => prime_filter[n-min_num] ^= true,
                            _ => (),
                    };},
                };
                n_3 += to_next_n_3;
                to_next_n_3 += 24;
            };
        };
        while{ //Do-while
            y_sq += to_next_y_sq;
            to_next_y_sq += 2;
            y_sq%6 == 0
        } {};
    };

    //Eliminate non-squarefree numbers

    let mut n_sq = 49; // 7^2
    let mut next_n_sq = 32; //9^2 - 7^2, skip even numbers.
    while n_sq < max_num {
        let mut non_sq_free =  n_sq * match min_num/n_sq{
            k if k%2 == 1 => k,
            k => k + 1,
        };
        while non_sq_free < max_num {
            if non_sq_free >= min_num {
                prime_filter[non_sq_free - min_num] = false;
            }
            while{ //Do-while
                non_sq_free += n_sq + n_sq;
                (non_sq_free%3==0) | (non_sq_free%5==0)
            } {};
        };
        while{ //Do-while
            n_sq += next_n_sq;
            next_n_sq += 8;
            (n_sq%3==0) | (n_sq%5 == 0)
        } {};
    }
    prime_filter
}

#[cfg(test)]
pub fn old_prime_filter(max_num: usize) -> std::vec::Vec<bool>{
    slow_prime_filter(max_num)
}

#[test]
fn private_filter_test(){
    assert_eq!(5, ceil_sqrt(24));
    assert_eq!(2, int_sqrt(4));
    assert_eq!(4, int_sqrt(24));
    assert_eq!(10, int_sqrt(101));
    assert_eq!(1, int_sqrt(1));
    assert_eq!(10, int_sqrt(100));
    assert_eq!(3, int_sqrt(13));
}

fn slow_prime_filter(max_num: usize) -> std::vec::Vec<bool>{
    if max_num < 5 {
         let mut ret = vec![false, false, true, true];
         ret.truncate(max_num);
         return ret
     }
    let mut prime_filter = vec![true; max_num];
    prime_filter[0] = false;
    prime_filter[1] = false;
    let mut cur_num = 2;
    'outer: loop{
        for i in (cur_num+1)..max_num{
            if 0 == i%cur_num { prime_filter[i] = false; }
        }
        cur_num += 1;
        while !prime_filter[cur_num]{
            if cur_num*cur_num > max_num {
                break 'outer
            }
            cur_num += 1;
        }
    };
    prime_filter
}
<|endoftext|>


<commit_before>use std;
use std::thread;
use std::sync::mpsc;

pub fn prime_filter(iter_size: usize) -> Vec<bool>{
    if iter_size<100{
        slow_prime_filter(iter_size)
    }else{
        prime_filter_section(0, iter_size)
    }
}

fn case_1(y_sq: usize, iter_size: usize)
    -> (usize, Vec<bool>) {
    //n_1 = 4x^2 + y^2 === 1 (mod 4)
    let (mut n_1, mut to_next_n_1) = (y_sq, 4);
    let offset = n_1;
    let mut temp_filter = vec![false; iter_size-offset];
    loop{
        n_1 += to_next_n_1;
        to_next_n_1 += 8;
        match n_1%60{
            _ if n_1 >= iter_size => break,
            1 | 13 | 17 | 29 | 37 | 41 | 49 | 53 => (),
            _ => continue,
        };

        // println!("1: {}", n_1);
        temp_filter[n_1-offset] ^= true;
    };
    (offset, temp_filter)
}
fn case_2(y_sq: usize, iter_size: usize)
    -> (usize, Vec<bool>) {
    //n_2 = 3x^2 + y^2 === 1 (mod 6)
    let (mut n_2, mut to_next_n_2) = (y_sq, 3);
    let offset = n_2.clone();

    let mut temp_filter = vec![false; iter_size-offset];
    loop {
        n_2 += to_next_n_2;
        to_next_n_2 += 6;
        match n_2%60{
            _ if n_2 >= iter_size => break,
            7 | 19 | 31 | 43 => (),
            _ => continue,
        };
        // println!("2: {}", n_2);
        temp_filter[n_2 - offset] ^= true;
    };
    (offset, temp_filter)
}
fn case_3(y_sq: usize, to_next_y_sq: usize, iter_size: usize)
    -> (usize, Vec<bool>) {
    //n_3 = 3x^2 - y^2 === 11 (mod 12)
    let (mut n_3, mut to_next_n_3) = (2*y_sq, 3*to_next_y_sq);
    let offset = n_3;
    let mut temp_filter = vec![false; iter_size - offset];
    loop {
        n_3 += to_next_n_3;
        to_next_n_3 += 6;
        match n_3%60{
            _ if n_3 >= iter_size => break,
            11 | 23 | 47 | 59 => (),
            _ => continue,
        };
        // println!("3: {}", n_3);
        temp_filter[n_3 - offset] ^= true;
    };
    (offset, temp_filter)
}

fn prime_filter_section(min:usize, max: usize) -> Vec<bool>{
    //Sieve of Atkin
    assert!(min<max);
    let mut prime_filter = vec![false; max-min];
    if (min <= 2) & (max > 2) {prime_filter[2-min] = true;}
    if (min <= 3) & (max > 3) {prime_filter[3-min] = true;}
    if (min <= 5) & (max > 5) {prime_filter[5-min] = true;}

    let (mut y_sq, mut to_next_y_sq) = match min {
        0|1 => (1, 3),
        _ => (min*min, 2*min + 1),
    };
    let mut spawned_threads = 0;
    let (tx, rx) = mpsc::channel();
    while y_sq<max {
        if y_sq%2 == 1 {
            spawned_threads += 1;
            let (y_sq, tx) = (y_sq, tx.clone());
            thread::spawn(move || {
                tx.send(case_1(y_sq, max)).unwrap();
            });
        };
        if y_sq%3 == 1 {

            spawned_threads += 1;
            let (y_sq, to_next_y_sq, tx1) = (y_sq, to_next_y_sq, tx.clone());
            thread::spawn(move || {
                tx1.send(case_2(y_sq, max)).unwrap();
            });
            if y_sq*2<iter_size{
                spawned_threads += 1;
                let tx2 = tx.clone();
                thread::spawn(move || {
                    tx2.send(case_3(y_sq, to_next_y_sq, max)).unwrap();
                });
            }
        };
        if spawned_threads>0{
            for mes in rx.try_iter(){
                spawned_threads -= 1;
                let (offset, temp_filter) = mes;
                for flip_i in temp_filter.into_iter().enumerate()
                                        .filter_map(|x| match x {
                                            (i, true) => Some(i),
                                            _ => None,
                                        }){
                    prime_filter[offset + flip_i - min] ^= true;
                }
            }
        }
        while{ //Do-while
            y_sq += to_next_y_sq;
            to_next_y_sq += 2;
            y_sq%6 == 0
        } {};
    };
    while spawned_threads>0{
        for mes in rx.try_iter(){
            spawned_threads -= 1;
            let (offset, temp_filter) = mes;
            for flip_i in temp_filter.into_iter().enumerate()
                                    .filter_map(|x| match x {
                                        (i, true) => Some(i),
                                        _ => None,
                                    }){
                prime_filter[offset + flip_i - min] ^= true;
            }
        }
    }
    //Eliminate non-squarefree numbers
    let mut n_sq = 49; // 7^2
    let mut next_n_sq = 32; //9^2 - 7^2, skip even numbers.
    while n_sq < iter_size {
        let mut non_sq_free = n_sq;
        while non_sq_free < max {
            if non_sq_free >= min {
                prime_filter[non_sq_free - min] = false;
            }
            while{ //Do-while
                non_sq_free += n_sq + n_sq;
                (non_sq_free%3==0) | (non_sq_free%5==0)
            } {};
        };
        while{ //Do-while
            non_sq_free += n_sq + n_sq;
            (non_sq_free%3==0) | (non_sq_free%5==0)
        } {};
    };
    while{ //Do-while
        n_sq += next_n_sq;
        next_n_sq += 8;
        (n_sq%3==0) | (n_sq%5 == 0)
    } {};
    prime_filter
}
#[cfg(test)]
pub fn old_prime_filter(iter_size: usize) -> std::vec::Vec<bool>{
    slow_prime_filter(iter_size)
}

fn slow_prime_filter(iter_size: usize) -> std::vec::Vec<bool>{
    if iter_size < 5 {
         let mut ret = vec![false, false, true, true];
         ret.truncate(iter_size);
         return ret
     }
    let mut prime_filter = vec![true; iter_size];
    prime_filter[0] = false;
    prime_filter[1] = false;
    let mut cur_num = 2;
    'outer: loop{
        for i in (cur_num+1)..iter_size{
            if 0 == i%cur_num { prime_filter[i] = false; }
        }
        cur_num += 1;
        while !prime_filter[cur_num]{
            if cur_num*cur_num > iter_size {
                break 'outer
            }
            cur_num += 1;
        }
    };
    prime_filter
}
<commit_msg>name resolution fix<commit_after>use std;
use std::thread;
use std::sync::mpsc;

pub fn prime_filter(iter_size: usize) -> Vec<bool>{
    if iter_size<100{
        slow_prime_filter(iter_size)
    }else{
        prime_filter_section(0, iter_size)
    }
}

fn case_1(y_sq: usize, iter_size: usize)
    -> (usize, Vec<bool>) {
    //n_1 = 4x^2 + y^2 === 1 (mod 4)
    let (mut n_1, mut to_next_n_1) = (y_sq, 4);
    let offset = n_1;
    let mut temp_filter = vec![false; iter_size-offset];
    loop{
        n_1 += to_next_n_1;
        to_next_n_1 += 8;
        match n_1%60{
            _ if n_1 >= iter_size => break,
            1 | 13 | 17 | 29 | 37 | 41 | 49 | 53 => (),
            _ => continue,
        };

        // println!("1: {}", n_1);
        temp_filter[n_1-offset] ^= true;
    };
    (offset, temp_filter)
}
fn case_2(y_sq: usize, iter_size: usize)
    -> (usize, Vec<bool>) {
    //n_2 = 3x^2 + y^2 === 1 (mod 6)
    let (mut n_2, mut to_next_n_2) = (y_sq, 3);
    let offset = n_2.clone();

    let mut temp_filter = vec![false; iter_size-offset];
    loop {
        n_2 += to_next_n_2;
        to_next_n_2 += 6;
        match n_2%60{
            _ if n_2 >= iter_size => break,
            7 | 19 | 31 | 43 => (),
            _ => continue,
        };
        // println!("2: {}", n_2);
        temp_filter[n_2 - offset] ^= true;
    };
    (offset, temp_filter)
}
fn case_3(y_sq: usize, to_next_y_sq: usize, iter_size: usize)
    -> (usize, Vec<bool>) {
    //n_3 = 3x^2 - y^2 === 11 (mod 12)
    let (mut n_3, mut to_next_n_3) = (2*y_sq, 3*to_next_y_sq);
    let offset = n_3;
    let mut temp_filter = vec![false; iter_size - offset];
    loop {
        n_3 += to_next_n_3;
        to_next_n_3 += 6;
        match n_3%60{
            _ if n_3 >= iter_size => break,
            11 | 23 | 47 | 59 => (),
            _ => continue,
        };
        // println!("3: {}", n_3);
        temp_filter[n_3 - offset] ^= true;
    };
    (offset, temp_filter)
}

fn prime_filter_section(min:usize, max: usize) -> Vec<bool>{
    //Sieve of Atkin
    assert!(min<max);
    let mut prime_filter = vec![false; max-min];
    if (min <= 2) & (max > 2) {prime_filter[2-min] = true;}
    if (min <= 3) & (max > 3) {prime_filter[3-min] = true;}
    if (min <= 5) & (max > 5) {prime_filter[5-min] = true;}

    let (mut y_sq, mut to_next_y_sq) = match min {
        0|1 => (1, 3),
        _ => (min*min, 2*min + 1),
    };
    let mut spawned_threads = 0;
    let (tx, rx) = mpsc::channel();
    while y_sq<max {
        if y_sq%2 == 1 {
            spawned_threads += 1;
            let (y_sq, tx) = (y_sq, tx.clone());
            thread::spawn(move || {
                tx.send(case_1(y_sq, max)).unwrap();
            });
        };
        if y_sq%3 == 1 {

            spawned_threads += 1;
            let (y_sq, to_next_y_sq, tx1) = (y_sq, to_next_y_sq, tx.clone());
            thread::spawn(move || {
                tx1.send(case_2(y_sq, max)).unwrap();
            });
            if y_sq*2<max{
                spawned_threads += 1;
                let tx2 = tx.clone();
                thread::spawn(move || {
                    tx2.send(case_3(y_sq, to_next_y_sq, max)).unwrap();
                });
            }
        };
        if spawned_threads>0{
            for mes in rx.try_iter(){
                spawned_threads -= 1;
                let (offset, temp_filter) = mes;
                for flip_i in temp_filter.into_iter().enumerate()
                                        .filter_map(|x| match x {
                                            (i, true) => Some(i),
                                            _ => None,
                                        }){
                    prime_filter[offset + flip_i - min] ^= true;
                }
            }
        }
        while{ //Do-while
            y_sq += to_next_y_sq;
            to_next_y_sq += 2;
            y_sq%6 == 0
        } {};
    };
    while spawned_threads>0{
        for mes in rx.try_iter(){
            spawned_threads -= 1;
            let (offset, temp_filter) = mes;
            for flip_i in temp_filter.into_iter().enumerate()
                                    .filter_map(|x| match x {
                                        (i, true) => Some(i),
                                        _ => None,
                                    }){
                prime_filter[offset + flip_i - min] ^= true;
            }
        }
    }
    //Eliminate non-squarefree numbers
    let mut n_sq = 49; // 7^2
    let mut next_n_sq = 32; //9^2 - 7^2, skip even numbers.
    while n_sq < max {
        let mut non_sq_free = n_sq;
        while non_sq_free < max {
            if non_sq_free >= min {
                prime_filter[non_sq_free - min] = false;
            }
            while{ //Do-while
                non_sq_free += n_sq + n_sq;
                (non_sq_free%3==0) | (non_sq_free%5==0)
            } {};
        };
        while{ //Do-while
            non_sq_free += n_sq + n_sq;
            (non_sq_free%3==0) | (non_sq_free%5==0)
        } {};
    };
    while{ //Do-while
        n_sq += next_n_sq;
        next_n_sq += 8;
        (n_sq%3==0) | (n_sq%5 == 0)
    } {};
    prime_filter
}
#[cfg(test)]
pub fn old_prime_filter(iter_size: usize) -> std::vec::Vec<bool>{
    slow_prime_filter(iter_size)
}

fn slow_prime_filter(iter_size: usize) -> std::vec::Vec<bool>{
    if iter_size < 5 {
         let mut ret = vec![false, false, true, true];
         ret.truncate(iter_size);
         return ret
     }
    let mut prime_filter = vec![true; iter_size];
    prime_filter[0] = false;
    prime_filter[1] = false;
    let mut cur_num = 2;
    'outer: loop{
        for i in (cur_num+1)..iter_size{
            if 0 == i%cur_num { prime_filter[i] = false; }
        }
        cur_num += 1;
        while !prime_filter[cur_num]{
            if cur_num*cur_num > iter_size {
                break 'outer
            }
            cur_num += 1;
        }
    };
    prime_filter
}
<|endoftext|>