<commit_before>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use std::time::Duration;
use futures::{Async,Poll,Sink,Stream,StartSend,Future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use tokio_timer::{Interval,Timer};
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec;

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length = buf.len();
      if length < 8192 {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(8192 - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  pub upstream: Framed<T,AMQPCodec>,
  pub heartbeat: Interval,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(upstream: Framed<T,AMQPCodec>, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_heartbeat(options.heartbeat);
    conn.connect();

    let mut t = AMQPTransport {
      upstream:  upstream,
      heartbeat: Timer::default().interval(Duration::from_secs(conn.configuration.heartbeat as u64)),
      conn:      conn,
    };

    t.send_and_handle_frames();
    t.poll_children();

    let mut connector = AMQPTransportConnector {
      transport: Some(t)
    };

    trace!("pre-poll");
    connector.poll();
    trace!("post-poll");

    Box::new(connector)
  }

  fn poll_heartbeat(&mut self) {
    if let Ok(Async::Ready(_)) = self.heartbeat.poll() {
      debug!("Heartbeat");
      self.start_send(Frame::Heartbeat(0));
      self.poll_complete();
    }
  }

  fn poll_upstream(&mut self) -> Poll<Option<()>, io::Error> {
    let value = match self.upstream.poll() {
      Ok(Async::Ready(t)) => t,
      Ok(Async::NotReady) => {
        trace!("NotReady");
        return Ok(Async::NotReady);
      },
      Err(e) => {
        error!("upstream poll got error: {:?}", e);
        return Err(From::from(e));
      },
    };

    if let Some(frame) = value {
      trace!("got frame: {:?}", frame);
      self.conn.handle_frame(frame);
      self.send_frames();
      self.upstream.poll_complete();
      Ok(Async::Ready(Some(())))
    } else {
      error!("did not get a frame");
      Ok(Async::Ready(None))
    }
  }

  pub fn poll_children(&mut self) {
    self.poll_heartbeat();
    self.poll_upstream();
  }

  pub fn send_and_handle_frames(&mut self) {
    self.send_frames();
    self.handle_frames();
  }

  pub fn send_frames(&mut self) {
    //FIXME: find a way to use a future here
    while let Some(f) = self.conn.next_frame() {
      self.upstream.start_send(f);
      self.upstream.poll_complete();
    }
    //self.upstream.poll_complete();
  }

  pub fn handle_frames(&mut self) {
    loop {
      match self.poll() {
        Ok(Async::Ready(Some(frame))) => {
          trace!("handle frames: AMQPTransport received frame: {:?}", frame);
          self.conn.handle_frame(frame);
        },
        Ok(Async::Ready(None)) => {
          trace!("handle frames: upstream poll gave Ready(None)");
          break;
        },
        Ok(Async::NotReady) => {
          trace!("handle frames: upstream poll gave NotReady");
          self.poll_children();
          break;
        },
        Err(e) => {
          error!("handle frames: upstream poll got error: {:?}", e);
          break;
        },
      };
    }
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.handle_frames();
    transport.send_frames();

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    transport.poll_heartbeat();
    transport.poll_upstream().map(|value| {
      match value {
        Async::Ready(Some(_)) => {
          if transport.conn.state == ConnectionState::Connected {
            // Upstream had frames available and we're connected, the transport is ready
            Async::Ready(transport)
          } else {
            // Upstream had frames but we're not yet connected, continue polling
            transport.poll_children();
            self.transport = Some(transport);
            Async::NotReady
          }
        },
        _ => {
          // Upstream had no frames
          self.transport = Some(transport);
          Async::NotReady
        },
      }
    })
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {
    type Item = Frame;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Frame>, io::Error> {
        self.poll_heartbeat();

        trace!("stream poll");
        // and Async::NotReady.
        match try_ready!(self.upstream.poll()) {
            Some(frame) => {
              debug!("AMQPTransport received frame: {:?}", frame);
              //try!(self.poll_complete());
              return Ok(Async::Ready(Some(frame)))
            },
            None => {
              trace!("AMQPTransport returned NotReady");
              return Ok(Async::NotReady)
            }
        }
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}

<commit_msg>futures: use poll_upstream in AMQPTrasport::poll and simplify handle_frames<commit_after>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use std::time::Duration;
use futures::{Async,Poll,Sink,Stream,StartSend,Future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use tokio_timer::{Interval,Timer};
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec;

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length = buf.len();
      if length < 8192 {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(8192 - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  pub upstream: Framed<T,AMQPCodec>,
  pub heartbeat: Interval,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(upstream: Framed<T,AMQPCodec>, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_heartbeat(options.heartbeat);
    conn.connect();

    let mut t = AMQPTransport {
      upstream:  upstream,
      heartbeat: Timer::default().interval(Duration::from_secs(conn.configuration.heartbeat as u64)),
      conn:      conn,
    };

    t.send_and_handle_frames();
    t.poll_children();

    let mut connector = AMQPTransportConnector {
      transport: Some(t)
    };

    trace!("pre-poll");
    connector.poll();
    trace!("post-poll");

    Box::new(connector)
  }

  fn poll_heartbeat(&mut self) {
    if let Ok(Async::Ready(_)) = self.heartbeat.poll() {
      debug!("Heartbeat");
      self.start_send(Frame::Heartbeat(0));
      self.poll_complete();
    }
  }

  fn poll_upstream(&mut self) -> Poll<Option<()>, io::Error> {
    let value = match self.upstream.poll() {
      Ok(Async::Ready(t)) => t,
      Ok(Async::NotReady) => {
        trace!("upstream poll gave NotReady");
        return Ok(Async::NotReady);
      },
      Err(e) => {
        error!("upstream poll gave error: {:?}", e);
        return Err(From::from(e));
      },
    };

    if let Some(frame) = value {
      trace!("upstream poll gave frame: {:?}", frame);
      self.conn.handle_frame(frame);
      self.send_frames();
      self.upstream.poll_complete();
      Ok(Async::Ready(Some(())))
    } else {
      error!("upstream poll gave Ready(None)");
      Ok(Async::Ready(None))
    }
  }

  pub fn poll_children(&mut self) {
    self.poll_heartbeat();
    self.poll_upstream();
  }

  pub fn send_and_handle_frames(&mut self) {
    self.send_frames();
    self.handle_frames();
  }

  pub fn send_frames(&mut self) {
    //FIXME: find a way to use a future here
    while let Some(f) = self.conn.next_frame() {
      self.upstream.start_send(f);
      self.upstream.poll_complete();
    }
    //self.upstream.poll_complete();
  }

  pub fn handle_frames(&mut self) {
    while let Ok(Async::Ready(Some(_))) = self.poll() {}
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.handle_frames();
    transport.send_frames();

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    transport.poll_heartbeat();
    transport.poll_upstream().map(|value| {
      match value {
        Async::Ready(Some(_)) => {
          if transport.conn.state == ConnectionState::Connected {
            // Upstream had frames available and we're connected, the transport is ready
            Async::Ready(transport)
          } else {
            // Upstream had frames but we're not yet connected, continue polling
            transport.poll_children();
            self.transport = Some(transport);
            Async::NotReady
          }
        },
        _ => {
          // Upstream had no frames
          self.transport = Some(transport);
          Async::NotReady
        },
      }
    })
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("stream poll");
      self.poll_heartbeat();
      self.poll_upstream()
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}

<|endoftext|>


<commit_before>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::cmp;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use futures::{Async,Poll,Sink,Stream,StartSend,Future,future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use channel::BasicProperties;
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec {
    pub frame_max: u32,
}

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length    = buf.len();
      // Ensure we at least allocate 8192 so that the buffer is big enough for the frame_max
      // negociation. Afterwards, use frame_max if > 8192.
      let frame_max = cmp::max(self.frame_max, 8192) as usize;
      if length < frame_max {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(frame_max - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream: Framed<T,AMQPCodec>,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: Sync+Send,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(stream: T, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_frame_max(options.frame_max);
    conn.set_heartbeat(options.heartbeat);
    if let Err(e) = conn.connect() {
      let err = format!("Failed to connect: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let codec = AMQPCodec {
      frame_max: conn.configuration.frame_max,
    };
    let mut t = AMQPTransport {
      upstream:  stream.framed(codec),
      conn:      conn,
    };

    if let Err(e) = t.send_and_handle_frames() {
      let err = format!("Failed to handle frames: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let mut connector = AMQPTransportConnector {
      transport: Some(t),
    };

    trace!("pre-poll");
    if let Err(e) = connector.poll() {
      let err = format!("Failed to handle frames: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }
    trace!("post-poll");

    Box::new(connector)
  }

  pub fn send_and_handle_frames(&mut self) -> Poll<Option<()>, io::Error> {
    self.send_frames()?;
    self.handle_frames()
  }

  pub fn send_and_handle_frames_with_payload(&mut self, channel_id: u16, payload: &[u8], properties: BasicProperties) -> Poll<Option<()>, io::Error> {
      self.send_and_handle_frames()?;
      self.conn.send_content_frames(channel_id, 60, payload, properties);
      self.send_and_handle_frames()
  }

  fn send_frames(&mut self) -> Result<(), io::Error> {
    //FIXME: find a way to use a future here
    trace!("send frames");
    while let Some(f) = self.conn.next_frame() {
      if let Err(e) = self.send_frame(f) {
        error!("Failed to send frame: {:?}", e);
        return Err(e);
      }
    }
    Ok(())
  }

  pub fn send_frame(&mut self, frame: Frame) -> Poll<(), io::Error> {
      self.start_send(frame).and_then(|_| self.poll_complete())
  }

  fn handle_frames(&mut self) -> Poll<Option<()>, io::Error> {
    trace!("handle frames");
    loop {
      // try_ready will return if we hit an error or NotReady.
      if try_ready!(self.poll()).is_none() {
        return Ok(Async::Ready(None));
      }
    }
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: Sync+Send,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.send_and_handle_frames()?;

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    match transport.poll()? {
      Async::Ready(Some(_)) => {
        if transport.conn.state == ConnectionState::Connected {
          // Upstream had frames available and we're connected, the transport is ready
          Ok(Async::Ready(transport))
        } else {
          // Upstream had frames but we're not yet connected, continue polling
          let poll_ret = transport.poll();
          self.transport = Some(transport);
          poll_ret?;
          Ok(Async::NotReady)
        }
      },
      _ => {
        // Upstream had no frames
        self.transport = Some(transport);
        Ok(Async::NotReady)
      },
    }
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: Sync+Send,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
        trace!("stream poll");
        let value = match self.upstream.poll() {
            Ok(Async::Ready(t)) => t,
            Ok(Async::NotReady) => {
                trace!("upstream poll gave NotReady");
                return Ok(Async::NotReady);
            },
            Err(e) => {
                error!("upstream poll gave error: {:?}", e);
                return Err(From::from(e));
            },
        };

        if let Some(frame) = value {
            trace!("upstream poll gave frame: {:?}", frame);
            if let Err(e) = self.conn.handle_frame(frame) {
                let err = format!("failed to handle frame: {:?}", e);
                return Err(io::Error::new(io::ErrorKind::Other, err));
            }
            self.send_frames()?;
            Ok(Async::Ready(Some(())))
        } else {
            error!("upstream poll gave Ready(None)");
            Ok(Async::Ready(None))
        }
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite,
          T: Sync+Send {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}
<commit_msg>futures: add comment about send_and_handle_frames return values<commit_after>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::cmp;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use futures::{Async,Poll,Sink,Stream,StartSend,Future,future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use channel::BasicProperties;
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec {
    pub frame_max: u32,
}

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length    = buf.len();
      // Ensure we at least allocate 8192 so that the buffer is big enough for the frame_max
      // negociation. Afterwards, use frame_max if > 8192.
      let frame_max = cmp::max(self.frame_max, 8192) as usize;
      if length < frame_max {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(frame_max - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream: Framed<T,AMQPCodec>,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: Sync+Send,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(stream: T, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_frame_max(options.frame_max);
    conn.set_heartbeat(options.heartbeat);
    if let Err(e) = conn.connect() {
      let err = format!("Failed to connect: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let codec = AMQPCodec {
      frame_max: conn.configuration.frame_max,
    };
    let mut t = AMQPTransport {
      upstream:  stream.framed(codec),
      conn:      conn,
    };

    if let Err(e) = t.send_and_handle_frames() {
      let err = format!("Failed to handle frames: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let mut connector = AMQPTransportConnector {
      transport: Some(t),
    };

    trace!("pre-poll");
    if let Err(e) = connector.poll() {
      let err = format!("Failed to handle frames: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }
    trace!("post-poll");

    Box::new(connector)
  }

  // Note that this can only return one of
  // - Error
  // - Async::NotReady
  // - Async::Ready(None)
  // All other results are handled until one of these three is reached.
  pub fn send_and_handle_frames(&mut self) -> Poll<Option<()>, io::Error> {
    self.send_frames()?;
    self.handle_frames()
  }

  pub fn send_and_handle_frames_with_payload(&mut self, channel_id: u16, payload: &[u8], properties: BasicProperties) -> Poll<Option<()>, io::Error> {
      self.send_and_handle_frames()?;
      self.conn.send_content_frames(channel_id, 60, payload, properties);
      self.send_and_handle_frames()
  }

  fn send_frames(&mut self) -> Result<(), io::Error> {
    //FIXME: find a way to use a future here
    trace!("send frames");
    while let Some(f) = self.conn.next_frame() {
      if let Err(e) = self.send_frame(f) {
        error!("Failed to send frame: {:?}", e);
        return Err(e);
      }
    }
    Ok(())
  }

  pub fn send_frame(&mut self, frame: Frame) -> Poll<(), io::Error> {
      self.start_send(frame).and_then(|_| self.poll_complete())
  }

  fn handle_frames(&mut self) -> Poll<Option<()>, io::Error> {
    trace!("handle frames");
    loop {
      // try_ready will return if we hit an error or NotReady.
      if try_ready!(self.poll()).is_none() {
        return Ok(Async::Ready(None));
      }
    }
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: Sync+Send,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.send_and_handle_frames()?;

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    match transport.poll()? {
      Async::Ready(Some(_)) => {
        if transport.conn.state == ConnectionState::Connected {
          // Upstream had frames available and we're connected, the transport is ready
          Ok(Async::Ready(transport))
        } else {
          // Upstream had frames but we're not yet connected, continue polling
          let poll_ret = transport.poll();
          self.transport = Some(transport);
          poll_ret?;
          Ok(Async::NotReady)
        }
      },
      _ => {
        // Upstream had no frames
        self.transport = Some(transport);
        Ok(Async::NotReady)
      },
    }
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: Sync+Send,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
        trace!("stream poll");
        let value = match self.upstream.poll() {
            Ok(Async::Ready(t)) => t,
            Ok(Async::NotReady) => {
                trace!("upstream poll gave NotReady");
                return Ok(Async::NotReady);
            },
            Err(e) => {
                error!("upstream poll gave error: {:?}", e);
                return Err(From::from(e));
            },
        };

        if let Some(frame) = value {
            trace!("upstream poll gave frame: {:?}", frame);
            if let Err(e) = self.conn.handle_frame(frame) {
                let err = format!("failed to handle frame: {:?}", e);
                return Err(io::Error::new(io::ErrorKind::Other, err));
            }
            self.send_frames()?;
            Ok(Async::Ready(Some(())))
        } else {
            error!("upstream poll gave Ready(None)");
            Ok(Async::Ready(None))
        }
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite,
          T: Sync+Send {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}
<|endoftext|>


<commit_before>///! Input will listens to user input, modify the query string, send special
///! keystrokes(such as Enter, Ctrl-p, Ctrl-n, etc) to the controller.
use crate::event::{parse_action, Event, EventArg};
use regex::Regex;
use std::collections::HashMap;
use tuikit::event::Event as TermEvent;
use tuikit::key::{from_keyname, Key};

#[derive(Debug, Clone, PartialEq)]
pub enum ActionArg {
    Char(char),
    String(String),
    None,
}

pub type Action = (Event, ActionArg);
pub type ActionChain = Vec<Action>;

pub struct Input {
    keymap: HashMap<Key, ActionChain>,
}

impl From<ActionArg> for EventArg {
    fn from(arg: ActionArg) -> Self {
        match arg {
            ActionArg::Char(ch) => Box::new(ch),
            ActionArg::String(string) => Box::new(Some(string)),
            ActionArg::None => Box::new(None as Option<String>),
        }
    }
}

impl Input {
    pub fn new() -> Self {
        Input {
            keymap: get_default_key_map(),
        }
    }

    pub fn translate_event(&self, event: TermEvent) -> Vec<(Event, EventArg)> {
        match event {
            // search event from keymap
            TermEvent::Key(key) => self
                .keymap
                .get(&key)
                .map(|chain| {
                    chain
                        .iter()
                        .map(|(action, arg)| (action.clone(), arg.clone().into()))
                        .collect()
                })
                .unwrap_or_else(|| {
                    if let Key::Char(ch) = key {
                        vec![(Event::EvActAddChar, Box::new(ch) as EventArg)]
                    } else {
                        vec![(Event::EvInputKey, Box::new(key) as EventArg)]
                    }
                }),
            TermEvent::Resize { .. } => vec![(Event::EvActRedraw, Box::new(true))],
            _ => vec![(Event::EvInputInvalid, Box::new(true))],
        }
    }

    pub fn bind(&mut self, key: &str, action_chain: ActionChain) {
        let key = from_keyname(key);
        if key == None || action_chain.is_empty() {
            return;
        }

        let key = key.unwrap();

        // remove the key for existing keymap;
        let _ = self.keymap.remove(&key);
        self.keymap.entry(key).or_insert(action_chain);
    }

    pub fn parse_keymaps(&mut self, maps: &[&str]) {
        for &map in maps {
            self.parse_keymap(map);
        }
    }

    // key_action is comma separated: 'ctrl-j:accept,ctrl-k:kill-line'
    pub fn parse_keymap(&mut self, key_action: &str) {
        for (key, action_chain) in parse_key_action(key_action).into_iter() {
            let action_chain = action_chain
                .into_iter()
                .filter_map(|(action, arg)| {
                    parse_action(action).map(|act| (act, arg.map(ActionArg::String).unwrap_or(ActionArg::None)))
                })
                .collect();
            self.bind(key, action_chain);
        }
    }

    pub fn parse_expect_keys(&mut self, keys: Option<&str>) {
        if let Some(keys) = keys {
            self.bind("enter", vec![(Event::EvActAccept, ActionArg::String("".to_string()))]);
            for key in keys.split(',') {
                self.bind(key, vec![(Event::EvActAccept, ActionArg::String(key.to_string()))]);
            }
        }
    }
}

/// parse key action string to `(key, action, argument)` tuple
/// key_action is comma separated: 'ctrl-j:accept,ctrl-k:kill-line'
pub fn parse_key_action(key_action: &str) -> Vec<(&str, Vec<(&str, Option<String>)>)> {
    lazy_static! {
        // match `key:action` or `key:action:arg` or `key:action(arg)` etc.
        static ref RE: Regex =
            Regex::new(r#"(?si)([^:]+?):((?:\+?[a-z-]+?(?:"[^"]*?"|'[^']*?'|\([^\)]*?\)|\[[^\]]*?\]|:[^:]*?)?\s*)+)(?:,|$)"#)
                .unwrap();
        // grab key, action and arg out.
        static ref RE_BIND: Regex = Regex::new(r#"(?si)([a-z-]+)(?:[:\(\["'](.+?)[\)"'\]]?)?(?:\+|$)"#).unwrap();
    }

    RE.captures_iter(&key_action)
        .map(|caps| {
            let key = caps.get(1).unwrap().as_str();
            let actions = RE_BIND
                .captures_iter(caps.get(2).unwrap().as_str())
                .map(|caps| {
                    (
                        caps.get(1).unwrap().as_str(),
                        caps.get(2).map(|s| s.as_str().to_string()),
                    )
                })
                .collect();
            (key, actions)
        })
        .collect()
}

/// e.g. execute(...) => Some(Event::EvActExecute, Box::new(Option("...")))
pub fn parse_action_arg(action_arg: &str) -> Option<(Event, EventArg)> {
    // construct a fake key_action: `fake_key:action(arg)`
    let fake_key_action = format!("fake_key:{}", action_arg);
    // get keys: [(key, [(action, arg), (action, arg)]), ...]
    let keys = parse_key_action(&fake_key_action);
    // only get the first key(since it is faked), and get the first action
    if keys.is_empty() || keys[0].1.is_empty() {
        None
    } else {
        // first action pair of key(keys[0].1) and first action (keys[0].1[0])
        let (action, new_arg) = keys[0].1[0].clone();
        parse_action(action).map(|act| (act, Box::new(new_arg) as EventArg))
    }
}

#[rustfmt::skip]
fn get_default_key_map() -> HashMap<Key, ActionChain> {
    use self::ActionArg::*;
    let mut ret = HashMap::new();
    ret.insert(Key::ESC,          vec![(Event::EvActAbort,              None)]);
    ret.insert(Key::Ctrl('c'),    vec![(Event::EvActAbort,              None)]);
    ret.insert(Key::Ctrl('g'),    vec![(Event::EvActAbort,              None)]);
    ret.insert(Key::Enter,        vec![(Event::EvActAccept,             None)]);
    ret.insert(Key::Left,         vec![(Event::EvActBackwardChar,       None)]);
    ret.insert(Key::Ctrl('b'),    vec![(Event::EvActBackwardChar,       None)]);
    ret.insert(Key::Ctrl('h'),    vec![(Event::EvActBackwardDeleteChar, None)]);
    ret.insert(Key::Backspace,    vec![(Event::EvActBackwardDeleteChar, None)]);
    ret.insert(Key::AltBackspace, vec![(Event::EvActBackwardKillWord,   None)]);
    ret.insert(Key::Alt('b'),     vec![(Event::EvActBackwardWord,       None)]);
    ret.insert(Key::ShiftLeft,    vec![(Event::EvActBackwardWord,       None)]);
    ret.insert(Key::Ctrl('a'),    vec![(Event::EvActBeginningOfLine,    None)]);
    ret.insert(Key::Ctrl('l'),    vec![(Event::EvActClearScreen,        None)]);
    ret.insert(Key::Delete,       vec![(Event::EvActDeleteChar,         None)]);
    ret.insert(Key::Ctrl('d'),    vec![(Event::EvActDeleteCharEOF,      None)]);
    ret.insert(Key::Ctrl('j'),    vec![(Event::EvActDown,               None)]);
    ret.insert(Key::Ctrl('n'),    vec![(Event::EvActDown,               None)]);
    ret.insert(Key::Down,         vec![(Event::EvActDown,               None)]);
    ret.insert(Key::Ctrl('e'),    vec![(Event::EvActEndOfLine,          None)]);
    ret.insert(Key::End,          vec![(Event::EvActEndOfLine,          None)]);
    ret.insert(Key::Ctrl('f'),    vec![(Event::EvActForwardChar,        None)]);
    ret.insert(Key::Right,        vec![(Event::EvActForwardChar,        None)]);
    ret.insert(Key::Alt('f'),     vec![(Event::EvActForwardWord,        None)]);
    ret.insert(Key::ShiftRight,   vec![(Event::EvActForwardWord,        None)]);
    ret.insert(Key::Alt('d'),     vec![(Event::EvActKillWord,           None)]);
    ret.insert(Key::ShiftUp,      vec![(Event::EvActPreviewPageUp,      None)]);
    ret.insert(Key::ShiftDown,    vec![(Event::EvActPreviewPageDown,    None)]);
    ret.insert(Key::PageDown,     vec![(Event::EvActPageDown,           None)]);
    ret.insert(Key::PageUp,       vec![(Event::EvActPageUp,             None)]);
    ret.insert(Key::Ctrl('r'),    vec![(Event::EvActRotateMode,         None)]);
    ret.insert(Key::Alt('h'),     vec![(Event::EvActScrollLeft,         None)]);
    ret.insert(Key::Alt('l'),     vec![(Event::EvActScrollRight,        None)]);
    ret.insert(Key::Tab,          vec![(Event::EvActToggle,             None), (Event::EvActDown, None)]);
    ret.insert(Key::Ctrl('q'),    vec![(Event::EvActToggleInteractive,  None)]);
    ret.insert(Key::BackTab,      vec![(Event::EvActToggle,             None), (Event::EvActUp,   None)]);
    ret.insert(Key::Ctrl('u'),    vec![(Event::EvActUnixLineDiscard,    None)]);
    ret.insert(Key::Ctrl('w'),    vec![(Event::EvActUnixWordRubout,     None)]);
    ret.insert(Key::Ctrl('p'),    vec![(Event::EvActUp,                 None)]);
    ret.insert(Key::Ctrl('k'),    vec![(Event::EvActUp,                 None)]);
    ret.insert(Key::Up,           vec![(Event::EvActUp,                 None)]);
    ret.insert(Key::Ctrl('y'),    vec![(Event::EvActYank,               None)]);
    ret.insert(Key::Null,         vec![(Event::EvActAbort,              None)]);
    ret
}

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

    #[test]
    fn execute_should_be_parsed_correctly() {
        // example from https://github.com/lotabout/skim/issues/73
        let cmd = "
        (grep -o '[a-f0-9]\\{7\\}' | head -1 |
        xargs -I % sh -c 'git show --color=always % | less -R') << 'FZF-EOF'
        {}
        FZF-EOF";

        let key_action_str = format!("ctrl-s:toggle-sort,ctrl-m:execute:{},ctrl-t:toggle", cmd);

        let key_action = parse_key_action(&key_action_str);
        assert_eq!(("ctrl-s", vec![("toggle-sort", None)]), key_action[0]);
        assert_eq!(("ctrl-m", vec![("execute", Some(cmd.to_string()))]), key_action[1]);
        assert_eq!(("ctrl-t", vec![("toggle", None)]), key_action[2]);

        let key_action_str = "f1:execute(less -f {}),ctrl-y:execute-silent(echo {} | pbcopy)";
        let key_action = parse_key_action(&key_action_str);
        assert_eq!(("f1", vec![("execute", Some("less -f {}".to_string()))]), key_action[0]);
        assert_eq!(
            ("ctrl-y", vec![("execute-silent", Some("echo {} | pbcopy".to_string()))]),
            key_action[1]
        );
    }

    #[test]
    fn action_chain_should_be_parsed() {
        let key_action = parse_key_action("ctrl-t:toggle+up");
        assert_eq!(("ctrl-t", vec![("toggle", None), ("up", None)]), key_action[0]);

        let key_action_str = "f1:execute(less -f {}),ctrl-y:execute-silent(echo {} | pbcopy)+abort";
        let key_action = parse_key_action(&key_action_str);
        assert_eq!(("f1", vec![("execute", Some("less -f {}".to_string()))]), key_action[0]);
        assert_eq!(
            (
                "ctrl-y",
                vec![
                    ("execute-silent", Some("echo {} | pbcopy".to_string())),
                    ("abort", None)
                ]
            ),
            key_action[1]
        );
    }
}
<commit_msg>fix #196: `+` in execute was eaten by skim<commit_after>///! Input will listens to user input, modify the query string, send special
///! keystrokes(such as Enter, Ctrl-p, Ctrl-n, etc) to the controller.
use crate::event::{parse_action, Event, EventArg};
use regex::Regex;
use std::collections::HashMap;
use tuikit::event::Event as TermEvent;
use tuikit::key::{from_keyname, Key};

#[derive(Debug, Clone, PartialEq)]
pub enum ActionArg {
    Char(char),
    String(String),
    None,
}

pub type Action = (Event, ActionArg);
pub type ActionChain = Vec<Action>;

pub struct Input {
    keymap: HashMap<Key, ActionChain>,
}

impl From<ActionArg> for EventArg {
    fn from(arg: ActionArg) -> Self {
        match arg {
            ActionArg::Char(ch) => Box::new(ch),
            ActionArg::String(string) => Box::new(Some(string)),
            ActionArg::None => Box::new(None as Option<String>),
        }
    }
}

impl Input {
    pub fn new() -> Self {
        Input {
            keymap: get_default_key_map(),
        }
    }

    pub fn translate_event(&self, event: TermEvent) -> Vec<(Event, EventArg)> {
        match event {
            // search event from keymap
            TermEvent::Key(key) => self
                .keymap
                .get(&key)
                .map(|chain| {
                    chain
                        .iter()
                        .map(|(action, arg)| (action.clone(), arg.clone().into()))
                        .collect()
                })
                .unwrap_or_else(|| {
                    if let Key::Char(ch) = key {
                        vec![(Event::EvActAddChar, Box::new(ch) as EventArg)]
                    } else {
                        vec![(Event::EvInputKey, Box::new(key) as EventArg)]
                    }
                }),
            TermEvent::Resize { .. } => vec![(Event::EvActRedraw, Box::new(true))],
            _ => vec![(Event::EvInputInvalid, Box::new(true))],
        }
    }

    pub fn bind(&mut self, key: &str, action_chain: ActionChain) {
        let key = from_keyname(key);
        if key == None || action_chain.is_empty() {
            return;
        }

        let key = key.unwrap();

        // remove the key for existing keymap;
        let _ = self.keymap.remove(&key);
        self.keymap.entry(key).or_insert(action_chain);
    }

    pub fn parse_keymaps(&mut self, maps: &[&str]) {
        for &map in maps {
            self.parse_keymap(map);
        }
    }

    // key_action is comma separated: 'ctrl-j:accept,ctrl-k:kill-line'
    pub fn parse_keymap(&mut self, key_action: &str) {
        debug!("got key_action: {:?}", key_action);
        for (key, action_chain) in parse_key_action(key_action).into_iter() {
            debug!("parsed key_action: {:?}: {:?}", key, action_chain);
            let action_chain = action_chain
                .into_iter()
                .filter_map(|(action, arg)| {
                    parse_action(action).map(|act| (act, arg.map(ActionArg::String).unwrap_or(ActionArg::None)))
                })
                .collect();
            self.bind(key, action_chain);
        }
    }

    pub fn parse_expect_keys(&mut self, keys: Option<&str>) {
        if let Some(keys) = keys {
            self.bind("enter", vec![(Event::EvActAccept, ActionArg::String("".to_string()))]);
            for key in keys.split(',') {
                self.bind(key, vec![(Event::EvActAccept, ActionArg::String(key.to_string()))]);
            }
        }
    }
}

/// parse key action string to `(key, action, argument)` tuple
/// key_action is comma separated: 'ctrl-j:accept,ctrl-k:kill-line'
pub fn parse_key_action(key_action: &str) -> Vec<(&str, Vec<(&str, Option<String>)>)> {
    lazy_static! {
        // match `key:action` or `key:action:arg` or `key:action(arg)` etc.
        static ref RE: Regex =
            Regex::new(r#"(?si)([^:]+?):((?:\+?[a-z-]+?(?:"[^"]*?"|'[^']*?'|\([^\)]*?\)|\[[^\]]*?\]|:[^:]*?)?\s*)+)(?:,|$)"#)
                .unwrap();
        // grab key, action and arg out.
        static ref RE_BIND: Regex = Regex::new(r#"(?si)([a-z-]+)("[^"]+?"|'[^']+?'|\([^\)]+?\)|\[[^\]]+?\]|:[^:]+?)?(?:\+|$)"#).unwrap();
    }

    RE.captures_iter(&key_action)
        .map(|caps| {
            debug!("RE: caps: {:?}", caps);
            let key = caps.get(1).unwrap().as_str();
            let actions = RE_BIND
                .captures_iter(caps.get(2).unwrap().as_str())
                .map(|caps| {
                    debug!("RE_BIND: caps: {:?}", caps);
                    (
                        caps.get(1).unwrap().as_str(),
                        caps.get(2).map(|s| {
                            // (arg) => arg, :end_arg => arg
                            let action = s.as_str();
                            if action.starts_with(":") {
                                return action[1..].to_string();
                            } else {
                                return action[1..action.len() - 1].to_string();
                            }
                        }),
                    )
                })
                .collect();
            (key, actions)
        })
        .collect()
}

/// e.g. execute(...) => Some(Event::EvActExecute, Box::new(Option("...")))
pub fn parse_action_arg(action_arg: &str) -> Option<(Event, EventArg)> {
    // construct a fake key_action: `fake_key:action(arg)`
    let fake_key_action = format!("fake_key:{}", action_arg);
    // get keys: [(key, [(action, arg), (action, arg)]), ...]
    let keys = parse_key_action(&fake_key_action);
    // only get the first key(since it is faked), and get the first action
    if keys.is_empty() || keys[0].1.is_empty() {
        None
    } else {
        // first action pair of key(keys[0].1) and first action (keys[0].1[0])
        let (action, new_arg) = keys[0].1[0].clone();
        parse_action(action).map(|act| (act, Box::new(new_arg) as EventArg))
    }
}

#[rustfmt::skip]
fn get_default_key_map() -> HashMap<Key, ActionChain> {
    use self::ActionArg::*;
    let mut ret = HashMap::new();
    ret.insert(Key::ESC,          vec![(Event::EvActAbort,              None)]);
    ret.insert(Key::Ctrl('c'),    vec![(Event::EvActAbort,              None)]);
    ret.insert(Key::Ctrl('g'),    vec![(Event::EvActAbort,              None)]);
    ret.insert(Key::Enter,        vec![(Event::EvActAccept,             None)]);
    ret.insert(Key::Left,         vec![(Event::EvActBackwardChar,       None)]);
    ret.insert(Key::Ctrl('b'),    vec![(Event::EvActBackwardChar,       None)]);
    ret.insert(Key::Ctrl('h'),    vec![(Event::EvActBackwardDeleteChar, None)]);
    ret.insert(Key::Backspace,    vec![(Event::EvActBackwardDeleteChar, None)]);
    ret.insert(Key::AltBackspace, vec![(Event::EvActBackwardKillWord,   None)]);
    ret.insert(Key::Alt('b'),     vec![(Event::EvActBackwardWord,       None)]);
    ret.insert(Key::ShiftLeft,    vec![(Event::EvActBackwardWord,       None)]);
    ret.insert(Key::Ctrl('a'),    vec![(Event::EvActBeginningOfLine,    None)]);
    ret.insert(Key::Ctrl('l'),    vec![(Event::EvActClearScreen,        None)]);
    ret.insert(Key::Delete,       vec![(Event::EvActDeleteChar,         None)]);
    ret.insert(Key::Ctrl('d'),    vec![(Event::EvActDeleteCharEOF,      None)]);
    ret.insert(Key::Ctrl('j'),    vec![(Event::EvActDown,               None)]);
    ret.insert(Key::Ctrl('n'),    vec![(Event::EvActDown,               None)]);
    ret.insert(Key::Down,         vec![(Event::EvActDown,               None)]);
    ret.insert(Key::Ctrl('e'),    vec![(Event::EvActEndOfLine,          None)]);
    ret.insert(Key::End,          vec![(Event::EvActEndOfLine,          None)]);
    ret.insert(Key::Ctrl('f'),    vec![(Event::EvActForwardChar,        None)]);
    ret.insert(Key::Right,        vec![(Event::EvActForwardChar,        None)]);
    ret.insert(Key::Alt('f'),     vec![(Event::EvActForwardWord,        None)]);
    ret.insert(Key::ShiftRight,   vec![(Event::EvActForwardWord,        None)]);
    ret.insert(Key::Alt('d'),     vec![(Event::EvActKillWord,           None)]);
    ret.insert(Key::ShiftUp,      vec![(Event::EvActPreviewPageUp,      None)]);
    ret.insert(Key::ShiftDown,    vec![(Event::EvActPreviewPageDown,    None)]);
    ret.insert(Key::PageDown,     vec![(Event::EvActPageDown,           None)]);
    ret.insert(Key::PageUp,       vec![(Event::EvActPageUp,             None)]);
    ret.insert(Key::Ctrl('r'),    vec![(Event::EvActRotateMode,         None)]);
    ret.insert(Key::Alt('h'),     vec![(Event::EvActScrollLeft,         None)]);
    ret.insert(Key::Alt('l'),     vec![(Event::EvActScrollRight,        None)]);
    ret.insert(Key::Tab,          vec![(Event::EvActToggle,             None), (Event::EvActDown, None)]);
    ret.insert(Key::Ctrl('q'),    vec![(Event::EvActToggleInteractive,  None)]);
    ret.insert(Key::BackTab,      vec![(Event::EvActToggle,             None), (Event::EvActUp,   None)]);
    ret.insert(Key::Ctrl('u'),    vec![(Event::EvActUnixLineDiscard,    None)]);
    ret.insert(Key::Ctrl('w'),    vec![(Event::EvActUnixWordRubout,     None)]);
    ret.insert(Key::Ctrl('p'),    vec![(Event::EvActUp,                 None)]);
    ret.insert(Key::Ctrl('k'),    vec![(Event::EvActUp,                 None)]);
    ret.insert(Key::Up,           vec![(Event::EvActUp,                 None)]);
    ret.insert(Key::Ctrl('y'),    vec![(Event::EvActYank,               None)]);
    ret.insert(Key::Null,         vec![(Event::EvActAbort,              None)]);
    ret
}

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

    #[test]
    fn execute_should_be_parsed_correctly() {
        // example from https://github.com/lotabout/skim/issues/73
        let cmd = "
        (grep -o '[a-f0-9]\\{7\\}' | head -1 |
        xargs -I % sh -c 'git show --color=always % | less -R') << 'FZF-EOF'
        {}
        FZF-EOF";

        let key_action_str = format!("ctrl-s:toggle-sort,ctrl-m:execute:{},ctrl-t:toggle", cmd);

        let key_action = parse_key_action(&key_action_str);
        assert_eq!(("ctrl-s", vec![("toggle-sort", None)]), key_action[0]);
        assert_eq!(("ctrl-m", vec![("execute", Some(cmd.to_string()))]), key_action[1]);
        assert_eq!(("ctrl-t", vec![("toggle", None)]), key_action[2]);

        let key_action_str = "f1:execute(less -f {}),ctrl-y:execute-silent(echo {} | pbcopy)";
        let key_action = parse_key_action(&key_action_str);
        assert_eq!(("f1", vec![("execute", Some("less -f {}".to_string()))]), key_action[0]);
        assert_eq!(
            ("ctrl-y", vec![("execute-silent", Some("echo {} | pbcopy".to_string()))]),
            key_action[1]
        );

        // #196
        let key_action_str = "enter:execute($EDITOR +{2} {1})";
        let key_action = parse_key_action(key_action_str);
        assert_eq!(
            ("enter", vec![("execute", Some("$EDITOR +{2} {1}".to_string()))]),
            key_action[0]
        );
    }

    #[test]
    fn action_chain_should_be_parsed() {
        let key_action = parse_key_action("ctrl-t:toggle+up");
        assert_eq!(("ctrl-t", vec![("toggle", None), ("up", None)]), key_action[0]);

        let key_action_str = "f1:execute(less -f {}),ctrl-y:execute-silent(echo {} | pbcopy)+abort";
        let key_action = parse_key_action(&key_action_str);
        assert_eq!(("f1", vec![("execute", Some("less -f {}".to_string()))]), key_action[0]);
        assert_eq!(
            (
                "ctrl-y",
                vec![
                    ("execute-silent", Some("echo {} | pbcopy".to_string())),
                    ("abort", None)
                ]
            ),
            key_action[1]
        );
    }
}
<|endoftext|>


<commit_before>#![allow(warnings)]

mod hlt;
use hlt::{ networking, types };
use std::collections::HashMap;

struct MoveFeatures {
    loc: types::Location,
    d: u8,
    owner_them: u8,
    distance: i32,
    friendly: bool,
    strength_us: i32,
    strength_them: i32,
    adjacent_strength_us: i32,
    assigned_strength: i32,
    production_them: i32,
    production_us: i32,
}

fn get_total_adjacent_strength(loc: types::Location, map: &types::GameMap, my_id: u8) -> i32 {
    // Check to see if we can use multiple moves to capture
    let total_adjacent_strength: i32 = types::CARDINALS.iter()
        .map(|d| {
            let site = map.get_site_ref(loc, *d);
            if site.owner == my_id { site.strength as i32 } else { 0 }
        })
        .sum();
    total_adjacent_strength
}

fn get_best_move_simple(
    loc: types::Location,
    map: &types::GameMap,
    my_id: u8,
    current_moves: &HashMap<types::Location, i32>) -> Vec<(types::Location, u8)>  {
    let mut moves = vec![];
    for d in &types::CARDINALS {
        let proposed_loc = map.get_location(loc, *d);
        let proposed = map.get_site_ref(loc, *d);
        let current = map.get_site_ref(loc, types::STILL);
        let already_assigned_strength: i32 = *current_moves.get(&proposed_loc)
            .unwrap_or(&0i32);
        moves.push(MoveFeatures {
            loc: loc,
            d: *d,
            distance: distance_to_border(loc, *d, &map, my_id),
            owner_them: proposed.owner,
            friendly: proposed.owner == my_id,
            strength_us: current.strength as i32,
            strength_them: proposed.strength as i32,
            adjacent_strength_us: get_total_adjacent_strength(proposed_loc, &map, my_id),
            assigned_strength: already_assigned_strength,
            production_us: current.production as i32,
            production_them: proposed.production as i32,
        });
    }
    // Sort by distance
    moves.sort_by(|a, b| a.distance.cmp(&b.distance));
    // Remove all but the shortest distances and filter losing battles and strength losses
    let shortest = moves[0].distance;
    let mut moves: Vec<MoveFeatures> = moves.into_iter()
        // Don't move weak pieces
        .filter(|a| a.strength_us > a.production_us * 5)
        // Only move towards the closest border
        .filter(|a| a.distance == shortest)
        // Don't allow losing battles
        .filter(|a| a.adjacent_strength_us > (a.strength_them + a.production_them))
        // Allow a small loss so full strength don't get stuck
        .filter(|a| !a.friendly || a.strength_us + a.strength_them <= 260)
        // Don't allow too many troops to move into the same space
        .filter(|a| a.strength_us + a.assigned_strength <= 260)
        .collect();
    match moves.len() {
        0 => return vec![(loc, types::STILL)],
        _ => {},
    };
    // If there's still more than one move available, go for production
    moves.sort_by(|a, b| a.production_them.cmp(&b.production_them));
    // Prefer losing less strength
    moves.sort_by(|a, b| b.strength_them.cmp(&a.strength_them));

    let m = moves.pop().unwrap();
    if m.strength_us > m.strength_them {
        vec![(m.loc, m.d)]
    } else {
        // We need to issue additional orders for other troups
        let mut r = vec![(m.loc, m.d)];
        let proposed = map.get_location(m.loc, m.d);
        for adj in &types::CARDINALS {
            let adj_loc = map.get_location(proposed, *adj);
            r.push((adj_loc, reverse(adj)));
        }
        r
    }
}

fn distance_to_border(loc: types::Location, dir: u8, map: &types::GameMap, my_id: u8) -> i32 {
    let mut l = loc.clone();
    let mut counter = 0;
    loop {
        if map.get_site_ref(l, types::STILL).owner != my_id {
            return counter
        }
        l = map.get_location(l, dir);
        if l == loc {
            return map.width as i32
        }
        counter += 1;
    }
}

fn get_units_of_player(id: u8, map: &types::GameMap) -> Vec<types::Location> {
    let mut result = Vec::new();
    for a in 0..map.height {
        for b in 0..map.width {
            let l = types::Location { x: b, y: a };
            let site = map.get_site_ref(l, types::STILL);
            if site.owner == id {
                result.push(l);
            }
        }
    }
    result
}

fn local_best_strategy(my_id: u8, game_map: &types::GameMap) -> HashMap<types::Location, u8> {
    let mut moves = HashMap::new();
    let mut planned_moves: HashMap<types::Location, i32> = HashMap::new();
    for l in get_units_of_player(my_id, game_map) {
        if moves.contains_key(&l) {
            continue;
        }

        let next = get_best_move_simple(l, game_map, my_id, &planned_moves);
        for (loc, d) in next {
            let site = game_map.get_site_ref(loc, types::STILL);
            let next_loc = game_map.get_location(loc, d);

            // Record the planned move
            let strength = site.strength as i32;
            *planned_moves.entry(loc).or_insert(strength) -= strength;
            *planned_moves.entry(next_loc).or_insert(0) += strength;

            moves.insert(loc, d);
        }
    }
    moves
}

fn max_capture_strategy(my_id: u8, game_map: &types::GameMap) -> HashMap<types::Location, u8> {
    // First identify the weakest target
    let my_units = get_units_of_player(my_id, &game_map);
    let mut possibilities = my_units
        .iter()
        .flat_map(|l| {
            types::CARDINALS.iter()
                .cloned()
                .map(move |d| (*l, d))
                .collect::<Vec<_>>()
        })
        .map(|(l, d)| {
            let proposed = game_map.get_site_ref(l, d);
            let current = game_map.get_site_ref(l, types::STILL);
            MoveFeatures {
                loc: l,
                d: d,
                distance: 1,
                owner_them: proposed.owner,
                friendly: proposed.owner == my_id,
                strength_us: current.strength as i32,
                strength_them: proposed.strength as i32,
                adjacent_strength_us: get_total_adjacent_strength(l, &game_map, my_id),
                assigned_strength: 0,
                production_us: current.production as i32,
                production_them: proposed.production as i32,
            }
        })
        // Don't move weak pieces
        .filter(|a| a.strength_us > a.production_us * 3)
        // Only consider moves that move us toward victory!
        .filter(|a| !a.friendly)
        .filter(|a| a.adjacent_strength_us + a.strength_us > a.strength_them)
        .collect::<Vec<_>>();
    possibilities.sort_by(|a, b| a.strength_us.cmp(&b.strength_us));
    possibilities.sort_by(|a, b| b.strength_them.cmp(&a.strength_them));

    let mut moves = HashMap::new();
    while possibilities.len() > 0 {
        let m = possibilities.pop().unwrap();
        if moves.contains_key(&m.loc) {
            continue
        }
        if m.strength_us > m.strength_them {
            // If we can capture, do so
            moves.insert(m.loc, m.d);
        } else {
            // Otherwise, move everything towards that point
            for adj in &types::CARDINALS {
                let adj_loc = game_map.get_location(m.loc, *adj);
                if !moves.contains_key(&adj_loc) {
                    moves.insert(adj_loc, reverse(adj));
                }
            }
            moves.insert(m.loc, types::STILL);
        }
    }
    for remaining in my_units {
        if !moves.contains_key(&remaining) {
            moves.insert(remaining, types::STILL);
        }
    }
    moves
}

fn reverse(d: &u8) -> u8 {
    match d {
        &types::NORTH => types::SOUTH,
        &types::SOUTH => types::NORTH,
        &types::WEST => types::EAST,
        &types::EAST => types::WEST,
        _ => panic!("unknown direction {}", d),
    }
}

fn find_poi(map: &types::GameMap, my_id: u8) -> Vec<types::Location> {
    let mut result = Vec::new();
    for a in 0..map.height {
        for b in 0..map.width {
            let l = types::Location { x: b, y: a };
            let site = map.get_site_ref(l, types::STILL);
            if site.production >= 5 && site.owner != my_id {
                result.push(l);
            }
        }
    }
    result
}

fn find_closest_poi(l: types::Location, map: &types::GameMap, poi: &Vec<types::Location>) -> types::Location {
    poi.iter()
    .min_by_key(|p| map.get_distance(l, **p)).unwrap_or(&types::Location{x: 0, y: 0}).clone()
}

fn poi_strategy(my_id: u8, game_map: &types::GameMap) -> HashMap<types::Location, u8> {
    let poi = find_poi(game_map, my_id);
    let my_units = get_units_of_player(my_id, &game_map);
    let unit_count = my_units.len();
    let mut moves = HashMap::new();
    let mut targets: HashMap<types::Location, ()> = HashMap::new();
    for l in my_units {
        let site = game_map.get_site_ref(l, types::STILL);
        if site.strength < site.production * 5 {
            moves.insert(l, types::STILL);
        } else {
            let closest = find_closest_poi(l, game_map, &poi);
            let d = game_map.get_direction(l, closest);
            let proposed_loc = game_map.get_location(l, d);
            if site.strength > game_map.get_site_ref(proposed_loc, types::STILL).strength {
                moves.insert(l, d);
            }
            targets.insert(closest, ());
        }
    }
    //log(
    //    format!("{} POI, {} units, {} distinct targets\n", poi.len(), unit_count, targets.len()),
    //    my_id
    //);
    moves
}

use std::io::prelude::*;
use std::fs::OpenOptions;
fn log(s: String, id: u8) {
    let mut file = OpenOptions::new()
         .append(true)
         .create(true)
         .open(format!("output-{}.log", id))
         .unwrap();
    file.write(s.as_bytes()).unwrap();
}

fn main() {
    let (my_id, mut game_map) = networking::get_init();
    networking::send_init(format!("{}{}", "Asp2Insp".to_string(), my_id.to_string()));
    let mut turn_counter = 0;
    loop {
        networking::get_frame(&mut game_map);
        let moves = poi_strategy(my_id, &game_map);

        //let my_count = get_units_of_player(my_id, &game_map).len();
        //let moves = if my_count < 20 {
        //    max_capture_strategy(my_id, &game_map)
        //} else {
        //    local_best_strategy(my_id, &game_map)
        //};
        networking::send_frame(moves);
        turn_counter += 1;
    }
}
<commit_msg>Got a working POI finder<commit_after>#![allow(warnings)]

mod hlt;
use hlt::{ networking, types };
use std::collections::HashMap;

struct MoveFeatures {
    loc: types::Location,
    d: u8,
    owner_them: u8,
    distance: i32,
    friendly: bool,
    strength_us: i32,
    strength_them: i32,
    adjacent_strength_us: i32,
    assigned_strength: i32,
    production_them: i32,
    production_us: i32,
}

fn get_total_adjacent_strength(loc: types::Location, map: &types::GameMap, my_id: u8) -> i32 {
    // Check to see if we can use multiple moves to capture
    let total_adjacent_strength: i32 = types::CARDINALS.iter()
        .map(|d| {
            let site = map.get_site_ref(loc, *d);
            if site.owner == my_id { site.strength as i32 } else { 0 }
        })
        .sum();
    total_adjacent_strength
}

fn get_best_move_simple(
    loc: types::Location,
    map: &types::GameMap,
    my_id: u8,
    current_moves: &HashMap<types::Location, i32>) -> Vec<(types::Location, u8)>  {
    let mut moves = vec![];
    for d in &types::CARDINALS {
        let proposed_loc = map.get_location(loc, *d);
        let proposed = map.get_site_ref(loc, *d);
        let current = map.get_site_ref(loc, types::STILL);
        let already_assigned_strength: i32 = *current_moves.get(&proposed_loc)
            .unwrap_or(&0i32);
        moves.push(MoveFeatures {
            loc: loc,
            d: *d,
            distance: distance_to_border(loc, *d, &map, my_id),
            owner_them: proposed.owner,
            friendly: proposed.owner == my_id,
            strength_us: current.strength as i32,
            strength_them: proposed.strength as i32,
            adjacent_strength_us: get_total_adjacent_strength(proposed_loc, &map, my_id),
            assigned_strength: already_assigned_strength,
            production_us: current.production as i32,
            production_them: proposed.production as i32,
        });
    }
    // Sort by distance
    moves.sort_by(|a, b| a.distance.cmp(&b.distance));
    // Remove all but the shortest distances and filter losing battles and strength losses
    let shortest = moves[0].distance;
    let mut moves: Vec<MoveFeatures> = moves.into_iter()
        // Don't move weak pieces
        .filter(|a| a.strength_us > a.production_us * 5)
        // Only move towards the closest border
        .filter(|a| a.distance == shortest)
        // Don't allow losing battles
        .filter(|a| a.adjacent_strength_us > (a.strength_them + a.production_them))
        // Allow a small loss so full strength don't get stuck
        .filter(|a| !a.friendly || a.strength_us + a.strength_them <= 260)
        // Don't allow too many troops to move into the same space
        .filter(|a| a.strength_us + a.assigned_strength <= 260)
        .collect();
    match moves.len() {
        0 => return vec![(loc, types::STILL)],
        _ => {},
    };
    // If there's still more than one move available, go for production
    moves.sort_by(|a, b| a.production_them.cmp(&b.production_them));
    // Prefer losing less strength
    moves.sort_by(|a, b| b.strength_them.cmp(&a.strength_them));

    let m = moves.pop().unwrap();
    if m.strength_us > m.strength_them {
        vec![(m.loc, m.d)]
    } else {
        // We need to issue additional orders for other troups
        let mut r = vec![(m.loc, m.d)];
        let proposed = map.get_location(m.loc, m.d);
        for adj in &types::CARDINALS {
            let adj_loc = map.get_location(proposed, *adj);
            r.push((adj_loc, reverse(adj)));
        }
        r
    }
}

fn distance_to_border(loc: types::Location, dir: u8, map: &types::GameMap, my_id: u8) -> i32 {
    let mut l = loc.clone();
    let mut counter = 0;
    loop {
        if map.get_site_ref(l, types::STILL).owner != my_id {
            return counter
        }
        l = map.get_location(l, dir);
        if l == loc {
            return map.width as i32
        }
        counter += 1;
    }
}

fn get_units_of_player(id: u8, map: &types::GameMap) -> Vec<types::Location> {
    let mut result = Vec::new();
    for a in 0..map.height {
        for b in 0..map.width {
            let l = types::Location { x: b, y: a };
            let site = map.get_site_ref(l, types::STILL);
            if site.owner == id {
                result.push(l);
            }
        }
    }
    result
}

fn local_best_strategy(my_id: u8, game_map: &types::GameMap) -> HashMap<types::Location, u8> {
    let mut moves = HashMap::new();
    let mut planned_moves: HashMap<types::Location, i32> = HashMap::new();
    for l in get_units_of_player(my_id, game_map) {
        if moves.contains_key(&l) {
            continue;
        }

        let next = get_best_move_simple(l, game_map, my_id, &planned_moves);
        for (loc, d) in next {
            let site = game_map.get_site_ref(loc, types::STILL);
            let next_loc = game_map.get_location(loc, d);

            // Record the planned move
            let strength = site.strength as i32;
            *planned_moves.entry(loc).or_insert(strength) -= strength;
            *planned_moves.entry(next_loc).or_insert(0) += strength;

            moves.insert(loc, d);
        }
    }
    moves
}

fn max_capture_strategy(my_id: u8, game_map: &types::GameMap) -> HashMap<types::Location, u8> {
    let my_units = get_units_of_player(my_id, &game_map);
    let mut possibilities = my_units
        .iter()
        .flat_map(|l| {
            types::CARDINALS.iter()
                .cloned()
                .map(move |d| (*l, d))
                .collect::<Vec<_>>()
        })
        .map(|(l, d)| {
            let proposed = game_map.get_site_ref(l, d);
            let current = game_map.get_site_ref(l, types::STILL);
            MoveFeatures {
                loc: l,
                d: d,
                distance: 1,
                owner_them: proposed.owner,
                friendly: proposed.owner == my_id,
                strength_us: current.strength as i32,
                strength_them: proposed.strength as i32,
                adjacent_strength_us: get_total_adjacent_strength(l, &game_map, my_id),
                assigned_strength: 0,
                production_us: current.production as i32,
                production_them: proposed.production as i32,
            }
        })
        // Don't move weak pieces
        .filter(|a| a.strength_us > a.production_us * 3)
        // Only consider moves that move us toward victory!
        .filter(|a| !a.friendly)
        .filter(|a| a.adjacent_strength_us + a.strength_us > a.strength_them)
        .collect::<Vec<_>>();
    possibilities.sort_by(|a, b| b.strength_them.cmp(&a.strength_them));
    possibilities.sort_by(|a, b| a.strength_us.cmp(&b.strength_us));
    possibilities.sort_by(|a, b| a.production_them.cmp(&b.production_them));

    let mut moves = HashMap::new();
    while possibilities.len() > 0 {
        let m = possibilities.pop().unwrap();
        if moves.contains_key(&m.loc) {
            continue
        }
        if m.strength_us > m.strength_them {
            // If we can capture, do so
            moves.insert(m.loc, m.d);
        } else {
            // Otherwise, move everything towards that point
            for adj in &types::CARDINALS {
                let adj_loc = game_map.get_location(m.loc, *adj);
                if !moves.contains_key(&adj_loc) {
                    moves.insert(adj_loc, reverse(adj));
                }
            }
            moves.insert(m.loc, types::STILL);
        }
    }
    for remaining in my_units {
        if !moves.contains_key(&remaining) {
            moves.insert(remaining, types::STILL);
        }
    }
    moves
}

fn reverse(d: &u8) -> u8 {
    match d {
        &types::NORTH => types::SOUTH,
        &types::SOUTH => types::NORTH,
        &types::WEST => types::EAST,
        &types::EAST => types::WEST,
        _ => panic!("unknown direction {}", d),
    }
}

fn find_poi(map: &types::GameMap, my_id: u8) -> Vec<types::Location> {
    let mut result = Vec::new();
    for a in 0..map.height {
        for b in 0..map.width {
            let l = types::Location { x: b, y: a };
            let site = map.get_site_ref(l, types::STILL);
            if site.production >= 5 && site.owner != my_id {
                result.push(l);
            }
        }
    }
    result
}

fn find_closest_poi(l: types::Location, map: &types::GameMap, poi: &Vec<types::Location>) -> types::Location {
    poi.iter()
    .min_by_key(|p| map.get_distance(l, **p)).unwrap_or(&types::Location{x: 0, y: 0}).clone()
}

fn poi_strategy(my_id: u8, game_map: &types::GameMap) -> HashMap<types::Location, u8> {
    let poi = find_poi(game_map, my_id);
    let my_units = get_units_of_player(my_id, &game_map);
    let unit_count = my_units.len();
    let mut moves = HashMap::new();
    let mut targets: HashMap<types::Location, ()> = HashMap::new();
    for l in my_units {
        let site = game_map.get_site_ref(l, types::STILL);
        if site.strength < site.production * 5 {
            moves.insert(l, types::STILL);
        } else {
            let closest = find_closest_poi(l, game_map, &poi);
            let d = game_map.get_direction(l, closest);
            let proposed_loc = game_map.get_location(l, d);
            if site.strength > game_map.get_site_ref(proposed_loc, types::STILL).strength {
                moves.insert(l, d);
            }
            targets.insert(closest, ());
        }
    }
    //log(
    //    format!("{} POI, {} units, {} distinct targets\n", poi.len(), unit_count, targets.len()),
    //    my_id
    //);
    moves
}

use std::io::prelude::*;
use std::fs::OpenOptions;
fn log(s: String, id: u8) {
    let mut file = OpenOptions::new()
         .append(true)
         .create(true)
         .open(format!("output-{}.log", id))
         .unwrap();
    file.write(s.as_bytes()).unwrap();
}

fn main() {
    let (my_id, mut game_map) = networking::get_init();
    networking::send_init(format!("{}{}", "Asp2Insp".to_string(), my_id.to_string()));
    let mut turn_counter = 0;
    loop {
        networking::get_frame(&mut game_map);
        //let moves = poi_strategy(my_id, &game_map);

        let my_count = get_units_of_player(my_id, &game_map).len();
        let moves = if my_count < 20 {
            max_capture_strategy(my_id, &game_map)
        } else {
            local_best_strategy(my_id, &game_map)
        };
        networking::send_frame(moves);
        turn_counter += 1;
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    account_address::AccountAddress,
    account_config::{
        type_tag_for_currency_code, AccountResource, AccountRole, BalanceResource, ChildVASP,
        ChildVASPRole, CurrencyInfoResource, EmptyRole, ParentVASP, ParentVASPRole, UnhostedRole,
        ACCOUNT_RECEIVED_EVENT_PATH, ACCOUNT_SENT_EVENT_PATH,
    },
    block_metadata::{LibraBlockResource, NEW_BLOCK_EVENT_PATH},
    event::EventHandle,
    libra_timestamp::LibraTimestampResource,
    on_chain_config::{ConfigurationResource, OnChainConfig, ValidatorSet},
    validator_config::ValidatorConfigResource,
};
use anyhow::{bail, Error, Result};
use move_core_types::{identifier::Identifier, move_resource::MoveResource};
use serde::{de::DeserializeOwned, export::Formatter, Deserialize, Serialize};
use std::{collections::btree_map::BTreeMap, convert::TryFrom, fmt};

#[derive(Default, Deserialize, PartialEq, Serialize)]
pub struct AccountState(BTreeMap<Vec<u8>, Vec<u8>>);

impl AccountState {
    // By design and do not remove
    pub fn get_account_address(&self) -> Result<Option<AccountAddress>> {
        self.get_account_resource()
            .map(|opt_ar| opt_ar.map(|ar| ar.sent_events().key().get_creator_address()))
    }

    pub fn get_account_resource(&self) -> Result<Option<AccountResource>> {
        self.get_resource(&AccountResource::resource_path())
    }

    pub fn get_balance_resources(
        &self,
        currency_codes: &[Identifier],
    ) -> Result<BTreeMap<Identifier, BalanceResource>> {
        currency_codes
            .iter()
            .filter_map(|currency_code| {
                let currency_type_tag = type_tag_for_currency_code(currency_code.to_owned());
                // TODO: update this to use BalanceResource::resource_path once that takes type
                // parameters
                self.get_resource(&BalanceResource::access_path_for(currency_type_tag))
                    .transpose()
                    .map(|balance| balance.map(|b| (currency_code.to_owned(), b)))
            })
            .collect()
    }

    pub fn get_configuration_resource(&self) -> Result<Option<ConfigurationResource>> {
        self.get_resource(&ConfigurationResource::resource_path())
    }

    pub fn get_libra_timestamp_resource(&self) -> Result<Option<LibraTimestampResource>> {
        self.get_resource(&LibraTimestampResource::resource_path())
    }

    pub fn get_validator_config_resource(&self) -> Result<Option<ValidatorConfigResource>> {
        self.get_resource(&ValidatorConfigResource::resource_path())
    }

    pub fn get_account_role(&self) -> Result<Option<AccountRole>> {
        if self
            .0
            .contains_key(&AccountRole::access_path_for::<ParentVASP>())
        {
            self.get_resource(&AccountRole::access_path_for::<ParentVASP>())
                .map(|role| role.map(|role: ParentVASPRole| AccountRole::ParentVASP(role.role)))
        } else if self
            .0
            .contains_key(&AccountRole::access_path_for::<ChildVASP>())
        {
            self.get_resource(&AccountRole::access_path_for::<ChildVASP>())
                .map(|x| x.map(|role: ChildVASPRole| AccountRole::ChildVASP(role.role)))
        } else if self
            .0
            .contains_key(&AccountRole::access_path_for::<UnhostedRole>())
        {
            self.get_resource(&AccountRole::access_path_for::<UnhostedRole>())
                .map(|x| x.map(|_: UnhostedRole| AccountRole::Unhosted))
        } else if self
            .0
            .contains_key(&AccountRole::access_path_for::<EmptyRole>())
        {
            self.get_resource(&AccountRole::access_path_for::<EmptyRole>())
                .map(|x| x.map(|_: EmptyRole| AccountRole::Empty))
        } else {
            Ok(Some(AccountRole::Unknown))
        }
    }

    pub fn get_currency_info_resource(
        &self,
        currency_code: Identifier,
    ) -> Result<Option<BalanceResource>> {
        // TODO: update this to use BalanceResource::resource_path once that takes type
        // parameters
        self.get_resource(&CurrencyInfoResource::access_path_for(currency_code))
    }

    pub fn get_validator_set(&self) -> Result<Option<ValidatorSet>> {
        self.get_resource(&ValidatorSet::CONFIG_ID.access_path().path)
    }

    pub fn get_libra_block_resource(&self) -> Result<Option<LibraBlockResource>> {
        self.get_resource(&LibraBlockResource::resource_path())
    }

    pub fn get_event_handle_by_query_path(&self, query_path: &[u8]) -> Result<Option<EventHandle>> {
        let event_handle = if *ACCOUNT_RECEIVED_EVENT_PATH == query_path {
            self.get_account_resource()?
                .map(|account_resource| account_resource.received_events().clone())
        } else if *ACCOUNT_SENT_EVENT_PATH == query_path {
            self.get_account_resource()?
                .map(|account_resource| account_resource.sent_events().clone())
        } else if *NEW_BLOCK_EVENT_PATH == query_path {
            self.get_libra_block_resource()?
                .map(|libra_block_resource| libra_block_resource.new_block_events().clone())
        } else {
            bail!("Unrecognized query path: {:?}", query_path);
        };

        Ok(event_handle)
    }

    pub fn get(&self, key: &[u8]) -> Option<&Vec<u8>> {
        self.0.get(key)
    }

    pub fn get_resource<T: DeserializeOwned>(&self, key: &[u8]) -> Result<Option<T>> {
        self.0
            .get(key)
            .map(|bytes| lcs::from_bytes(bytes))
            .transpose()
            .map_err(Into::into)
    }

    pub fn insert(&mut self, key: Vec<u8>, value: Vec<u8>) -> Option<Vec<u8>> {
        self.0.insert(key, value)
    }

    pub fn remove(&mut self, key: &[u8]) -> Option<Vec<u8>> {
        self.0.remove(key)
    }

    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    pub fn iter(&self) -> impl std::iter::Iterator<Item = (&Vec<u8>, &Vec<u8>)> {
        self.0.iter()
    }
}

impl fmt::Debug for AccountState {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        // TODO: add support for other types of resources
        let account_resource_str = self
            .get_account_resource()
            .map(|account_resource_opt| format!("{:#?}", account_resource_opt))
            .unwrap_or_else(|e| format!("parse error: {:#?}", e));

        let libra_timestamp_str = self
            .get_libra_timestamp_resource()
            .map(|libra_timestamp_opt| format!("{:#?}", libra_timestamp_opt))
            .unwrap_or_else(|e| format!("parse: {:#?}", e));

        let validator_config_str = self
            .get_validator_config_resource()
            .map(|validator_config_opt| format!("{:#?}", validator_config_opt))
            .unwrap_or_else(|e| format!("parse error: {:#?}", e));

        let validator_set_str = self
            .get_validator_set()
            .map(|validator_set_opt| format!("{:#?}", validator_set_opt))
            .unwrap_or_else(|e| format!("parse error: {:#?}", e));

        write!(
            f,
            "{{ \n \
             AccountResource {{ {} }} \n \
             LibraTimestamp {{ {} }} \n \
             ValidatorConfig {{ {} }} \n \
             ValidatorSet {{ {} }} \n \
             }}",
            account_resource_str, libra_timestamp_str, validator_config_str, validator_set_str,
        )
    }
}

impl TryFrom<(&AccountResource, &BalanceResource)> for AccountState {
    type Error = Error;

    fn try_from(
        (account_resource, balance_resource): (&AccountResource, &BalanceResource),
    ) -> Result<Self> {
        let mut btree_map: BTreeMap<Vec<u8>, Vec<u8>> = BTreeMap::new();
        btree_map.insert(
            AccountResource::resource_path(),
            lcs::to_bytes(account_resource)?,
        );
        btree_map.insert(
            BalanceResource::resource_path(),
            lcs::to_bytes(balance_resource)?,
        );

        Ok(Self(btree_map))
    }
}
<commit_msg>Removing get_currency_info_resource from types/src/account_state.rs<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    account_address::AccountAddress,
    account_config::{
        type_tag_for_currency_code, AccountResource, AccountRole, BalanceResource, ChildVASP,
        ChildVASPRole, CurrencyInfoResource, EmptyRole, ParentVASP, ParentVASPRole, UnhostedRole,
        ACCOUNT_RECEIVED_EVENT_PATH, ACCOUNT_SENT_EVENT_PATH,
    },
    block_metadata::{LibraBlockResource, NEW_BLOCK_EVENT_PATH},
    event::EventHandle,
    libra_timestamp::LibraTimestampResource,
    on_chain_config::{ConfigurationResource, OnChainConfig, ValidatorSet},
    validator_config::ValidatorConfigResource,
};
use anyhow::{bail, Error, Result};
use move_core_types::{identifier::Identifier, move_resource::MoveResource};
use serde::{de::DeserializeOwned, export::Formatter, Deserialize, Serialize};
use std::{collections::btree_map::BTreeMap, convert::TryFrom, fmt};

#[derive(Default, Deserialize, PartialEq, Serialize)]
pub struct AccountState(BTreeMap<Vec<u8>, Vec<u8>>);

impl AccountState {
    // By design and do not remove
    pub fn get_account_address(&self) -> Result<Option<AccountAddress>> {
        self.get_account_resource()
            .map(|opt_ar| opt_ar.map(|ar| ar.sent_events().key().get_creator_address()))
    }

    pub fn get_account_resource(&self) -> Result<Option<AccountResource>> {
        self.get_resource(&AccountResource::resource_path())
    }

    pub fn get_balance_resources(
        &self,
        currency_codes: &[Identifier],
    ) -> Result<BTreeMap<Identifier, BalanceResource>> {
        currency_codes
            .iter()
            .filter_map(|currency_code| {
                let currency_type_tag = type_tag_for_currency_code(currency_code.to_owned());
                // TODO: update this to use BalanceResource::resource_path once that takes type
                // parameters
                self.get_resource(&BalanceResource::access_path_for(currency_type_tag))
                    .transpose()
                    .map(|balance| balance.map(|b| (currency_code.to_owned(), b)))
            })
            .collect()
    }

    pub fn get_configuration_resource(&self) -> Result<Option<ConfigurationResource>> {
        self.get_resource(&ConfigurationResource::resource_path())
    }

    pub fn get_libra_timestamp_resource(&self) -> Result<Option<LibraTimestampResource>> {
        self.get_resource(&LibraTimestampResource::resource_path())
    }

    pub fn get_validator_config_resource(&self) -> Result<Option<ValidatorConfigResource>> {
        self.get_resource(&ValidatorConfigResource::resource_path())
    }

    pub fn get_account_role(&self) -> Result<Option<AccountRole>> {
        if self
            .0
            .contains_key(&AccountRole::access_path_for::<ParentVASP>())
        {
            self.get_resource(&AccountRole::access_path_for::<ParentVASP>())
                .map(|role| role.map(|role: ParentVASPRole| AccountRole::ParentVASP(role.role)))
        } else if self
            .0
            .contains_key(&AccountRole::access_path_for::<ChildVASP>())
        {
            self.get_resource(&AccountRole::access_path_for::<ChildVASP>())
                .map(|x| x.map(|role: ChildVASPRole| AccountRole::ChildVASP(role.role)))
        } else if self
            .0
            .contains_key(&AccountRole::access_path_for::<UnhostedRole>())
        {
            self.get_resource(&AccountRole::access_path_for::<UnhostedRole>())
                .map(|x| x.map(|_: UnhostedRole| AccountRole::Unhosted))
        } else if self
            .0
            .contains_key(&AccountRole::access_path_for::<EmptyRole>())
        {
            self.get_resource(&AccountRole::access_path_for::<EmptyRole>())
                .map(|x| x.map(|_: EmptyRole| AccountRole::Empty))
        } else {
            Ok(Some(AccountRole::Unknown))
        }
    }

    pub fn get_validator_set(&self) -> Result<Option<ValidatorSet>> {
        self.get_resource(&ValidatorSet::CONFIG_ID.access_path().path)
    }

    pub fn get_libra_block_resource(&self) -> Result<Option<LibraBlockResource>> {
        self.get_resource(&LibraBlockResource::resource_path())
    }

    pub fn get_event_handle_by_query_path(&self, query_path: &[u8]) -> Result<Option<EventHandle>> {
        let event_handle = if *ACCOUNT_RECEIVED_EVENT_PATH == query_path {
            self.get_account_resource()?
                .map(|account_resource| account_resource.received_events().clone())
        } else if *ACCOUNT_SENT_EVENT_PATH == query_path {
            self.get_account_resource()?
                .map(|account_resource| account_resource.sent_events().clone())
        } else if *NEW_BLOCK_EVENT_PATH == query_path {
            self.get_libra_block_resource()?
                .map(|libra_block_resource| libra_block_resource.new_block_events().clone())
        } else {
            bail!("Unrecognized query path: {:?}", query_path);
        };

        Ok(event_handle)
    }

    pub fn get(&self, key: &[u8]) -> Option<&Vec<u8>> {
        self.0.get(key)
    }

    pub fn get_resource<T: DeserializeOwned>(&self, key: &[u8]) -> Result<Option<T>> {
        self.0
            .get(key)
            .map(|bytes| lcs::from_bytes(bytes))
            .transpose()
            .map_err(Into::into)
    }

    pub fn insert(&mut self, key: Vec<u8>, value: Vec<u8>) -> Option<Vec<u8>> {
        self.0.insert(key, value)
    }

    pub fn remove(&mut self, key: &[u8]) -> Option<Vec<u8>> {
        self.0.remove(key)
    }

    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    pub fn iter(&self) -> impl std::iter::Iterator<Item = (&Vec<u8>, &Vec<u8>)> {
        self.0.iter()
    }
}

impl fmt::Debug for AccountState {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        // TODO: add support for other types of resources
        let account_resource_str = self
            .get_account_resource()
            .map(|account_resource_opt| format!("{:#?}", account_resource_opt))
            .unwrap_or_else(|e| format!("parse error: {:#?}", e));

        let libra_timestamp_str = self
            .get_libra_timestamp_resource()
            .map(|libra_timestamp_opt| format!("{:#?}", libra_timestamp_opt))
            .unwrap_or_else(|e| format!("parse: {:#?}", e));

        let validator_config_str = self
            .get_validator_config_resource()
            .map(|validator_config_opt| format!("{:#?}", validator_config_opt))
            .unwrap_or_else(|e| format!("parse error: {:#?}", e));

        let validator_set_str = self
            .get_validator_set()
            .map(|validator_set_opt| format!("{:#?}", validator_set_opt))
            .unwrap_or_else(|e| format!("parse error: {:#?}", e));

        write!(
            f,
            "{{ \n \
             AccountResource {{ {} }} \n \
             LibraTimestamp {{ {} }} \n \
             ValidatorConfig {{ {} }} \n \
             ValidatorSet {{ {} }} \n \
             }}",
            account_resource_str, libra_timestamp_str, validator_config_str, validator_set_str,
        )
    }
}

impl TryFrom<(&AccountResource, &BalanceResource)> for AccountState {
    type Error = Error;

    fn try_from(
        (account_resource, balance_resource): (&AccountResource, &BalanceResource),
    ) -> Result<Self> {
        let mut btree_map: BTreeMap<Vec<u8>, Vec<u8>> = BTreeMap::new();
        btree_map.insert(
            AccountResource::resource_path(),
            lcs::to_bytes(account_resource)?,
        );
        btree_map.insert(
            BalanceResource::resource_path(),
            lcs::to_bytes(balance_resource)?,
        );

        Ok(Self(btree_map))
    }
}
<|endoftext|>