<commit_before>use std::old_io::MemReader;

use tls::TLS_VERSION;
use tls_result::TlsResult;
use tls_result::TlsErrorKind::{InternalError, UnexpectedMessage};
use tls_item::{TlsItem, DummyItem, ObscureData};
use signature::SignatureAndHashAlgorithmVec;
use cipher::CipherSuite;

// This is actually `struct { gmt_unix_time: u32, random_bytes: [u8, ..28] }`
// cf: http://tools.ietf.org/html/draft-mathewson-no-gmtunixtime-00
tls_array!(Random = [u8, ..32]);

tls_vec!(CipherSuiteVec = CipherSuite(2, (1 << 16) - 2));

tls_enum!(u8, enum CompressionMethod {
    null(0),
    DEFLATE(1) // RFC 3749
});
tls_vec!(CompressionMethodVec = CompressionMethod(1, (1 << 8) - 1));

tls_struct!(struct ProtocolVersion { major: u8, minor: u8 });

tls_vec!(SessionId = u8(0, 32));

tls_vec!(Asn1Cert = u8(1, (1 << 24) - 1));

// RFC 4492

tls_enum!(u16, enum NamedCurve {
    sect163k1 (1), sect163r1 (2), sect163r2 (3),
    sect193r1 (4), sect193r2 (5), sect233k1 (6),
    sect233r1 (7), sect239k1 (8), sect283k1 (9),
    sect283r1 (10), sect409k1 (11), sect409r1 (12),
    sect571k1 (13), sect571r1 (14), secp160k1 (15),
    secp160r1 (16), secp160r2 (17), secp192k1 (18),
    secp192r1 (19), secp224k1 (20), secp224r1 (21),
    secp256k1 (22), secp256r1 (23), secp384r1 (24),
    secp521r1 (25),
    arbitrary_explicit_prime_curves(0xFF01),
    arbitrary_explicit_char2_curves(0xFF02)
});
tls_vec!(EllipticCurveList = NamedCurve(1, (1 << 16) - 1));

tls_enum!(u8, enum ECPointFormat {
    uncompressed (0), ansiX962_compressed_prime (1),
    ansiX962_compressed_char2 (2)
});
tls_vec!(ECPointFormatList = ECPointFormat(1, (1 << 8) - 1));

// FIXME: Extension has the following structure:
// struct Extension {
//     extension_type: u8,
//     extension_data: opaque<1..2^16-1>,
// }
// and actual structure of `extension_data` depends on `extension_type`.
// Note that this is not exactly what `tls_enum_struct` wants!
// It's why we define horrible structs here.
tls_vec!(EllipticCurveListList = EllipticCurveList(1, (1 << 16) - 1));
tls_vec!(ECPointFormatListList = ECPointFormatList(1, (1 << 16) - 1));

tls_enum_struct!(u16, enum Extension {
    // RFC 6066
    //server_name(0),
    //max_fragment_length(1),
    //client_certificate_url(2),
    //trusted_ca_keys(3),
    //truncated_hmac(4),
    //status_request(5),
    // RFC 4492
    elliptic_curves(EllipticCurveListList) = 10,
    ec_point_formats(ECPointFormatListList) = 11
    // RFC 5246
    //signature_algorithms(13)
});
impl Extension {
    pub fn new_elliptic_curve_list(list: Vec<NamedCurve>) -> TlsResult<Extension> {
        let list = try!(EllipticCurveList::new(list));
        let list = try!(EllipticCurveListList::new(vec!(list)));
        let list = Extension::elliptic_curves(list);
        Ok(list)
    }

    pub fn new_ec_point_formats(list: Vec<ECPointFormat>) -> TlsResult<Extension> {
        let list = try!(ECPointFormatList::new(list));
        let list = try!(ECPointFormatListList::new(vec!(list)));
        let list = Extension::ec_point_formats(list);
        Ok(list)
    }
}

tls_vec!(ExtensionVec = Extension(0, (1 << 16) - 1));
tls_option!(ExtensionVec);

// struct Handshake {
//     msg_type: u8,
//     len: u24,
//     data: <depends on msg_type>
// }
macro_rules! tls_handshake(
    (
        $(
            $name:ident($body_ty:ty) = $num:tt, // $num: integer literal
        )+
    ) => (
        #[allow(non_camel_case_types)]
        pub enum Handshake {
            $(
                $name($body_ty),
            )+
        }

        impl TlsItem for Handshake {
            fn tls_write<W: Writer>(&self, writer: &mut W) -> TlsResult<()> {
                match *self {
                    $(
                        Handshake::$name(ref body) => {
                            try!(writer.write_u8(tt_to_expr!($num)));

                            let len = body.tls_size();
                            try!(writer.write_u8((len >> 16) as u8));
                            try!(writer.write_u8((len >> 8) as u8));
                            try!(writer.write_u8(len as u8));

                            try!(body.tls_write(writer));
                        }
                    )+
                }

                Ok(())
            }

            fn tls_read<R: Reader>(reader: &mut R) -> TlsResult<Handshake> {
                let ty = try!(reader.read_u8());

                // HandshakeBuffer already checked validity of length
                let _len = {
                    let n1 = try!(reader.read_u8()) as u32;
                    let n2 = try!(reader.read_u8()) as u32;
                    let n3 = try!(reader.read_u8()) as u32;
                    (n1 << 16) | (n2 << 8) | n3
                };

                let ret = match ty {
                    $(
                        tt_to_pat!($num) => {
                            let body: $body_ty = try!(TlsItem::tls_read(reader));
                            Handshake::$name(body)
                        }
                    )+
                    _ => return tls_err!(UnexpectedMessage,
                                         "unexpected Handshake message: type {}",
                                         ty),
                };

                let should_be_err = reader.read_u8();
                match should_be_err {
                    Err(_) => {},
                    Ok(_) => return tls_err!(InternalError, "expected EOF but found not"),
                }

                Ok(ret)
            }

            fn tls_size(&self) -> u64 {
                let body_len = match *self {
                    $(
                        Handshake::$name(ref body) => body.tls_size(),
                    )+
                };
                // msg_type 1 byte, length 3 bytes
                1 + 3 + body_len
            }
        }
    )
);

tls_handshake!(
    hello_request(DummyItem) = 0,
    client_hello(ClientHello) = 1,
    server_hello(ServerHello) = 2,
    // hello_verify_request(..) = 3, RFC 6347: DTLS
    // NewSessionTicket(..) = 4, RFC 5077: session resumption w/o server-side state
    certificate(CertificateList) = 11,
    server_key_exchange(ObscureData) = 12,
    certificate_request(CertificateRequest) = 13,
    server_hello_done(DummyItem) = 14,
    // certificate_verify = 15,
    client_key_exchange(ObscureData) = 16,
    finished(VerifyData) = 20,
);

tls_struct!(struct ClientHello {
    client_version: ProtocolVersion,
    random: Random,
    session_id: SessionId,
    cipher_suites: CipherSuiteVec,
    compression_methods: CompressionMethodVec,
    extensions: Option<ExtensionVec>
});

tls_struct!(struct ServerHello {
    server_version: ProtocolVersion,
    random: Random,
    session_id: SessionId,
    cipher_suite: CipherSuite,
    compression_method: CompressionMethod,
    extensions: Option<ExtensionVec>
});

tls_vec!(CertificateList = Asn1Cert(0, (1 << 24) - 1));

tls_enum!(u8, enum ClientCertificateType {
      rsa_sign(1), dss_sign(2), rsa_fixed_dh(3), dss_fixed_dh(4),
      rsa_ephemeral_dh_RESERVED(5), dss_ephemeral_dh_RESERVED(6),
      fortezza_dms_RESERVED(20)
});
tls_vec!(CertificiateTypeVec = ClientCertificateType(1, (1 << 8) - 1));

tls_vec!(DistinguishedName = u8(1, (1 << 16) - 1));
tls_vec!(DistinguishedNameVec = DistinguishedName(0, (1 << 16) - 1));

tls_struct!(struct CertificateRequest {
    certificate_types: CertificiateTypeVec,
    supported_signature_algorithms: SignatureAndHashAlgorithmVec,
    certificate_authorities: DistinguishedNameVec
});

// FIXME TLS 1.2 says the length can be longer for future ciphe suites.
tls_array!(VerifyData = [u8, ..12]);

// buffer for handshake protocol
pub struct HandshakeBuffer {
    buf: Vec<u8>,
}

impl HandshakeBuffer {
    pub fn new() -> HandshakeBuffer {
        HandshakeBuffer { buf: Vec::new() }
    }

    pub fn add_record(&mut self, fragment: Vec<u8>) {
        self.buf.push_all(&fragment);
    }

    // if message is arrived but has unknown type, the message is discarded and returns error.
    pub fn get_message(&mut self) -> TlsResult<Option<Handshake>> {
        let len = self.buf.len();
        // we need to read at least ty and length
        if len < 4 {
            return Ok(None);
        }

        let n1 = self.buf[1] as usize;
        let n2 = self.buf[2] as usize;
        let n3 = self.buf[3] as usize;
        let wanted_len: usize = (n1 << 16) | (n2 << 8) | n3;
        let wanted_len = wanted_len + 4;

        if len < wanted_len {
            return Ok(None);
        }

        // FIXME bad clone?
        let (message, remaining) = {
            let (message, remaining) = self.buf.split_at_mut(wanted_len);
            let message = message.to_vec();
            let remaining = remaining.to_vec();
            (message, remaining)
        };
        self.buf = remaining;

        let mut reader = MemReader::new(message);
        let message: Handshake = try!(TlsItem::tls_read(&mut reader));
        let ret = Ok(Some(message));

        ret
    }
}

impl Handshake {
    pub fn new_client_hello(random: Random,
                            cipher_suite: CipherSuite,
                            extensions: Vec<Extension>) -> TlsResult<Handshake> {
        let client_hello_body = {
            let client_version = {
                let (major, minor) = TLS_VERSION;

                ProtocolVersion {
                    major: major,
                    minor: minor,
                }
            };

            // TODO support session resumption
            let session_id = {
                let data = Vec::new();
                try!(SessionId::new(data))
            };

            let cipher_suites = {
                let data = vec!(cipher_suite);
                try!(CipherSuiteVec::new(data))
            };

            let compression_methods = {
                let data = vec!(CompressionMethod::null);
                try!(CompressionMethodVec::new(data))
            };

            let extensions = if extensions.len() == 0 {
                None
            } else {
                let ext = try!(ExtensionVec::new(extensions));
                Some(ext)
            };
            ClientHello {
                client_version: client_version,
                random: random,
                session_id: session_id,
                cipher_suites: cipher_suites,
                compression_methods: compression_methods,
                extensions: extensions,
            }
        };

        Ok(Handshake::client_hello(client_hello_body))
    }

    pub fn new_client_key_exchange(data: Vec<u8>) -> TlsResult<Handshake> {
        let data = ObscureData::new(data);
        Ok(Handshake::client_key_exchange(data))
    }

    pub fn new_finished(data: Vec<u8>) -> TlsResult<Handshake> {
        let data = try!(VerifyData::new(data));
        Ok(Handshake::finished(data))
    }
}

#[cfg(test)]
mod test {
    use std::old_io::MemReader;
    use tls_item::TlsItem;
    use cipher::CipherSuite;

    use super::{ProtocolVersion, SessionId, CipherSuiteVec, CompressionMethod,
                CompressionMethodVec, ClientHello, Handshake, Random};

    #[test]
    fn test_parse_client_hello() {
        let client_hello_msg = {
            let client_hello_body = {
                let client_version = {
                    let (major, minor) = (3, 3);

                    ProtocolVersion {
                        major: major,
                        minor: minor,
                    }
                };

                let random = {
                    let random_bytes = [0u8; 32].to_vec();
                    Random::new(random_bytes).unwrap()
                };

                let session_id = {
                    let data = Vec::new();
                    SessionId::new(data).unwrap()
                };

                let cipher_suites = {
                    let data = vec!(CipherSuite::TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256);
                    CipherSuiteVec::new(data).unwrap()
                };

                let compression_methods = {
                    let data = vec!(CompressionMethod::null);
                    CompressionMethodVec::new(data).unwrap()
                };

                ClientHello {
                    client_version: client_version,
                    random: random,
                    session_id: session_id,
                    cipher_suites: cipher_suites,
                    compression_methods: compression_methods,
                    extensions: None,
                }
            };

            Handshake::client_hello(client_hello_body)
        };

        let mut packet = Vec::new();
        client_hello_msg.tls_write(&mut packet).unwrap();

        let mut reader = MemReader::new(packet.clone());
        let client_hello_msg_2: Handshake = TlsItem::tls_read(&mut reader).unwrap();

        let mut packet_2 = Vec::new();
        client_hello_msg_2.tls_write(&mut packet_2).unwrap();

        assert_eq!(packet, packet_2);
    }
}
<commit_msg>Fix for overflow semantics change of `as`<commit_after>use std::old_io::MemReader;

use tls::TLS_VERSION;
use tls_result::TlsResult;
use tls_result::TlsErrorKind::{InternalError, UnexpectedMessage};
use tls_item::{TlsItem, DummyItem, ObscureData};
use signature::SignatureAndHashAlgorithmVec;
use cipher::CipherSuite;

// This is actually `struct { gmt_unix_time: u32, random_bytes: [u8, ..28] }`
// cf: http://tools.ietf.org/html/draft-mathewson-no-gmtunixtime-00
tls_array!(Random = [u8, ..32]);

tls_vec!(CipherSuiteVec = CipherSuite(2, (1 << 16) - 2));

tls_enum!(u8, enum CompressionMethod {
    null(0),
    DEFLATE(1) // RFC 3749
});
tls_vec!(CompressionMethodVec = CompressionMethod(1, (1 << 8) - 1));

tls_struct!(struct ProtocolVersion { major: u8, minor: u8 });

tls_vec!(SessionId = u8(0, 32));

tls_vec!(Asn1Cert = u8(1, (1 << 24) - 1));

// RFC 4492

tls_enum!(u16, enum NamedCurve {
    sect163k1 (1), sect163r1 (2), sect163r2 (3),
    sect193r1 (4), sect193r2 (5), sect233k1 (6),
    sect233r1 (7), sect239k1 (8), sect283k1 (9),
    sect283r1 (10), sect409k1 (11), sect409r1 (12),
    sect571k1 (13), sect571r1 (14), secp160k1 (15),
    secp160r1 (16), secp160r2 (17), secp192k1 (18),
    secp192r1 (19), secp224k1 (20), secp224r1 (21),
    secp256k1 (22), secp256r1 (23), secp384r1 (24),
    secp521r1 (25),
    arbitrary_explicit_prime_curves(0xFF01),
    arbitrary_explicit_char2_curves(0xFF02)
});
tls_vec!(EllipticCurveList = NamedCurve(1, (1 << 16) - 1));

tls_enum!(u8, enum ECPointFormat {
    uncompressed (0), ansiX962_compressed_prime (1),
    ansiX962_compressed_char2 (2)
});
tls_vec!(ECPointFormatList = ECPointFormat(1, (1 << 8) - 1));

// FIXME: Extension has the following structure:
// struct Extension {
//     extension_type: u8,
//     extension_data: opaque<1..2^16-1>,
// }
// and actual structure of `extension_data` depends on `extension_type`.
// Note that this is not exactly what `tls_enum_struct` wants!
// It's why we define horrible structs here.
tls_vec!(EllipticCurveListList = EllipticCurveList(1, (1 << 16) - 1));
tls_vec!(ECPointFormatListList = ECPointFormatList(1, (1 << 16) - 1));

tls_enum_struct!(u16, enum Extension {
    // RFC 6066
    //server_name(0),
    //max_fragment_length(1),
    //client_certificate_url(2),
    //trusted_ca_keys(3),
    //truncated_hmac(4),
    //status_request(5),
    // RFC 4492
    elliptic_curves(EllipticCurveListList) = 10,
    ec_point_formats(ECPointFormatListList) = 11
    // RFC 5246
    //signature_algorithms(13)
});
impl Extension {
    pub fn new_elliptic_curve_list(list: Vec<NamedCurve>) -> TlsResult<Extension> {
        let list = try!(EllipticCurveList::new(list));
        let list = try!(EllipticCurveListList::new(vec!(list)));
        let list = Extension::elliptic_curves(list);
        Ok(list)
    }

    pub fn new_ec_point_formats(list: Vec<ECPointFormat>) -> TlsResult<Extension> {
        let list = try!(ECPointFormatList::new(list));
        let list = try!(ECPointFormatListList::new(vec!(list)));
        let list = Extension::ec_point_formats(list);
        Ok(list)
    }
}

tls_vec!(ExtensionVec = Extension(0, (1 << 16) - 1));
tls_option!(ExtensionVec);

// struct Handshake {
//     msg_type: u8,
//     len: u24,
//     data: <depends on msg_type>
// }
macro_rules! tls_handshake(
    (
        $(
            $name:ident($body_ty:ty) = $num:tt, // $num: integer literal
        )+
    ) => (
        #[allow(non_camel_case_types)]
        pub enum Handshake {
            $(
                $name($body_ty),
            )+
        }

        impl TlsItem for Handshake {
            fn tls_write<W: Writer>(&self, writer: &mut W) -> TlsResult<()> {
                match *self {
                    $(
                        Handshake::$name(ref body) => {
                            try!(writer.write_u8(tt_to_expr!($num)));

                            let len = body.tls_size();
                            try!(writer.write_u8(((len >> 16) & 0xff) as u8));
                            try!(writer.write_u8(((len >> 8) & 0xff) as u8));
                            try!(writer.write_u8((len & 0xff) as u8));

                            try!(body.tls_write(writer));
                        }
                    )+
                }

                Ok(())
            }

            fn tls_read<R: Reader>(reader: &mut R) -> TlsResult<Handshake> {
                let ty = try!(reader.read_u8());

                // HandshakeBuffer already checked validity of length
                let _len = {
                    let n1 = try!(reader.read_u8()) as u32;
                    let n2 = try!(reader.read_u8()) as u32;
                    let n3 = try!(reader.read_u8()) as u32;
                    (n1 << 16) | (n2 << 8) | n3
                };

                let ret = match ty {
                    $(
                        tt_to_pat!($num) => {
                            let body: $body_ty = try!(TlsItem::tls_read(reader));
                            Handshake::$name(body)
                        }
                    )+
                    _ => return tls_err!(UnexpectedMessage,
                                         "unexpected Handshake message: type {}",
                                         ty),
                };

                let should_be_err = reader.read_u8();
                match should_be_err {
                    Err(_) => {},
                    Ok(_) => return tls_err!(InternalError, "expected EOF but found not"),
                }

                Ok(ret)
            }

            fn tls_size(&self) -> u64 {
                let body_len = match *self {
                    $(
                        Handshake::$name(ref body) => body.tls_size(),
                    )+
                };
                // msg_type 1 byte, length 3 bytes
                1 + 3 + body_len
            }
        }
    )
);

tls_handshake!(
    hello_request(DummyItem) = 0,
    client_hello(ClientHello) = 1,
    server_hello(ServerHello) = 2,
    // hello_verify_request(..) = 3, RFC 6347: DTLS
    // NewSessionTicket(..) = 4, RFC 5077: session resumption w/o server-side state
    certificate(CertificateList) = 11,
    server_key_exchange(ObscureData) = 12,
    certificate_request(CertificateRequest) = 13,
    server_hello_done(DummyItem) = 14,
    // certificate_verify = 15,
    client_key_exchange(ObscureData) = 16,
    finished(VerifyData) = 20,
);

tls_struct!(struct ClientHello {
    client_version: ProtocolVersion,
    random: Random,
    session_id: SessionId,
    cipher_suites: CipherSuiteVec,
    compression_methods: CompressionMethodVec,
    extensions: Option<ExtensionVec>
});

tls_struct!(struct ServerHello {
    server_version: ProtocolVersion,
    random: Random,
    session_id: SessionId,
    cipher_suite: CipherSuite,
    compression_method: CompressionMethod,
    extensions: Option<ExtensionVec>
});

tls_vec!(CertificateList = Asn1Cert(0, (1 << 24) - 1));

tls_enum!(u8, enum ClientCertificateType {
      rsa_sign(1), dss_sign(2), rsa_fixed_dh(3), dss_fixed_dh(4),
      rsa_ephemeral_dh_RESERVED(5), dss_ephemeral_dh_RESERVED(6),
      fortezza_dms_RESERVED(20)
});
tls_vec!(CertificiateTypeVec = ClientCertificateType(1, (1 << 8) - 1));

tls_vec!(DistinguishedName = u8(1, (1 << 16) - 1));
tls_vec!(DistinguishedNameVec = DistinguishedName(0, (1 << 16) - 1));

tls_struct!(struct CertificateRequest {
    certificate_types: CertificiateTypeVec,
    supported_signature_algorithms: SignatureAndHashAlgorithmVec,
    certificate_authorities: DistinguishedNameVec
});

// FIXME TLS 1.2 says the length can be longer for future ciphe suites.
tls_array!(VerifyData = [u8, ..12]);

// buffer for handshake protocol
pub struct HandshakeBuffer {
    buf: Vec<u8>,
}

impl HandshakeBuffer {
    pub fn new() -> HandshakeBuffer {
        HandshakeBuffer { buf: Vec::new() }
    }

    pub fn add_record(&mut self, fragment: Vec<u8>) {
        self.buf.push_all(&fragment);
    }

    // if message is arrived but has unknown type, the message is discarded and returns error.
    pub fn get_message(&mut self) -> TlsResult<Option<Handshake>> {
        let len = self.buf.len();
        // we need to read at least ty and length
        if len < 4 {
            return Ok(None);
        }

        let n1 = self.buf[1] as usize;
        let n2 = self.buf[2] as usize;
        let n3 = self.buf[3] as usize;
        let wanted_len: usize = (n1 << 16) | (n2 << 8) | n3;
        let wanted_len = wanted_len + 4;

        if len < wanted_len {
            return Ok(None);
        }

        // FIXME bad clone?
        let (message, remaining) = {
            let (message, remaining) = self.buf.split_at_mut(wanted_len);
            let message = message.to_vec();
            let remaining = remaining.to_vec();
            (message, remaining)
        };
        self.buf = remaining;

        let mut reader = MemReader::new(message);
        let message: Handshake = try!(TlsItem::tls_read(&mut reader));
        let ret = Ok(Some(message));

        ret
    }
}

impl Handshake {
    pub fn new_client_hello(random: Random,
                            cipher_suite: CipherSuite,
                            extensions: Vec<Extension>) -> TlsResult<Handshake> {
        let client_hello_body = {
            let client_version = {
                let (major, minor) = TLS_VERSION;

                ProtocolVersion {
                    major: major,
                    minor: minor,
                }
            };

            // TODO support session resumption
            let session_id = {
                let data = Vec::new();
                try!(SessionId::new(data))
            };

            let cipher_suites = {
                let data = vec!(cipher_suite);
                try!(CipherSuiteVec::new(data))
            };

            let compression_methods = {
                let data = vec!(CompressionMethod::null);
                try!(CompressionMethodVec::new(data))
            };

            let extensions = if extensions.len() == 0 {
                None
            } else {
                let ext = try!(ExtensionVec::new(extensions));
                Some(ext)
            };
            ClientHello {
                client_version: client_version,
                random: random,
                session_id: session_id,
                cipher_suites: cipher_suites,
                compression_methods: compression_methods,
                extensions: extensions,
            }
        };

        Ok(Handshake::client_hello(client_hello_body))
    }

    pub fn new_client_key_exchange(data: Vec<u8>) -> TlsResult<Handshake> {
        let data = ObscureData::new(data);
        Ok(Handshake::client_key_exchange(data))
    }

    pub fn new_finished(data: Vec<u8>) -> TlsResult<Handshake> {
        let data = try!(VerifyData::new(data));
        Ok(Handshake::finished(data))
    }
}

#[cfg(test)]
mod test {
    use std::old_io::MemReader;
    use tls_item::TlsItem;
    use cipher::CipherSuite;

    use super::{ProtocolVersion, SessionId, CipherSuiteVec, CompressionMethod,
                CompressionMethodVec, ClientHello, Handshake, Random};

    #[test]
    fn test_parse_client_hello() {
        let client_hello_msg = {
            let client_hello_body = {
                let client_version = {
                    let (major, minor) = (3, 3);

                    ProtocolVersion {
                        major: major,
                        minor: minor,
                    }
                };

                let random = {
                    let random_bytes = [0u8; 32].to_vec();
                    Random::new(random_bytes).unwrap()
                };

                let session_id = {
                    let data = Vec::new();
                    SessionId::new(data).unwrap()
                };

                let cipher_suites = {
                    let data = vec!(CipherSuite::TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256);
                    CipherSuiteVec::new(data).unwrap()
                };

                let compression_methods = {
                    let data = vec!(CompressionMethod::null);
                    CompressionMethodVec::new(data).unwrap()
                };

                ClientHello {
                    client_version: client_version,
                    random: random,
                    session_id: session_id,
                    cipher_suites: cipher_suites,
                    compression_methods: compression_methods,
                    extensions: None,
                }
            };

            Handshake::client_hello(client_hello_body)
        };

        let mut packet = Vec::new();
        client_hello_msg.tls_write(&mut packet).unwrap();

        let mut reader = MemReader::new(packet.clone());
        let client_hello_msg_2: Handshake = TlsItem::tls_read(&mut reader).unwrap();

        let mut packet_2 = Vec::new();
        client_hello_msg_2.tls_write(&mut packet_2).unwrap();

        assert_eq!(packet, packet_2);
    }
}
<|endoftext|>


<commit_before>extern crate getopts;
extern crate rusqlite;
extern crate pibq;

use std::env;
use std::time::{Duration, Instant};
use getopts::Options;

use pibq::bluetherm;
use pibq::sql;
use pibq::models::{ConnectionStatus, Reading};

// interval between sending query packtets, in ms
const QUERY_INTERVAL: u64 = 7000;

// if no packets are returned, amount of time to wait before creating a timeout error, in ms
const TIMEOUT_INTERVAL: u64 = 15000;

const HEARTBEAT_INTERVAL: u64 = 1000;

struct Harvester {
    sql_conn: rusqlite::Connection,
    bt_conn: Option<bluetherm::Connection>,
    serial: String,
    disconnected: bool,
    disconnect_reason: Option<bluetherm::ConnectionEvent>,
    error_count: i64,
    send_interval: Duration,
    timeout_interval: Duration,
    last_send: Option<Instant>,
    last_receive: Option<Instant>
}

impl Harvester {
    fn new(sql_conn: rusqlite::Connection, serial: &str) -> Harvester {
        Harvester {
            sql_conn: sql_conn,
            bt_conn: Some(Harvester::connect_bluetherm(serial)),
            serial: serial.to_string(),
            disconnected: true,
            disconnect_reason: None,
            error_count: 0,
            send_interval: Duration::from_millis(QUERY_INTERVAL),
            timeout_interval: Duration::from_millis(TIMEOUT_INTERVAL),
            last_send: None,
            last_receive: None
        }
    }

    fn connect_bluetherm(serial: &str) -> bluetherm::Connection {
        bluetherm::Connection::new(serial, Some(HEARTBEAT_INTERVAL)).unwrap()
    }

    fn start(&mut self) {
        loop {
            match (self.last_send, self.last_receive) {
                (Some(sent), _) if sent.elapsed() < self.send_interval => {},
                //(Some(sent), Some(received)) if sent > received => {},
                _ => {
                    self.send_packet();
                    self.last_send = Some(Instant::now());
                }
            }

            let event = self.get_bt_conn().wait().unwrap();

            println!("{}", event);
            match event {
                bluetherm::ConnectionEvent::Packet(p) => {
                    self.record_packet(p);
                    self.last_send = Some(Instant::now());
                    self.last_receive = Some(Instant::now());
                },
                e @ bluetherm::ConnectionEvent::InvalidPacket(_) => { self.bt_error(e); },
                e @ bluetherm::ConnectionEvent::ReadError(_) => { self.bt_error(e); },
                e @ bluetherm::ConnectionEvent::WriteError(_) => { self.bt_error(e); },
                e @ bluetherm::ConnectionEvent::Heartbeat => {
                    match self.last_receive {
                        None => {
                            match self.last_send {
                                None => {},
                                Some(sent) => {
                                    if sent.elapsed() > self.timeout_interval {
                                        self.bt_error(e);
                                    }
                                }
                            }
                        },
                        Some(rec) => {
                            if rec.elapsed() > self.timeout_interval {
                                self.bt_error(e);
                            }
                        }
                    }
                }
            }
        }
    }

    fn get_bt_conn(&mut self) -> &mut bluetherm::Connection {
        self.bt_conn.as_mut().unwrap()
    }

    fn record_packet(&mut self, packet: bluetherm::Packet) {
        let mut reading = Reading::new();
        reading.value1 = packet.get_sensor1_reading();
        reading.value2 = packet.get_sensor2_reading();
        sql::insert_reading(&self.sql_conn, &mut reading).unwrap();
    }

    fn send_packet(&mut self) {
        println!("sending packet...");
        let p = bluetherm::Packet::temp_packet();
        match self.get_bt_conn().send(p) {
            Err(e) => {
                self.bt_error(bluetherm::ConnectionEvent::ReadError(e));
            },
            Ok(_) => {}
        }
        println!("done sending packet.");
    }

    fn bt_error(&mut self, evt: bluetherm::ConnectionEvent) {
        let mut report_error = true;

        if self.disconnected {
            match &self.disconnect_reason {
                &None => {},
                &Some(ref reason) => {
                    match (reason, &evt) {
                        (&bluetherm::ConnectionEvent::InvalidPacket(_), &bluetherm::ConnectionEvent::InvalidPacket(_)) |
                        (&bluetherm::ConnectionEvent::ReadError(_), &bluetherm::ConnectionEvent::ReadError(_)) |
                        (&bluetherm::ConnectionEvent::WriteError(_), &bluetherm::ConnectionEvent::WriteError(_)) |
                        (&bluetherm::ConnectionEvent::Heartbeat, &bluetherm::ConnectionEvent::Heartbeat) => {
                            report_error = false;
                        },
                        _ => {},
                    }
                }
            }
        }

        if report_error {
            let msg = match evt {
                bluetherm::ConnectionEvent::InvalidPacket(ref p) => { format!("Invalid Packet: [{}]", p) },
                bluetherm::ConnectionEvent::ReadError(ref err) => { format!("Read Error: {}", err) },
                bluetherm::ConnectionEvent::WriteError(ref err) => { format!("Write Error: {}", err) },
                bluetherm::ConnectionEvent::Heartbeat => { "Timeout".to_string() },
                _ => "Unknown Error".to_string()
            };

            let mut s = ConnectionStatus::new();
            s.is_disconnect = true;
            s.info = Some(msg);
            sql::insert_connection_status(&self.sql_conn, &mut s).unwrap();
        }

        println!("error: {}", evt);

        self.disconnected = true;
        self.disconnect_reason = Some(evt);
        self.last_send = None;
        self.last_receive = None;

        self.error_count = self.error_count + 1;

        if self.error_count > 3 {
            self.error_count = 0;
            println!("killing old connection.");

            let old = self.bt_conn.take();
            drop(old.unwrap());

            println!("old conneciton dropped; making new");

            self.bt_conn = Some(Harvester::connect_bluetherm(&self.serial));

            println!("new connection made");
        }
    }
}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = Options::new();
    opts.reqopt("s", "serial", "tty serial device", "DEV");
    opts.optopt("d", "dbfile", "sqlite DB file", "FILE");
    opts.optopt("m", "migrations", "migration folder", "DIR");
    opts.optflag("h", "help", "print this help menu");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => {
            println!("{}", f.to_string());
            print_usage(&program, opts);
            return;
        }
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    let serial = matches.opt_str("s").unwrap();
    let dbfile = matches.opt_str("d").unwrap_or("pibq.sqlite".to_string());
    let migrations = matches.opt_str("m").unwrap_or("migrations".to_string());

    let db = sql::get_connection(&dbfile, Some(migrations)).unwrap();

    let mut h = Harvester::new(db, &serial);
    h.start();
}
<commit_msg>debug<commit_after>extern crate getopts;
extern crate rusqlite;
extern crate pibq;

use std::env;
use std::time::{Duration, Instant};
use getopts::Options;

use pibq::bluetherm;
use pibq::sql;
use pibq::models::{ConnectionStatus, Reading};

// interval between sending query packtets, in ms
const QUERY_INTERVAL: u64 = 7000;

// if no packets are returned, amount of time to wait before creating a timeout error, in ms
const TIMEOUT_INTERVAL: u64 = 15000;

const HEARTBEAT_INTERVAL: u64 = 1000;

struct Harvester {
    sql_conn: rusqlite::Connection,
    bt_conn: Option<bluetherm::Connection>,
    serial: String,
    disconnected: bool,
    disconnect_reason: Option<bluetherm::ConnectionEvent>,
    error_count: i64,
    send_interval: Duration,
    timeout_interval: Duration,
    last_send: Option<Instant>,
    last_receive: Option<Instant>
}

impl Harvester {
    fn new(sql_conn: rusqlite::Connection, serial: &str) -> Harvester {
        Harvester {
            sql_conn: sql_conn,
            bt_conn: Some(Harvester::connect_bluetherm(serial)),
            serial: serial.to_string(),
            disconnected: true,
            disconnect_reason: None,
            error_count: 0,
            send_interval: Duration::from_millis(QUERY_INTERVAL),
            timeout_interval: Duration::from_millis(TIMEOUT_INTERVAL),
            last_send: None,
            last_receive: None
        }
    }

    fn connect_bluetherm(serial: &str) -> bluetherm::Connection {
        bluetherm::Connection::new(serial, Some(HEARTBEAT_INTERVAL)).unwrap()
    }

    fn start(&mut self) {
        loop {
            match (self.last_send, self.last_receive) {
                (Some(sent), _) if sent.elapsed() < self.send_interval => {},
                //(Some(sent), Some(received)) if sent > received => {},
                _ => {
                    self.send_packet();
                    self.last_send = Some(Instant::now());
                }
            }

            let event = self.get_bt_conn().wait().unwrap();

            println!("{}", event);
            match event {
                bluetherm::ConnectionEvent::Packet(p) => {
                    self.record_packet(p);
                    self.last_send = Some(Instant::now());
                    self.last_receive = Some(Instant::now());
                },
                e @ bluetherm::ConnectionEvent::InvalidPacket(_) => { self.bt_error(e); },
                e @ bluetherm::ConnectionEvent::ReadError(_) => { self.bt_error(e); },
                e @ bluetherm::ConnectionEvent::WriteError(_) => { self.bt_error(e); },
                e @ bluetherm::ConnectionEvent::Heartbeat => {
                    match self.last_receive {
                        None => {
                            match self.last_send {
                                None => {},
                                Some(sent) => {
                                    if sent.elapsed() > self.timeout_interval {
                                        self.bt_error(e);
                                    }
                                }
                            }
                        },
                        Some(rec) => {
                            if rec.elapsed() > self.timeout_interval {
                                self.bt_error(e);
                            }
                        }
                    }
                }
            }
        }
    }

    fn get_bt_conn(&mut self) -> &mut bluetherm::Connection {
        self.bt_conn.as_mut().unwrap()
    }

    fn record_packet(&mut self, packet: bluetherm::Packet) {
        let mut reading = Reading::new();
        reading.value1 = packet.get_sensor1_reading();
        reading.value2 = packet.get_sensor2_reading();
        sql::insert_reading(&self.sql_conn, &mut reading).unwrap();
    }

    fn send_packet(&mut self) {
        println!("sending packet...");
        let p = bluetherm::Packet::temp_packet();
        match self.get_bt_conn().send(p) {
            Err(e) => {
                self.bt_error(bluetherm::ConnectionEvent::ReadError(e));
            },
            Ok(_) => {}
        }
        println!("done sending packet.");
    }

    fn bt_error(&mut self, evt: bluetherm::ConnectionEvent) {
        let mut report_error = true;

        if self.disconnected {
            match &self.disconnect_reason {
                &None => {},
                &Some(ref reason) => {
                    match (reason, &evt) {
                        (&bluetherm::ConnectionEvent::InvalidPacket(_), &bluetherm::ConnectionEvent::InvalidPacket(_)) |
                        (&bluetherm::ConnectionEvent::ReadError(_), &bluetherm::ConnectionEvent::ReadError(_)) |
                        (&bluetherm::ConnectionEvent::WriteError(_), &bluetherm::ConnectionEvent::WriteError(_)) |
                        (&bluetherm::ConnectionEvent::Heartbeat, &bluetherm::ConnectionEvent::Heartbeat) => {
                            report_error = false;
                        },
                        _ => {},
                    }
                }
            }
        }

        if report_error {
            let msg = match evt {
                bluetherm::ConnectionEvent::InvalidPacket(ref p) => { format!("Invalid Packet: [{}]", p) },
                bluetherm::ConnectionEvent::ReadError(ref err) => { format!("Read Error: {}", err) },
                bluetherm::ConnectionEvent::WriteError(ref err) => { format!("Write Error: {}", err) },
                bluetherm::ConnectionEvent::Heartbeat => { "Timeout".to_string() },
                _ => "Unknown Error".to_string()
            };

            let mut s = ConnectionStatus::new();
            s.is_disconnect = true;
            s.info = Some(msg);
            sql::insert_connection_status(&self.sql_conn, &mut s).unwrap();
        }

        println!("error: {}", evt);

        self.disconnected = true;
        self.disconnect_reason = Some(evt);

        self.error_count = self.error_count + 1;

        if self.error_count > 3 {
            self.error_count = 0;
            println!("killing old connection.");

            self.last_send = None;
            self.last_receive = None;

            let old = self.bt_conn.take();
            drop(old.unwrap());

            println!("old conneciton dropped; making new");

            self.bt_conn = Some(Harvester::connect_bluetherm(&self.serial));

            println!("new connection made");
        }
    }
}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = Options::new();
    opts.reqopt("s", "serial", "tty serial device", "DEV");
    opts.optopt("d", "dbfile", "sqlite DB file", "FILE");
    opts.optopt("m", "migrations", "migration folder", "DIR");
    opts.optflag("h", "help", "print this help menu");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => {
            println!("{}", f.to_string());
            print_usage(&program, opts);
            return;
        }
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    let serial = matches.opt_str("s").unwrap();
    let dbfile = matches.opt_str("d").unwrap_or("pibq.sqlite".to_string());
    let migrations = matches.opt_str("m").unwrap_or("migrations".to_string());

    let db = sql::get_connection(&dbfile, Some(migrations)).unwrap();

    let mut h = Harvester::new(db, &serial);
    h.start();
}
<|endoftext|>


<commit_before>//! Functions to perform numerical integration using Shapes
//!
//! # Definitions
//!
//! * `n_integ_point` -- number of integration (Gauss) points
//!
//! When performing numerical integrations, we use the following notation:
//! |J| is the determinant of the Jacobian, ||J|| is the norm of the Jacobian vector
//! for line in multi-dimensions, `n_integ_point` is the number of integration points,
//! `ιp := ξp` is the reference coordinate of the integration point,
//! and `wp` is the weight of the p-th integration point.

mod analytical_tet4;
mod analytical_tri3;
mod integ_points;
mod matrix_cases;
mod vector_cases;
pub use crate::integ::analytical_tet4::*;
pub use crate::integ::analytical_tri3::*;
pub use crate::integ::integ_points::*;
pub use crate::integ::matrix_cases::*;
pub use crate::integ::vector_cases::*;
<commit_msg>Improve doc<commit_after>//! Functions to perform numerical integration using Shapes
//!
//! # Definitions
//!
//! * `n_integ_point` -- number of integration (Gauss) points
//! * `|J|` --is the determinant of the Jacobian
//! * `||J||` -- is the norm of the Jacobian vector for lines in multi-dimensions
//! * `ιp := ξp` -- iota-p gets ksi-p (or xi-p) -- is the reference coordinate of the integration point
//! * `wp` -- is the weight of the p-th integration point

mod analytical_tet4;
mod analytical_tri3;
mod integ_points;
mod matrix_cases;
mod vector_cases;
pub use crate::integ::analytical_tet4::*;
pub use crate::integ::analytical_tri3::*;
pub use crate::integ::integ_points::*;
pub use crate::integ::matrix_cases::*;
pub use crate::integ::vector_cases::*;
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{self, BufRead};
use std::path::Path;
use regex::Regex;

#[macro_use]
extern crate lazy_static;
pub fn get_number<P: AsRef<Path>>(filename: P) -> io::Result<u32> {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"NUMBER=(\d+)").unwrap();
    }
    if let Ok(lines) = read_lines(filename) {
        for line in lines {
            if let Some(n) = RE.captures(&line.unwrap()) {
                return Ok(n.get(1).unwrap().as_str().parse().unwrap());
            }
        }
    }
    Ok(0)
}

fn read_lines<P: AsRef<Path>>(filename: P) -> io::Result<io::Lines<io::BufReader<File>>> {
    let file = File::open(filename)?;
    Ok(io::BufReader::new(file).lines())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn get_test_number() {
        assert_eq!(get_number("../data/test.txt").unwrap() , 3);
    }
}
<commit_msg>Added own Iterator<commit_after>use std::fs::File;
use std::io::{self, BufRead};
use std::path::Path;
use regex::Regex;

#[macro_use]
extern crate lazy_static;
pub fn get_number<P: AsRef<Path>>(filename: P) -> io::Result<u32> {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"NUMBER=(\d+)").unwrap();
    }
    if let Ok(lines) = read_lines(filename) {
        for line in lines {
            if let Some(n) = RE.captures(&line) {
                return Ok(n.get(1).unwrap().as_str().parse().unwrap());
            }
        }
    }
    Ok(0)
}

/*
fn read_lines<P: AsRef<Path>>(filename: P) -> io::Result<io::Lines<io::BufReader<File>>> {
    let file = File::open(filename)?;
    Ok(io::BufReader::new(file).lines())
}
*/

fn read_lines<P: AsRef<Path>>(filename: P) -> io::Result<impl Iterator<Item = String>> {
    let file = File::open(filename)?;
    Ok(FileReader::new(file))
}

struct FileReader {
    lines: io::Lines<io::BufReader<File>>
}

impl FileReader {
    fn new(file: File) -> FileReader {
        FileReader{ lines: io::BufReader::new(file).lines() }
    }
}

impl Iterator for FileReader {
    type Item = String;
    fn next(&mut self) -> Option<Self::Item> {
        if let Some(line) = self.lines.next() {
            return Some(line.unwrap());
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn get_test_number() {
        assert_eq!(get_number("../data/test.txt").unwrap() , 3);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2016 Sandstorm Development Group, Inc. and contributors
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use futures::{Future, Stream};
use futures::sync::oneshot;

pub struct TaskSet<T, E> {
    sender: ::futures::sync::mpsc::UnboundedSender<Box<Future<Item=T,Error=E>>>,
    _canceler: ::futures::sync::oneshot::Sender<()>, // when dropped, the tasks get canceled
}

impl<T, E> TaskSet<T, E> {
    pub fn new(_reaper: Box<TaskReaper<T, E>>, handle: &::tokio_core::reactor::Handle)
               -> TaskSet<T, E>
        where E: 'static, T: 'static, E: ::std::fmt::Debug,
    {
        let (tx, rx) = ::futures::sync::mpsc::unbounded::<Box<Future<Item=T,Error=E>>>();
        let stream = rx.map_err(|()| unreachable!())
            .buffer_unordered(1000); // XXX hack that should basically work in small cases.

        let (fulfiller, dropped) = oneshot::channel::<()>();
        let dropped = dropped.map_err(|_| ());

        let f = dropped.join(
            stream.for_each(|_t| Ok(()) ).map_err(|e| { println!("error {:?}", e); ()})).map(|_| {println!("task set done"); ()});

        handle.spawn(f);

        TaskSet {
            sender: tx,
            _canceler: fulfiller,
        }
    }

    pub fn add<F>(&mut self, promise: F)
        where F: Future<Item=T, Error=E> + 'static
    {
        self.sender.send(Box::new(promise)).unwrap();
    }
}

pub trait TaskReaper<T, E> where T: 'static, E: 'static
{
    fn task_succeeded(&mut self, _value: T) {}
    fn task_failed(&mut self, error: E);
}

<commit_msg>actually hook up the reaper<commit_after>// Copyright (c) 2013-2016 Sandstorm Development Group, Inc. and contributors
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use futures::{Future, Stream};
use futures::sync::oneshot;

use std::cell::{RefCell};
use std::rc::{Rc};

pub struct TaskSet<T, E> {
    sender: ::futures::sync::mpsc::UnboundedSender<Box<Future<Item=(),Error=()>>>,
    _canceler: ::futures::sync::oneshot::Sender<()>, // when dropped, the tasks get canceled
    reaper: Rc<RefCell<Box<TaskReaper<T, E>>>>,
}

impl<T, E> TaskSet<T, E> where T: 'static, E: 'static {
    pub fn new(reaper: Box<TaskReaper<T, E>>, handle: &::tokio_core::reactor::Handle)
               -> TaskSet<T, E>
        where E: 'static, T: 'static, E: ::std::fmt::Debug,
    {
        let (tx, rx) = ::futures::sync::mpsc::unbounded::<Box<Future<Item=(),Error=()>>>();
        let stream = rx.map_err(|()| unreachable!())
            .buffer_unordered(1000); // XXX hack that should basically work in small cases.

        let (fulfiller, dropped) = oneshot::channel::<()>();
        let dropped = dropped.map_err(|_| ());

        let f = dropped.join(
            stream.for_each(|()| Ok(()) ).map_err(|e| { println!("error {:?}", e); ()})).map(|_| {println!("task set done"); ()});

        handle.spawn(f);

        TaskSet {
            sender: tx,
            _canceler: fulfiller,
            reaper: Rc::new(RefCell::new(reaper)),
        }
    }

    pub fn add<F>(&mut self, promise: F)
        where F: Future<Item=T, Error=E> + 'static
    {
        let reaper = self.reaper.clone();
        let task = promise.then(move |r| {
            match r {
                Ok(v) => reaper.borrow_mut().task_succeeded(v),
                Err(e) => reaper.borrow_mut().task_failed(e),
            }
            Ok(())
        });
        self.sender.send(Box::new(task)).unwrap();
    }
}

pub trait TaskReaper<T, E> where T: 'static, E: 'static
{
    fn task_succeeded(&mut self, _value: T) {}
    fn task_failed(&mut self, error: E);
}

<|endoftext|>


<commit_before>#![feature(proc_macro_hygiene)]
#![feature(decl_macro)]
#[macro_use] extern crate rocket;

use rocket::request::{Outcome, FromRequest};
use rocket::Outcome::{Success};
use rocket::Request;

pub trait UseCase {
    fn run(&self) -> &'static str;
}

struct HelloWorldCase {
}

impl UseCase for HelloWorldCase {
    fn run(&self) -> &'static str {
        "Hello world"
    }
}

impl<'a, 'r> FromRequest<'a, 'r> for HelloWorldCase {
    type Error = ();
    fn from_request(request: &'a Request<'r>) -> Outcome<Self, Self::Error> {
        Success(HelloWorldCase{})
    }
}

#[get("/")]
fn hello() -> &'static str {
    HelloWorldCase{}.run()
}

fn main() {
    println!("Hello, world!");
    rocket::ignite().mount("/", routes![hello]).launch();
}

#[test]
fn test1() {
    use rocket::local::Client;
    let use_case = HelloWorldCase{};
    let rocket = rocket::ignite().mount("/", routes![hello]);
    let client = Client::new(rocket).unwrap();
    let mut response = client.get("/").dispatch();
    assert_eq!(response.body_string(), Some("Hello world".into()));
}
<commit_msg>Added db classes<commit_after>#![feature(proc_macro_hygiene)]
#![feature(decl_macro)]
#[macro_use] extern crate rocket;
#[macro_use] extern crate diesel;
extern crate dotenv;

use rocket::request::{Outcome, FromRequest};
use rocket::Outcome::{Success};
use rocket::Request;

use diesel::prelude::*;
use diesel::pg::PgConnection;
use dotenv::dotenv;
use std::env;

pub mod schema;

// DATABASE classes
pub fn establish_connection() -> PgConnection {
    dotenv().ok();
    let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
    PgConnection::establish(&database_url).expect(&format!("Error connecting to {}", database_url))
}

#[derive(Queryable)]
pub struct Post {
    pub id: i32,
    pub title: String,
    pub body: String,
    pub publihed: bool,
}

// Entities classes
pub trait UseCase {
    fn run(&self) -> &'static str;
}

struct HelloWorldCase {
}

impl UseCase for HelloWorldCase {
    fn run(&self) -> &'static str {
        "Hello world"
    }
}

impl<'a, 'r> FromRequest<'a, 'r> for HelloWorldCase {
    type Error = ();
    fn from_request(request: &'a Request<'r>) -> Outcome<Self, Self::Error> {
        Success(HelloWorldCase{})
    }
}

#[get("/")]
fn hello() -> &'static str {
    HelloWorldCase{}.run()
}

fn main() {
    println!("Hello, world!");
    rocket::ignite().mount("/", routes![hello]).launch();
}

#[test]
fn test1() {
    use rocket::local::Client;
    let use_case = HelloWorldCase{};
    let rocket = rocket::ignite().mount("/", routes![hello]);
    let client = Client::new(rocket).unwrap();
    let mut response = client.get("/").dispatch();
    assert_eq!(response.body_string(), Some("Hello world".into()));
}
<|endoftext|>


<commit_before>use super::stacker::{Addr, MemoryArena, TermHashMap};

use crate::fieldnorm::FieldNormReaders;
use crate::postings::recorder::{
    BufferLender, NothingRecorder, Recorder, TfAndPositionRecorder, TermFrequencyRecorder,
};
use crate::postings::UnorderedTermId;
use crate::postings::{FieldSerializer, InvertedIndexSerializer};
use crate::schema::IndexRecordOption;
use crate::schema::{Field, FieldEntry, FieldType, Schema, Term};
use crate::termdict::TermOrdinal;
use crate::tokenizer::TokenStream;
use crate::tokenizer::{Token, MAX_TOKEN_LEN};
use crate::DocId;
use fnv::FnvHashMap;
use std::collections::HashMap;
use std::io;
use std::marker::PhantomData;
use std::ops::{DerefMut, Range};

fn posting_from_field_entry(field_entry: &FieldEntry) -> Box<dyn PostingsWriter> {
    match *field_entry.field_type() {
        FieldType::Str(ref text_options) => text_options
            .get_indexing_options()
            .map(|indexing_options| match indexing_options.index_option() {
                IndexRecordOption::Basic => {
                    SpecializedPostingsWriter::<NothingRecorder>::new_boxed()
                }
                IndexRecordOption::WithFreqs => {
                    SpecializedPostingsWriter::<TermFrequencyRecorder>::new_boxed()
                }
                IndexRecordOption::WithFreqsAndPositions => {
                    SpecializedPostingsWriter::<TfAndPositionRecorder>::new_boxed()
                }
            })
            .unwrap_or_else(|| SpecializedPostingsWriter::<NothingRecorder>::new_boxed()),
        FieldType::U64(_)
        | FieldType::I64(_)
        | FieldType::F64(_)
        | FieldType::Date(_)
        | FieldType::Bytes(_)
        | FieldType::HierarchicalFacet(_) => {
            SpecializedPostingsWriter::<NothingRecorder>::new_boxed()
        }
    }
}

pub struct MultiFieldPostingsWriter {
    heap: MemoryArena,
    schema: Schema,
    term_index: TermHashMap,
    per_field_postings_writers: Vec<Box<dyn PostingsWriter>>,
}

fn make_field_partition(
    term_offsets: &[(&[u8], Addr, UnorderedTermId)],
) -> Vec<(Field, Range<usize>)> {
    let term_offsets_it = term_offsets
        .iter()
        .map(|(key, _, _)| Term::wrap(key).field())
        .enumerate();
    let mut prev_field_opt = None;
    let mut fields = vec![];
    let mut offsets = vec![];
    for (offset, field) in term_offsets_it {
        if Some(field) != prev_field_opt {
            prev_field_opt = Some(field);
            fields.push(field);
            offsets.push(offset);
        }
    }
    offsets.push(term_offsets.len());
    let mut field_offsets = vec![];
    for i in 0..fields.len() {
        field_offsets.push((fields[i], offsets[i]..offsets[i + 1]));
    }
    field_offsets
}

impl MultiFieldPostingsWriter {
    /// Create a new `MultiFieldPostingsWriter` given
    /// a schema and a heap.
    pub fn new(schema: &Schema, table_bits: usize) -> MultiFieldPostingsWriter {
        let term_index = TermHashMap::new(table_bits);
        let per_field_postings_writers: Vec<_> = schema
            .fields()
            .map(|(_, field_entry)| posting_from_field_entry(field_entry))
            .collect();
        MultiFieldPostingsWriter {
            heap: MemoryArena::new(),
            schema: schema.clone(),
            term_index,
            per_field_postings_writers,
        }
    }

    pub fn mem_usage(&self) -> usize {
        self.term_index.mem_usage() + self.heap.mem_usage()
    }

    pub fn index_text(
        &mut self,
        doc: DocId,
        field: Field,
        token_stream: &mut dyn TokenStream,
        term_buffer: &mut Term,
    ) -> u32 {
        let postings_writer =
            self.per_field_postings_writers[field.field_id() as usize].deref_mut();
        postings_writer.index_text(
            &mut self.term_index,
            doc,
            field,
            token_stream,
            &mut self.heap,
            term_buffer,
        )
    }

    pub fn subscribe(&mut self, doc: DocId, term: &Term) -> UnorderedTermId {
        let postings_writer =
            self.per_field_postings_writers[term.field().field_id() as usize].deref_mut();
        postings_writer.subscribe(&mut self.term_index, doc, 0u32, term, &mut self.heap)
    }

    /// Serialize the inverted index.
    /// It pushes all term, one field at a time, towards the
    /// postings serializer.
    pub fn serialize(
        &self,
        serializer: &mut InvertedIndexSerializer,
        fieldnorm_readers: FieldNormReaders,
    ) -> crate::Result<HashMap<Field, FnvHashMap<UnorderedTermId, TermOrdinal>>> {
        let mut term_offsets: Vec<(&[u8], Addr, UnorderedTermId)> =
            self.term_index.iter().collect();
        term_offsets.sort_unstable_by_key(|&(k, _, _)| k);

        let mut unordered_term_mappings: HashMap<Field, FnvHashMap<UnorderedTermId, TermOrdinal>> =
            HashMap::new();

        let field_offsets = make_field_partition(&term_offsets);

        for (field, byte_offsets) in field_offsets {
            let field_entry = self.schema.get_field_entry(field);

            match *field_entry.field_type() {
                FieldType::Str(_) | FieldType::HierarchicalFacet(_) => {
                    // populating the (unordered term ord) -> (ordered term ord) mapping
                    // for the field.
                    let unordered_term_ids = term_offsets[byte_offsets.clone()]
                        .iter()
                        .map(|&(_, _, bucket)| bucket);
                    let mapping: FnvHashMap<UnorderedTermId, TermOrdinal> = unordered_term_ids
                        .enumerate()
                        .map(|(term_ord, unord_term_id)| {
                            (unord_term_id as UnorderedTermId, term_ord as TermOrdinal)
                        })
                        .collect();
                    unordered_term_mappings.insert(field, mapping);
                }
                FieldType::U64(_) | FieldType::I64(_) | FieldType::F64(_) | FieldType::Date(_) => {}
                FieldType::Bytes(_) => {}
            }

            let postings_writer =
                self.per_field_postings_writers[field.field_id() as usize].as_ref();
            let fieldnorm_reader = fieldnorm_readers.get_field(field)?;
            let mut field_serializer = serializer.new_field(
                field,
                postings_writer.total_num_tokens(),
                fieldnorm_reader,
            )?;
            postings_writer.serialize(
                &term_offsets[byte_offsets],
                &mut field_serializer,
                &self.term_index.heap,
                &self.heap,
            )?;
            field_serializer.close()?;
        }
        Ok(unordered_term_mappings)
    }
}

/// The `PostingsWriter` is in charge of receiving documenting
/// and building a `Segment` in anonymous memory.
///
/// `PostingsWriter` writes in a `MemoryArena`.
pub trait PostingsWriter {
    /// Record that a document contains a term at a given position.
    ///
    /// * doc  - the document id
    /// * pos  - the term position (expressed in tokens)
    /// * term - the term
    /// * heap - heap used to store the postings informations as well as the terms
    /// in the hashmap.
    fn subscribe(
        &mut self,
        term_index: &mut TermHashMap,
        doc: DocId,
        pos: u32,
        term: &Term,
        heap: &mut MemoryArena,
    ) -> UnorderedTermId;

    /// Serializes the postings on disk.
    /// The actual serialization format is handled by the `PostingsSerializer`.
    fn serialize(
        &self,
        term_addrs: &[(&[u8], Addr, UnorderedTermId)],
        serializer: &mut FieldSerializer<'_>,
        term_heap: &MemoryArena,
        heap: &MemoryArena,
    ) -> io::Result<()>;

    /// Tokenize a text and subscribe all of its token.
    fn index_text(
        &mut self,
        term_index: &mut TermHashMap,
        doc_id: DocId,
        field: Field,
        token_stream: &mut dyn TokenStream,
        heap: &mut MemoryArena,
        term_buffer: &mut Term,
    ) -> u32 {
        term_buffer.set_field(field);
        let mut sink = |token: &Token| {
            // We skip all tokens with a len greater than u16.
            if token.text.len() <= MAX_TOKEN_LEN {
                term_buffer.set_text(token.text.as_str());
                self.subscribe(
                    term_index,
                    doc_id,
                    token.position as u32,
                    &term_buffer,
                    heap,
                );
            } else {
                info!(
                    "A token exceeding MAX_TOKEN_LEN ({}>{}) was dropped. Search for \
                     MAX_TOKEN_LEN in the documentation for more information.",
                    token.text.len(),
                    MAX_TOKEN_LEN
                );
            }
        };
        token_stream.process(&mut sink)
    }

    fn total_num_tokens(&self) -> u64;
}

/// The `SpecializedPostingsWriter` is just here to remove dynamic
/// dispatch to the recorder information.
pub(crate) struct SpecializedPostingsWriter<Rec: Recorder + 'static> {
    total_num_tokens: u64,
    _recorder_type: PhantomData<Rec>,
}

impl<Rec: Recorder + 'static> SpecializedPostingsWriter<Rec> {
    /// constructor
    pub fn new() -> SpecializedPostingsWriter<Rec> {
        SpecializedPostingsWriter {
            total_num_tokens: 0u64,
            _recorder_type: PhantomData,
        }
    }

    /// Builds a `SpecializedPostingsWriter` storing its data in a heap.
    pub fn new_boxed() -> Box<dyn PostingsWriter> {
        Box::new(SpecializedPostingsWriter::<Rec>::new())
    }
}

impl<Rec: Recorder + 'static> PostingsWriter for SpecializedPostingsWriter<Rec> {
    fn subscribe(
        &mut self,
        term_index: &mut TermHashMap,
        doc: DocId,
        position: u32,
        term: &Term,
        heap: &mut MemoryArena,
    ) -> UnorderedTermId {
        debug_assert!(term.as_slice().len() >= 4);
        self.total_num_tokens += 1;
        term_index.mutate_or_create(term, |opt_recorder: Option<Rec>| {
            if let Some(mut recorder) = opt_recorder {
                let current_doc = recorder.current_doc();
                if current_doc != doc {
                    recorder.close_doc(heap);
                    recorder.new_doc(doc, heap);
                }
                recorder.record_position(position, heap);
                recorder
            } else {
                let mut recorder = Rec::new();
                recorder.new_doc(doc, heap);
                recorder.record_position(position, heap);
                recorder
            }
        }) as UnorderedTermId
    }

    fn serialize(
        &self,
        term_addrs: &[(&[u8], Addr, UnorderedTermId)],
        serializer: &mut FieldSerializer<'_>,
        termdict_heap: &MemoryArena,
        heap: &MemoryArena,
    ) -> io::Result<()> {
        let mut buffer_lender = BufferLender::default();
        for &(term_bytes, addr, _) in term_addrs {
            let recorder: Rec = termdict_heap.read(addr);
            let term_doc_freq = recorder.term_doc_freq().unwrap_or(0u32);
            serializer.new_term(&term_bytes[4..], term_doc_freq)?;
            recorder.serialize(&mut buffer_lender, serializer, heap)?;
            serializer.close_term()?;
        }
        Ok(())
    }

    fn total_num_tokens(&self) -> u64 {
        self.total_num_tokens
    }
}
<commit_msg>Cargo fmt<commit_after>use super::stacker::{Addr, MemoryArena, TermHashMap};

use crate::fieldnorm::FieldNormReaders;
use crate::postings::recorder::{
    BufferLender, NothingRecorder, Recorder, TermFrequencyRecorder, TfAndPositionRecorder,
};
use crate::postings::UnorderedTermId;
use crate::postings::{FieldSerializer, InvertedIndexSerializer};
use crate::schema::IndexRecordOption;
use crate::schema::{Field, FieldEntry, FieldType, Schema, Term};
use crate::termdict::TermOrdinal;
use crate::tokenizer::TokenStream;
use crate::tokenizer::{Token, MAX_TOKEN_LEN};
use crate::DocId;
use fnv::FnvHashMap;
use std::collections::HashMap;
use std::io;
use std::marker::PhantomData;
use std::ops::{DerefMut, Range};

fn posting_from_field_entry(field_entry: &FieldEntry) -> Box<dyn PostingsWriter> {
    match *field_entry.field_type() {
        FieldType::Str(ref text_options) => text_options
            .get_indexing_options()
            .map(|indexing_options| match indexing_options.index_option() {
                IndexRecordOption::Basic => {
                    SpecializedPostingsWriter::<NothingRecorder>::new_boxed()
                }
                IndexRecordOption::WithFreqs => {
                    SpecializedPostingsWriter::<TermFrequencyRecorder>::new_boxed()
                }
                IndexRecordOption::WithFreqsAndPositions => {
                    SpecializedPostingsWriter::<TfAndPositionRecorder>::new_boxed()
                }
            })
            .unwrap_or_else(|| SpecializedPostingsWriter::<NothingRecorder>::new_boxed()),
        FieldType::U64(_)
        | FieldType::I64(_)
        | FieldType::F64(_)
        | FieldType::Date(_)
        | FieldType::Bytes(_)
        | FieldType::HierarchicalFacet(_) => {
            SpecializedPostingsWriter::<NothingRecorder>::new_boxed()
        }
    }
}

pub struct MultiFieldPostingsWriter {
    heap: MemoryArena,
    schema: Schema,
    term_index: TermHashMap,
    per_field_postings_writers: Vec<Box<dyn PostingsWriter>>,
}

fn make_field_partition(
    term_offsets: &[(&[u8], Addr, UnorderedTermId)],
) -> Vec<(Field, Range<usize>)> {
    let term_offsets_it = term_offsets
        .iter()
        .map(|(key, _, _)| Term::wrap(key).field())
        .enumerate();
    let mut prev_field_opt = None;
    let mut fields = vec![];
    let mut offsets = vec![];
    for (offset, field) in term_offsets_it {
        if Some(field) != prev_field_opt {
            prev_field_opt = Some(field);
            fields.push(field);
            offsets.push(offset);
        }
    }
    offsets.push(term_offsets.len());
    let mut field_offsets = vec![];
    for i in 0..fields.len() {
        field_offsets.push((fields[i], offsets[i]..offsets[i + 1]));
    }
    field_offsets
}

impl MultiFieldPostingsWriter {
    /// Create a new `MultiFieldPostingsWriter` given
    /// a schema and a heap.
    pub fn new(schema: &Schema, table_bits: usize) -> MultiFieldPostingsWriter {
        let term_index = TermHashMap::new(table_bits);
        let per_field_postings_writers: Vec<_> = schema
            .fields()
            .map(|(_, field_entry)| posting_from_field_entry(field_entry))
            .collect();
        MultiFieldPostingsWriter {
            heap: MemoryArena::new(),
            schema: schema.clone(),
            term_index,
            per_field_postings_writers,
        }
    }

    pub fn mem_usage(&self) -> usize {
        self.term_index.mem_usage() + self.heap.mem_usage()
    }

    pub fn index_text(
        &mut self,
        doc: DocId,
        field: Field,
        token_stream: &mut dyn TokenStream,
        term_buffer: &mut Term,
    ) -> u32 {
        let postings_writer =
            self.per_field_postings_writers[field.field_id() as usize].deref_mut();
        postings_writer.index_text(
            &mut self.term_index,
            doc,
            field,
            token_stream,
            &mut self.heap,
            term_buffer,
        )
    }

    pub fn subscribe(&mut self, doc: DocId, term: &Term) -> UnorderedTermId {
        let postings_writer =
            self.per_field_postings_writers[term.field().field_id() as usize].deref_mut();
        postings_writer.subscribe(&mut self.term_index, doc, 0u32, term, &mut self.heap)
    }

    /// Serialize the inverted index.
    /// It pushes all term, one field at a time, towards the
    /// postings serializer.
    pub fn serialize(
        &self,
        serializer: &mut InvertedIndexSerializer,
        fieldnorm_readers: FieldNormReaders,
    ) -> crate::Result<HashMap<Field, FnvHashMap<UnorderedTermId, TermOrdinal>>> {
        let mut term_offsets: Vec<(&[u8], Addr, UnorderedTermId)> =
            self.term_index.iter().collect();
        term_offsets.sort_unstable_by_key(|&(k, _, _)| k);

        let mut unordered_term_mappings: HashMap<Field, FnvHashMap<UnorderedTermId, TermOrdinal>> =
            HashMap::new();

        let field_offsets = make_field_partition(&term_offsets);

        for (field, byte_offsets) in field_offsets {
            let field_entry = self.schema.get_field_entry(field);

            match *field_entry.field_type() {
                FieldType::Str(_) | FieldType::HierarchicalFacet(_) => {
                    // populating the (unordered term ord) -> (ordered term ord) mapping
                    // for the field.
                    let unordered_term_ids = term_offsets[byte_offsets.clone()]
                        .iter()
                        .map(|&(_, _, bucket)| bucket);
                    let mapping: FnvHashMap<UnorderedTermId, TermOrdinal> = unordered_term_ids
                        .enumerate()
                        .map(|(term_ord, unord_term_id)| {
                            (unord_term_id as UnorderedTermId, term_ord as TermOrdinal)
                        })
                        .collect();
                    unordered_term_mappings.insert(field, mapping);
                }
                FieldType::U64(_) | FieldType::I64(_) | FieldType::F64(_) | FieldType::Date(_) => {}
                FieldType::Bytes(_) => {}
            }

            let postings_writer =
                self.per_field_postings_writers[field.field_id() as usize].as_ref();
            let fieldnorm_reader = fieldnorm_readers.get_field(field)?;
            let mut field_serializer = serializer.new_field(
                field,
                postings_writer.total_num_tokens(),
                fieldnorm_reader,
            )?;
            postings_writer.serialize(
                &term_offsets[byte_offsets],
                &mut field_serializer,
                &self.term_index.heap,
                &self.heap,
            )?;
            field_serializer.close()?;
        }
        Ok(unordered_term_mappings)
    }
}

/// The `PostingsWriter` is in charge of receiving documenting
/// and building a `Segment` in anonymous memory.
///
/// `PostingsWriter` writes in a `MemoryArena`.
pub trait PostingsWriter {
    /// Record that a document contains a term at a given position.
    ///
    /// * doc  - the document id
    /// * pos  - the term position (expressed in tokens)
    /// * term - the term
    /// * heap - heap used to store the postings informations as well as the terms
    /// in the hashmap.
    fn subscribe(
        &mut self,
        term_index: &mut TermHashMap,
        doc: DocId,
        pos: u32,
        term: &Term,
        heap: &mut MemoryArena,
    ) -> UnorderedTermId;

    /// Serializes the postings on disk.
    /// The actual serialization format is handled by the `PostingsSerializer`.
    fn serialize(
        &self,
        term_addrs: &[(&[u8], Addr, UnorderedTermId)],
        serializer: &mut FieldSerializer<'_>,
        term_heap: &MemoryArena,
        heap: &MemoryArena,
    ) -> io::Result<()>;

    /// Tokenize a text and subscribe all of its token.
    fn index_text(
        &mut self,
        term_index: &mut TermHashMap,
        doc_id: DocId,
        field: Field,
        token_stream: &mut dyn TokenStream,
        heap: &mut MemoryArena,
        term_buffer: &mut Term,
    ) -> u32 {
        term_buffer.set_field(field);
        let mut sink = |token: &Token| {
            // We skip all tokens with a len greater than u16.
            if token.text.len() <= MAX_TOKEN_LEN {
                term_buffer.set_text(token.text.as_str());
                self.subscribe(
                    term_index,
                    doc_id,
                    token.position as u32,
                    &term_buffer,
                    heap,
                );
            } else {
                info!(
                    "A token exceeding MAX_TOKEN_LEN ({}>{}) was dropped. Search for \
                     MAX_TOKEN_LEN in the documentation for more information.",
                    token.text.len(),
                    MAX_TOKEN_LEN
                );
            }
        };
        token_stream.process(&mut sink)
    }

    fn total_num_tokens(&self) -> u64;
}

/// The `SpecializedPostingsWriter` is just here to remove dynamic
/// dispatch to the recorder information.
pub(crate) struct SpecializedPostingsWriter<Rec: Recorder + 'static> {
    total_num_tokens: u64,
    _recorder_type: PhantomData<Rec>,
}

impl<Rec: Recorder + 'static> SpecializedPostingsWriter<Rec> {
    /// constructor
    pub fn new() -> SpecializedPostingsWriter<Rec> {
        SpecializedPostingsWriter {
            total_num_tokens: 0u64,
            _recorder_type: PhantomData,
        }
    }

    /// Builds a `SpecializedPostingsWriter` storing its data in a heap.
    pub fn new_boxed() -> Box<dyn PostingsWriter> {
        Box::new(SpecializedPostingsWriter::<Rec>::new())
    }
}

impl<Rec: Recorder + 'static> PostingsWriter for SpecializedPostingsWriter<Rec> {
    fn subscribe(
        &mut self,
        term_index: &mut TermHashMap,
        doc: DocId,
        position: u32,
        term: &Term,
        heap: &mut MemoryArena,
    ) -> UnorderedTermId {
        debug_assert!(term.as_slice().len() >= 4);
        self.total_num_tokens += 1;
        term_index.mutate_or_create(term, |opt_recorder: Option<Rec>| {
            if let Some(mut recorder) = opt_recorder {
                let current_doc = recorder.current_doc();
                if current_doc != doc {
                    recorder.close_doc(heap);
                    recorder.new_doc(doc, heap);
                }
                recorder.record_position(position, heap);
                recorder
            } else {
                let mut recorder = Rec::new();
                recorder.new_doc(doc, heap);
                recorder.record_position(position, heap);
                recorder
            }
        }) as UnorderedTermId
    }

    fn serialize(
        &self,
        term_addrs: &[(&[u8], Addr, UnorderedTermId)],
        serializer: &mut FieldSerializer<'_>,
        termdict_heap: &MemoryArena,
        heap: &MemoryArena,
    ) -> io::Result<()> {
        let mut buffer_lender = BufferLender::default();
        for &(term_bytes, addr, _) in term_addrs {
            let recorder: Rec = termdict_heap.read(addr);
            let term_doc_freq = recorder.term_doc_freq().unwrap_or(0u32);
            serializer.new_term(&term_bytes[4..], term_doc_freq)?;
            recorder.serialize(&mut buffer_lender, serializer, heap)?;
            serializer.close_term()?;
        }
        Ok(())
    }

    fn total_num_tokens(&self) -> u64 {
        self.total_num_tokens
    }
}
<|endoftext|>


<commit_before>// #![feature(scoped)]
// #![feature(collections)]

extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate dataflow_join;

extern crate docopt;
use docopt::Docopt;

use std::thread;

use dataflow_join::graph::{GraphTrait, GraphMMap};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_static::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

// static USAGE: &'static str = "
// Usage: pagerank <source> <workers>
// ";

static USAGE: &'static str = "
Usage: scc <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      list of host:port for workers
";

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    // let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
    //               else { panic!("invalid setting for workers: {}", args.get_str("<workers>")) };
    // println!("starting pagerank dataflow with {:?} worker{}", workers, if workers == 1 { "" } else { "s" });
    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_multi(communicators, source);
    }
    else if workers > 1 {
        println!("Initializing ProcessCommunicator");
        pagerank_multi(ProcessCommunicator::new_vector(workers), source);
    }
    else {
        println!("Initializing ThreadCommunicator");
        pagerank_multi(vec![ThreadCommunicator], source);
    };
}

fn pagerank_multi<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

fn pagerank<C>(communicator: C, filename: String, workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    {   // new scope avoids long borrow on root
        let mut builder = root.new_subgraph();

        // establish the beginnings of a loop,
        // 20 iterations, each time around += 1.
        let (helper, stream) = builder.loop_variable::<(u32, f32)>(RootTimestamp::new(20), Local(1));

        let graph = GraphMMap::<u32>::new(&filename);

        let nodes = graph.nodes();

        let mut src = vec![1.0; nodes / peers as usize];    // local rank accumulation
        let mut dst = vec![0.0; nodes];                     // local rank accumulation

        let mut start = time::precise_time_s();

        // from feedback, place an operator that
        // aggregates and broadcasts ranks along edges.
        let ranks = stream.enable(builder).unary_notify(

            Exchange::new(|x: &(u32, f32)| x.0 as u64),     // 1. how data should be exchanged
            format!("PageRank"),                            // 2. a tasteful, descriptive name
            vec![RootTimestamp::new(0)],                    // 3. indicate an initial capability
            move |input, output, iterator| {                // 4. provide the operator logic

                while let Some((iter, _)) = iterator.next() {
                    // /---- should look familiar! ----\
                    for node in 0..src.len() {
                        src[node] = 0.15 + 0.85 * src[node];
                    }

                    for node in 0..src.len() {
                        let edges = graph.edges(index + peers * node);
                        let value = src[node] / edges.len() as f32;
                        for &b in edges {
                            dst[b as usize] += value;
                        }
                    }
                    // \------ end familiar part ------/
                    output.give_at(&iter, dst.drain_temp()
                                             .enumerate()
                                             .filter(|&(_,f)| f != 0.0)
                                             .map(|(u,f)| (u as u32, f)));

                    // dst.resize(graph.nodes(), 0.0);
                    for _ in 0..graph.nodes() { dst.push(0.0); }

                    println!("{}s", time::precise_time_s() - start);
                    start = time::precise_time_s();
                }

                while let Some((iter, data)) = input.pull() {
                    iterator.notify_at(&iter);
                    for (node, rank) in data.drain_temp() {
                        src[node as usize / peers] += rank;
                    }
                }
            }
        );

        let local_index = index as usize % workers;

        let mut acc = vec![0.0; nodes / workers];

        ranks
        .unary_notify(
            Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
            format!("Aggregation"),
            vec![],
            move |input, output, iterator| {
                while let Some((iter, data)) = input.pull() {
                    iterator.notify_at(&iter);
                    for (node, rank) in data.drain_temp() {
                        acc[node as usize / workers] += rank;
                    }
                }

                while let Some((item, _)) = iterator.next() {

                    let mut count = 0;
                    for x in &acc { if *x != 0.0 { count += 1; } }
                    println!("about to send {} non-zero elements of {}: ({} bytes)", count, acc.len(), 8 * count);

                    output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
                                             .map(|(u,f)| (((u * workers + local_index) as u32), f)));

                    for _ in 0..(nodes/workers) { acc.push(0.0); }
                    assert!(acc.len() == nodes/workers);
                }
            }
        )
        .connect_loop(helper);
    }

    while root.step() { }
}
<commit_msg>diagnostics<commit_after>// #![feature(scoped)]
// #![feature(collections)]

extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate dataflow_join;

extern crate docopt;
use docopt::Docopt;

use std::thread;

use dataflow_join::graph::{GraphTrait, GraphMMap};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_static::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

// static USAGE: &'static str = "
// Usage: pagerank <source> <workers>
// ";

static USAGE: &'static str = "
Usage: scc <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      list of host:port for workers
";

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    // let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
    //               else { panic!("invalid setting for workers: {}", args.get_str("<workers>")) };
    // println!("starting pagerank dataflow with {:?} worker{}", workers, if workers == 1 { "" } else { "s" });
    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_multi(communicators, source);
    }
    else if workers > 1 {
        println!("Initializing ProcessCommunicator");
        pagerank_multi(ProcessCommunicator::new_vector(workers), source);
    }
    else {
        println!("Initializing ThreadCommunicator");
        pagerank_multi(vec![ThreadCommunicator], source);
    };
}

fn pagerank_multi<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

fn pagerank<C>(communicator: C, filename: String, workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    {   // new scope avoids long borrow on root
        let mut builder = root.new_subgraph();

        // establish the beginnings of a loop,
        // 20 iterations, each time around += 1.
        let (helper, stream) = builder.loop_variable::<(u32, f32)>(RootTimestamp::new(20), Local(1));

        let graph = GraphMMap::<u32>::new(&filename);

        let nodes = graph.nodes();

        let mut src = vec![1.0; nodes / peers as usize];    // local rank accumulation
        let mut dst = vec![0.0; nodes];                     // local rank accumulation

        let mut start = time::precise_time_s();

        // from feedback, place an operator that
        // aggregates and broadcasts ranks along edges.
        let ranks = stream.enable(builder).unary_notify(

            Exchange::new(|x: &(u32, f32)| x.0 as u64),     // 1. how data should be exchanged
            format!("PageRank"),                            // 2. a tasteful, descriptive name
            vec![RootTimestamp::new(0)],                    // 3. indicate an initial capability
            move |input, output, iterator| {                // 4. provide the operator logic

                while let Some((iter, _)) = iterator.next() {
                    // /---- should look familiar! ----\
                    for node in 0..src.len() {
                        src[node] = 0.15 + 0.85 * src[node];
                    }

                    for node in 0..src.len() {
                        let edges = graph.edges(index + peers * node);
                        let value = src[node] / edges.len() as f32;
                        for &b in edges {
                            dst[b as usize] += value;
                        }
                    }

                    let mut count = 0;
                    for x in &dst { if *x != 0.0 { count += 1; } }
                    println!("about to send {} non-zero elements of {}: ({} bytes)", count, dst.len(), 8 * count);

                    // \------ end familiar part ------/
                    output.give_at(&iter, dst.drain_temp()
                                             .enumerate()
                                             .filter(|&(_,f)| f != 0.0)
                                             .map(|(u,f)| (u as u32, f)));

                    // dst.resize(graph.nodes(), 0.0);
                    for _ in 0..graph.nodes() { dst.push(0.0); }

                    println!("{}s", time::precise_time_s() - start);
                    start = time::precise_time_s();
                }

                while let Some((iter, data)) = input.pull() {
                    iterator.notify_at(&iter);
                    for (node, rank) in data.drain_temp() {
                        src[node as usize / peers] += rank;
                    }
                }
            }
        );

        let local_index = index as usize % workers;

        let mut acc = vec![0.0; nodes / workers];

        ranks
        .unary_notify(
            Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
            format!("Aggregation"),
            vec![],
            move |input, output, iterator| {
                while let Some((iter, data)) = input.pull() {
                    iterator.notify_at(&iter);
                    for (node, rank) in data.drain_temp() {
                        acc[node as usize / workers] += rank;
                    }
                }

                while let Some((item, _)) = iterator.next() {

                    output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
                                             .map(|(u,f)| (((u * workers + local_index) as u32), f)));

                    for _ in 0..(nodes/workers) { acc.push(0.0); }
                    assert!(acc.len() == nodes/workers);
                }
            }
        )
        .connect_loop(helper);
    }

    while root.step() { }
}
<|endoftext|>