<commit_before>use config::Config;
use clockpro_cache::*;
use dns;
use dns::{NormalizedQuestion, NormalizedQuestionKey, DNS_CLASS_IN, DNS_RCODE_NXDOMAIN};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

#[derive(Clone, Debug)]
pub struct CacheEntry {
    pub expiration: Instant,
    pub packet: Vec<u8>,
}

impl CacheEntry {
    pub fn is_expired(&self) -> bool {
        let now = Instant::now();
        now > self.expiration
    }
}

#[derive(Clone)]
pub struct Cache {
    config: Config,
    arc_mx: Arc<Mutex<ClockProCache<NormalizedQuestionKey, CacheEntry>>>,
}

pub struct CacheStats {
    pub frequent_len: usize,
    pub recent_len: usize,
    pub test_len: usize,
    pub inserted: u64,
    pub evicted: u64,
}

impl Cache {
    pub fn new(config: Config) -> Cache {
        let arc = ClockProCache::new(config.cache_size).unwrap();
        let arc_mx = Arc::new(Mutex::new(arc));
        Cache {
            config: config,
            arc_mx: arc_mx,
        }
    }

    pub fn stats(&self) -> CacheStats {
        let cache = self.arc_mx.lock().unwrap();
        CacheStats {
            frequent_len: cache.frequent_len(),
            recent_len: cache.recent_len(),
            test_len: cache.test_len(),
            inserted: cache.inserted(),
            evicted: cache.evicted(),
        }
    }

    pub fn insert(&mut self,
                  normalized_question_key: NormalizedQuestionKey,
                  packet: Vec<u8>,
                  ttl: u32)
                  -> bool {
        debug_assert!(packet.len() >= dns::DNS_HEADER_SIZE);
        if packet.len() < dns::DNS_HEADER_SIZE {
            return false;
        }
        let now = Instant::now();
        let duration = Duration::from_secs(ttl as u64);
        let expiration = now + duration;
        let mut cache = self.arc_mx.lock().unwrap();
        let cache_entry = CacheEntry {
            expiration: expiration,
            packet: packet,
        };
        cache.insert(normalized_question_key, cache_entry)
    }

    pub fn get(&mut self, normalized_question_key: &NormalizedQuestionKey) -> Option<CacheEntry> {
        let mut cache = self.arc_mx.lock().unwrap();
        cache.get_mut(normalized_question_key).and_then(|res| Some(res.clone()))
    }

    pub fn get2(&mut self, normalized_question: &NormalizedQuestion) -> Option<CacheEntry> {
        if let Some(special_packet) = self.handle_special_queries(normalized_question) {
            Some(CacheEntry {
                expiration: Instant::now() + Duration::from_secs(self.config.max_ttl as u64),
                packet: special_packet,
            })
        } else if normalized_question.qclass != DNS_CLASS_IN {
            Some(CacheEntry {
                expiration: Instant::now() + Duration::from_secs(self.config.max_ttl as u64),
                packet: dns::build_refused_packet(normalized_question).unwrap(),
            })
        } else {
            let normalized_question_key = normalized_question.key();
            let cache_entry = self.get(&normalized_question_key);
            if let Some(mut cache_entry) = cache_entry {
                if self.config.decrement_ttl {
                    let now = Instant::now();
                    if now <= cache_entry.expiration {
                        let remaining_ttl = cache_entry.expiration.duration_since(now).as_secs();
                        let _ = dns::set_ttl(&mut cache_entry.packet, remaining_ttl as u32);
                    }
                }
                return Some(cache_entry);
            }
            if !normalized_question_key.dnssec {
                let qname = normalized_question_key.qname_lc;
                if let Some(qname_shifted) = dns::qname_shift(&qname) {
                    let mut normalized_question_key = normalized_question.key();
                    normalized_question_key.qname_lc = qname_shifted.to_owned();
                    let shifted_cache_entry = self.get(&normalized_question_key);
                    if let Some(shifted_cache_entry) = shifted_cache_entry {
                        debug!("Shifted query cached");
                        let shifted_packet = shifted_cache_entry.packet;
                        if shifted_packet.len() >= dns::DNS_HEADER_SIZE &&
                           dns::rcode(&shifted_packet) == DNS_RCODE_NXDOMAIN {
                            debug!("Shifted query returned NXDOMAIN");
                            return Some(CacheEntry {
                                expiration: shifted_cache_entry.expiration,
                                packet: dns::build_nxdomain_packet(normalized_question).unwrap(),
                            });
                        }
                    }
                }
            }
            None
        }
    }

    fn handle_special_queries(&self, normalized_question: &NormalizedQuestion) -> Option<Vec<u8>> {
        if normalized_question.qclass == dns::DNS_CLASS_IN &&
           normalized_question.qtype == dns::DNS_TYPE_ANY {
            debug!("ANY query");
            let packet = dns::build_any_packet(normalized_question, self.config.max_ttl).unwrap();
            return Some(packet);
        }
        if normalized_question.qclass == dns::DNS_CLASS_CH &&
           normalized_question.qtype == dns::DNS_TYPE_TXT {
            debug!("CHAOS TXT");
            let packet = dns::build_version_packet(normalized_question, self.config.max_ttl)
                .unwrap();
            return Some(packet);
        }
        None
    }
}
<commit_msg>Keep the lock close to the operation requiring it<commit_after>use config::Config;
use clockpro_cache::*;
use dns;
use dns::{NormalizedQuestion, NormalizedQuestionKey, DNS_CLASS_IN, DNS_RCODE_NXDOMAIN};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

#[derive(Clone, Debug)]
pub struct CacheEntry {
    pub expiration: Instant,
    pub packet: Vec<u8>,
}

impl CacheEntry {
    pub fn is_expired(&self) -> bool {
        let now = Instant::now();
        now > self.expiration
    }
}

#[derive(Clone)]
pub struct Cache {
    config: Config,
    arc_mx: Arc<Mutex<ClockProCache<NormalizedQuestionKey, CacheEntry>>>,
}

pub struct CacheStats {
    pub frequent_len: usize,
    pub recent_len: usize,
    pub test_len: usize,
    pub inserted: u64,
    pub evicted: u64,
}

impl Cache {
    pub fn new(config: Config) -> Cache {
        let arc = ClockProCache::new(config.cache_size).unwrap();
        let arc_mx = Arc::new(Mutex::new(arc));
        Cache {
            config: config,
            arc_mx: arc_mx,
        }
    }

    pub fn stats(&self) -> CacheStats {
        let cache = self.arc_mx.lock().unwrap();
        CacheStats {
            frequent_len: cache.frequent_len(),
            recent_len: cache.recent_len(),
            test_len: cache.test_len(),
            inserted: cache.inserted(),
            evicted: cache.evicted(),
        }
    }

    pub fn insert(&mut self,
                  normalized_question_key: NormalizedQuestionKey,
                  packet: Vec<u8>,
                  ttl: u32)
                  -> bool {
        debug_assert!(packet.len() >= dns::DNS_HEADER_SIZE);
        if packet.len() < dns::DNS_HEADER_SIZE {
            return false;
        }
        let now = Instant::now();
        let duration = Duration::from_secs(ttl as u64);
        let expiration = now + duration;
        let cache_entry = CacheEntry {
            expiration: expiration,
            packet: packet,
        };
        let mut cache = self.arc_mx.lock().unwrap();
        cache.insert(normalized_question_key, cache_entry)
    }

    pub fn get(&mut self, normalized_question_key: &NormalizedQuestionKey) -> Option<CacheEntry> {
        let mut cache = self.arc_mx.lock().unwrap();
        cache.get_mut(normalized_question_key).and_then(|res| Some(res.clone()))
    }

    pub fn get2(&mut self, normalized_question: &NormalizedQuestion) -> Option<CacheEntry> {
        if let Some(special_packet) = self.handle_special_queries(normalized_question) {
            Some(CacheEntry {
                expiration: Instant::now() + Duration::from_secs(self.config.max_ttl as u64),
                packet: special_packet,
            })
        } else if normalized_question.qclass != DNS_CLASS_IN {
            Some(CacheEntry {
                expiration: Instant::now() + Duration::from_secs(self.config.max_ttl as u64),
                packet: dns::build_refused_packet(normalized_question).unwrap(),
            })
        } else {
            let normalized_question_key = normalized_question.key();
            let cache_entry = self.get(&normalized_question_key);
            if let Some(mut cache_entry) = cache_entry {
                if self.config.decrement_ttl {
                    let now = Instant::now();
                    if now <= cache_entry.expiration {
                        let remaining_ttl = cache_entry.expiration.duration_since(now).as_secs();
                        let _ = dns::set_ttl(&mut cache_entry.packet, remaining_ttl as u32);
                    }
                }
                return Some(cache_entry);
            }
            if !normalized_question_key.dnssec {
                let qname = normalized_question_key.qname_lc;
                if let Some(qname_shifted) = dns::qname_shift(&qname) {
                    let mut normalized_question_key = normalized_question.key();
                    normalized_question_key.qname_lc = qname_shifted.to_owned();
                    let shifted_cache_entry = self.get(&normalized_question_key);
                    if let Some(shifted_cache_entry) = shifted_cache_entry {
                        debug!("Shifted query cached");
                        let shifted_packet = shifted_cache_entry.packet;
                        if shifted_packet.len() >= dns::DNS_HEADER_SIZE &&
                           dns::rcode(&shifted_packet) == DNS_RCODE_NXDOMAIN {
                            debug!("Shifted query returned NXDOMAIN");
                            return Some(CacheEntry {
                                expiration: shifted_cache_entry.expiration,
                                packet: dns::build_nxdomain_packet(normalized_question).unwrap(),
                            });
                        }
                    }
                }
            }
            None
        }
    }

    fn handle_special_queries(&self, normalized_question: &NormalizedQuestion) -> Option<Vec<u8>> {
        if normalized_question.qclass == dns::DNS_CLASS_IN &&
           normalized_question.qtype == dns::DNS_TYPE_ANY {
            debug!("ANY query");
            let packet = dns::build_any_packet(normalized_question, self.config.max_ttl).unwrap();
            return Some(packet);
        }
        if normalized_question.qclass == dns::DNS_CLASS_CH &&
           normalized_question.qtype == dns::DNS_TYPE_TXT {
            debug!("CHAOS TXT");
            let packet = dns::build_version_packet(normalized_question, self.config.max_ttl)
                .unwrap();
            return Some(packet);
        }
        None
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate helix;
extern crate csv;

use helix::sys;
use helix::sys::VALUE;
use helix::{UncheckedValue, CheckResult, CheckedValue, ToRust};

struct VecWrap<T>(Vec<T>);

impl<T> UncheckedValue<VecWrap<T>> for VALUE
    where VALUE: UncheckedValue<T>
{
    fn to_checked(self) -> CheckResult<VecWrap<T>> {
        if unsafe { sys::RB_TYPE_P(self, sys::T_ARRAY) } {
            // TODO: Make sure we can actually do the conversions for the values.
            let len = unsafe { sys::RARRAY_LEN(self) };
            let ptr = unsafe { sys::RARRAY_PTR(self) };
            for i in 0..len {
                let val = unsafe { *ptr.offset(i) };
                if let Err(error) = val.to_checked() {
                    return Err(format!("Failed to convert value for Vec<T>: {}", error));
                }
            }
            Ok(unsafe { CheckedValue::<VecWrap<T>>::new(self) })
        } else {
            let val = unsafe { CheckedValue::<String>::new(sys::rb_inspect(self)) };
            Err(format!("No implicit conversion of {} into Vec<String>",
                        val.to_rust()))
        }
    }
}

impl ToRust<VecWrap<String>> for CheckedValue<VecWrap<String>>
    where VALUE: UncheckedValue<String>,
          CheckedValue<String>: ToRust<String>
{
    fn to_rust(self) -> VecWrap<String> {
        let mut vec: Vec<String> = Vec::new();
        let len = unsafe { sys::RARRAY_LEN(self.inner) };
        let ptr = unsafe { sys::RARRAY_PTR(self.inner) };
        for i in 0..len {
            let val = unsafe { *ptr.offset(i) };
            let checked = val.to_checked().unwrap();
            vec.push(checked.to_rust());
        }
        return VecWrap(vec);
    }
}

ruby! {
    class RscsvWriter {
        def generate_line(row: VecWrap<String>) -> String {
            let mut writer = csv::Writer::from_memory();
            writer.write(row.0.into_iter()).unwrap();
            let result = writer.as_string();
            return result.to_owned();
        }
    }
}
<commit_msg>Remove comment<commit_after>#[macro_use]
extern crate helix;
extern crate csv;

use helix::sys;
use helix::sys::VALUE;
use helix::{UncheckedValue, CheckResult, CheckedValue, ToRust};

struct VecWrap<T>(Vec<T>);

impl<T> UncheckedValue<VecWrap<T>> for VALUE
    where VALUE: UncheckedValue<T>
{
    fn to_checked(self) -> CheckResult<VecWrap<T>> {
        if unsafe { sys::RB_TYPE_P(self, sys::T_ARRAY) } {
            let len = unsafe { sys::RARRAY_LEN(self) };
            let ptr = unsafe { sys::RARRAY_PTR(self) };
            for i in 0..len {
                let val = unsafe { *ptr.offset(i) };
                if let Err(error) = val.to_checked() {
                    return Err(format!("Failed to convert value for Vec<T>: {}", error));
                }
            }
            Ok(unsafe { CheckedValue::<VecWrap<T>>::new(self) })
        } else {
            let val = unsafe { CheckedValue::<String>::new(sys::rb_inspect(self)) };
            Err(format!("No implicit conversion of {} into Vec<String>",
                        val.to_rust()))
        }
    }
}

impl ToRust<VecWrap<String>> for CheckedValue<VecWrap<String>>
    where VALUE: UncheckedValue<String>,
          CheckedValue<String>: ToRust<String>
{
    fn to_rust(self) -> VecWrap<String> {
        let mut vec: Vec<String> = Vec::new();
        let len = unsafe { sys::RARRAY_LEN(self.inner) };
        let ptr = unsafe { sys::RARRAY_PTR(self.inner) };
        for i in 0..len {
            let val = unsafe { *ptr.offset(i) };
            let checked = val.to_checked().unwrap();
            vec.push(checked.to_rust());
        }
        return VecWrap(vec);
    }
}

ruby! {
    class RscsvWriter {
        def generate_line(row: VecWrap<String>) -> String {
            let mut writer = csv::Writer::from_memory();
            writer.write(row.0.into_iter()).unwrap();
            let result = writer.as_string();
            return result.to_owned();
        }
    }
}
<|endoftext|>


<commit_before>use crate::error::ZmqError;
use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::Display;
use tokio_util::codec::{Decoder, Encoder};

use crate::SocketType;

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqMechanism {
    NULL,
    PLAIN,
    CURVE,
}

impl Display for ZmqMechanism {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return match self {
            ZmqMechanism::NULL => write!(f, "NULL"),
            ZmqMechanism::PLAIN => write!(f, "PLAIN"),
            ZmqMechanism::CURVE => write!(f, "CURVE"),
        };
    }
}

impl TryFrom<Vec<u8>> for ZmqMechanism {
    type Error = ZmqError;

    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
        let mech = value.split(|x| *x == 0x0).next().unwrap_or(b"");
        // mechanism-char = "A"-"Z" | DIGIT
        //                  | "-" | "_" | "." | "+" | %x0
        // according to https://rfc.zeromq.org/spec:23/ZMTP/
        let mechanism = unsafe { String::from_utf8_unchecked(mech.to_vec()) };
        match mechanism.as_str() {
            "NULL" => Ok(ZmqMechanism::NULL),
            "PLAIN" => Ok(ZmqMechanism::PLAIN),
            "CURVE" => Ok(ZmqMechanism::CURVE),
            _ => Err(ZmqError::OTHER("Failed to parse ZmqMechanism")),
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct ZmqGreeting {
    pub version: (u8, u8),
    pub mechanism: ZmqMechanism,
    pub as_server: bool,
}

impl Default for ZmqGreeting {
    fn default() -> Self {
        Self {
            version: (3, 0),
            mechanism: ZmqMechanism::NULL,
            as_server: false,
        }
    }
}

impl TryFrom<Bytes> for ZmqGreeting {
    type Error = ZmqError;

    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
        if !(value[0] == 0xff && value[9] == 0x7f) {
            return Err(ZmqError::CODEC("Failed to parse greeting"));
        }
        Ok(ZmqGreeting {
            version: (value[10], value[11]),
            mechanism: ZmqMechanism::try_from(value[12..32].to_vec())?,
            as_server: value[32] == 0x01,
        })
    }
}

impl From<ZmqGreeting> for BytesMut {
    fn from(greet: ZmqGreeting) -> Self {
        let mut data: [u8; 64] = [0; 64];
        data[0] = 0xff;
        data[9] = 0x7f;
        data[10] = greet.version.0;
        data[11] = greet.version.1;
        let mech = format!("{}", greet.mechanism);
        data[12..12 + mech.len()].copy_from_slice(mech.as_bytes());
        data[32] = greet.as_server.into();
        let mut bytes = BytesMut::new();
        bytes.extend_from_slice(&data);
        bytes
    }
}

#[derive(Debug, Clone)]
pub struct ZmqMessage {
    pub data: Bytes,
    pub more: bool,
}

#[derive(Debug, Clone)]
pub(crate) enum Message {
    Greeting(ZmqGreeting),
    Command(ZmqCommand),
    Message(ZmqMessage),
    MultipartMessage(Vec<ZmqMessage>),
}

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqCommandName {
    READY,
}

impl From<ZmqCommandName> for String {
    fn from(c_name: ZmqCommandName) -> Self {
        match c_name {
            ZmqCommandName::READY => "READY".into(),
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct ZmqCommand {
    pub name: ZmqCommandName,
    pub properties: HashMap<String, String>,
}

impl ZmqCommand {
    pub fn ready(socket: SocketType) -> Self {
        let mut properties = HashMap::new();
        properties.insert("Socket-Type".into(), format!("{}", socket));
        Self {
            name: ZmqCommandName::READY,
            properties,
        }
    }
}

impl TryFrom<BytesMut> for ZmqCommand {
    type Error = ZmqError;

    fn try_from(mut buf: BytesMut) -> Result<Self, Self::Error> {
        let command_len = buf[0] as usize;
        buf.advance(1);
        // command-name-char = ALPHA according to https://rfc.zeromq.org/spec:23/ZMTP/
        let command_name =
            unsafe { String::from_utf8_unchecked(buf.split_to(command_len).to_vec()) };
        let command = match command_name.as_str() {
            "READY" => ZmqCommandName::READY,
            _ => return Err(ZmqError::CODEC("Uknown command received")),
        };
        let mut properties = HashMap::new();

        while !buf.is_empty() {
            // Collect command properties
            let prop_len = buf[0] as usize;
            buf.advance(1);
            let property = unsafe { String::from_utf8_unchecked(buf.split_to(prop_len).to_vec()) };

            use std::u32;
            let prop_val_len = unsafe {
                u32::from_be_bytes(*(buf.split_to(4).as_ptr() as *const [u8; 4])) as usize
            };
            let prop_value =
                unsafe { String::from_utf8_unchecked(buf.split_to(prop_val_len).to_vec()) };
            properties.insert(property, prop_value);
        }
        Ok(Self {
            name: command,
            properties,
        })
    }
}

impl From<ZmqCommand> for BytesMut {
    fn from(command: ZmqCommand) -> Self {
        let mut message_len = 0;

        let command_name: String = command.name.into();
        message_len += command_name.len() + 1;
        for (prop, val) in command.properties.iter() {
            message_len += prop.len() + 1;
            message_len += val.len() + 4;
        }

        let long_message = message_len > 255;

        let mut bytes = BytesMut::new();
        if long_message {
            bytes.reserve(message_len + 9);
            bytes.put_u8(0x06);
            bytes.put_u64(message_len as u64);
        } else {
            bytes.reserve(message_len + 2);
            bytes.put_u8(0x04);
            bytes.put_u8(message_len as u8);
        };
        bytes.put_u8(command_name.len() as u8);
        bytes.extend_from_slice(command_name.as_ref());
        for (prop, val) in command.properties.iter() {
            bytes.put_u8(prop.len() as u8);
            bytes.extend_from_slice(prop.as_ref());
            bytes.put_u32(val.len() as u32);
            bytes.extend_from_slice(val.as_ref());
        }
        bytes
    }
}

#[derive(Debug)]
enum DecoderState {
    Greeting,
    FrameHeader,
    Frame(bool, usize, bool),
    End,
}

pub(crate) struct ZmqCodec {
    state: DecoderState,
}

impl ZmqCodec {
    pub fn new() -> Self {
        Self {
            state: DecoderState::Greeting,
        }
    }
}

impl Decoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        if src.is_empty() {
            return Ok(None);
        }
        match self.state {
            DecoderState::Greeting => {
                if src.len() >= 64 {
                    self.state = DecoderState::FrameHeader;
                    Ok(Some(Message::Greeting(ZmqGreeting::try_from(
                        src.split_to(64).freeze(),
                    )?)))
                } else {
                    if src[0] == 0xff {
                        src.reserve(64);
                        Ok(None)
                    } else {
                        Err(ZmqError::CODEC("Bad first byte of greeting"))
                    }
                }
            }
            DecoderState::FrameHeader => {
                let flags = src[0];
                let command = (flags & 0b0000_0100) != 0;
                let long = (flags & 0b0000_0010) != 0;
                let more = ((flags & 0b0000_0001) != 0) | command;

                let frame_len = if !long && src.len() >= 2 {
                    src.advance(1);
                    src.get_u8() as usize
                } else if long && src.len() >= 9 {
                    src.advance(1);
                    src.get_u64() as usize
                } else {
                    return Ok(None);
                };
                self.state = DecoderState::Frame(command, frame_len, more);
                return self.decode(src);
            }
            DecoderState::Frame(command, frame_len, more) => {
                if src.len() < frame_len {
                    src.reserve(frame_len);
                    return Ok(None);
                }
                self.state = DecoderState::FrameHeader;
                let message = if command {
                    Message::Command(ZmqCommand::try_from(src.split_to(frame_len))?)
                } else {
                    Message::Message(ZmqMessage {
                        data: src.split_to(frame_len).freeze(),
                        more,
                    })
                };
                return Ok(Some(message));
            }
            DecoderState::End => Err(ZmqError::NO_MESSAGE),
        }
    }
}

impl ZmqCodec {
    fn _encode_message(&mut self, message: ZmqMessage, dst: &mut BytesMut) {
        let mut flags: u8 = 0;
        if message.more {
            flags |= 0b0000_0001;
        }
        let len = message.data.len();
        if len > 255 {
            flags |= 0b0000_0010;
            dst.reserve(len + 9);
        } else {
            dst.reserve(len + 2);
        }
        dst.put_u8(flags);
        if len > 255 {
            dst.put_u64(len as u64);
        } else {
            dst.put_u8(len as u8);
        }
        dst.extend_from_slice(message.data.as_ref());
    }
}

impl Encoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn encode(&mut self, message: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match message {
            Message::Greeting(payload) => dst.unsplit(payload.into()),
            Message::Message(message) => self._encode_message(message, dst),
            Message::Command(command) => dst.unsplit(command.into()),
            Message::MultipartMessage(parts) => {
                for part in parts {
                    self._encode_message(part, dst);
                }
            }
        }
        Ok(())
    }
}
<commit_msg>Unsorted refactoring<commit_after>use crate::error::ZmqError;
use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::Display;
use tokio_util::codec::{Decoder, Encoder};

use crate::SocketType;

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqMechanism {
    NULL,
    PLAIN,
    CURVE,
}

impl Display for ZmqMechanism {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return match self {
            ZmqMechanism::NULL => write!(f, "NULL"),
            ZmqMechanism::PLAIN => write!(f, "PLAIN"),
            ZmqMechanism::CURVE => write!(f, "CURVE"),
        };
    }
}

impl TryFrom<Vec<u8>> for ZmqMechanism {
    type Error = ZmqError;

    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
        let mech = value.split(|x| *x == 0x0).next().unwrap_or(b"");
        // mechanism-char = "A"-"Z" | DIGIT
        //                  | "-" | "_" | "." | "+" | %x0
        // according to https://rfc.zeromq.org/spec:23/ZMTP/
        let mechanism = unsafe { String::from_utf8_unchecked(mech.to_vec()) };
        match mechanism.as_str() {
            "NULL" => Ok(ZmqMechanism::NULL),
            "PLAIN" => Ok(ZmqMechanism::PLAIN),
            "CURVE" => Ok(ZmqMechanism::CURVE),
            _ => Err(ZmqError::OTHER("Failed to parse ZmqMechanism")),
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct ZmqGreeting {
    pub version: (u8, u8),
    pub mechanism: ZmqMechanism,
    pub as_server: bool,
}

impl Default for ZmqGreeting {
    fn default() -> Self {
        Self {
            version: (3, 0),
            mechanism: ZmqMechanism::NULL,
            as_server: false,
        }
    }
}

impl TryFrom<Bytes> for ZmqGreeting {
    type Error = ZmqError;

    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
        if !(value[0] == 0xff && value[9] == 0x7f) {
            return Err(ZmqError::CODEC("Failed to parse greeting"));
        }
        Ok(ZmqGreeting {
            version: (value[10], value[11]),
            mechanism: ZmqMechanism::try_from(value[12..32].to_vec())?,
            as_server: value[32] == 0x01,
        })
    }
}

impl From<ZmqGreeting> for BytesMut {
    fn from(greet: ZmqGreeting) -> Self {
        let mut data: [u8; 64] = [0; 64];
        data[0] = 0xff;
        data[9] = 0x7f;
        data[10] = greet.version.0;
        data[11] = greet.version.1;
        let mech = format!("{}", greet.mechanism);
        data[12..12 + mech.len()].copy_from_slice(mech.as_bytes());
        data[32] = greet.as_server.into();
        let mut bytes = BytesMut::new();
        bytes.extend_from_slice(&data);
        bytes
    }
}

#[derive(Debug, Clone)]
pub struct ZmqMessage {
    pub data: Bytes,
    pub more: bool,
}

#[derive(Debug, Clone)]
pub(crate) enum Message {
    Greeting(ZmqGreeting),
    Command(ZmqCommand),
    Message(ZmqMessage),
    MultipartMessage(Vec<ZmqMessage>),
}

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqCommandName {
    READY,
}

impl From<ZmqCommandName> for String {
    fn from(c_name: ZmqCommandName) -> Self {
        match c_name {
            ZmqCommandName::READY => "READY".into(),
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct ZmqCommand {
    pub name: ZmqCommandName,
    pub properties: HashMap<String, String>,
}

impl ZmqCommand {
    pub fn ready(socket: SocketType) -> Self {
        let mut properties = HashMap::new();
        properties.insert("Socket-Type".into(), format!("{}", socket));
        Self {
            name: ZmqCommandName::READY,
            properties,
        }
    }
}

impl TryFrom<BytesMut> for ZmqCommand {
    type Error = ZmqError;

    fn try_from(mut buf: BytesMut) -> Result<Self, Self::Error> {
        let command_len = buf.get_u8() as usize;
        // command-name-char = ALPHA according to https://rfc.zeromq.org/spec:23/ZMTP/
        let command_name =
            unsafe { String::from_utf8_unchecked(buf.split_to(command_len).to_vec()) };
        let command = match command_name.as_str() {
            "READY" => ZmqCommandName::READY,
            _ => return Err(ZmqError::CODEC("Uknown command received")),
        };
        let mut properties = HashMap::new();

        while !buf.is_empty() {
            // Collect command properties
            let prop_len = buf.get_u8() as usize;
            let property = unsafe { String::from_utf8_unchecked(buf.split_to(prop_len).to_vec()) };

            let prop_val_len = buf.get_u32() as usize;
            let prop_value =
                unsafe { String::from_utf8_unchecked(buf.split_to(prop_val_len).to_vec()) };
            properties.insert(property, prop_value);
        }
        Ok(Self {
            name: command,
            properties,
        })
    }
}

impl From<ZmqCommand> for BytesMut {
    fn from(command: ZmqCommand) -> Self {
        let mut message_len = 0;

        let command_name: String = command.name.into();
        message_len += command_name.len() + 1;
        for (prop, val) in command.properties.iter() {
            message_len += prop.len() + 1;
            message_len += val.len() + 4;
        }

        let long_message = message_len > 255;

        let mut bytes = BytesMut::new();
        if long_message {
            bytes.reserve(message_len + 9);
            bytes.put_u8(0x06);
            bytes.put_u64(message_len as u64);
        } else {
            bytes.reserve(message_len + 2);
            bytes.put_u8(0x04);
            bytes.put_u8(message_len as u8);
        };
        bytes.put_u8(command_name.len() as u8);
        bytes.extend_from_slice(command_name.as_ref());
        for (prop, val) in command.properties.iter() {
            bytes.put_u8(prop.len() as u8);
            bytes.extend_from_slice(prop.as_ref());
            bytes.put_u32(val.len() as u32);
            bytes.extend_from_slice(val.as_ref());
        }
        bytes
    }
}

#[derive(Debug)]
enum DecoderState {
    Greeting,
    FrameHeader,
    Frame(bool, usize, bool),
    End,
}

pub(crate) struct ZmqCodec {
    state: DecoderState,
}

impl ZmqCodec {
    pub fn new() -> Self {
        Self {
            state: DecoderState::Greeting,
        }
    }
}

impl Decoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        if src.is_empty() {
            return Ok(None);
        }
        match self.state {
            DecoderState::Greeting => {
                if src.len() >= 64 {
                    self.state = DecoderState::FrameHeader;
                    Ok(Some(Message::Greeting(ZmqGreeting::try_from(
                        src.split_to(64).freeze(),
                    )?)))
                } else {
                    if src[0] == 0xff {
                        src.reserve(64);
                        Ok(None)
                    } else {
                        Err(ZmqError::CODEC("Bad first byte of greeting"))
                    }
                }
            }
            DecoderState::FrameHeader => {
                let flags = src[0];
                let command = (flags & 0b0000_0100) != 0;
                let long = (flags & 0b0000_0010) != 0;
                let more = ((flags & 0b0000_0001) != 0) | command;

                let frame_len = if !long && src.len() >= 2 {
                    src.advance(1);
                    src.get_u8() as usize
                } else if long && src.len() >= 9 {
                    src.advance(1);
                    src.get_u64() as usize
                } else {
                    return Ok(None);
                };
                self.state = DecoderState::Frame(command, frame_len, more);
                return self.decode(src);
            }
            DecoderState::Frame(command, frame_len, more) => {
                if src.len() < frame_len {
                    src.reserve(frame_len);
                    return Ok(None);
                }
                self.state = DecoderState::FrameHeader;
                let message = if command {
                    Message::Command(ZmqCommand::try_from(src.split_to(frame_len))?)
                } else {
                    Message::Message(ZmqMessage {
                        data: src.split_to(frame_len).freeze(),
                        more,
                    })
                };
                return Ok(Some(message));
            }
            DecoderState::End => Err(ZmqError::NO_MESSAGE),
        }
    }
}

impl ZmqCodec {
    fn _encode_message(&mut self, message: ZmqMessage, dst: &mut BytesMut) {
        let mut flags: u8 = 0;
        if message.more {
            flags |= 0b0000_0001;
        }
        let len = message.data.len();
        if len > 255 {
            flags |= 0b0000_0010;
            dst.reserve(len + 9);
        } else {
            dst.reserve(len + 2);
        }
        dst.put_u8(flags);
        if len > 255 {
            dst.put_u64(len as u64);
        } else {
            dst.put_u8(len as u8);
        }
        dst.extend_from_slice(message.data.as_ref());
    }
}

impl Encoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn encode(&mut self, message: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match message {
            Message::Greeting(payload) => dst.unsplit(payload.into()),
            Message::Message(message) => self._encode_message(message, dst),
            Message::Command(command) => dst.unsplit(command.into()),
            Message::MultipartMessage(parts) => {
                for part in parts {
                    self._encode_message(part, dst);
                }
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use crate::error::ZmqError;
use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::Display;
use tokio_util::codec::{Decoder, Encoder};

use crate::message::*;
use crate::SocketType;

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqMechanism {
    NULL,
    PLAIN,
    CURVE,
}

impl Display for ZmqMechanism {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return match self {
            ZmqMechanism::NULL => write!(f, "NULL"),
            ZmqMechanism::PLAIN => write!(f, "PLAIN"),
            ZmqMechanism::CURVE => write!(f, "CURVE"),
        };
    }
}

impl TryFrom<Vec<u8>> for ZmqMechanism {
    type Error = ZmqError;

    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
        let mech = value.split(|x| *x == 0x0).next().unwrap_or(b"");
        // mechanism-char = "A"-"Z" | DIGIT
        //                  | "-" | "_" | "." | "+" | %x0
        // according to https://rfc.zeromq.org/spec:23/ZMTP/
        let mechanism = unsafe { String::from_utf8_unchecked(mech.to_vec()) };
        match mechanism.as_str() {
            "NULL" => Ok(ZmqMechanism::NULL),
            "PLAIN" => Ok(ZmqMechanism::PLAIN),
            "CURVE" => Ok(ZmqMechanism::CURVE),
            _ => Err(ZmqError::Other("Failed to parse ZmqMechanism")),
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct ZmqGreeting {
    pub version: (u8, u8),
    pub mechanism: ZmqMechanism,
    pub as_server: bool,
}

impl Default for ZmqGreeting {
    fn default() -> Self {
        Self {
            version: (3, 0),
            mechanism: ZmqMechanism::NULL,
            as_server: false,
        }
    }
}

impl TryFrom<Bytes> for ZmqGreeting {
    type Error = ZmqError;

    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
        if !(value[0] == 0xff && value[9] == 0x7f) {
            return Err(ZmqError::Codec("Failed to parse greeting"));
        }
        Ok(ZmqGreeting {
            version: (value[10], value[11]),
            mechanism: ZmqMechanism::try_from(value[12..32].to_vec())?,
            as_server: value[32] == 0x01,
        })
    }
}

impl From<ZmqGreeting> for BytesMut {
    fn from(greet: ZmqGreeting) -> Self {
        let mut data: [u8; 64] = [0; 64];
        data[0] = 0xff;
        data[9] = 0x7f;
        data[10] = greet.version.0;
        data[11] = greet.version.1;
        let mech = format!("{}", greet.mechanism);
        data[12..12 + mech.len()].copy_from_slice(mech.as_bytes());
        data[32] = greet.as_server.into();
        let mut bytes = BytesMut::new();
        bytes.extend_from_slice(&data);
        bytes
    }
}

#[derive(Debug, Clone)]
pub(crate) enum Message {
    Greeting(ZmqGreeting),
    Command(ZmqCommand),
    Message(ZmqMessage),
    Multipart(Vec<ZmqMessage>),
}

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqCommandName {
    READY,
}

impl From<ZmqCommandName> for String {
    fn from(c_name: ZmqCommandName) -> Self {
        match c_name {
            ZmqCommandName::READY => "READY".into(),
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct ZmqCommand {
    pub name: ZmqCommandName,
    pub properties: HashMap<String, String>,
}

impl ZmqCommand {
    pub fn ready(socket: SocketType) -> Self {
        let mut properties = HashMap::new();
        properties.insert("Socket-Type".into(), format!("{}", socket));
        Self {
            name: ZmqCommandName::READY,
            properties,
        }
    }
}

impl TryFrom<BytesMut> for ZmqCommand {
    type Error = ZmqError;

    fn try_from(mut buf: BytesMut) -> Result<Self, Self::Error> {
        let command_len = buf.get_u8() as usize;
        // command-name-char = ALPHA according to https://rfc.zeromq.org/spec:23/ZMTP/
        let command_name =
            unsafe { String::from_utf8_unchecked(buf.split_to(command_len).to_vec()) };
        let command = match command_name.as_str() {
            "READY" => ZmqCommandName::READY,
            _ => return Err(ZmqError::Codec("Uknown command received")),
        };
        let mut properties = HashMap::new();

        while !buf.is_empty() {
            // Collect command properties
            let prop_len = buf.get_u8() as usize;
            let property = unsafe { String::from_utf8_unchecked(buf.split_to(prop_len).to_vec()) };

            let prop_val_len = buf.get_u32() as usize;
            let prop_value =
                unsafe { String::from_utf8_unchecked(buf.split_to(prop_val_len).to_vec()) };
            properties.insert(property, prop_value);
        }
        Ok(Self {
            name: command,
            properties,
        })
    }
}

impl From<ZmqCommand> for BytesMut {
    fn from(command: ZmqCommand) -> Self {
        let mut message_len = 0;

        let command_name: String = command.name.into();
        message_len += command_name.len() + 1;
        for (prop, val) in command.properties.iter() {
            message_len += prop.len() + 1;
            message_len += val.len() + 4;
        }

        let long_message = message_len > 255;

        let mut bytes = BytesMut::new();
        if long_message {
            bytes.reserve(message_len + 9);
            bytes.put_u8(0x06);
            bytes.put_u64(message_len as u64);
        } else {
            bytes.reserve(message_len + 2);
            bytes.put_u8(0x04);
            bytes.put_u8(message_len as u8);
        };
        bytes.put_u8(command_name.len() as u8);
        bytes.extend_from_slice(command_name.as_ref());
        for (prop, val) in command.properties.iter() {
            bytes.put_u8(prop.len() as u8);
            bytes.extend_from_slice(prop.as_ref());
            bytes.put_u32(val.len() as u32);
            bytes.extend_from_slice(val.as_ref());
        }
        bytes
    }
}

#[derive(Debug, Clone, Copy)]
struct Frame {
    command: bool,
    long: bool,
    more: bool,
}

#[derive(Debug)]
enum DecoderState {
    Greeting,
    FrameHeader,
    FrameLen(Frame),
    Frame(Frame),
}

#[derive(Debug)]
pub(crate) struct ZmqCodec {
    state: DecoderState,
    waiting_for: usize, // Number of bytes needed to decode frame
    // Needed to store incoming multipart message
    // This allows to incapsulate it's processing inside codec and not expose
    // internal details to higher levels
    buffered_message: Option<Message>,
}

impl ZmqCodec {
    pub fn new() -> Self {
        Self {
            state: DecoderState::Greeting,
            waiting_for: 64, // len of the greeting frame,
            buffered_message: None,
        }
    }
}

impl Decoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        if src.len() < self.waiting_for {
            src.reserve(self.waiting_for - src.len());
            return Ok(None);
        }
        match self.state {
            DecoderState::Greeting => {
                if src[0] != 0xff {
                    return Err(ZmqError::Codec("Bad first byte of greeting"));
                }
                self.state = DecoderState::FrameHeader;
                self.waiting_for = 1;
                Ok(Some(Message::Greeting(ZmqGreeting::try_from(
                    src.split_to(64).freeze(),
                )?)))
            }
            DecoderState::FrameHeader => {
                let flags = src.get_u8();
                let command = (flags & 0b0000_0100) != 0;
                let long = (flags & 0b0000_0010) != 0;
                let more = (flags & 0b0000_0001) != 0;

                if more && self.buffered_message.is_none() {
                    self.buffered_message = Some(Message::Multipart(Vec::new()));
                }
                let frame = Frame {
                    command,
                    long,
                    more,
                };
                self.state = DecoderState::FrameLen(frame);
                self.waiting_for = if frame.long { 8 } else { 1 };
                return self.decode(src);
            }
            DecoderState::FrameLen(frame) => {
                self.state = DecoderState::Frame(frame);
                self.waiting_for = if frame.long {
                    src.get_u64() as usize
                } else {
                    src.get_u8() as usize
                };
                return self.decode(src);
            }
            DecoderState::Frame(frame) => {
                let data = src.split_to(self.waiting_for);
                self.state = DecoderState::FrameHeader;
                self.waiting_for = 1;
                if frame.command {
                    Ok(Some(Message::Command(ZmqCommand::try_from(data)?)))
                } else if frame.more {
                    // cache incoming multipart message
                    match &mut self.buffered_message {
                        Some(Message::Multipart(message)) => message.push(data.into()),
                        _ => panic!("Corrupted decoder state"),
                    }
                    return self.decode(src);
                } else {
                    if let Some(Message::Multipart(mut message)) = self.buffered_message.take() {
                        message.push(data.into());
                        Ok(Some(Message::Multipart(message)))
                    } else {
                        Ok(Some(Message::Message(data.into())))
                    }
                }
            }
        }
    }
}

impl ZmqCodec {
    fn _encode_message(&mut self, message: ZmqMessage, dst: &mut BytesMut, more: bool) {
        let mut flags: u8 = 0;
        if more {
            flags |= 0b0000_0001;
        }
        let len = message.data.len();
        if len > 255 {
            flags |= 0b0000_0010;
            dst.reserve(len + 9);
        } else {
            dst.reserve(len + 2);
        }
        dst.put_u8(flags);
        if len > 255 {
            dst.put_u64(len as u64);
        } else {
            dst.put_u8(len as u8);
        }
        dst.extend_from_slice(message.data.as_ref());
    }
}

impl Encoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn encode(&mut self, message: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match message {
            Message::Greeting(payload) => dst.unsplit(payload.into()),
            Message::Message(message) => self._encode_message(message, dst, false),
            Message::Command(command) => dst.unsplit(command.into()),
            Message::Multipart(parts) => {
                let last_element = parts.len() - 1;
                for (idx, part) in parts.into_iter().enumerate() {
                    self._encode_message(part, dst, idx != last_element);
                }
            }
        }
        Ok(())
    }
}
<commit_msg>Collapsed nested if<commit_after>use crate::error::ZmqError;
use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::Display;
use tokio_util::codec::{Decoder, Encoder};

use crate::message::*;
use crate::SocketType;

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqMechanism {
    NULL,
    PLAIN,
    CURVE,
}

impl Display for ZmqMechanism {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return match self {
            ZmqMechanism::NULL => write!(f, "NULL"),
            ZmqMechanism::PLAIN => write!(f, "PLAIN"),
            ZmqMechanism::CURVE => write!(f, "CURVE"),
        };
    }
}

impl TryFrom<Vec<u8>> for ZmqMechanism {
    type Error = ZmqError;

    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
        let mech = value.split(|x| *x == 0x0).next().unwrap_or(b"");
        // mechanism-char = "A"-"Z" | DIGIT
        //                  | "-" | "_" | "." | "+" | %x0
        // according to https://rfc.zeromq.org/spec:23/ZMTP/
        let mechanism = unsafe { String::from_utf8_unchecked(mech.to_vec()) };
        match mechanism.as_str() {
            "NULL" => Ok(ZmqMechanism::NULL),
            "PLAIN" => Ok(ZmqMechanism::PLAIN),
            "CURVE" => Ok(ZmqMechanism::CURVE),
            _ => Err(ZmqError::Other("Failed to parse ZmqMechanism")),
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct ZmqGreeting {
    pub version: (u8, u8),
    pub mechanism: ZmqMechanism,
    pub as_server: bool,
}

impl Default for ZmqGreeting {
    fn default() -> Self {
        Self {
            version: (3, 0),
            mechanism: ZmqMechanism::NULL,
            as_server: false,
        }
    }
}

impl TryFrom<Bytes> for ZmqGreeting {
    type Error = ZmqError;

    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
        if !(value[0] == 0xff && value[9] == 0x7f) {
            return Err(ZmqError::Codec("Failed to parse greeting"));
        }
        Ok(ZmqGreeting {
            version: (value[10], value[11]),
            mechanism: ZmqMechanism::try_from(value[12..32].to_vec())?,
            as_server: value[32] == 0x01,
        })
    }
}

impl From<ZmqGreeting> for BytesMut {
    fn from(greet: ZmqGreeting) -> Self {
        let mut data: [u8; 64] = [0; 64];
        data[0] = 0xff;
        data[9] = 0x7f;
        data[10] = greet.version.0;
        data[11] = greet.version.1;
        let mech = format!("{}", greet.mechanism);
        data[12..12 + mech.len()].copy_from_slice(mech.as_bytes());
        data[32] = greet.as_server.into();
        let mut bytes = BytesMut::new();
        bytes.extend_from_slice(&data);
        bytes
    }
}

#[derive(Debug, Clone)]
pub(crate) enum Message {
    Greeting(ZmqGreeting),
    Command(ZmqCommand),
    Message(ZmqMessage),
    Multipart(Vec<ZmqMessage>),
}

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqCommandName {
    READY,
}

impl From<ZmqCommandName> for String {
    fn from(c_name: ZmqCommandName) -> Self {
        match c_name {
            ZmqCommandName::READY => "READY".into(),
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct ZmqCommand {
    pub name: ZmqCommandName,
    pub properties: HashMap<String, String>,
}

impl ZmqCommand {
    pub fn ready(socket: SocketType) -> Self {
        let mut properties = HashMap::new();
        properties.insert("Socket-Type".into(), format!("{}", socket));
        Self {
            name: ZmqCommandName::READY,
            properties,
        }
    }
}

impl TryFrom<BytesMut> for ZmqCommand {
    type Error = ZmqError;

    fn try_from(mut buf: BytesMut) -> Result<Self, Self::Error> {
        let command_len = buf.get_u8() as usize;
        // command-name-char = ALPHA according to https://rfc.zeromq.org/spec:23/ZMTP/
        let command_name =
            unsafe { String::from_utf8_unchecked(buf.split_to(command_len).to_vec()) };
        let command = match command_name.as_str() {
            "READY" => ZmqCommandName::READY,
            _ => return Err(ZmqError::Codec("Uknown command received")),
        };
        let mut properties = HashMap::new();

        while !buf.is_empty() {
            // Collect command properties
            let prop_len = buf.get_u8() as usize;
            let property = unsafe { String::from_utf8_unchecked(buf.split_to(prop_len).to_vec()) };

            let prop_val_len = buf.get_u32() as usize;
            let prop_value =
                unsafe { String::from_utf8_unchecked(buf.split_to(prop_val_len).to_vec()) };
            properties.insert(property, prop_value);
        }
        Ok(Self {
            name: command,
            properties,
        })
    }
}

impl From<ZmqCommand> for BytesMut {
    fn from(command: ZmqCommand) -> Self {
        let mut message_len = 0;

        let command_name: String = command.name.into();
        message_len += command_name.len() + 1;
        for (prop, val) in command.properties.iter() {
            message_len += prop.len() + 1;
            message_len += val.len() + 4;
        }

        let long_message = message_len > 255;

        let mut bytes = BytesMut::new();
        if long_message {
            bytes.reserve(message_len + 9);
            bytes.put_u8(0x06);
            bytes.put_u64(message_len as u64);
        } else {
            bytes.reserve(message_len + 2);
            bytes.put_u8(0x04);
            bytes.put_u8(message_len as u8);
        };
        bytes.put_u8(command_name.len() as u8);
        bytes.extend_from_slice(command_name.as_ref());
        for (prop, val) in command.properties.iter() {
            bytes.put_u8(prop.len() as u8);
            bytes.extend_from_slice(prop.as_ref());
            bytes.put_u32(val.len() as u32);
            bytes.extend_from_slice(val.as_ref());
        }
        bytes
    }
}

#[derive(Debug, Clone, Copy)]
struct Frame {
    command: bool,
    long: bool,
    more: bool,
}

#[derive(Debug)]
enum DecoderState {
    Greeting,
    FrameHeader,
    FrameLen(Frame),
    Frame(Frame),
}

#[derive(Debug)]
pub(crate) struct ZmqCodec {
    state: DecoderState,
    waiting_for: usize, // Number of bytes needed to decode frame
    // Needed to store incoming multipart message
    // This allows to incapsulate it's processing inside codec and not expose
    // internal details to higher levels
    buffered_message: Option<Message>,
}

impl ZmqCodec {
    pub fn new() -> Self {
        Self {
            state: DecoderState::Greeting,
            waiting_for: 64, // len of the greeting frame,
            buffered_message: None,
        }
    }
}

impl Decoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        if src.len() < self.waiting_for {
            src.reserve(self.waiting_for - src.len());
            return Ok(None);
        }
        match self.state {
            DecoderState::Greeting => {
                if src[0] != 0xff {
                    return Err(ZmqError::Codec("Bad first byte of greeting"));
                }
                self.state = DecoderState::FrameHeader;
                self.waiting_for = 1;
                Ok(Some(Message::Greeting(ZmqGreeting::try_from(
                    src.split_to(64).freeze(),
                )?)))
            }
            DecoderState::FrameHeader => {
                let flags = src.get_u8();
                let command = (flags & 0b0000_0100) != 0;
                let long = (flags & 0b0000_0010) != 0;
                let more = (flags & 0b0000_0001) != 0;

                if more && self.buffered_message.is_none() {
                    self.buffered_message = Some(Message::Multipart(Vec::new()));
                }
                let frame = Frame {
                    command,
                    long,
                    more,
                };
                self.state = DecoderState::FrameLen(frame);
                self.waiting_for = if frame.long { 8 } else { 1 };
                return self.decode(src);
            }
            DecoderState::FrameLen(frame) => {
                self.state = DecoderState::Frame(frame);
                self.waiting_for = if frame.long {
                    src.get_u64() as usize
                } else {
                    src.get_u8() as usize
                };
                return self.decode(src);
            }
            DecoderState::Frame(frame) => {
                let data = src.split_to(self.waiting_for);
                self.state = DecoderState::FrameHeader;
                self.waiting_for = 1;
                if frame.command {
                    Ok(Some(Message::Command(ZmqCommand::try_from(data)?)))
                } else if frame.more {
                    // cache incoming multipart message
                    match &mut self.buffered_message {
                        Some(Message::Multipart(message)) => message.push(data.into()),
                        _ => panic!("Corrupted decoder state"),
                    }
                    return self.decode(src);
                } else if let Some(Message::Multipart(mut message)) = self.buffered_message.take() {
                    message.push(data.into());
                    Ok(Some(Message::Multipart(message)))
                } else {
                    Ok(Some(Message::Message(data.into())))
                }
            }
        }
    }
}

impl ZmqCodec {
    fn _encode_message(&mut self, message: ZmqMessage, dst: &mut BytesMut, more: bool) {
        let mut flags: u8 = 0;
        if more {
            flags |= 0b0000_0001;
        }
        let len = message.data.len();
        if len > 255 {
            flags |= 0b0000_0010;
            dst.reserve(len + 9);
        } else {
            dst.reserve(len + 2);
        }
        dst.put_u8(flags);
        if len > 255 {
            dst.put_u64(len as u64);
        } else {
            dst.put_u8(len as u8);
        }
        dst.extend_from_slice(message.data.as_ref());
    }
}

impl Encoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn encode(&mut self, message: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match message {
            Message::Greeting(payload) => dst.unsplit(payload.into()),
            Message::Message(message) => self._encode_message(message, dst, false),
            Message::Command(command) => dst.unsplit(command.into()),
            Message::Multipart(parts) => {
                let last_element = parts.len() - 1;
                for (idx, part) in parts.into_iter().enumerate() {
                    self._encode_message(part, dst, idx != last_element);
                }
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//! A type-checked scaling factor between units.

use num::One;

use num_traits::NumCast;
#[cfg(feature = "serde")]
use serde;
use core::fmt;
use core::ops::{Add, Div, Mul, Neg, Sub};
use core::marker::PhantomData;
use {Point2D, Rect, Size2D, Vector2D};

/// A scaling factor between two different units of measurement.
///
/// This is effectively a type-safe float, intended to be used in combination with other types like
/// `length::Length` to enforce conversion between systems of measurement at compile time.
///
/// `Src` and `Dst` represent the units before and after multiplying a value by a `Scale`. They
/// may be types without values, such as empty enums.  For example:
///
/// ```rust
/// use euclid::Scale;
/// use euclid::Length;
/// enum Mm {};
/// enum Inch {};
///
/// let mm_per_inch: Scale<f32, Inch, Mm> = Scale::new(25.4);
///
/// let one_foot: Length<f32, Inch> = Length::new(12.0);
/// let one_foot_in_mm: Length<f32, Mm> = one_foot * mm_per_inch;
/// ```
#[repr(C)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", serde(bound(serialize = "T: serde::Serialize", deserialize = "T: serde::Deserialize<'de>")))]
pub struct Scale<T, Src, Dst>(pub T, #[doc(hidden)] pub PhantomData<(Src, Dst)>);

impl<T, Src, Dst> Scale<T, Src, Dst> {
    #[inline]
    pub const fn new(x: T) -> Self {
        Scale(x, PhantomData)
    }
}

impl<T: Clone, Src, Dst> Scale<T, Src, Dst> {
    #[inline]
    pub fn get(&self) -> T {
        self.0.clone()
    }
}

impl<Src, Dst> Scale<f32, Src, Dst> {
    /// Identity scaling, could be used to safely transit from one space to another.
    pub const ONE: Self = Scale(1.0, PhantomData);
}

impl<T: Clone + One + Div<T, Output = T>, Src, Dst> Scale<T, Src, Dst> {
    /// The inverse Scale (1.0 / self).
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::Scale;
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let cm_per_mm: Scale<f32, Cm, Mm> = Scale::new(0.1);
    ///
    /// assert_eq!(cm_per_mm.inv(), Scale::new(10.0));
    /// ```
    pub fn inv(&self) -> Scale<T, Dst, Src> {
        let one: T = One::one();
        Scale::new(one / self.get())
    }
}

// scale0 * scale1
impl<T: Mul<T, Output = T>, A, B, C> Mul<Scale<T, B, C>> for Scale<T, A, B> {
    type Output = Scale<T, A, C>;
    #[inline]
    fn mul(self, other: Scale<T, B, C>) -> Scale<T, A, C> {
        Scale::new(self.0 * other.0)
    }
}

// scale0 + scale1
impl<T: Add<T, Output = T>, Src, Dst> Add for Scale<T, Src, Dst> {
    type Output = Scale<T, Src, Dst>;
    #[inline]
    fn add(self, other: Scale<T, Src, Dst>) -> Scale<T, Src, Dst> {
        Scale::new(self.0 + other.0)
    }
}

// scale0 - scale1
impl<T: Sub<T, Output = T>, Src, Dst> Sub for Scale<T, Src, Dst> {
    type Output = Scale<T, Src, Dst>;
    #[inline]
    fn sub(self, other: Scale<T, Src, Dst>) -> Scale<T, Src, Dst> {
        Scale::new(self.0 - other.0)
    }
}

impl<T: NumCast + Clone, Src, Dst> Scale<T, Src, Dst> {
    /// Cast from one numeric representation to another, preserving the units.
    ///
    /// # Panics
    ///
    /// If the source value cannot be represented by the target type `NewT`, then
    /// method panics. Use `try_cast` if that must be case.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::Scale;
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    ///
    /// assert_eq!(to_mm.cast::<f32>(), Scale::new(10.0));
    /// ```
    /// That conversion will panic, because `i32` not enough to store such big numbers:
    /// ```rust,should_panic
    /// use euclid::Scale;
    /// enum Mm {};// millimeter = 10^-2 meters
    /// enum Em {};// exameter   = 10^18 meters
    ///
    /// // Panics
    /// let to_em: Scale<i32, Mm, Em> = Scale::new(10e20).cast();
    /// ```
    #[inline]
    pub fn cast<NewT: NumCast>(&self) -> Scale<NewT, Src, Dst> {
        self.try_cast().unwrap()
    }

    /// Fallible cast from one numeric representation to another, preserving the units.
    /// If the source value cannot be represented by the target type `NewT`, then `None`
    /// is returned.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::Scale;
    /// enum Mm {};
    /// enum Cm {};
    /// enum Em {};// Exameter = 10^18 meters
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    /// let to_em: Scale<f32, Mm, Em> = Scale::new(10e20);
    ///
    /// assert_eq!(to_mm.try_cast::<f32>(), Some(Scale::new(10.0)));
    /// // Integer to small to store that number
    /// assert_eq!(to_em.try_cast::<i32>(), None);
    /// ```
    pub fn try_cast<NewT: NumCast>(&self) -> Option<Scale<NewT, Src, Dst>> {
        NumCast::from(self.get()).map(Scale::new)
    }
}

impl<T, Src, Dst> Scale<T, Src, Dst>
where
    T: Copy + Mul<T, Output = T> + Neg<Output = T> + PartialEq + One,
{
    /// Returns the given point transformed by this scale.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::{Scale, point2};
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    ///
    /// assert_eq!(to_mm.transform_point(point2(42, -42)), point2(420, -420));
    /// ```
    #[inline]
    pub fn transform_point(&self, point: Point2D<T, Src>) -> Point2D<T, Dst> {
        Point2D::new(point.x * self.get(), point.y * self.get())
    }

    /// Returns the given vector transformed by this scale.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::{Scale, vec2};
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    ///
    /// assert_eq!(to_mm.transform_vector(vec2(42, -42)), vec2(420, -420));
    /// ```
    #[inline]
    pub fn transform_vector(&self, vec: Vector2D<T, Src>) -> Vector2D<T, Dst> {
        Vector2D::new(vec.x * self.get(), vec.y * self.get())
    }

    /// Returns the given vector transformed by this scale.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::{Scale, size2};
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    ///
    /// assert_eq!(to_mm.transform_size(size2(42, -42)), size2(420, -420));
    /// ```
    #[inline]
    pub fn transform_size(&self, size: Size2D<T, Src>) -> Size2D<T, Dst> {
        Size2D::new(size.width * self.get(), size.height * self.get())
    }

    /// Returns the given rect transformed by this scale.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::{Scale, rect};
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    ///
    /// assert_eq!(to_mm.transform_rect(&rect(1, 2, 42, -42)), rect(10, 20, 420, -420));
    /// ```
    #[inline]
    pub fn transform_rect(&self, rect: &Rect<T, Src>) -> Rect<T, Dst> {
        Rect::new(
            self.transform_point(rect.origin),
            self.transform_size(rect.size),
        )
    }

    /// Returns the inverse of this scale.
    #[inline]
    pub fn inverse(&self) -> Scale<T, Dst, Src> {
        Scale::new(-self.get())
    }

    /// Returns `true` if this scale has no effect.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::Scale;
    /// use euclid::num::One;
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let cm_per_mm: Scale<f32, Mm, Cm> = Scale::new(0.1);
    /// let mm_per_mm: Scale<f32, Mm, Mm> = Scale::new(1.0);
    ///
    /// assert_eq!(cm_per_mm.is_identity(), false);
    /// assert_eq!(mm_per_mm.is_identity(), true);
    /// ```
    #[inline]
    pub fn is_identity(&self) -> bool {
        self.0 == T::one()
    }
}

// FIXME: Switch to `derive(PartialEq, Clone)` after this Rust issue is fixed:
// https://github.com/rust-lang/rust/issues/26925

impl<T: PartialEq, Src, Dst> PartialEq for Scale<T, Src, Dst> {
    fn eq(&self, other: &Scale<T, Src, Dst>) -> bool {
        self.0 == other.0
    }
}

impl<T: Clone, Src, Dst> Clone for Scale<T, Src, Dst> {
    fn clone(&self) -> Scale<T, Src, Dst> {
        Scale::new(self.get())
    }
}

impl<T: Copy, Src, Dst> Copy for Scale<T, Src, Dst> {}

impl<T: fmt::Debug, Src, Dst> fmt::Debug for Scale<T, Src, Dst> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl<T: fmt::Display, Src, Dst> fmt::Display for Scale<T, Src, Dst> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

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

    enum Inch {}
    enum Cm {}
    enum Mm {}

    #[test]
    fn test_scale() {
        let mm_per_inch: Scale<f32, Inch, Mm> = Scale::new(25.4);
        let cm_per_mm: Scale<f32, Mm, Cm> = Scale::new(0.1);

        let mm_per_cm: Scale<f32, Cm, Mm> = cm_per_mm.inv();
        assert_eq!(mm_per_cm.get(), 10.0);

        let cm_per_inch: Scale<f32, Inch, Cm> = mm_per_inch * cm_per_mm;
        assert_eq!(cm_per_inch, Scale::new(2.54));

        let a: Scale<isize, Inch, Inch> = Scale::new(2);
        let b: Scale<isize, Inch, Inch> = Scale::new(3);
        assert!(a != b);
        assert_eq!(a, a.clone());
        assert_eq!(a.clone() + b.clone(), Scale::new(5));
        assert_eq!(a - b, Scale::new(-1));
    }
}
<commit_msg>Auto merge of #410 - Mingun:scale, r=nical<commit_after>// Copyright 2014 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//! A type-checked scaling factor between units.

use num::One;

use num_traits::NumCast;
#[cfg(feature = "serde")]
use serde;
use core::fmt;
use core::ops::{Add, Div, Mul, Neg, Sub};
use core::hash::{Hash, Hasher};
use core::marker::PhantomData;
use core::cmp::Ordering;
use {Point2D, Rect, Size2D, Vector2D};

/// A scaling factor between two different units of measurement.
///
/// This is effectively a type-safe float, intended to be used in combination with other types like
/// `length::Length` to enforce conversion between systems of measurement at compile time.
///
/// `Src` and `Dst` represent the units before and after multiplying a value by a `Scale`. They
/// may be types without values, such as empty enums.  For example:
///
/// ```rust
/// use euclid::Scale;
/// use euclid::Length;
/// enum Mm {};
/// enum Inch {};
///
/// let mm_per_inch: Scale<f32, Inch, Mm> = Scale::new(25.4);
///
/// let one_foot: Length<f32, Inch> = Length::new(12.0);
/// let one_foot_in_mm: Length<f32, Mm> = one_foot * mm_per_inch;
/// ```
#[repr(C)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", serde(bound(serialize = "T: serde::Serialize", deserialize = "T: serde::Deserialize<'de>")))]
pub struct Scale<T, Src, Dst>(pub T, #[doc(hidden)] pub PhantomData<(Src, Dst)>);

impl<T, Src, Dst> Scale<T, Src, Dst> {
    #[inline]
    pub const fn new(x: T) -> Self {
        Scale(x, PhantomData)
    }

    /// Returns the given point transformed by this scale.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::{Scale, point2};
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    ///
    /// assert_eq!(to_mm.transform_point(point2(42, -42)), point2(420, -420));
    /// ```
    #[inline]
    pub fn transform_point(&self, point: Point2D<T, Src>) -> Point2D<T::Output, Dst>
        where T: Clone + Mul
    {
        Point2D::new(point.x * self.get(), point.y * self.get())
    }

    /// Returns the given vector transformed by this scale.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::{Scale, vec2};
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    ///
    /// assert_eq!(to_mm.transform_vector(vec2(42, -42)), vec2(420, -420));
    /// ```
    #[inline]
    pub fn transform_vector(&self, vec: Vector2D<T, Src>) -> Vector2D<T::Output, Dst>
        where T: Clone + Mul
    {
        Vector2D::new(vec.x * self.get(), vec.y * self.get())
    }

    /// Returns the given vector transformed by this scale.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::{Scale, size2};
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    ///
    /// assert_eq!(to_mm.transform_size(size2(42, -42)), size2(420, -420));
    /// ```
    #[inline]
    pub fn transform_size(&self, size: Size2D<T, Src>) -> Size2D<T::Output, Dst>
        where T: Clone + Mul
    {
        Size2D::new(size.width * self.get(), size.height * self.get())
    }

    /// Returns the given rect transformed by this scale.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::{Scale, rect};
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    ///
    /// assert_eq!(to_mm.transform_rect(&rect(1, 2, 42, -42)), rect(10, 20, 420, -420));
    /// ```
    #[inline]
    pub fn transform_rect(&self, rect: &Rect<T, Src>) -> Rect<T::Output, Dst>
        where T: Copy + Mul
    {
        Rect::new(
            self.transform_point(rect.origin),
            self.transform_size(rect.size),
        )
    }

    /// Returns the inverse of this scale.
    #[inline]
    pub fn inverse(&self) -> Scale<T::Output, Dst, Src>
        where T: Clone + Neg
    {
        Scale::new(-self.get())
    }

    /// Returns `true` if this scale has no effect.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::Scale;
    /// use euclid::num::One;
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let cm_per_mm: Scale<f32, Mm, Cm> = Scale::new(0.1);
    /// let mm_per_mm: Scale<f32, Mm, Mm> = Scale::new(1.0);
    ///
    /// assert_eq!(cm_per_mm.is_identity(), false);
    /// assert_eq!(mm_per_mm.is_identity(), true);
    /// assert_eq!(mm_per_mm, Scale::one());
    /// ```
    #[inline]
    pub fn is_identity(&self) -> bool
        where T: PartialEq + One
    {
        self.0 == T::one()
    }
}

impl<T: Clone, Src, Dst> Scale<T, Src, Dst> {
    #[inline]
    pub fn get(&self) -> T {
        self.0.clone()
    }

    /// The inverse Scale (1.0 / self).
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::Scale;
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let cm_per_mm: Scale<f32, Cm, Mm> = Scale::new(0.1);
    ///
    /// assert_eq!(cm_per_mm.inv(), Scale::new(10.0));
    /// ```
    pub fn inv(&self) -> Scale<T::Output, Dst, Src>
    where
        T: One + Div
    {
        let one: T = One::one();
        Scale::new(one / self.0.clone())
    }
}

impl<T: NumCast + Clone, Src, Dst> Scale<T, Src, Dst> {
    /// Cast from one numeric representation to another, preserving the units.
    ///
    /// # Panics
    ///
    /// If the source value cannot be represented by the target type `NewT`, then
    /// method panics. Use `try_cast` if that must be case.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::Scale;
    /// enum Mm {};
    /// enum Cm {};
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    ///
    /// assert_eq!(to_mm.cast::<f32>(), Scale::new(10.0));
    /// ```
    /// That conversion will panic, because `i32` not enough to store such big numbers:
    /// ```rust,should_panic
    /// use euclid::Scale;
    /// enum Mm {};// millimeter = 10^-2 meters
    /// enum Em {};// exameter   = 10^18 meters
    ///
    /// // Panics
    /// let to_em: Scale<i32, Mm, Em> = Scale::new(10e20).cast();
    /// ```
    #[inline]
    pub fn cast<NewT: NumCast>(&self) -> Scale<NewT, Src, Dst> {
        self.try_cast().unwrap()
    }

    /// Fallible cast from one numeric representation to another, preserving the units.
    /// If the source value cannot be represented by the target type `NewT`, then `None`
    /// is returned.
    ///
    /// # Example
    ///
    /// ```rust
    /// use euclid::Scale;
    /// enum Mm {};
    /// enum Cm {};
    /// enum Em {};// Exameter = 10^18 meters
    ///
    /// let to_mm: Scale<i32, Cm, Mm> = Scale::new(10);
    /// let to_em: Scale<f32, Mm, Em> = Scale::new(10e20);
    ///
    /// assert_eq!(to_mm.try_cast::<f32>(), Some(Scale::new(10.0)));
    /// // Integer to small to store that number
    /// assert_eq!(to_em.try_cast::<i32>(), None);
    /// ```
    pub fn try_cast<NewT: NumCast>(&self) -> Option<Scale<NewT, Src, Dst>> {
        NumCast::from(self.get()).map(Scale::new)
    }
}

impl<Src, Dst> Scale<f32, Src, Dst> {
    /// Identity scaling, could be used to safely transit from one space to another.
    pub const ONE: Self = Scale(1.0, PhantomData);
}

// scale0 * scale1
// (A,B) * (B,C) = (A,C)
impl<T: Mul, A, B, C> Mul<Scale<T, B, C>> for Scale<T, A, B> {
    type Output = Scale<T::Output, A, C>;

    #[inline]
    fn mul(self, other: Scale<T, B, C>) -> Self::Output {
        Scale::new(self.0 * other.0)
    }
}

// scale0 + scale1
impl<T: Add, Src, Dst> Add for Scale<T, Src, Dst> {
    type Output = Scale<T::Output, Src, Dst>;

    #[inline]
    fn add(self, other: Scale<T, Src, Dst>) -> Self::Output {
        Scale::new(self.0 + other.0)
    }
}

// scale0 - scale1
impl<T: Sub, Src, Dst> Sub for Scale<T, Src, Dst> {
    type Output = Scale<T::Output, Src, Dst>;

    #[inline]
    fn sub(self, other: Scale<T, Src, Dst>) -> Self::Output {
        Scale::new(self.0 - other.0)
    }
}

// FIXME: Switch to `derive(PartialEq, Clone)` after this Rust issue is fixed:
// https://github.com/rust-lang/rust/issues/26925

impl<T: PartialEq, Src, Dst> PartialEq for Scale<T, Src, Dst> {
    fn eq(&self, other: &Scale<T, Src, Dst>) -> bool {
        self.0 == other.0
    }
}

impl<T: Eq, Src, Dst> Eq for Scale<T, Src, Dst> {}

impl<T: PartialOrd, Src, Dst> PartialOrd for Scale<T, Src, Dst> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.0.partial_cmp(&other.0)
    }
}

impl<T: Ord, Src, Dst> Ord for Scale<T, Src, Dst> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.cmp(&other.0)
    }
}

impl<T: Clone, Src, Dst> Clone for Scale<T, Src, Dst> {
    fn clone(&self) -> Scale<T, Src, Dst> {
        Scale::new(self.get())
    }
}

impl<T: Copy, Src, Dst> Copy for Scale<T, Src, Dst> {}

impl<T: fmt::Debug, Src, Dst> fmt::Debug for Scale<T, Src, Dst> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl<T: fmt::Display, Src, Dst> fmt::Display for Scale<T, Src, Dst> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl<T: Default, Src, Dst> Default for Scale<T, Src, Dst> {
    fn default() -> Self {
        Self::new(T::default())
    }
}

impl<T: Hash, Src, Dst> Hash for Scale<T, Src, Dst> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.hash(state)
    }
}

impl<T: One, Src, Dst> One for Scale<T, Src, Dst> {
    #[inline]
    fn one() -> Self {
        Scale::new(T::one())
    }
}

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

    enum Inch {}
    enum Cm {}
    enum Mm {}

    #[test]
    fn test_scale() {
        let mm_per_inch: Scale<f32, Inch, Mm> = Scale::new(25.4);
        let cm_per_mm: Scale<f32, Mm, Cm> = Scale::new(0.1);

        let mm_per_cm: Scale<f32, Cm, Mm> = cm_per_mm.inv();
        assert_eq!(mm_per_cm.get(), 10.0);

        let one: Scale<f32, Mm, Mm> = cm_per_mm * mm_per_cm;
        assert_eq!(one.get(), 1.0);

        let one: Scale<f32, Cm, Cm> = mm_per_cm * cm_per_mm;
        assert_eq!(one.get(), 1.0);

        let cm_per_inch: Scale<f32, Inch, Cm> = mm_per_inch * cm_per_mm;
        //  mm     cm     cm
        // ---- x ---- = ----
        // inch    mm    inch
        assert_eq!(cm_per_inch, Scale::new(2.54));

        let a: Scale<isize, Inch, Inch> = Scale::new(2);
        let b: Scale<isize, Inch, Inch> = Scale::new(3);
        assert_ne!(a, b);
        assert_eq!(a, a.clone());
        assert_eq!(a.clone() + b.clone(), Scale::new(5));
        assert_eq!(a - b, Scale::new(-1));
    }
}
<|endoftext|>


<commit_before>use std::path::Path;

use cache::{Cache, Get};
use resource::{Load, Reload};

/// The scene type.
///
/// This type gathers all the required objects a scene needs to correctly handle and render all
/// visual effects.
pub struct Scene<'a> {
  /// Cache.
  cache: Cache<'a>
}

impl<'a> Scene<'a> {
  pub fn new<P>(root: P) -> Self where P: AsRef<Path>{
    Scene {
      cache: Cache::new(root),
    }
  }

  pub fn get_id<T>(&mut self, name: &str, args: <T as Load>::Args) -> Option<<Cache<'a> as Get<T>>::Id> where Cache<'a>: Get<T>, T: Reload {
    self.cache.get_id(name, args)
  }

  pub fn get_by_id<T>(&mut self, id: &<Cache<'a> as Get<T>>::Id) -> Option<&T> where Cache<'a>: Get<T>, T: Reload {
    self.cache.get_by_id(id)
  }

  pub fn get<T>(&mut self, name: &str, args: <T as Load>::Args) -> Option<&T> where Cache<'a>: Get<T>, T: Reload {
    self.cache.get(name, args)
  }
}
<commit_msg>Added a get_id! macro to scene.rs.<commit_after>use std::path::Path;

use cache::{Cache, Get};
use resource::{Load, Reload};

/// The scene type.
///
/// This type gathers all the required objects a scene needs to correctly handle and render all
/// visual effects.
pub struct Scene<'a> {
  /// Cache.
  cache: Cache<'a>
}

impl<'a> Scene<'a> {
  pub fn new<P>(root: P) -> Self where P: AsRef<Path>{
    Scene {
      cache: Cache::new(root),
    }
  }

  pub fn get_id<T>(&mut self, name: &str, args: <T as Load>::Args) -> Option<<Cache<'a> as Get<T>>::Id> where Cache<'a>: Get<T>, T: Reload {
    self.cache.get_id(name, args)
  }

  pub fn get_by_id<T>(&mut self, id: &<Cache<'a> as Get<T>>::Id) -> Option<&T> where Cache<'a>: Get<T>, T: Reload {
    self.cache.get_by_id(id)
  }

  pub fn get<T>(&mut self, name: &str, args: <T as Load>::Args) -> Option<&T> where Cache<'a>: Get<T>, T: Reload {
    self.cache.get(name, args)
  }
}

#[macro_export]
macro_rules! get_id {
  ($scene:ident, $name:ident) => {
    $scene.get_id($name, ())
  };

  ($scene:ident, $name:ident, $($arg:ident),*) => {
    $scene.get_id($name, ($($arg),*))
  }
}
<|endoftext|>