<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::io::{Read, Write};

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use log::warn;

use fleetspeak_proto::common::{Message};
use fleetspeak_proto::channel::{StartupData};

use super::{ReadError, WriteError};

/// A Fleetspeak client communication packet.
///
/// This structure represents incoming or outgoing packet objects delivered by
/// Fleetspeak. This is a simplified version of the underlying Protocol Buffers
/// message that exposes too much irrelevant fields and makes the protocol easy
/// to misuse.
pub struct Packet<M> {
    /// A name of the server-side service that sent or should receive the data.
    pub service: String,
    /// An optional message type that can be used by the server-side service.
    pub kind: Option<String>,
    /// A message to sent to the specified service.
    pub data: M,
}

/// Executes the handshake procedure.
///
/// The handshake procedure consists of writing and reading magic numbers from
/// the connection buffers. This validates that the communication between the
/// Fleetspeak client and the service daemon is working as expected.
///
/// All Fleetspeak connection buffers are required to perform the handshake
/// before they became usable for sending and receiving messages.
pub fn handshake<R, W>(input: &mut R, output: &mut W) -> std::io::Result<()>
where
    R: Read,
    W: Write,
{
    write_magic(output)?;
    output.flush()?;
    read_magic(input)?;

    Ok(())
}

/// Sends a heartbeat information through this connection.
///
/// All client services should heartbeat from time to time. Otherwise, from
/// the Fleetspeak perspective, the service is unresponsive and should be
/// restarted.
///
/// The exact frequency of the required heartbeat is defined in the service
/// configuration file.
pub fn heartbeat<W>(output: &mut W) -> Result<(), WriteError>
where
    W: Write,
{
    let mut msg = Message::new();
    msg.set_message_type(String::from("Heartbeat"));
    msg.mut_destination().set_service_name(String::from("system"));

    write_raw(output, msg)
}

/// Sends the startup information through this connection.
///
/// All clients are required to send this information on startup. If the
/// client does not receive this information quickly enough, the service
/// will be killed.
///
/// The `version` string should contain a self-reported version of the
/// service. This data is used primarily for statistics.
pub fn startup<W>(output: &mut W, version: &str) -> Result<(), WriteError>
where
    W: Write,
{
    use protobuf::Message as _;

    let mut data = StartupData::new();
    data.set_pid(i64::from(std::process::id()));
    data.set_version(String::from(version));

    let mut msg = Message::new();
    msg.set_message_type(String::from("StartupData"));
    msg.mut_destination().set_service_name(String::from("system"));
    msg.mut_data().set_type_url(type_url(&data));
    msg.mut_data().set_value(data.write_to_bytes()?);

    write_raw(output, msg)
}

/// Sends the message to the Fleetspeak server through the output buffer.
///
/// The message is sent to the server-side `service` and tagged with the
/// `kind` type. Note that this message type is rather irrelevant for
/// Fleetspeak and it is up to the service what to do with this information.
pub fn send<W, M>(output: &mut W, packet: Packet<M>) -> Result<(), WriteError>
where
    W: Write,
    M: protobuf::Message,
{
    let mut msg = Message::new();
    msg.set_message_type(packet.kind.unwrap_or_else(String::new));
    msg.mut_destination().set_service_name(packet.service);
    msg.mut_data().set_type_url(type_url(&packet.data));
    msg.mut_data().set_value(packet.data.write_to_bytes()?);

    write_raw(output, msg)
}

/// Receives the message from the Fleetspeak server through the input buffer.
///
/// This function will block until there is a message to be read in the
/// input. Errors are reported in case of any I/O failure or if the read
/// message was malformed (e.g. it cannot be parsed to the expected type).
pub fn receive<R, M>(input: &mut R) -> Result<Packet<M>, ReadError>
where
    R: Read,
    M: protobuf::Message,
{
    let mut msg = read_raw(input)?;

    // While missing source address might not be consider a critical error
    // in most cases, for our own sanity we just disregard such messages.
    // Allowing such behaviour might indicate a more severe problem with
    // Fleetspeak and ignoring it simply masks the issue. This might be
    // reconsidered in the future.
    let service = if msg.has_source() {
        msg.take_source().take_service_name()
    } else {
        return Err(ReadError::malformed("missing source address"))
    };

    // It is not clear what is the best approach here. If there is no data,
    // should we error-out or return a default value? For the time being we
    // stick to the default approach, but if this proves to be not working
    // well in practice, it might be reconsidered.
    let data = if msg.has_data() {
        msg.take_data()
    } else {
        warn!(target: "fleetspeak", "empty message from '{}'", service);
        Default::default()
    };

    Ok(Packet {
        service: service,
        kind: Some(msg.message_type),
        data: protobuf::Message::parse_from_bytes(&data.value[..])?,
    })
}

/// Writes a raw Fleetspeak message to the output buffer.
///
/// This method does not perform any validation of the message being emitted
/// and assumes that all the required fields are present.
///
/// Note that this call will fail only if the message cannot be written to
/// the output or cannot be properly encoded but will succeed even if the
/// message is not what the server expects.
fn write_raw<W>(output: &mut W, msg: Message) -> Result<(), WriteError>
where
    W: Write,
{
    use protobuf::Message as _;

    output.write_u32::<LittleEndian>(msg.compute_size())?;
    msg.write_to_writer(output)?;
    write_magic(output)?;
    output.flush()?;

    Ok(())
}

/// Reads a raw Fleetspeeak message from the input buffer.
///
/// This function will block until there is a message to be read from the
/// input. It will fail in case of any I/O error or if the message cannot
/// be parsed as a Fleetspeak message.
fn read_raw<R>(input: &mut R) -> Result<Message, ReadError>
where
    R: Read,
{
    let len = input.read_u32::<LittleEndian>()? as usize;
    let mut buf = vec!(0; len);

    input.read_exact(&mut buf[..])?;
    read_magic(input)?;

    Ok(protobuf::Message::parse_from_bytes(&buf[..])?)
}

/// Writes the Fleetspeak magic to the output buffer.
fn write_magic<W>(output: &mut W) -> Result<(), WriteError>
where
    W: Write,
{
    output.write_u32::<LittleEndian>(MAGIC)?;

    Ok(())
}

/// Reads the Fleetspeak magic from the input buffer.
fn read_magic<R>(input: &mut R) -> Result<(), ReadError>
where
    R: Read,
{
    let magic = input.read_u32::<LittleEndian>()?;
    if magic != MAGIC {
        return Err(ReadError::Magic(magic));
    }

    Ok(())
}

const MAGIC: u32 = 0xf1ee1001;

// Computes a type URL of the given Protocol Buffers message.
//
// This function should probably be part of the `protobuf` package but for some
// reason it is not and we have to implement it ourselves.
fn type_url<M: protobuf::Message>(message: &M) -> String {
    format!("{}/{}", TYPE_URL_PREFIX, message.descriptor().full_name())
}

const TYPE_URL_PREFIX: &'static str = "type.googleapis.com";

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

    #[test]
    fn handshake_good_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(MAGIC).is_ok());

        let mut cur_in = Cursor::new(&mut buf_in[..]);
        let mut cur_out = Cursor::new(&mut buf_out[..]);
        assert!(handshake(&mut cur_in, &mut cur_out).is_ok());

        let mut cur = Cursor::new(&mut buf_out[..]);
        assert_eq!(cur.read_u32::<LittleEndian>().unwrap(), MAGIC);
    }

    #[test]
    fn handshake_bad_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(0xf1ee1337).is_ok());

        let mut cur_in = Cursor::new(&mut buf_in[..]);
        let mut cur_out = Cursor::new(&mut buf_out[..]);
        assert!(handshake(&mut cur_in, &mut cur_out).is_err());
    }

    #[test]
    fn type_url_startup_data() {
        assert_eq! {
            type_url(&StartupData::new()),
            "type.googleapis.com/fleetspeak.channel.StartupData"
        };
    }
}
<commit_msg>Remove unnecessary top-level imports.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::io::{Read, Write};

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

use super::{ReadError, WriteError};

/// A Fleetspeak client communication packet.
///
/// This structure represents incoming or outgoing packet objects delivered by
/// Fleetspeak. This is a simplified version of the underlying Protocol Buffers
/// message that exposes too much irrelevant fields and makes the protocol easy
/// to misuse.
pub struct Packet<M> {
    /// A name of the server-side service that sent or should receive the data.
    pub service: String,
    /// An optional message type that can be used by the server-side service.
    pub kind: Option<String>,
    /// A message to sent to the specified service.
    pub data: M,
}

/// Executes the handshake procedure.
///
/// The handshake procedure consists of writing and reading magic numbers from
/// the connection buffers. This validates that the communication between the
/// Fleetspeak client and the service daemon is working as expected.
///
/// All Fleetspeak connection buffers are required to perform the handshake
/// before they became usable for sending and receiving messages.
pub fn handshake<R, W>(input: &mut R, output: &mut W) -> std::io::Result<()>
where
    R: Read,
    W: Write,
{
    write_magic(output)?;
    output.flush()?;
    read_magic(input)?;

    Ok(())
}

/// Sends a heartbeat information through this connection.
///
/// All client services should heartbeat from time to time. Otherwise, from
/// the Fleetspeak perspective, the service is unresponsive and should be
/// restarted.
///
/// The exact frequency of the required heartbeat is defined in the service
/// configuration file.
pub fn heartbeat<W>(output: &mut W) -> Result<(), WriteError>
where
    W: Write,
{
    let mut msg = fleetspeak_proto::common::Message::new();
    msg.set_message_type(String::from("Heartbeat"));
    msg.mut_destination().set_service_name(String::from("system"));

    write_raw(output, msg)
}

/// Sends the startup information through this connection.
///
/// All clients are required to send this information on startup. If the
/// client does not receive this information quickly enough, the service
/// will be killed.
///
/// The `version` string should contain a self-reported version of the
/// service. This data is used primarily for statistics.
pub fn startup<W>(output: &mut W, version: &str) -> Result<(), WriteError>
where
    W: Write,
{
    use protobuf::Message as _;

    let mut data = fleetspeak_proto::channel::StartupData::new();
    data.set_pid(i64::from(std::process::id()));
    data.set_version(String::from(version));

    let mut msg = fleetspeak_proto::common::Message::new();
    msg.set_message_type(String::from("StartupData"));
    msg.mut_destination().set_service_name(String::from("system"));
    msg.mut_data().set_type_url(type_url(&data));
    msg.mut_data().set_value(data.write_to_bytes()?);

    write_raw(output, msg)
}

/// Sends the message to the Fleetspeak server through the output buffer.
///
/// The message is sent to the server-side `service` and tagged with the
/// `kind` type. Note that this message type is rather irrelevant for
/// Fleetspeak and it is up to the service what to do with this information.
pub fn send<W, M>(output: &mut W, packet: Packet<M>) -> Result<(), WriteError>
where
    W: Write,
    M: protobuf::Message,
{
    let mut msg = fleetspeak_proto::common::Message::new();
    msg.set_message_type(packet.kind.unwrap_or_else(String::new));
    msg.mut_destination().set_service_name(packet.service);
    msg.mut_data().set_type_url(type_url(&packet.data));
    msg.mut_data().set_value(packet.data.write_to_bytes()?);

    write_raw(output, msg)
}

/// Receives the message from the Fleetspeak server through the input buffer.
///
/// This function will block until there is a message to be read in the
/// input. Errors are reported in case of any I/O failure or if the read
/// message was malformed (e.g. it cannot be parsed to the expected type).
pub fn receive<R, M>(input: &mut R) -> Result<Packet<M>, ReadError>
where
    R: Read,
    M: protobuf::Message,
{
    let mut msg = read_raw(input)?;

    // While missing source address might not be consider a critical error
    // in most cases, for our own sanity we just disregard such messages.
    // Allowing such behaviour might indicate a more severe problem with
    // Fleetspeak and ignoring it simply masks the issue. This might be
    // reconsidered in the future.
    let service = if msg.has_source() {
        msg.take_source().take_service_name()
    } else {
        return Err(ReadError::malformed("missing source address"))
    };

    // It is not clear what is the best approach here. If there is no data,
    // should we error-out or return a default value? For the time being we
    // stick to the default approach, but if this proves to be not working
    // well in practice, it might be reconsidered.
    let data = if msg.has_data() {
        msg.take_data()
    } else {
        log::warn!(target: "fleetspeak", "empty message from '{}'", service);
        Default::default()
    };

    Ok(Packet {
        service: service,
        kind: Some(msg.message_type),
        data: protobuf::Message::parse_from_bytes(&data.value[..])?,
    })
}

/// Writes a raw Fleetspeak message to the output buffer.
///
/// This method does not perform any validation of the message being emitted
/// and assumes that all the required fields are present.
///
/// Note that this call will fail only if the message cannot be written to
/// the output or cannot be properly encoded but will succeed even if the
/// message is not what the server expects.
fn write_raw<W>(output: &mut W, msg: fleetspeak_proto::common::Message) -> Result<(), WriteError>
where
    W: Write,
{
    use protobuf::Message as _;

    output.write_u32::<LittleEndian>(msg.compute_size())?;
    msg.write_to_writer(output)?;
    write_magic(output)?;
    output.flush()?;

    Ok(())
}

/// Reads a raw Fleetspeeak message from the input buffer.
///
/// This function will block until there is a message to be read from the
/// input. It will fail in case of any I/O error or if the message cannot
/// be parsed as a Fleetspeak message.
fn read_raw<R>(input: &mut R) -> Result<fleetspeak_proto::common::Message, ReadError>
where
    R: Read,
{
    let len = input.read_u32::<LittleEndian>()? as usize;
    let mut buf = vec!(0; len);

    input.read_exact(&mut buf[..])?;
    read_magic(input)?;

    Ok(protobuf::Message::parse_from_bytes(&buf[..])?)
}

/// Writes the Fleetspeak magic to the output buffer.
fn write_magic<W>(output: &mut W) -> Result<(), WriteError>
where
    W: Write,
{
    output.write_u32::<LittleEndian>(MAGIC)?;

    Ok(())
}

/// Reads the Fleetspeak magic from the input buffer.
fn read_magic<R>(input: &mut R) -> Result<(), ReadError>
where
    R: Read,
{
    let magic = input.read_u32::<LittleEndian>()?;
    if magic != MAGIC {
        return Err(ReadError::Magic(magic));
    }

    Ok(())
}

const MAGIC: u32 = 0xf1ee1001;

// Computes a type URL of the given Protocol Buffers message.
//
// This function should probably be part of the `protobuf` package but for some
// reason it is not and we have to implement it ourselves.
fn type_url<M: protobuf::Message>(message: &M) -> String {
    format!("{}/{}", TYPE_URL_PREFIX, message.descriptor().full_name())
}

const TYPE_URL_PREFIX: &'static str = "type.googleapis.com";

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

    #[test]
    fn handshake_good_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(MAGIC).is_ok());

        let mut cur_in = Cursor::new(&mut buf_in[..]);
        let mut cur_out = Cursor::new(&mut buf_out[..]);
        assert!(handshake(&mut cur_in, &mut cur_out).is_ok());

        let mut cur = Cursor::new(&mut buf_out[..]);
        assert_eq!(cur.read_u32::<LittleEndian>().unwrap(), MAGIC);
    }

    #[test]
    fn handshake_bad_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(0xf1ee1337).is_ok());

        let mut cur_in = Cursor::new(&mut buf_in[..]);
        let mut cur_out = Cursor::new(&mut buf_out[..]);
        assert!(handshake(&mut cur_in, &mut cur_out).is_err());
    }

    #[test]
    fn type_url_startup_data() {
        assert_eq! {
            type_url(&fleetspeak_proto::channel::StartupData::new()),
            "type.googleapis.com/fleetspeak.channel.StartupData"
        };
    }
}
<|endoftext|>


<commit_before>use std::io::{Error,ErrorKind,Read,Result,Write};
use std::str;
use std::iter::repeat;
use std::collections::HashMap;
use amq_protocol_types::AMQPValue;
use nom::{HexDisplay,IResult,Offset};
use sasl::{ChannelBinding, Credentials, Secret, Mechanism};
use sasl::mechanisms::Plain;

use format::frame::*;
use format::field::*;
use channel::{Channel,ChannelState};
use buffer::Buffer;
use generated::*;

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ConnectionState {
  Initial,
  Connecting(ConnectingState),
  Connected,
  Closing(ClosingState),
  Closed,
  Error,
}

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ConnectingState {
  Initial,
  SentProtocolHeader,
  ReceivedStart,
  SentStartOk,
  ReceivedSecure,
  SentSecure,
  ReceivedSecondSecure,
  ReceivedTune,
  SentTuneOk,
  SentOpen,
  Error,
}

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ClosingState {
  Initial,
  SentClose,
  ReceivedClose,
  SentCloseOk,
  ReceivedCloseOk,
  Error,
}

#[derive(Clone,Debug,PartialEq)]
pub struct Configuration {
  channel_max: u16,
  frame_max:   u32,
  heartbeat:   u16,
}

#[derive(Clone,Debug,PartialEq)]
pub struct Connection {
  pub state:            ConnectionState,
  pub channels:         HashMap<u16, Channel>,
  pub send_buffer:      Buffer,
  pub receive_buffer:   Buffer,
  pub configuration:    Configuration,
}

impl Connection {
  pub fn new() -> Connection {
    let mut h = HashMap::new();
    h.insert(0, Channel::global());

    let configuration = Configuration {
      channel_max: 0,
      frame_max:   8192,
      heartbeat:   60,
    };

    Connection {
      state:          ConnectionState::Initial,
      channels:       h,
      send_buffer:    Buffer::with_capacity(configuration.frame_max as usize),
      receive_buffer: Buffer::with_capacity(configuration.frame_max as usize),
      configuration:  configuration,
    }
  }

  pub fn connect(&mut self, writer: &mut Write) -> Result<ConnectionState> {
    if self.state != ConnectionState::Initial {
      self.state = ConnectionState::Error;
      return Err(Error::new(ErrorKind::Other, "invalid state"))
    }

    let res = gen_protocol_header((self.send_buffer.space(), 0)).map(|tup| tup.1);
    if let Ok(sz) = res {
      self.send_buffer.fill(sz);
      match writer.write(&mut self.send_buffer.data()) {
        Ok(sz2) => {
          self.send_buffer.consume(sz2);
          self.state = ConnectionState::Connecting(ConnectingState::SentProtocolHeader);
          Ok(self.state)
        },
        Err(e) => Err(e),
      }
    } else {
      Err(Error::new(ErrorKind::WouldBlock, "could not write protocol header"))
    }
  }

  pub fn write(&mut self, writer: &mut Write) -> Result<ConnectionState> {
    println!("will write:\n{}", (&self.send_buffer.data()).to_hex(16));
    match writer.write(&mut self.send_buffer.data()) {
      Ok(sz) => {
        println!("wrote {} bytes", sz);
        self.send_buffer.consume(sz);
        Ok(self.state)
      },
      Err(e) => Err(e),
    }
  }

  pub fn read(&mut self, reader: &mut Read) -> Result<ConnectionState> {
    if self.state == ConnectionState::Initial || self.state == ConnectionState::Error {
      self.state = ConnectionState::Error;
      return Err(Error::new(ErrorKind::Other, "invalid state"))
    }

    match reader.read(&mut self.receive_buffer.space()) {
      Ok(sz) => {
        println!("read {} bytes", sz);
        self.receive_buffer.fill(sz);
      },
      Err(e) => return Err(e),
    }
    println!("will parse:\n{}", (&self.receive_buffer.data()).to_hex(16));
    let (channel_id, method, consumed) = {
      let parsed_frame = raw_frame(&self.receive_buffer.data());
      match parsed_frame {
        IResult::Done(i,_)     => {},
        IResult::Incomplete(_) => {
          return Ok(self.state);
        },
        IResult::Error(e) => {
          //FIXME: should probably disconnect on error here
          let err = format!("parse error: {:?}", e);
          self.state = ConnectionState::Error;
          return Err(Error::new(ErrorKind::Other, err))
        }
      }

      let (i, f) = parsed_frame.unwrap();

      //println!("parsed frame: {:?}", f);
      //FIXME: what happens if we fail to parse a packet in a channel?
      // do we continue?
      let consumed = self.receive_buffer.data().offset(i);

      match f.frame_type {
        FrameType::Method => {
          let parsed = parse_class(f.payload);
          println!("parsed method: {:?}", parsed);
          match parsed {
            IResult::Done(b"", m) => {
              (f.channel_id, Some(m), consumed)
            },
            e => {
              //we should not get an incomplete here
              //FIXME: should probably disconnect channel on error here
              let err = format!("parse error: {:?}", e);
              if f.channel_id == 0 {
                self.state = ConnectionState::Error;
              } else {
                self.channels.get_mut(&f.channel_id).map(|channel| channel.state = ChannelState::Error);
              }
              return Err(Error::new(ErrorKind::Other, err))
            }
          }
        },
        FrameType::Heartbeat => {
          match gen_heartbeat_frame((&mut self.send_buffer.space(), 0)).map(|tup| tup.1) {
            Ok(sz) => {
              self.send_buffer.fill(sz);
            }
            Err(e) => {
              println!("error generating start-ok frame: {:?}", e);
              self.state = ConnectionState::Error;
            },
          };

          (f.channel_id, None, consumed)
        },
        t => {
          println!("frame type: {:?} -> unknown payload:\n{}", t, f.payload.to_hex(16));
          let err = format!("parse error: {:?}", t);
          return Err(Error::new(ErrorKind::Other, err))
        }
      }
    };

    self.receive_buffer.consume(consumed);

    if let Some(method) = method {
      if channel_id == 0 {
        self.handle_global_method(method);
      } else {
        self.channels.get_mut(&channel_id).map(|channel| channel.received_method(method));
      }
    }

    return Ok(self.state);

    unreachable!();
  }

  pub fn handle_global_method(&mut self, c: Class) {
    match self.state {
      ConnectionState::Initial | ConnectionState::Closed | ConnectionState::Error => {
        self.state = ConnectionState::Error
      },
      ConnectionState::Connecting(connecting_state) => {
        match connecting_state {
          ConnectingState::Initial | ConnectingState::Error => {
            self.state = ConnectionState::Error
          },
          ConnectingState::SentProtocolHeader => {
            if let Class::Connection(connection::Methods::Start(s)) = c {
              println!("Server sent Connection::Start: {:?}", s);
              self.state = ConnectionState::Connecting(ConnectingState::ReceivedStart);

              let mut h = HashMap::new();
              h.insert("product".to_string(), AMQPValue::LongString("lapin".to_string()));

              let creds = Credentials {
                username: Some("guest".to_owned()),
                secret: Secret::Password("guest".to_owned()),
                channel_binding: ChannelBinding::None,
              };

              let mut mechanism = Plain::from_credentials(creds).unwrap();

              let initial_data = mechanism.initial().unwrap();
              let s = str::from_utf8(&initial_data).unwrap();

              //FIXME: fill with user configured data
              let start_ok = Class::Connection(connection::Methods::StartOk(
                connection::StartOk {
                  client_properties: h,
                  mechanism: "PLAIN".to_string(),
                  locale:    "en_US".to_string(), // FIXME: comes from the server
                  response:  s.to_string(),     //FIXME: implement SASL
                }
              ));

              println!("client sending Connection::StartOk: {:?}", start_ok);

              match gen_method_frame((&mut self.send_buffer.space(), 0), 0, &start_ok).map(|tup| tup.1) {
                Ok(sz) => {
                  self.send_buffer.fill(sz);
                  self.state = ConnectionState::Connecting(ConnectingState::SentStartOk);
                }
                Err(e) => {
                  println!("error generating start-ok frame: {:?}", e);
                  self.state = ConnectionState::Error;
                },
              }
            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          /*ConnectingState::ReceivedStart => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },*/
          ConnectingState::SentStartOk => {
            if let Class::Connection(connection::Methods::Tune(t)) = c {
              println!("Server sent Connection::Tune: {:?}", t);
              self.state = ConnectionState::Connecting(ConnectingState::ReceivedTune);

              self.configuration.channel_max = t.channel_max;
              self.configuration.heartbeat   = t.heartbeat;

              if t.frame_max > self.configuration.frame_max {
                self.send_buffer.grow(t.frame_max as usize);
                self.receive_buffer.grow(t.frame_max as usize);
              }

              self.configuration.frame_max = t.frame_max;

              let tune_ok = Class::Connection(connection::Methods::TuneOk(
                connection::TuneOk {
                  channel_max : self.configuration.channel_max,
                  frame_max   : self.configuration.frame_max,
                  heartbeat   : self.configuration.heartbeat,
                }
              ));

              println!("client sending Connection::TuneOk: {:?}", tune_ok);

              match gen_method_frame((&mut self.send_buffer.space(), 0), 0, &tune_ok).map(|tup| tup.1) {
                Ok(sz) => {
                  self.send_buffer.fill(sz);
                  self.state = ConnectionState::Connecting(ConnectingState::SentTuneOk);

                  let open = Class::Connection(connection::Methods::Open(
                    connection::Open {
                      virtual_host: "/".to_string(),
                      capabilities: "".to_string(),
                      insist:       false,
                    }
                  ));

                  println!("client sending Connection::Open: {:?}", open);

                  match gen_method_frame((&mut self.send_buffer.space(), 0), 0, &open).map(|tup| tup.1) {
                    Ok(sz) => {
                      self.send_buffer.fill(sz);
                      self.state = ConnectionState::Connecting(ConnectingState::SentOpen);
                    }
                    Err(e) => {
                      println!("error generating start-ok frame: {:?}", e);
                      self.state = ConnectionState::Error;
                    },
                  }
                }
                Err(e) => {
                  println!("error generating start-ok frame: {:?}", e);
                  self.state = ConnectionState::Error;
                },
              }
            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          ConnectingState::ReceivedSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::SentSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::ReceivedSecondSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::ReceivedTune => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::SentOpen => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
            if let Class::Connection(connection::Methods::OpenOk(o)) = c {
              println!("Server sent Connection::OpenOk: {:?}, client now connected", o);
              self.state = ConnectionState::Connected;
            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          ConnectingState::Error => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          s => {
            println!("invalid state {:?}", s);
            self.state = ConnectionState::Error;
          }
        }
      },
      ConnectionState::Connected => {},
      ConnectionState::Closing(ClosingState) => {},
    };
  }

  //FIXME: no chunking for now
  pub fn send_content_frames(&mut self, channel_id: u16, class_id: u16, slice: &[u8]) {
    match gen_content_header_frame((&mut self.send_buffer.space(), 0), channel_id, class_id, slice.len() as u64).map(|tup| tup.1) {
      Ok(sz) => {
        self.send_buffer.fill(sz);
        match gen_content_body_frame((&mut self.send_buffer.space(), 0), channel_id, slice).map(|tup| tup.1) {
          Ok(sz) => {
            self.send_buffer.fill(sz);
          },
          Err(e) => {
            println!("error generating frame: {:?}", e);
            self.state = ConnectionState::Error;
          }
        }
      },
      Err(e) => {
        println!("error generating frame: {:?}", e);
        self.state = ConnectionState::Error;
      }
    }
  }

  pub fn send_method_frame(&mut self, channel: u16, method: &Class) {
    match gen_method_frame((&mut self.send_buffer.space(), 0), channel, &method).map(|tup| tup.1) {
      Ok(sz) => {
        self.send_buffer.fill(sz);
      },
      Err(e) => {
        println!("error generating frame: {:?}", e);
        self.state = ConnectionState::Error;
      }
    }
  }
}
<commit_msg>transform generator errors to the error type<commit_after>use std::io::{Error,ErrorKind,Read,Result,Write};
use std::{result,str};
use std::iter::repeat;
use std::collections::HashMap;
use amq_protocol_types::AMQPValue;
use nom::{HexDisplay,IResult,Offset};
use sasl::{ChannelBinding, Credentials, Secret, Mechanism};
use sasl::mechanisms::Plain;
use cookie_factory::GenError;

use format::frame::*;
use format::field::*;
use channel::Channel;
use api::ChannelState;
use buffer::Buffer;
use generated::*;
use error;

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ConnectionState {
  Initial,
  Connecting(ConnectingState),
  Connected,
  Closing(ClosingState),
  Closed,
  Error,
}

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ConnectingState {
  Initial,
  SentProtocolHeader,
  ReceivedStart,
  SentStartOk,
  ReceivedSecure,
  SentSecure,
  ReceivedSecondSecure,
  ReceivedTune,
  SentTuneOk,
  SentOpen,
  Error,
}

#[derive(Clone,Copy,Debug,PartialEq,Eq)]
pub enum ClosingState {
  Initial,
  SentClose,
  ReceivedClose,
  SentCloseOk,
  ReceivedCloseOk,
  Error,
}

#[derive(Clone,Debug,PartialEq)]
pub struct Configuration {
  channel_max: u16,
  frame_max:   u32,
  heartbeat:   u16,
}

#[derive(Clone,Debug,PartialEq)]
pub struct Connection {
  pub state:            ConnectionState,
  pub channels:         HashMap<u16, Channel>,
  pub send_buffer:      Buffer,
  pub receive_buffer:   Buffer,
  pub configuration:    Configuration,
}

impl Connection {
  pub fn new() -> Connection {
    let mut h = HashMap::new();
    h.insert(0, Channel::global());

    let configuration = Configuration {
      channel_max: 0,
      frame_max:   8192,
      heartbeat:   60,
    };

    Connection {
      state:          ConnectionState::Initial,
      channels:       h,
      send_buffer:    Buffer::with_capacity(configuration.frame_max as usize),
      receive_buffer: Buffer::with_capacity(configuration.frame_max as usize),
      configuration:  configuration,
    }
  }

  pub fn connect(&mut self, writer: &mut Write) -> Result<ConnectionState> {
    if self.state != ConnectionState::Initial {
      self.state = ConnectionState::Error;
      return Err(Error::new(ErrorKind::Other, "invalid state"))
    }

    let res = gen_protocol_header((self.send_buffer.space(), 0)).map(|tup| tup.1);
    if let Ok(sz) = res {
      self.send_buffer.fill(sz);
      match writer.write(&mut self.send_buffer.data()) {
        Ok(sz2) => {
          self.send_buffer.consume(sz2);
          self.state = ConnectionState::Connecting(ConnectingState::SentProtocolHeader);
          Ok(self.state)
        },
        Err(e) => Err(e),
      }
    } else {
      Err(Error::new(ErrorKind::WouldBlock, "could not write protocol header"))
    }
  }

  pub fn write(&mut self, writer: &mut Write) -> Result<ConnectionState> {
    println!("will write:\n{}", (&self.send_buffer.data()).to_hex(16));
    match writer.write(&mut self.send_buffer.data()) {
      Ok(sz) => {
        println!("wrote {} bytes", sz);
        self.send_buffer.consume(sz);
        Ok(self.state)
      },
      Err(e) => Err(e),
    }
  }

  pub fn read(&mut self, reader: &mut Read) -> Result<ConnectionState> {
    if self.state == ConnectionState::Initial || self.state == ConnectionState::Error {
      self.state = ConnectionState::Error;
      return Err(Error::new(ErrorKind::Other, "invalid state"))
    }

    match reader.read(&mut self.receive_buffer.space()) {
      Ok(sz) => {
        println!("read {} bytes", sz);
        self.receive_buffer.fill(sz);
      },
      Err(e) => return Err(e),
    }
    println!("will parse:\n{}", (&self.receive_buffer.data()).to_hex(16));
    let (channel_id, method, consumed) = {
      let parsed_frame = raw_frame(&self.receive_buffer.data());
      match parsed_frame {
        IResult::Done(i,_)     => {},
        IResult::Incomplete(_) => {
          return Ok(self.state);
        },
        IResult::Error(e) => {
          //FIXME: should probably disconnect on error here
          let err = format!("parse error: {:?}", e);
          self.state = ConnectionState::Error;
          return Err(Error::new(ErrorKind::Other, err))
        }
      }

      let (i, f) = parsed_frame.unwrap();

      //println!("parsed frame: {:?}", f);
      //FIXME: what happens if we fail to parse a packet in a channel?
      // do we continue?
      let consumed = self.receive_buffer.data().offset(i);

      match f.frame_type {
        FrameType::Method => {
          let parsed = parse_class(f.payload);
          println!("parsed method: {:?}", parsed);
          match parsed {
            IResult::Done(b"", m) => {
              (f.channel_id, Some(m), consumed)
            },
            e => {
              //we should not get an incomplete here
              //FIXME: should probably disconnect channel on error here
              let err = format!("parse error: {:?}", e);
              if f.channel_id == 0 {
                self.state = ConnectionState::Error;
              } else {
                self.channels.get_mut(&f.channel_id).map(|channel| channel.state = ChannelState::Error);
              }
              return Err(Error::new(ErrorKind::Other, err))
            }
          }
        },
        FrameType::Heartbeat => {
          match gen_heartbeat_frame((&mut self.send_buffer.space(), 0)).map(|tup| tup.1) {
            Ok(sz) => {
              self.send_buffer.fill(sz);
            }
            Err(e) => {
              println!("error generating start-ok frame: {:?}", e);
              self.state = ConnectionState::Error;
            },
          };

          (f.channel_id, None, consumed)
        },
        t => {
          println!("frame type: {:?} -> unknown payload:\n{}", t, f.payload.to_hex(16));
          let err = format!("parse error: {:?}", t);
          return Err(Error::new(ErrorKind::Other, err))
        }
      }
    };

    self.receive_buffer.consume(consumed);

    if let Some(method) = method {
      if channel_id == 0 {
        self.handle_global_method(method);
      } else {
        self.channels.get_mut(&channel_id).map(|channel| channel.received_method(method));
      }
    }

    return Ok(self.state);

    unreachable!();
  }

  pub fn handle_global_method(&mut self, c: Class) {
    match self.state {
      ConnectionState::Initial | ConnectionState::Closed | ConnectionState::Error => {
        self.state = ConnectionState::Error
      },
      ConnectionState::Connecting(connecting_state) => {
        match connecting_state {
          ConnectingState::Initial | ConnectingState::Error => {
            self.state = ConnectionState::Error
          },
          ConnectingState::SentProtocolHeader => {
            if let Class::Connection(connection::Methods::Start(s)) = c {
              println!("Server sent Connection::Start: {:?}", s);
              self.state = ConnectionState::Connecting(ConnectingState::ReceivedStart);

              let mut h = HashMap::new();
              h.insert("product".to_string(), AMQPValue::LongString("lapin".to_string()));

              let creds = Credentials {
                username: Some("guest".to_owned()),
                secret: Secret::Password("guest".to_owned()),
                channel_binding: ChannelBinding::None,
              };

              let mut mechanism = Plain::from_credentials(creds).unwrap();

              let initial_data = mechanism.initial().unwrap();
              let s = str::from_utf8(&initial_data).unwrap();

              //FIXME: fill with user configured data
              let start_ok = Class::Connection(connection::Methods::StartOk(
                connection::StartOk {
                  client_properties: h,
                  mechanism: "PLAIN".to_string(),
                  locale:    "en_US".to_string(), // FIXME: comes from the server
                  response:  s.to_string(),     //FIXME: implement SASL
                }
              ));

              println!("client sending Connection::StartOk: {:?}", start_ok);

              match gen_method_frame((&mut self.send_buffer.space(), 0), 0, &start_ok).map(|tup| tup.1) {
                Ok(sz) => {
                  self.send_buffer.fill(sz);
                  self.state = ConnectionState::Connecting(ConnectingState::SentStartOk);
                }
                Err(e) => {
                  println!("error generating start-ok frame: {:?}", e);
                  self.state = ConnectionState::Error;
                },
              }
            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          /*ConnectingState::ReceivedStart => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },*/
          ConnectingState::SentStartOk => {
            if let Class::Connection(connection::Methods::Tune(t)) = c {
              println!("Server sent Connection::Tune: {:?}", t);
              self.state = ConnectionState::Connecting(ConnectingState::ReceivedTune);

              self.configuration.channel_max = t.channel_max;
              self.configuration.heartbeat   = t.heartbeat;

              if t.frame_max > self.configuration.frame_max {
                self.send_buffer.grow(t.frame_max as usize);
                self.receive_buffer.grow(t.frame_max as usize);
              }

              self.configuration.frame_max = t.frame_max;

              let tune_ok = Class::Connection(connection::Methods::TuneOk(
                connection::TuneOk {
                  channel_max : self.configuration.channel_max,
                  frame_max   : self.configuration.frame_max,
                  heartbeat   : self.configuration.heartbeat,
                }
              ));

              println!("client sending Connection::TuneOk: {:?}", tune_ok);

              match gen_method_frame((&mut self.send_buffer.space(), 0), 0, &tune_ok).map(|tup| tup.1) {
                Ok(sz) => {
                  self.send_buffer.fill(sz);
                  self.state = ConnectionState::Connecting(ConnectingState::SentTuneOk);

                  let open = Class::Connection(connection::Methods::Open(
                    connection::Open {
                      virtual_host: "/".to_string(),
                      capabilities: "".to_string(),
                      insist:       false,
                    }
                  ));

                  println!("client sending Connection::Open: {:?}", open);

                  match gen_method_frame((&mut self.send_buffer.space(), 0), 0, &open).map(|tup| tup.1) {
                    Ok(sz) => {
                      self.send_buffer.fill(sz);
                      self.state = ConnectionState::Connecting(ConnectingState::SentOpen);
                    }
                    Err(e) => {
                      println!("error generating start-ok frame: {:?}", e);
                      self.state = ConnectionState::Error;
                    },
                  }
                }
                Err(e) => {
                  println!("error generating start-ok frame: {:?}", e);
                  self.state = ConnectionState::Error;
                },
              }
            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          ConnectingState::ReceivedSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::SentSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::ReceivedSecondSecure => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::ReceivedTune => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          ConnectingState::SentOpen => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
            if let Class::Connection(connection::Methods::OpenOk(o)) = c {
              println!("Server sent Connection::OpenOk: {:?}, client now connected", o);
              self.state = ConnectionState::Connected;
            } else {
              println!("waiting for class Connection method Start, got {:?}", c);
              self.state = ConnectionState::Error;
            }
          },
          ConnectingState::Error => {
            println!("state {:?}\treceived\t{:?}", self.state, c);
          },
          s => {
            println!("invalid state {:?}", s);
            self.state = ConnectionState::Error;
          }
        }
      },
      ConnectionState::Connected => {},
      ConnectionState::Closing(ClosingState) => {},
    };
  }

  //FIXME: no chunking for now
  pub fn send_content_frames(&mut self, channel_id: u16, class_id: u16, slice: &[u8]) {
    match gen_content_header_frame((&mut self.send_buffer.space(), 0), channel_id, class_id, slice.len() as u64).map(|tup| tup.1) {
      Ok(sz) => {
        self.send_buffer.fill(sz);
        match gen_content_body_frame((&mut self.send_buffer.space(), 0), channel_id, slice).map(|tup| tup.1) {
          Ok(sz) => {
            self.send_buffer.fill(sz);
          },
          Err(e) => {
            println!("error generating frame: {:?}", e);
            self.state = ConnectionState::Error;
          }
        }
      },
      Err(e) => {
        println!("error generating frame: {:?}", e);
        self.state = ConnectionState::Error;
      }
    }
  }

  pub fn send_method_frame(&mut self, channel: u16, method: &Class) -> result::Result<(), error::Error> {
    match gen_method_frame((&mut self.send_buffer.space(), 0), channel, &method).map(|tup| tup.1) {
      Ok(sz) => {
        self.send_buffer.fill(sz);
        Ok(())
      },
      Err(e) => {
        println!("error generating frame: {:?}", e);
        self.state = ConnectionState::Error;
        match e {
          GenError::BufferTooSmall(_) => Err(error::Error::SendBufferTooSmall),
          GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
            //self.channels.get_mut(&channel).map(|c| c.state = ChannelState::Error);
            Err(error::Error::GeneratorError)
          }
        }
      }
    }
  }
}
<|endoftext|>


<commit_before>use std::mem;

use rustc::ty::{self, layout::{self, Size, Align}};
use rustc::hir::def_id::{DefId, CRATE_DEF_INDEX};
use rustc::mir;

use rand::RngCore;

use crate::*;

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}

pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    /// Gets an instance for a path.
    fn resolve_path(&self, path: &[&str]) -> InterpResult<'tcx, ty::Instance<'tcx>> {
        let this = self.eval_context_ref();
        this.tcx
            .crates()
            .iter()
            .find(|&&krate| this.tcx.original_crate_name(krate).as_str() == path[0])
            .and_then(|krate| {
                let krate = DefId {
                    krate: *krate,
                    index: CRATE_DEF_INDEX,
                };
                let mut items = this.tcx.item_children(krate);
                let mut path_it = path.iter().skip(1).peekable();

                while let Some(segment) = path_it.next() {
                    for item in mem::replace(&mut items, Default::default()).iter() {
                        if item.ident.name.as_str() == *segment {
                            if path_it.peek().is_none() {
                                return Some(ty::Instance::mono(this.tcx.tcx, item.res.def_id()));
                            }

                            items = this.tcx.item_children(item.res.def_id());
                            break;
                        }
                    }
                }
                None
            })
            .ok_or_else(|| {
                let path = path.iter().map(|&s| s.to_owned()).collect();
                err_unsup!(PathNotFound(path)).into()
            })
    }

    /// Write a 0 of the appropriate size to `dest`.
    fn write_null(&mut self, dest: PlaceTy<'tcx, Tag>) -> InterpResult<'tcx> {
        self.eval_context_mut().write_scalar(Scalar::from_int(0, dest.layout.size), dest)
    }

    /// Test if this immediate equals 0.
    fn is_null(&self, val: Scalar<Tag>) -> InterpResult<'tcx, bool> {
        let this = self.eval_context_ref();
        let null = Scalar::from_int(0, this.memory().pointer_size());
        this.ptr_eq(val, null)
    }

    /// Turn a Scalar into an Option<NonNullScalar>
    fn test_null(&self, val: Scalar<Tag>) -> InterpResult<'tcx, Option<Scalar<Tag>>> {
        let this = self.eval_context_ref();
        Ok(if this.is_null(val)? {
            None
        } else {
            Some(val)
        })
    }

    /// Get the `Place` for a local
    fn local_place(&mut self, local: mir::Local) -> InterpResult<'tcx, PlaceTy<'tcx, Tag>> {
        let this = self.eval_context_mut();
        let place = mir::Place { base: mir::PlaceBase::Local(local), projection: None };
        this.eval_place(&place)
    }

    /// Generate some random bytes, and write them to `dest`.
    fn gen_random(
        &mut self,
        ptr: Scalar<Tag>,
        len: usize,
    ) -> InterpResult<'tcx>  {
        // Some programs pass in a null pointer and a length of 0
        // to their platform's random-generation function (e.g. getrandom())
        // on Linux. For compatibility with these programs, we don't perform
        // any additional checks - it's okay if the pointer is invalid,
        // since we wouldn't actually be writing to it.
        if len == 0 {
            return Ok(());
        }
        let this = self.eval_context_mut();

        let ptr = this.memory().check_ptr_access(
            ptr,
            Size::from_bytes(len as u64),
            Align::from_bytes(1).unwrap()
        )?.expect("we already checked for size 0");

        let mut data = vec![0; len];

        if this.machine.communicate {
            // Fill the buffer using the host's rng.
            getrandom::getrandom(&mut data)
                .map_err(|err| err_unsup_format!("getrandom failed: {}", err))?;
        }
        else {
            let rng = this.memory_mut().extra.rng.get_mut();
            rng.fill_bytes(&mut data);
        }

        let tcx = &{this.tcx.tcx};
        this.memory_mut().get_mut(ptr.alloc_id)?.write_bytes(tcx, ptr, &data)
    }

    /// Visits the memory covered by `place`, sensitive to freezing: the 3rd parameter
    /// will be true if this is frozen, false if this is in an `UnsafeCell`.
    fn visit_freeze_sensitive(
        &self,
        place: MPlaceTy<'tcx, Tag>,
        size: Size,
        mut action: impl FnMut(Pointer<Tag>, Size, bool) -> InterpResult<'tcx>,
    ) -> InterpResult<'tcx> {
        let this = self.eval_context_ref();
        trace!("visit_frozen(place={:?}, size={:?})", *place, size);
        debug_assert_eq!(size,
            this.size_and_align_of_mplace(place)?
            .map(|(size, _)| size)
            .unwrap_or_else(|| place.layout.size)
        );
        // Store how far we proceeded into the place so far. Everything to the left of
        // this offset has already been handled, in the sense that the frozen parts
        // have had `action` called on them.
        let mut end_ptr = place.ptr.assert_ptr();
        // Called when we detected an `UnsafeCell` at the given offset and size.
        // Calls `action` and advances `end_ptr`.
        let mut unsafe_cell_action = |unsafe_cell_ptr: Scalar<Tag>, unsafe_cell_size: Size| {
            let unsafe_cell_ptr = unsafe_cell_ptr.assert_ptr();
            debug_assert_eq!(unsafe_cell_ptr.alloc_id, end_ptr.alloc_id);
            debug_assert_eq!(unsafe_cell_ptr.tag, end_ptr.tag);
            // We assume that we are given the fields in increasing offset order,
            // and nothing else changes.
            let unsafe_cell_offset = unsafe_cell_ptr.offset;
            let end_offset = end_ptr.offset;
            assert!(unsafe_cell_offset >= end_offset);
            let frozen_size = unsafe_cell_offset - end_offset;
            // Everything between the end_ptr and this `UnsafeCell` is frozen.
            if frozen_size != Size::ZERO {
                action(end_ptr, frozen_size, /*frozen*/true)?;
            }
            // This `UnsafeCell` is NOT frozen.
            if unsafe_cell_size != Size::ZERO {
                action(unsafe_cell_ptr, unsafe_cell_size, /*frozen*/false)?;
            }
            // Update end end_ptr.
            end_ptr = unsafe_cell_ptr.wrapping_offset(unsafe_cell_size, this);
            // Done
            Ok(())
        };
        // Run a visitor
        {
            let mut visitor = UnsafeCellVisitor {
                ecx: this,
                unsafe_cell_action: |place| {
                    trace!("unsafe_cell_action on {:?}", place.ptr);
                    // We need a size to go on.
                    let unsafe_cell_size = this.size_and_align_of_mplace(place)?
                        .map(|(size, _)| size)
                        // for extern types, just cover what we can
                        .unwrap_or_else(|| place.layout.size);
                    // Now handle this `UnsafeCell`, unless it is empty.
                    if unsafe_cell_size != Size::ZERO {
                        unsafe_cell_action(place.ptr, unsafe_cell_size)
                    } else {
                        Ok(())
                    }
                },
            };
            visitor.visit_value(place)?;
        }
        // The part between the end_ptr and the end of the place is also frozen.
        // So pretend there is a 0-sized `UnsafeCell` at the end.
        unsafe_cell_action(place.ptr.ptr_wrapping_offset(size, this), Size::ZERO)?;
        // Done!
        return Ok(());

        /// Visiting the memory covered by a `MemPlace`, being aware of
        /// whether we are inside an `UnsafeCell` or not.
        struct UnsafeCellVisitor<'ecx, 'mir, 'tcx, F>
            where F: FnMut(MPlaceTy<'tcx, Tag>) -> InterpResult<'tcx>
        {
            ecx: &'ecx MiriEvalContext<'mir, 'tcx>,
            unsafe_cell_action: F,
        }

        impl<'ecx, 'mir, 'tcx, F>
            ValueVisitor<'mir, 'tcx, Evaluator<'tcx>>
        for
            UnsafeCellVisitor<'ecx, 'mir, 'tcx, F>
        where
            F: FnMut(MPlaceTy<'tcx, Tag>) -> InterpResult<'tcx>
        {
            type V = MPlaceTy<'tcx, Tag>;

            #[inline(always)]
            fn ecx(&self) -> &MiriEvalContext<'mir, 'tcx> {
                &self.ecx
            }

            // Hook to detect `UnsafeCell`.
            fn visit_value(&mut self, v: MPlaceTy<'tcx, Tag>) -> InterpResult<'tcx>
            {
                trace!("UnsafeCellVisitor: {:?} {:?}", *v, v.layout.ty);
                let is_unsafe_cell = match v.layout.ty.sty {
                    ty::Adt(adt, _) => Some(adt.did) == self.ecx.tcx.lang_items().unsafe_cell_type(),
                    _ => false,
                };
                if is_unsafe_cell {
                    // We do not have to recurse further, this is an `UnsafeCell`.
                    (self.unsafe_cell_action)(v)
                } else if self.ecx.type_is_freeze(v.layout.ty) {
                    // This is `Freeze`, there cannot be an `UnsafeCell`
                    Ok(())
                } else {
                    // Proceed further
                    self.walk_value(v)
                }
            }

            // Make sure we visit aggregrates in increasing offset order.
            fn visit_aggregate(
                &mut self,
                place: MPlaceTy<'tcx, Tag>,
                fields: impl Iterator<Item=InterpResult<'tcx, MPlaceTy<'tcx, Tag>>>,
            ) -> InterpResult<'tcx> {
                match place.layout.fields {
                    layout::FieldPlacement::Array { .. } => {
                        // For the array layout, we know the iterator will yield sorted elements so
                        // we can avoid the allocation.
                        self.walk_aggregate(place, fields)
                    }
                    layout::FieldPlacement::Arbitrary { .. } => {
                        // Gather the subplaces and sort them before visiting.
                        let mut places = fields.collect::<InterpResult<'tcx, Vec<MPlaceTy<'tcx, Tag>>>>()?;
                        places.sort_by_key(|place| place.ptr.assert_ptr().offset);
                        self.walk_aggregate(place, places.into_iter().map(Ok))
                    }
                    layout::FieldPlacement::Union { .. } => {
                        // Uh, what?
                        bug!("a union is not an aggregate we should ever visit")
                    }
                }
            }

            // We have to do *something* for unions.
            fn visit_union(&mut self, v: MPlaceTy<'tcx, Tag>) -> InterpResult<'tcx>
            {
                // With unions, we fall back to whatever the type says, to hopefully be consistent
                // with LLVM IR.
                // FIXME: are we consistent, and is this really the behavior we want?
                let frozen = self.ecx.type_is_freeze(v.layout.ty);
                if frozen {
                    Ok(())
                } else {
                    (self.unsafe_cell_action)(v)
                }
            }

            // We should never get to a primitive, but always short-circuit somewhere above.
            fn visit_primitive(&mut self, _v: MPlaceTy<'tcx, Tag>) -> InterpResult<'tcx>
            {
                bug!("we should always short-circuit before coming to a primitive")
            }
        }
    }
}
<commit_msg>Auto merge of #931 - RalfJung:retag-makes-my-head-hurt, r=RalfJung<commit_after>use std::mem;

use rustc::ty::{self, layout::{self, Size, Align}};
use rustc::hir::def_id::{DefId, CRATE_DEF_INDEX};
use rustc::mir;

use rand::RngCore;

use crate::*;

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}

pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    /// Gets an instance for a path.
    fn resolve_path(&self, path: &[&str]) -> InterpResult<'tcx, ty::Instance<'tcx>> {
        let this = self.eval_context_ref();
        this.tcx
            .crates()
            .iter()
            .find(|&&krate| this.tcx.original_crate_name(krate).as_str() == path[0])
            .and_then(|krate| {
                let krate = DefId {
                    krate: *krate,
                    index: CRATE_DEF_INDEX,
                };
                let mut items = this.tcx.item_children(krate);
                let mut path_it = path.iter().skip(1).peekable();

                while let Some(segment) = path_it.next() {
                    for item in mem::replace(&mut items, Default::default()).iter() {
                        if item.ident.name.as_str() == *segment {
                            if path_it.peek().is_none() {
                                return Some(ty::Instance::mono(this.tcx.tcx, item.res.def_id()));
                            }

                            items = this.tcx.item_children(item.res.def_id());
                            break;
                        }
                    }
                }
                None
            })
            .ok_or_else(|| {
                let path = path.iter().map(|&s| s.to_owned()).collect();
                err_unsup!(PathNotFound(path)).into()
            })
    }

    /// Write a 0 of the appropriate size to `dest`.
    fn write_null(&mut self, dest: PlaceTy<'tcx, Tag>) -> InterpResult<'tcx> {
        self.eval_context_mut().write_scalar(Scalar::from_int(0, dest.layout.size), dest)
    }

    /// Test if this immediate equals 0.
    fn is_null(&self, val: Scalar<Tag>) -> InterpResult<'tcx, bool> {
        let this = self.eval_context_ref();
        let null = Scalar::from_int(0, this.memory().pointer_size());
        this.ptr_eq(val, null)
    }

    /// Turn a Scalar into an Option<NonNullScalar>
    fn test_null(&self, val: Scalar<Tag>) -> InterpResult<'tcx, Option<Scalar<Tag>>> {
        let this = self.eval_context_ref();
        Ok(if this.is_null(val)? {
            None
        } else {
            Some(val)
        })
    }

    /// Get the `Place` for a local
    fn local_place(&mut self, local: mir::Local) -> InterpResult<'tcx, PlaceTy<'tcx, Tag>> {
        let this = self.eval_context_mut();
        let place = mir::Place { base: mir::PlaceBase::Local(local), projection: None };
        this.eval_place(&place)
    }

    /// Generate some random bytes, and write them to `dest`.
    fn gen_random(
        &mut self,
        ptr: Scalar<Tag>,
        len: usize,
    ) -> InterpResult<'tcx>  {
        // Some programs pass in a null pointer and a length of 0
        // to their platform's random-generation function (e.g. getrandom())
        // on Linux. For compatibility with these programs, we don't perform
        // any additional checks - it's okay if the pointer is invalid,
        // since we wouldn't actually be writing to it.
        if len == 0 {
            return Ok(());
        }
        let this = self.eval_context_mut();

        let ptr = this.memory().check_ptr_access(
            ptr,
            Size::from_bytes(len as u64),
            Align::from_bytes(1).unwrap()
        )?.expect("we already checked for size 0");

        let mut data = vec![0; len];

        if this.machine.communicate {
            // Fill the buffer using the host's rng.
            getrandom::getrandom(&mut data)
                .map_err(|err| err_unsup_format!("getrandom failed: {}", err))?;
        }
        else {
            let rng = this.memory_mut().extra.rng.get_mut();
            rng.fill_bytes(&mut data);
        }

        let tcx = &{this.tcx.tcx};
        this.memory_mut().get_mut(ptr.alloc_id)?.write_bytes(tcx, ptr, &data)
    }

    /// Visits the memory covered by `place`, sensitive to freezing: the 3rd parameter
    /// will be true if this is frozen, false if this is in an `UnsafeCell`.
    fn visit_freeze_sensitive(
        &self,
        place: MPlaceTy<'tcx, Tag>,
        size: Size,
        mut action: impl FnMut(Pointer<Tag>, Size, bool) -> InterpResult<'tcx>,
    ) -> InterpResult<'tcx> {
        let this = self.eval_context_ref();
        trace!("visit_frozen(place={:?}, size={:?})", *place, size);
        debug_assert_eq!(size,
            this.size_and_align_of_mplace(place)?
            .map(|(size, _)| size)
            .unwrap_or_else(|| place.layout.size)
        );
        // Store how far we proceeded into the place so far. Everything to the left of
        // this offset has already been handled, in the sense that the frozen parts
        // have had `action` called on them.
        let mut end_ptr = place.ptr.assert_ptr();
        // Called when we detected an `UnsafeCell` at the given offset and size.
        // Calls `action` and advances `end_ptr`.
        let mut unsafe_cell_action = |unsafe_cell_ptr: Scalar<Tag>, unsafe_cell_size: Size| {
            let unsafe_cell_ptr = unsafe_cell_ptr.assert_ptr();
            debug_assert_eq!(unsafe_cell_ptr.alloc_id, end_ptr.alloc_id);
            debug_assert_eq!(unsafe_cell_ptr.tag, end_ptr.tag);
            // We assume that we are given the fields in increasing offset order,
            // and nothing else changes.
            let unsafe_cell_offset = unsafe_cell_ptr.offset;
            let end_offset = end_ptr.offset;
            assert!(unsafe_cell_offset >= end_offset);
            let frozen_size = unsafe_cell_offset - end_offset;
            // Everything between the end_ptr and this `UnsafeCell` is frozen.
            if frozen_size != Size::ZERO {
                action(end_ptr, frozen_size, /*frozen*/true)?;
            }
            // This `UnsafeCell` is NOT frozen.
            if unsafe_cell_size != Size::ZERO {
                action(unsafe_cell_ptr, unsafe_cell_size, /*frozen*/false)?;
            }
            // Update end end_ptr.
            end_ptr = unsafe_cell_ptr.wrapping_offset(unsafe_cell_size, this);
            // Done
            Ok(())
        };
        // Run a visitor
        {
            let mut visitor = UnsafeCellVisitor {
                ecx: this,
                unsafe_cell_action: |place| {
                    trace!("unsafe_cell_action on {:?}", place.ptr);
                    // We need a size to go on.
                    let unsafe_cell_size = this.size_and_align_of_mplace(place)?
                        .map(|(size, _)| size)
                        // for extern types, just cover what we can
                        .unwrap_or_else(|| place.layout.size);
                    // Now handle this `UnsafeCell`, unless it is empty.
                    if unsafe_cell_size != Size::ZERO {
                        unsafe_cell_action(place.ptr, unsafe_cell_size)
                    } else {
                        Ok(())
                    }
                },
            };
            visitor.visit_value(place)?;
        }
        // The part between the end_ptr and the end of the place is also frozen.
        // So pretend there is a 0-sized `UnsafeCell` at the end.
        unsafe_cell_action(place.ptr.ptr_wrapping_offset(size, this), Size::ZERO)?;
        // Done!
        return Ok(());

        /// Visiting the memory covered by a `MemPlace`, being aware of
        /// whether we are inside an `UnsafeCell` or not.
        struct UnsafeCellVisitor<'ecx, 'mir, 'tcx, F>
            where F: FnMut(MPlaceTy<'tcx, Tag>) -> InterpResult<'tcx>
        {
            ecx: &'ecx MiriEvalContext<'mir, 'tcx>,
            unsafe_cell_action: F,
        }

        impl<'ecx, 'mir, 'tcx, F>
            ValueVisitor<'mir, 'tcx, Evaluator<'tcx>>
        for
            UnsafeCellVisitor<'ecx, 'mir, 'tcx, F>
        where
            F: FnMut(MPlaceTy<'tcx, Tag>) -> InterpResult<'tcx>
        {
            type V = MPlaceTy<'tcx, Tag>;

            #[inline(always)]
            fn ecx(&self) -> &MiriEvalContext<'mir, 'tcx> {
                &self.ecx
            }

            // Hook to detect `UnsafeCell`.
            fn visit_value(&mut self, v: MPlaceTy<'tcx, Tag>) -> InterpResult<'tcx>
            {
                trace!("UnsafeCellVisitor: {:?} {:?}", *v, v.layout.ty);
                let is_unsafe_cell = match v.layout.ty.sty {
                    ty::Adt(adt, _) => Some(adt.did) == self.ecx.tcx.lang_items().unsafe_cell_type(),
                    _ => false,
                };
                if is_unsafe_cell {
                    // We do not have to recurse further, this is an `UnsafeCell`.
                    (self.unsafe_cell_action)(v)
                } else if self.ecx.type_is_freeze(v.layout.ty) {
                    // This is `Freeze`, there cannot be an `UnsafeCell`
                    Ok(())
                } else {
                    // We want to not actually read from memory for this visit. So, before
                    // walking this value, we have to make sure it is not a
                    // `Variants::Multiple`.
                    match v.layout.variants {
                        layout::Variants::Multiple { .. } => {
                            // A multi-variant enum, or generator, or so.
                            // Treat this like a union: without reading from memory,
                            // we cannot determine the variant we are in. Reading from
                            // memory would be subject to Stacked Borrows rules, leading
                            // to all sorts of "funny" recursion.
                            // We only end up here if the type is *not* freeze, so we just call the
                            // `UnsafeCell` action.
                            (self.unsafe_cell_action)(v)
                        }
                        layout::Variants::Single { .. } => {
                            // Proceed further, try to find where exactly that `UnsafeCell`
                            // is hiding.
                            self.walk_value(v)
                        }
                    }
                }
            }

            // Make sure we visit aggregrates in increasing offset order.
            fn visit_aggregate(
                &mut self,
                place: MPlaceTy<'tcx, Tag>,
                fields: impl Iterator<Item=InterpResult<'tcx, MPlaceTy<'tcx, Tag>>>,
            ) -> InterpResult<'tcx> {
                match place.layout.fields {
                    layout::FieldPlacement::Array { .. } => {
                        // For the array layout, we know the iterator will yield sorted elements so
                        // we can avoid the allocation.
                        self.walk_aggregate(place, fields)
                    }
                    layout::FieldPlacement::Arbitrary { .. } => {
                        // Gather the subplaces and sort them before visiting.
                        let mut places = fields.collect::<InterpResult<'tcx, Vec<MPlaceTy<'tcx, Tag>>>>()?;
                        places.sort_by_key(|place| place.ptr.assert_ptr().offset);
                        self.walk_aggregate(place, places.into_iter().map(Ok))
                    }
                    layout::FieldPlacement::Union { .. } => {
                        // Uh, what?
                        bug!("a union is not an aggregate we should ever visit")
                    }
                }
            }

            // We have to do *something* for unions.
            fn visit_union(&mut self, v: MPlaceTy<'tcx, Tag>) -> InterpResult<'tcx>
            {
                // With unions, we fall back to whatever the type says, to hopefully be consistent
                // with LLVM IR.
                // FIXME: are we consistent, and is this really the behavior we want?
                let frozen = self.ecx.type_is_freeze(v.layout.ty);
                if frozen {
                    Ok(())
                } else {
                    (self.unsafe_cell_action)(v)
                }
            }

            // We should never get to a primitive, but always short-circuit somewhere above.
            fn visit_primitive(&mut self, _v: MPlaceTy<'tcx, Tag>) -> InterpResult<'tcx>
            {
                bug!("we should always short-circuit before coming to a primitive")
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 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.

use point::{Point2D, Point3D};
use vector::{Vector2D, Vector3D};

use num::{One, Zero};

use core::fmt;
use core::marker::PhantomData;
use core::ops::Div;
use core::cmp::{Eq, PartialEq};
use core::hash::{Hash};
#[cfg(feature = "serde")]
use serde;

/// Homogeneous vector in 3D space.
#[repr(C)]
pub struct HomogeneousVector<T, U> {
    pub x: T,
    pub y: T,
    pub z: T,
    pub w: T,
    #[doc(hidden)]
    pub _unit: PhantomData<U>,
}

impl<T: Copy, U> Copy for HomogeneousVector<T, U> {}

impl<T: Clone, U> Clone for HomogeneousVector<T, U> {
    fn clone(&self) -> Self {
        HomogeneousVector {
            x: self.x.clone(),
            y: self.y.clone(),
            z: self.z.clone(),
            w: self.w.clone(),
            _unit: PhantomData,
        }
    }
}

#[cfg(feature = "serde")]
impl<'de, T, U> serde::Deserialize<'de> for HomogeneousVector<T, U>
    where T: serde::Deserialize<'de>
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where D: serde::Deserializer<'de>
    {
        let (x, y, z, w) = serde::Deserialize::deserialize(deserializer)?;
        Ok(HomogeneousVector { x, y, z, w, _unit: PhantomData })
    }
}

#[cfg(feature = "serde")]
impl<T, U> serde::Serialize for HomogeneousVector<T, U>
    where T: serde::Serialize
{
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: serde::Serializer
    {
        (&self.x, &self.y, &self.z, &self.w).serialize(serializer)
    }
}

impl<T, U> Eq for HomogeneousVector<T, U> where T: Eq {}

impl<T, U> PartialEq for HomogeneousVector<T, U>
    where T: PartialEq
{
    fn eq(&self, other: &Self) -> bool {
        self.x == other.x && self.y == other.y && self.z == other.z && self.w == other.w
    }
}

impl<T, U> Hash for HomogeneousVector<T, U>
    where T: Hash
{
    fn hash<H: ::core::hash::Hasher>(&self, h: &mut H) {
        self.x.hash(h);
        self.y.hash(h);
        self.z.hash(h);
        self.w.hash(h);
    }
}

impl<T, U> HomogeneousVector<T, U> {
    /// Constructor taking scalar values directly.
    #[inline]
    pub const fn new(x: T, y: T, z: T, w: T) -> Self {
        HomogeneousVector { x, y, z, w, _unit: PhantomData }
    }
}

impl<T: Copy + Div<T, Output=T> + Zero + PartialOrd, U> HomogeneousVector<T, U> {
    /// Convert into Cartesian 2D point.
    ///
    /// Returns None if the point is on or behind the W=0 hemisphere.
    #[inline]
    pub fn to_point2d(&self) -> Option<Point2D<T, U>> {
        if self.w > T::zero() {
            Some(Point2D::new(self.x / self.w, self.y / self.w))
        } else {
            None
        }
    }

    /// Convert into Cartesian 3D point.
    ///
    /// Returns None if the point is on or behind the W=0 hemisphere.
    #[inline]
    pub fn to_point3d(&self) -> Option<Point3D<T, U>> {
        if self.w > T::zero() {
            Some(Point3D::new(self.x / self.w, self.y / self.w, self.z / self.w))
        } else {
            None
        }
    }
}

impl<T: Zero, U> From<Vector2D<T, U>> for HomogeneousVector<T, U> {
    #[inline]
    fn from(v: Vector2D<T, U>) -> Self {
        HomogeneousVector::new(v.x, v.y, T::zero(), T::zero())
    }
}

impl<T: Zero, U> From<Vector3D<T, U>> for HomogeneousVector<T, U> {
    #[inline]
    fn from(v: Vector3D<T, U>) -> Self {
        HomogeneousVector::new(v.x, v.y, v.z, T::zero())
    }
}

impl<T: Zero + One, U> From<Point2D<T, U>> for HomogeneousVector<T, U> {
    #[inline]
    fn from(p: Point2D<T, U>) -> Self {
        HomogeneousVector::new(p.x, p.y, T::zero(), T::one())
    }
}

impl<T: One, U> From<Point3D<T, U>> for HomogeneousVector<T, U> {
    #[inline]
    fn from(p: Point3D<T, U>) -> Self {
        HomogeneousVector::new(p.x, p.y, p.z, T::one())
    }
}

impl<T: fmt::Debug, U> fmt::Debug for HomogeneousVector<T, U> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "({:?},{:?},{:?},{:?})", self.x, self.y, self.z, self.w)
    }
}

impl<T: fmt::Display, U> fmt::Display for HomogeneousVector<T, U> {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "({},{},{},{})", self.x, self.y, self.z, self.w)
    }
}

#[cfg(test)]
mod homogeneous {
    use super::HomogeneousVector;
    use default::{Point2D, Point3D};

    #[test]
    fn roundtrip() {
        assert_eq!(Some(Point2D::new(1.0, 2.0)), HomogeneousVector::from(Point2D::new(1.0, 2.0)).to_point2d());
        assert_eq!(Some(Point3D::new(1.0, -2.0, 0.1)), HomogeneousVector::from(Point3D::new(1.0, -2.0, 0.1)).to_point3d());
    }

    #[test]
    fn negative() {
        assert_eq!(None, HomogeneousVector::<f32, ()>::new(1.0, 2.0, 3.0, 0.0).to_point2d());
        assert_eq!(None, HomogeneousVector::<f32, ()>::new(1.0, -2.0, -3.0, -2.0).to_point3d());
    }
}
<commit_msg>Propagate formatting parameters to components for 4d vectors.<commit_after>// Copyright 2018 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.

use point::{Point2D, Point3D};
use vector::{Vector2D, Vector3D};

use num::{One, Zero};

use core::fmt;
use core::marker::PhantomData;
use core::ops::Div;
use core::cmp::{Eq, PartialEq};
use core::hash::{Hash};
#[cfg(feature = "serde")]
use serde;

/// Homogeneous vector in 3D space.
#[repr(C)]
pub struct HomogeneousVector<T, U> {
    pub x: T,
    pub y: T,
    pub z: T,
    pub w: T,
    #[doc(hidden)]
    pub _unit: PhantomData<U>,
}

impl<T: Copy, U> Copy for HomogeneousVector<T, U> {}

impl<T: Clone, U> Clone for HomogeneousVector<T, U> {
    fn clone(&self) -> Self {
        HomogeneousVector {
            x: self.x.clone(),
            y: self.y.clone(),
            z: self.z.clone(),
            w: self.w.clone(),
            _unit: PhantomData,
        }
    }
}

#[cfg(feature = "serde")]
impl<'de, T, U> serde::Deserialize<'de> for HomogeneousVector<T, U>
    where T: serde::Deserialize<'de>
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where D: serde::Deserializer<'de>
    {
        let (x, y, z, w) = serde::Deserialize::deserialize(deserializer)?;
        Ok(HomogeneousVector { x, y, z, w, _unit: PhantomData })
    }
}

#[cfg(feature = "serde")]
impl<T, U> serde::Serialize for HomogeneousVector<T, U>
    where T: serde::Serialize
{
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: serde::Serializer
    {
        (&self.x, &self.y, &self.z, &self.w).serialize(serializer)
    }
}

impl<T, U> Eq for HomogeneousVector<T, U> where T: Eq {}

impl<T, U> PartialEq for HomogeneousVector<T, U>
    where T: PartialEq
{
    fn eq(&self, other: &Self) -> bool {
        self.x == other.x && self.y == other.y && self.z == other.z && self.w == other.w
    }
}

impl<T, U> Hash for HomogeneousVector<T, U>
    where T: Hash
{
    fn hash<H: ::core::hash::Hasher>(&self, h: &mut H) {
        self.x.hash(h);
        self.y.hash(h);
        self.z.hash(h);
        self.w.hash(h);
    }
}

impl<T, U> HomogeneousVector<T, U> {
    /// Constructor taking scalar values directly.
    #[inline]
    pub const fn new(x: T, y: T, z: T, w: T) -> Self {
        HomogeneousVector { x, y, z, w, _unit: PhantomData }
    }
}

impl<T: Copy + Div<T, Output=T> + Zero + PartialOrd, U> HomogeneousVector<T, U> {
    /// Convert into Cartesian 2D point.
    ///
    /// Returns None if the point is on or behind the W=0 hemisphere.
    #[inline]
    pub fn to_point2d(&self) -> Option<Point2D<T, U>> {
        if self.w > T::zero() {
            Some(Point2D::new(self.x / self.w, self.y / self.w))
        } else {
            None
        }
    }

    /// Convert into Cartesian 3D point.
    ///
    /// Returns None if the point is on or behind the W=0 hemisphere.
    #[inline]
    pub fn to_point3d(&self) -> Option<Point3D<T, U>> {
        if self.w > T::zero() {
            Some(Point3D::new(self.x / self.w, self.y / self.w, self.z / self.w))
        } else {
            None
        }
    }
}

impl<T: Zero, U> From<Vector2D<T, U>> for HomogeneousVector<T, U> {
    #[inline]
    fn from(v: Vector2D<T, U>) -> Self {
        HomogeneousVector::new(v.x, v.y, T::zero(), T::zero())
    }
}

impl<T: Zero, U> From<Vector3D<T, U>> for HomogeneousVector<T, U> {
    #[inline]
    fn from(v: Vector3D<T, U>) -> Self {
        HomogeneousVector::new(v.x, v.y, v.z, T::zero())
    }
}

impl<T: Zero + One, U> From<Point2D<T, U>> for HomogeneousVector<T, U> {
    #[inline]
    fn from(p: Point2D<T, U>) -> Self {
        HomogeneousVector::new(p.x, p.y, T::zero(), T::one())
    }
}

impl<T: One, U> From<Point3D<T, U>> for HomogeneousVector<T, U> {
    #[inline]
    fn from(p: Point3D<T, U>) -> Self {
        HomogeneousVector::new(p.x, p.y, p.z, T::one())
    }
}

impl<T: fmt::Debug, U> fmt::Debug for HomogeneousVector<T, U> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "(")?;
        fmt::Debug::fmt(&self.x, f)?;
        write!(f, ",")?;
        fmt::Debug::fmt(&self.y, f)?;
        write!(f, ",")?;
        fmt::Debug::fmt(&self.z, f)?;
        write!(f, ",")?;
        fmt::Debug::fmt(&self.w, f)?;
        write!(f, ")")?;

        Ok(())
    }
}

impl<T: fmt::Display, U> fmt::Display for HomogeneousVector<T, U> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "(")?;
        fmt::Display::fmt(&self.x, f)?;
        write!(f, ",")?;
        fmt::Display::fmt(&self.y, f)?;
        write!(f, ",")?;
        fmt::Display::fmt(&self.z, f)?;
        write!(f, ",")?;
        fmt::Display::fmt(&self.w, f)?;
        write!(f, ")")?;

        Ok(())
    }
}

#[cfg(test)]
mod homogeneous {
    use super::HomogeneousVector;
    use default::{Point2D, Point3D};

    #[test]
    fn roundtrip() {
        assert_eq!(Some(Point2D::new(1.0, 2.0)), HomogeneousVector::from(Point2D::new(1.0, 2.0)).to_point2d());
        assert_eq!(Some(Point3D::new(1.0, -2.0, 0.1)), HomogeneousVector::from(Point3D::new(1.0, -2.0, 0.1)).to_point3d());
    }

    #[test]
    fn negative() {
        assert_eq!(None, HomogeneousVector::<f32, ()>::new(1.0, 2.0, 3.0, 0.0).to_point2d());
        assert_eq!(None, HomogeneousVector::<f32, ()>::new(1.0, -2.0, -3.0, -2.0).to_point3d());
    }
}
<|endoftext|>


<commit_before>#![feature(specialization)]
#![feature(conservative_impl_trait)]
#![feature(cfg_target_feature)]
#![feature(asm)]

#![allow(unused_imports)]

extern crate core;
extern crate typenum;
extern crate momms;

use std::time::{Instant};
use typenum::{U1, B0, UInt, Unsigned};

use momms::kern::{KernelNM,KernelMN};
use momms::matrix::{Scalar, Mat, ColumnPanelMatrix, RowPanelMatrix, Matrix, Hierarch};
use momms::composables::{GemmNode, AlgorithmStep, PartM, PartN, PartK, PackA, PackB, SpawnThreads, ParallelM, ParallelN, TheRest};
use momms::thread_comm::ThreadInfo;
use momms::util;

fn test_algorithm<T: Scalar, Mr: Unsigned, Nr: Unsigned, Kc:Unsigned, CLRS: Unsigned, CLCS: Unsigned,
    S: GemmNode<T, Hierarch<T, Mr, Kc, U1, Mr>, Hierarch<T, Kc, Nr, Nr, U1>, Hierarch<T, Mr, Nr, CLRS, CLCS>>>
    ( m:usize, n: usize, k: usize, algo: &mut S, flusher: &mut Vec<f64>, n_reps: usize ) -> (f64, T)
{
    let algo_desc = S::hierarchy_description();
    let mut best_time: f64 = 9999999999.0;
    let mut worst_err: T = T::zero();

    for _ in 0..n_reps {
        //Create matrices.
        let mut a : Hierarch<T, Mr, Kc, U1, Mr> = Hierarch::new(m, k, &algo_desc, AlgorithmStep::M{bsz: 0}, AlgorithmStep::K{bsz: 0});
        let mut b : Hierarch<T, Kc, Nr, Nr, U1> = Hierarch::new(k, n, &algo_desc, AlgorithmStep::K{bsz: 0}, AlgorithmStep::N{bsz: 0});
        let mut c : Hierarch<T, Mr, Nr, CLRS, CLCS> = Hierarch::new(m, n, &algo_desc, AlgorithmStep::M{bsz: 0}, AlgorithmStep::N{bsz: 0});

        //Fill the matrices
        a.fill_rand(); c.fill_zero(); b.fill_rand();

        //Read a buffer so that A, B, and C are cold in cache.
        for i in flusher.iter_mut() { *i += 1.0; }

        //Time and run algorithm
        let start = Instant::now();
        unsafe{ algo.run( &mut a, &mut b, &mut c, &ThreadInfo::single_thread() ); }
        best_time = best_time.min(util::dur_seconds(start));
        let err = util::test_c_eq_a_b( &mut a, &mut b, &mut c);
        worst_err = worst_err.max(err);
    }
    (best_time, worst_err)
}

fn test(m_selector: isize, n_selector: isize, k_selector: isize) {
    //Goto's algorithm
    type U3000 = UInt<UInt<typenum::U750, B0>, B0>;
    type Nc = U3000;
    type Kc = typenum::U192;
    type Mc = typenum::U120;
    type Mr = typenum::U4;
    type Nr = typenum::U12;

    type Goto<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, Nc,
          PartK<T, MTA, MTB, MTC, Kc,
          PartM<T, MTA, MTB, MTC, Mc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>;

    type RootS3 = typenum::U768;
    //Resident A algorithm
    type NcL2 = typenum::U120;
    type L3A<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartM<T, MTA, MTB, MTC, RootS3,
          PartK<T, MTA, MTB, MTC, RootS3,
          PartN<T, MTA, MTB, MTC, NcL2,
          PartK<T, MTA, MTB, MTC, Kc,
          ParallelM<T, MTA, MTB, MTC, Nr, TheRest, //This algorithm uses a 12x4 ukernel instead of 4x12 like the other two
          KernelMN<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    //Resident B algorithm
    type McL2 = typenum::U120;
    type L3B<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, RootS3,
          PartK<T, MTA, MTB, MTC, RootS3,
          PartM<T, MTA, MTB, MTC, McL2,
          PartK<T, MTA, MTB, MTC, Kc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    //Resident C algorithm
    type L3C<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, RootS3,
          PartM<T, MTA, MTB, MTC, typenum::U720, //Use 720 as blocksize as it is divisible by Mc=120
          PartK<T, MTA, MTB, MTC, Kc,
          PartM<T, MTA, MTB, MTC, McL2,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    //We can use the same matrix type for all three algorithms
    type HierA<T> = Hierarch<T, Mr, Kc, U1, Mr>;
    type HierB<T> = Hierarch<T, Kc, Nr, Nr, U1>;
    type HierC<T> = Hierarch<T, Mr, Nr, Nr, U1>;

    type HierAL3a<T> = Hierarch<T, Nr, Kc, U1, Nr>;
    type HierBL3a<T> = Hierarch<T, Kc, Mr, Mr, U1>;
    type HierCL3a<T> = Hierarch<T, Nr, Mr, U1, Nr>;

    let mut goto = <Goto<f64, HierA<f64>, HierB<f64>, HierC<f64>>>::new();
    let mut l3a = <L3A<f64, HierAL3a<f64>, HierBL3a<f64>, HierCL3a<f64>>>::new();
    let mut l3b = <L3B<f64, HierA<f64>, HierB<f64>, HierC<f64>>>::new();
    let mut l3c = <L3C<f64, HierA<f64>, HierB<f64>, HierC<f64>>>::new();

    goto.set_n_threads(4);
    l3a.set_n_threads(4);
    l3b.set_n_threads(4);
    l3c.set_n_threads(4);

    let flusher_len = 2*1024*1024; //16MB
    let mut flusher : Vec<f64> = Vec::with_capacity(flusher_len);
    for _ in 0..flusher_len {
        flusher.push(0.0);
    }

    println!("m\tn\tk\t{: <13}{: <13}{: <13}{: <13}{: <15}{: <15}{: <15}{: <15}", "goto", "l3a", "l3b", "l3c", "goto", "l3a", "l3b", "l3c");
    for index in 01..81 {
        let size = index*50;
//        let (m, n, k) = (size, size, size);
        let m = if m_selector < 0 { size * isize::abs(m_selector) } else { m_selector } as usize;
        let n = if n_selector < 0 { size * isize::abs(n_selector) } else { n_selector } as usize;
        let k = if k_selector < 0 { size * isize::abs(k_selector) } else { k_selector } as usize;

        let n_reps = 5;
        let (goto_time, goto_err) = test_algorithm(m, n, k, &mut goto, &mut flusher, n_reps);
        let (l3a_time, l3a_err) = test_algorithm(m, n, k, &mut l3a, &mut flusher, n_reps);
        let (l3b_time, l3b_err) = test_algorithm(m, n, k, &mut l3b, &mut flusher, n_reps);
        let (l3c_time, l3c_err) = test_algorithm(m, n, k, &mut l3c, &mut flusher, n_reps);

        println!("{}\t{}\t{}\t{}{}{}{}{}{}{}{}",
                 m, n, k,
                 format!("{: <13.5}", util::gflops(m,n,k,goto_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3a_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3b_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3c_time)),
                 format!("{: <15.5e}", goto_err.sqrt()),
                 format!("{: <15.5e}", l3a_err.sqrt()),
                 format!("{: <15.5e}", l3b_err.sqrt()),
                 format!("{: <15.5e}", l3c_err.sqrt()));
    }

    let mut sum = 0.0;
    for a in flusher.iter() {
        sum += *a;
    }
    println!("Flush value {}", sum);
}

fn main() {
    //Best for Resident A
    test(768,-1,768);
    //Best for Resident B
    test(-1,768,768);
    //Best for Resident C
    test(768,768,-1);
}
<commit_msg>Adjusting blocksizes for L3C algorithm<commit_after>#![feature(specialization)]
#![feature(conservative_impl_trait)]
#![feature(cfg_target_feature)]
#![feature(asm)]

#![allow(unused_imports)]

extern crate core;
extern crate typenum;
extern crate momms;

use std::time::{Instant};
use typenum::{U1, B0, UInt, Unsigned};

use momms::kern::{KernelNM,KernelMN};
use momms::matrix::{Scalar, Mat, ColumnPanelMatrix, RowPanelMatrix, Matrix, Hierarch};
use momms::composables::{GemmNode, AlgorithmStep, PartM, PartN, PartK, PackA, PackB, SpawnThreads, ParallelM, ParallelN, TheRest};
use momms::thread_comm::ThreadInfo;
use momms::util;

fn test_algorithm<T: Scalar, Mr: Unsigned, Nr: Unsigned, Kc:Unsigned, CLRS: Unsigned, CLCS: Unsigned,
    S: GemmNode<T, Hierarch<T, Mr, Kc, U1, Mr>, Hierarch<T, Kc, Nr, Nr, U1>, Hierarch<T, Mr, Nr, CLRS, CLCS>>>
    ( m:usize, n: usize, k: usize, algo: &mut S, flusher: &mut Vec<f64>, n_reps: usize ) -> (f64, T)
{
    let algo_desc = S::hierarchy_description();
    let mut best_time: f64 = 9999999999.0;
    let mut worst_err: T = T::zero();

    for _ in 0..n_reps {
        //Create matrices.
        let mut a : Hierarch<T, Mr, Kc, U1, Mr> = Hierarch::new(m, k, &algo_desc, AlgorithmStep::M{bsz: 0}, AlgorithmStep::K{bsz: 0});
        let mut b : Hierarch<T, Kc, Nr, Nr, U1> = Hierarch::new(k, n, &algo_desc, AlgorithmStep::K{bsz: 0}, AlgorithmStep::N{bsz: 0});
        let mut c : Hierarch<T, Mr, Nr, CLRS, CLCS> = Hierarch::new(m, n, &algo_desc, AlgorithmStep::M{bsz: 0}, AlgorithmStep::N{bsz: 0});

        //Fill the matrices
        a.fill_rand(); c.fill_zero(); b.fill_rand();

        //Read a buffer so that A, B, and C are cold in cache.
        for i in flusher.iter_mut() { *i += 1.0; }

        //Time and run algorithm
        let start = Instant::now();
        unsafe{ algo.run( &mut a, &mut b, &mut c, &ThreadInfo::single_thread() ); }
        best_time = best_time.min(util::dur_seconds(start));
        let err = util::test_c_eq_a_b( &mut a, &mut b, &mut c);
        worst_err = worst_err.max(err);
    }
    (best_time, worst_err)
}

fn test(m_selector: isize, n_selector: isize, k_selector: isize) {
    //Goto's algorithm
    type U3000 = UInt<UInt<typenum::U750, B0>, B0>;
    type Nc = U3000;
    type Kc = typenum::U192;
    type Mc = typenum::U120;
    type Mr = typenum::U4;
    type Nr = typenum::U12;

    type Goto<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, Nc,
          PartK<T, MTA, MTB, MTC, Kc,
          PartM<T, MTA, MTB, MTC, Mc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>;

    type RootS3 = typenum::U768;
    //Resident A algorithm
    type NcL2 = typenum::U120;
    type L3A<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartM<T, MTA, MTB, MTC, RootS3,
          PartK<T, MTA, MTB, MTC, RootS3,
          PartN<T, MTA, MTB, MTC, NcL2,
          PartK<T, MTA, MTB, MTC, Kc,
          ParallelM<T, MTA, MTB, MTC, Nr, TheRest, //This algorithm uses a 12x4 ukernel instead of 4x12 like the other two
          KernelMN<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    //Resident B algorithm
    type McL2 = typenum::U120;
    type L3B<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, RootS3,
          PartK<T, MTA, MTB, MTC, RootS3,
          PartM<T, MTA, MTB, MTC, McL2,
          PartK<T, MTA, MTB, MTC, Kc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    type L3CNc = typenum::U780;
    type L3CKc = typenum::U156;
    type L3CMc = typenum::U156;
    //Resident C algorithm
    type L3C<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, L3CNc,
          PartM<T, MTA, MTB, MTC, L3CKc,
          PartK<T, MTA, MTB, MTC, L3CKc,
          PartM<T, MTA, MTB, MTC, L3CMc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    //We can use the same matrix type for all three algorithms
    type HierA<T> = Hierarch<T, Mr, Kc, U1, Mr>;
    type HierB<T> = Hierarch<T, Kc, Nr, Nr, U1>;
    type HierC<T> = Hierarch<T, Mr, Nr, Nr, U1>;

    type HierAL3a<T> = Hierarch<T, Nr, Kc, U1, Nr>;
    type HierBL3a<T> = Hierarch<T, Kc, Mr, Mr, U1>;
    type HierCL3a<T> = Hierarch<T, Nr, Mr, U1, Nr>;

    type HierAL3c<T> = Hierarch<T, Mr, L3CKc, U1, Mr>;
    type HierBL3c<T> = Hierarch<T, L3CKc, Nr, Nr, U1>;
    type HierCL3c<T> = Hierarch<T, Mr, Nr, Nr, U1>;

    let mut goto = <Goto<f64, HierA<f64>, HierB<f64>, HierC<f64>>>::new();
    let mut l3a = <L3A<f64, HierAL3a<f64>, HierBL3a<f64>, HierCL3a<f64>>>::new();
    let mut l3b = <L3B<f64, HierA<f64>, HierB<f64>, HierC<f64>>>::new();
    let mut l3c = <L3C<f64, HierAL3c<f64>, HierBL3c<f64>, HierCL3c<f64>>>::new();

    goto.set_n_threads(4);
    l3a.set_n_threads(4);
    l3b.set_n_threads(4);
    l3c.set_n_threads(4);

    let flusher_len = 2*1024*1024; //16MB
    let mut flusher : Vec<f64> = Vec::with_capacity(flusher_len);
    for _ in 0..flusher_len {
        flusher.push(0.0);
    }

    println!("m\tn\tk\t{: <13}{: <13}{: <13}{: <13}{: <15}{: <15}{: <15}{: <15}", "goto", "l3a", "l3b", "l3c", "goto", "l3a", "l3b", "l3c");
    for index in 01..81 {
        let size = index*50;
//        let (m, n, k) = (size, size, size);
        let m = if m_selector < 0 { size * isize::abs(m_selector) } else { m_selector } as usize;
        let n = if n_selector < 0 { size * isize::abs(n_selector) } else { n_selector } as usize;
        let k = if k_selector < 0 { size * isize::abs(k_selector) } else { k_selector } as usize;

        let n_reps = 5;
        let (goto_time, goto_err) = test_algorithm(m, n, k, &mut goto, &mut flusher, n_reps);
        let (l3a_time, l3a_err) = test_algorithm(m, n, k, &mut l3a, &mut flusher, n_reps);
        let (l3b_time, l3b_err) = test_algorithm(m, n, k, &mut l3b, &mut flusher, n_reps);
        let (l3c_time, l3c_err) = test_algorithm(m, n, k, &mut l3c, &mut flusher, n_reps);

        println!("{}\t{}\t{}\t{}{}{}{}{}{}{}{}",
                 m, n, k,
                 format!("{: <13.5}", util::gflops(m,n,k,goto_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3a_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3b_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3c_time)),
                 format!("{: <15.5e}", goto_err.sqrt()),
                 format!("{: <15.5e}", l3a_err.sqrt()),
                 format!("{: <15.5e}", l3b_err.sqrt()),
                 format!("{: <15.5e}", l3c_err.sqrt()));
    }

    let mut sum = 0.0;
    for a in flusher.iter() {
        sum += *a;
    }
    println!("Flush value {}", sum);
}

fn main() {
    //Best for Resident A
    test(768,-1,768);
    //Best for Resident B
    test(-1,768,768);
    //Best for Resident C
    test(768,768,-1);
}
<|endoftext|>