<commit_before>
// language: Rust master branch

use persistent::list::List;

pub mod persistent {
pub mod list {

// Is reference-counting the best choice for the shared immutable data?
use std::rc::Rc;

// Rc's Eq/Ord compares the contained data, not the pointer.
#[deriving(Clone, DeepClone, Eq, Ord, TotalEq, TotalOrd)]
pub struct List<T> {
  priv node : Rc<Node<T>>
}

#[deriving(Clone, DeepClone, Eq, Ord, TotalEq, TotalOrd)]
pub enum Node<T> {
  Nil,
  Cons(T, List<T>)
}

impl<'self, T> Iterator<&'self T> for &'self List<T> {
  fn next(&mut self) -> Option<&'self T> {
    match *self.node.borrow() {
      Nil => None,
      Cons(ref x, ref xs) => {
        *self = xs;
        Some(x)
      }
    }
  }
}

impl<T> List<T> {
  pub fn iter<'t>(&'t self) -> &'t List<T> {
    self
  }
  pub fn node<'t>(&'t self) -> &'t Node<T> {
    self.node.borrow()
  }
}

// Ought Freeze really be required for members of persistent lists?
// Generally, yes, because there's shared data; but what if you want
// a list of Cells or RefCells that deliberately have shared mutable identity?
// Freeze is forced by Rc.
impl<T: Freeze> List<T> {
  pub fn new(node: Node<T>) -> List<T> {
    List{node: Rc::new(node)}
  }
  pub fn nil() -> List<T> {
    List::new(Nil)
  }
  pub fn cons(x:T, xs:List<T>) -> List<T> {
    List::new(Cons(x, xs))
  }
}
impl<T: Clone+Freeze> List<T> {
  fn reverse_impl(&self, acc : List<T>) -> List<T> {
    match *self.node.borrow() {
      Nil => acc,
      Cons(ref x, ref xs) => xs.reverse_impl(List::cons(x.clone(), acc))
    }
  }
  pub fn reverse(&self) -> List<T> {
    self.reverse_impl(List::nil())
  }
}

impl<T> Container for List<T> {
  fn len(&self) -> uint {
    let mut result = 0;
    for _ in self.iter() { result += 1; }
    result
  }
  fn is_empty(&self) -> bool {
    match *self.node.borrow() {
      Nil => true,
      Cons(_, _) => false
    }
  }
}

impl<T: Freeze> Default for List<T> {
  fn default() -> List<T> {
    List::nil()
  }
}

/* Does this even make sense for an immutable container?
impl<A> Extendable<A> for List<A> {
  fn extend<T: Iterator<A>>(&mut self, iter: &mut T) {
  }
}
*/

impl<A: Freeze> FromIterator<A> for List<A> {
  // Is it possible to write this function without a
  // stack (implicit via recursion, or explicit),
  // and without 'unsafe' code?
  fn from_iterator<T: Iterator<A>>(iter: &mut T) -> List<A> {
    match iter.next() {
      None => List::nil(),
      Some(a) => List::cons(a, FromIterator::from_iterator(iter))
    }
  }
}

#[cfg(test)]
mod test {
use super::List;
//use std::cell::RefCell;
#[test]
fn test() {
  let p0 = List::nil();
  let p1 : List<int> = List::cons(1, p0.clone());
  let p2a = List::cons(2, p1.clone());
  let p2b = List::cons(2, p1.clone());
  let p2c = List::cons(3, p1.clone());
  assert!(p0 == p0);
  assert!(p1 == p1);
  assert!(p0 != p1);
  assert!(p2a == p2b);
  assert!(p2a < p2c);
  assert!(p1 < p2c);
  assert!(p0 == p0.reverse());
  assert!(p1 == p1.reverse());
  assert!(p2a > p2a.reverse());
  assert!(p2a == p2a.reverse().reverse());
  let mut sum = 0;
  for i in p2c.iter() {
    sum += *i;
  }
  assert!(sum == 4);
  let seq = ~[1,3,2]; //why did I have to allocate here in order to get .move_iter()?
  let mut digits = 0;
  let seql : List<int> = seq.move_iter().collect();
  for i in seql.iter() {
    digits = digits * 10 + *i;
  }
  assert!(digits == 132);
  assert!(p0 == Default::default());
  // doesn't meet Freeze requirement:
  //let sdf : List<RefCell<int>> = List::nil();
}
}

}
}

fn main() {
  // Trying out printing some stuff
  let p0 = List::nil();
  let p1 : List<int> = List::cons(1, p0);
  let p2 = List::cons(2, p1.clone());
  let p2a = List::cons(2, p1.clone());
  let p2b = List::cons(3, p1.clone());
  println(format!("Successyays:\n{}, {}", p2 == p2a, p2 == p2b));
  for i in p2b.iter() {
    println(format!("{}", *i))
  }
  // Is there a way not to use the temporary var name here?
  // "for i in p2b.reverse().iter()" didn't work (I guess the
  // lifetime of p2b.reverse() there didn't include the body of the
  // for loop).
  let x = p2b.reverse();
  for i in x.iter() {
    println(format!("{}", *i))
  }
}

<commit_msg>Document persistent.rs<commit_after>// (language: Rust master branch)

#[link(name="persistent")];

/*! Persistent data structures

This is a collection of persistent data structures
modeled after purely functional programming languages
( https://en.wikipedia.org/wiki/Persistent_data_structure ).

Currently, only an inductive linked list is implemented.

*/

use persistent::list::List;

pub mod persistent {
pub mod list {

// Is reference-counting the best choice for the shared immutable data?
use std::rc::Rc;

/// Persistent cons/nil list.
/// O(1) access to the head of the list.
/// Modified versions of the list with shared tails can be created efficiently.
/// List members cannot be moved from,
/// because a List might be sharing data with other Lists.

// Rc's Eq/Ord compares the contained data, not the pointer.
#[deriving(Clone, DeepClone, Eq, Ord, TotalEq, TotalOrd)]
pub struct List<T> {
  priv node : Rc<Node<T>>
}

#[deriving(Clone, DeepClone, Eq, Ord, TotalEq, TotalOrd)]
pub enum Node<T> {
  Nil,
  Cons(T, List<T>)
}

impl<'self, T> Iterator<&'self T> for &'self List<T> {
  fn next(&mut self) -> Option<&'self T> {
    match *self.node.borrow() {
      Nil => None,
      Cons(ref x, ref xs) => {
        *self = xs;
        Some(x)
      }
    }
  }
}

impl<T> List<T> {
  /// Lists are iterable.
  pub fn iter<'t>(&'t self) -> &'t List<T> {
    self
  }
  /// Use this to pattern match on Nil vs Cons.
  pub fn node<'t>(&'t self) -> &'t Node<T> {
    self.node.borrow()
  }
}

// Ought Freeze really be required for members of persistent lists?
// Generally, yes, because there's shared data; but what if you want
// a list of Cells or RefCells that deliberately have shared mutable identity?
// Freeze is forced by Rc.
impl<T: Freeze> List<T> {
  /// Create an empty list
  pub fn nil() -> List<T> {
    List::new(Nil)
  }
  /// Create a list from the head and the rest of the list
  pub fn cons(x:T, xs:List<T>) -> List<T> {
    List::new(Cons(x, xs))
  }
  /// Create a list from a node (you probably won't need this function).
  pub fn new(node: Node<T>) -> List<T> {
    List{node: Rc::new(node)}
  }
}
impl<T: Clone+Freeze> List<T> {
  fn reverse_impl(&self, acc : List<T>) -> List<T> {
    match *self.node.borrow() {
      Nil => acc,
      Cons(ref x, ref xs) => xs.reverse_impl(List::cons(x.clone(), acc))
    }
  }
  /// Create a copy of this list in reverse order.
  /// It cannot move the members because it can't know whether it has
  /// the only reference to them.  (Rc doesn't even provide a runtime
  /// feature to find out whether you have the last reference.)
  pub fn reverse(&self) -> List<T> {
    self.reverse_impl(List::nil())
  }
}

impl<T> Container for List<T> {
  fn len(&self) -> uint {
    let mut result = 0;
    for _ in self.iter() { result += 1; }
    result
  }
  fn is_empty(&self) -> bool {
    match *self.node.borrow() {
      Nil => true,
      Cons(_, _) => false
    }
  }
}

impl<T: Freeze> Default for List<T> {
  fn default() -> List<T> {
    List::nil()
  }
}

/* Does this even make sense for an immutable container?
impl<A> Extendable<A> for List<A> {
  fn extend<T: Iterator<A>>(&mut self, iter: &mut T) {
  }
}
*/

impl<A: Freeze> FromIterator<A> for List<A> {
  // Is it possible to write this function without a
  // stack (implicit via recursion, or explicit),
  // and without 'unsafe' code?
  fn from_iterator<T: Iterator<A>>(iter: &mut T) -> List<A> {
    match iter.next() {
      None => List::nil(),
      Some(a) => List::cons(a, FromIterator::from_iterator(iter))
    }
  }
}

#[cfg(test)]
mod test {
use super::List;
//use std::cell::RefCell;
#[test]
fn test() {
  let p0 = List::nil();
  let p1 : List<int> = List::cons(1, p0.clone());
  let p2a = List::cons(2, p1.clone());
  let p2b = List::cons(2, p1.clone());
  let p2c = List::cons(3, p1.clone());
  assert!(p0 == p0);
  assert!(p1 == p1);
  assert!(p0 != p1);
  assert!(p2a == p2b);
  assert!(p2a < p2c);
  assert!(p1 < p2c);
  assert!(p0 == p0.reverse());
  assert!(p1 == p1.reverse());
  assert!(p2a > p2a.reverse());
  assert!(p2a == p2a.reverse().reverse());
  let mut sum = 0;
  for i in p2c.iter() {
    sum += *i;
  }
  assert!(sum == 4);
  let seq = ~[1,3,2]; //why did I have to allocate here in order to get .move_iter()?
  let mut digits = 0;
  let seql : List<int> = seq.move_iter().collect();
  for i in seql.iter() {
    digits = digits * 10 + *i;
  }
  assert!(digits == 132);
  assert!(p0 == Default::default());
  // doesn't meet Freeze requirement:
  //let sdf : List<RefCell<int>> = List::nil();
}
}

}
}

// Trying out printing some stuff
fn main() {
  let p0 = List::nil();
  let p1 : List<int> = List::cons(1, p0);
  let p2 = List::cons(2, p1.clone());
  let p2a = List::cons(2, p1.clone());
  let p2b = List::cons(3, p1.clone());
  println(format!("Successyays:\n{}, {}", p2 == p2a, p2 == p2b));
  for i in p2b.iter() {
    println(format!("{}", *i))
  }
  // Is there a way not to use the temporary var name here?
  // "for i in p2b.reverse().iter()" didn't work (I guess the
  // lifetime of p2b.reverse() there didn't include the body of the
  // for loop).
  let x = p2b.reverse();
  for i in x.iter() {
    println(format!("{}", *i))
  }
}

<|endoftext|>


<commit_before>#![feature(slicing_syntax)]

extern crate test;
extern crate roaring;

use std::{ u32 };
use test::Bencher;

use roaring::RoaringBitmap;

#[bench]
fn create(b: &mut Bencher) {
    b.iter(|| {
        let mut bitmap = RoaringBitmap::new();
        bitmap
    })
}

#[bench]
fn insert1(b: &mut Bencher) {
    b.iter(|| {
        let mut bitmap = RoaringBitmap::new();
        bitmap.insert(1);
        bitmap
    })
}

#[bench]
fn insert2(b: &mut Bencher) {
    b.iter(|| {
        let mut bitmap = RoaringBitmap::new();
        bitmap.insert(1);
        bitmap.insert(2);
        bitmap
    })
}

<commit_msg>Make benchmarks work<commit_after>#![feature(slicing_syntax)]

extern crate test;
extern crate roaring;

use std::{ u32 };
use test::Bencher;

use roaring::RoaringBitmap;

#[bench]
fn create(b: &mut Bencher) {
    b.iter(|| {
        let mut bitmap: RoaringBitmap<u32> = RoaringBitmap::new();
        bitmap
    })
}

#[bench]
fn insert1(b: &mut Bencher) {
    b.iter(|| {
        let mut bitmap: RoaringBitmap<u32> = RoaringBitmap::new();
        bitmap.insert(1);
        bitmap
    })
}

#[bench]
fn insert2(b: &mut Bencher) {
    b.iter(|| {
        let mut bitmap: RoaringBitmap<u32> = RoaringBitmap::new();
        bitmap.insert(1);
        bitmap.insert(2);
        bitmap
    })
}

<|endoftext|>


<commit_before>use core::{mem, str};

use arch::interrupt::irq::{ACKS, COUNTS, acknowledge};
use context;
use syscall::{Error, Result};
use super::Scheme;

pub struct IrqScheme;

impl Scheme for IrqScheme {
    fn open(&mut self, path: &[u8], _flags: usize) -> Result<usize> {
        let path_str = str::from_utf8(path).or(Err(Error::NoEntry))?;

        let id = path_str.parse::<usize>().or(Err(Error::NoEntry))?;

        if id < COUNTS.lock().len() {
            Ok(id)
        } else {
            Err(Error::NoEntry)
        }
    }

    fn dup(&mut self, file: usize) -> Result<usize> {
        Err(Error::NotPermitted)
    }

    fn read(&mut self, file: usize, buffer: &mut [u8]) -> Result<usize> {
        // Ensures that the length of the buffer is larger than the size of a usize
        if buffer.len() >= mem::size_of::<usize>() {
            let ack = ACKS.lock()[file];
            let current = COUNTS.lock()[file];
            if ack != current {
                // Safe if the length of the buffer is larger than the size of a usize
                assert!(buffer.len() >= mem::size_of::<usize>());
                unsafe { *(buffer.as_mut_ptr() as *mut usize) = current; }
                return Ok(mem::size_of::<usize>());
            } else {
                return Ok(0);
            }
        } else {
            Err(Error::InvalidValue)
        }
    }

    fn write(&mut self, file: usize, buffer: &[u8]) -> Result<usize> {
        if buffer.len() >= mem::size_of::<usize>() {
            assert!(buffer.len() >= mem::size_of::<usize>());
            let ack = unsafe { *(buffer.as_ptr() as *const usize) };
            let current = COUNTS.lock()[file];
            if ack == current {
                ACKS.lock()[file] = ack;
                unsafe { acknowledge(file); }
                return Ok(mem::size_of::<usize>());
            } else {
                return Ok(0);
            }
        } else {
            Err(Error::InvalidValue)
        }
    }

    fn fsync(&mut self, _file: usize) -> Result<()> {
        Ok(())
    }

    fn close(&mut self, file: usize) -> Result<()> {
        Ok(())
    }
}
<commit_msg>PS/2 driver convert to char<commit_after>use core::{mem, str};

use arch::interrupt::irq::{ACKS, COUNTS, acknowledge};
use syscall::{Error, Result};
use super::Scheme;

pub struct IrqScheme;

impl Scheme for IrqScheme {
    fn open(&mut self, path: &[u8], _flags: usize) -> Result<usize> {
        let path_str = str::from_utf8(path).or(Err(Error::NoEntry))?;

        let id = path_str.parse::<usize>().or(Err(Error::NoEntry))?;

        if id < COUNTS.lock().len() {
            Ok(id)
        } else {
            Err(Error::NoEntry)
        }
    }

    fn dup(&mut self, file: usize) -> Result<usize> {
        Err(Error::NotPermitted)
    }

    fn read(&mut self, file: usize, buffer: &mut [u8]) -> Result<usize> {
        // Ensures that the length of the buffer is larger than the size of a usize
        if buffer.len() >= mem::size_of::<usize>() {
            let ack = ACKS.lock()[file];
            let current = COUNTS.lock()[file];
            if ack != current {
                // Safe if the length of the buffer is larger than the size of a usize
                assert!(buffer.len() >= mem::size_of::<usize>());
                unsafe { *(buffer.as_mut_ptr() as *mut usize) = current; }
                return Ok(mem::size_of::<usize>());
            } else {
                return Ok(0);
            }
        } else {
            Err(Error::InvalidValue)
        }
    }

    fn write(&mut self, file: usize, buffer: &[u8]) -> Result<usize> {
        if buffer.len() >= mem::size_of::<usize>() {
            assert!(buffer.len() >= mem::size_of::<usize>());
            let ack = unsafe { *(buffer.as_ptr() as *const usize) };
            let current = COUNTS.lock()[file];
            if ack == current {
                ACKS.lock()[file] = ack;
                unsafe { acknowledge(file); }
                return Ok(mem::size_of::<usize>());
            } else {
                return Ok(0);
            }
        } else {
            Err(Error::InvalidValue)
        }
    }

    fn fsync(&mut self, _file: usize) -> Result<()> {
        Ok(())
    }

    fn close(&mut self, file: usize) -> Result<()> {
        Ok(())
    }
}
<|endoftext|>


<commit_before>extern crate gcc;

use std::process::Command;
use std::io::{ stderr, Write };

fn main() {
    let perl = std::env::var("PERL").unwrap();

    Command::new(perl).arg("build/regen.pl").status().unwrap();

    gcc::compile_library("libouroboros.a", &[ "ouroboros/libouroboros.c" ]);
}
<commit_msg>Bootstrap build from just a path to perl.<commit_after>extern crate gcc;

use std::path::Path;
use std::process::Command;

struct Perl {
    bin: String,
}

impl Perl {
    fn new() -> Perl {
        Perl {
            bin: std::env::var("PERL").unwrap(),
        }
    }

    fn cfg(&self, key: &str) -> String {
        let out = Command::new(&self.bin)
            .arg("-MConfig")
            .arg("-e")
            .arg(format!("print $Config::Config{{{}}}", key))
            .output()
            .unwrap();

        String::from_utf8(out.stdout).unwrap()
    }

    fn run(&self, script: &str) {
        Command::new(&self.bin)
            .arg(script)
            .status()
            .unwrap();
    }
}

fn main() {
    let perl = Perl::new();

    let prefix = perl.cfg("archlibexp");
    let perl_multi = perl.cfg("usemultiplicity");

    let perl_inc = Path::new(&prefix).join("CORE");

    perl.run("build/regen.pl");

    gcc::Config::new()
        .file("ouroboros/libouroboros.c")
        .include(&perl_inc)
        .compile("libouroboros.a");

    if perl_multi == "define" {
        println!("cargo:rustc-cfg=perl_multiplicity");
    }
}
<|endoftext|>


<commit_before>#![feature(phase)]

extern crate getopts;
#[phase(plugin, link)]
extern crate log;

use getopts::{optopt, optflag, getopts, usage, Matches};

use std::os;

use std::io::net::udp::UdpSocket;
use std::io::net::tcp::TcpStream;
use std::io::net::ip::SocketAddr;
use std::io::net::ip::{IpAddr, Ipv4Addr, Ipv6Addr, Port};
use std::io::net::addrinfo::get_host_addresses;
use std::io::{MemWriter, BufferedStream};
use std::io::stdio::{stdin, stdout};
use std::io::{IoResult, Reader, Writer};

pub const SOCKS5_VERSION : u8 = 0x05;

pub const SOCKS5_AUTH_METHOD_NONE            : u8 = 0x00;
pub const SOCKS5_AUTH_METHOD_GSSAPI          : u8 = 0x01;
pub const SOCKS5_AUTH_METHOD_PASSWORD        : u8 = 0x02;
pub const SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE  : u8 = 0xff;

pub const SOCKS5_CMD_TCP_CONNECT   : u8 = 0x01;
pub const SOCKS5_CMD_TCP_BIND      : u8 = 0x02;
pub const SOCKS5_CMD_UDP_ASSOCIATE : u8 = 0x03;

pub const SOCKS5_ADDR_TYPE_IPV4        : u8 = 0x01;
pub const SOCKS5_ADDR_TYPE_DOMAIN_NAME : u8 = 0x03;
pub const SOCKS5_ADDR_TYPE_IPV6        : u8 = 0x04;

pub const SOCKS5_REPLY_SUCCEEDED                     : u8 = 0x00;
pub const SOCKS5_REPLY_GENERAL_FAILURE               : u8 = 0x01;
pub const SOCKS5_REPLY_CONNECTION_NOT_ALLOWED        : u8 = 0x02;
pub const SOCKS5_REPLY_NETWORK_UNREACHABLE           : u8 = 0x03;
pub const SOCKS5_REPLY_HOST_UNREACHABLE              : u8 = 0x04;
pub const SOCKS5_REPLY_CONNECTION_REFUSED            : u8 = 0x05;
pub const SOCKS5_REPLY_TTL_EXPIRED                   : u8 = 0x06;
pub const SOCKS5_REPLY_COMMAND_NOT_SUPPORTED         : u8 = 0x07;
pub const SOCKS5_REPLY_ADDRESS_TYPE_NOT_SUPPORTED    : u8 = 0x08;

#[deriving(Show, Clone)]
pub struct Error {
    pub code: u8,
    pub message: String,
}

impl Error {
    pub fn new(code: u8, message: &str) -> Error {
        Error {
            code: code,
            message: message.to_string(),
        }
    }
}

#[deriving(Clone, PartialEq, Eq, Hash)]
pub struct DomainNameAddr {
    pub domain_name: String,
    pub port: Port,
}

#[deriving(Clone, PartialEq, Eq, Hash)]
pub enum AddressType {
    SocketAddress(SocketAddr),
    DomainNameAddress(DomainNameAddr),
}

fn write_addr(addr: &AddressType, writer: &mut Writer) -> IoResult<()> {
    match addr {
        &SocketAddress(ref ip) => {
            match ip.ip {
                Ipv4Addr(v1, v2, v3, v4) => {
                    try!(writer.write([SOCKS5_ADDR_TYPE_IPV4, v1, v2, v3, v4]));
                },
                Ipv6Addr(v1, v2, v3, v4, v5, v6, v7, v8) => {
                    try!(writer.write_u8(SOCKS5_ADDR_TYPE_IPV6));
                    try!(writer.write_be_u16(v1));
                    try!(writer.write_be_u16(v2));
                    try!(writer.write_be_u16(v3));
                    try!(writer.write_be_u16(v4));
                    try!(writer.write_be_u16(v5));
                    try!(writer.write_be_u16(v6));
                    try!(writer.write_be_u16(v7));
                    try!(writer.write_be_u16(v8));
                }
            }
            try!(writer.write_be_u16(ip.port));
        },
        &DomainNameAddress(ref domain_name) => {
            try!(writer.write_u8(SOCKS5_ADDR_TYPE_DOMAIN_NAME));
            try!(writer.write_u8(domain_name.domain_name.len() as u8));
            try!(writer.write(domain_name.domain_name.as_slice().as_bytes()));
            try!(writer.write_be_u16(domain_name.port));
        }
    }

    Ok(())
}

pub fn parse_request_header(stream: &mut Reader) -> Result<(uint, AddressType), Error> {
    let atyp = match stream.read_byte() {
        Ok(atyp) => atyp,
        Err(_) => return Err(Error::new(SOCKS5_REPLY_GENERAL_FAILURE, "Error while reading address type"))
    };

    match atyp {
        SOCKS5_ADDR_TYPE_IPV4 => {
            let wrapper = || {
                let v4addr = Ipv4Addr(try!(stream.read_byte()),
                                      try!(stream.read_byte()),
                                      try!(stream.read_byte()),
                                      try!(stream.read_byte()));
                let port = try!(stream.read_be_u16());
                Ok((v4addr, port))
            };
            match wrapper() {
                Ok((v4addr, port)) => Ok((7u, SocketAddress(SocketAddr{ip: v4addr, port: port}))),
                Err(_) =>
                    Err(Error::new(SOCKS5_REPLY_GENERAL_FAILURE, "Error while parsing IPv4 address"))
            }
        },
        SOCKS5_ADDR_TYPE_IPV6 => {
            let wrapper = || {
                let v6addr = Ipv6Addr(try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()));
                let port = try!(stream.read_be_u16());
                Ok((v6addr, port))
            };

            match wrapper() {
                Ok((v6addr, port)) =>
                    Ok((19u, SocketAddress(SocketAddr{ip: v6addr, port: port}))),
                Err(_) =>
                    Err(Error::new(SOCKS5_REPLY_GENERAL_FAILURE, "Error while parsing IPv6 address"))
            }
        },
        SOCKS5_ADDR_TYPE_DOMAIN_NAME => {
            let wrapper = || {
                let addr_len = try!(stream.read_byte()) as uint;
                let raw_addr = try!(stream.read_exact(addr_len));
                let port = try!(stream.read_be_u16());

                Ok((addr_len, raw_addr, port))
            };

            match wrapper() {
                Ok((addr_len, raw_addr, port)) =>
                    Ok((4 + addr_len, DomainNameAddress(DomainNameAddr{
                                                            domain_name: String::from_utf8(raw_addr).unwrap(),
                                                            port: port,
                                                        }))),
                Err(_) => {
                    Err(Error::new(SOCKS5_REPLY_GENERAL_FAILURE, "Error while parsing domain name"))
                }
            }
        },
        _ => {
            // Address type not supported
            Err(Error::new(SOCKS5_REPLY_ADDRESS_TYPE_NOT_SUPPORTED, "Not supported address type"))
        }
    }
}

fn get_addr_len(response: &[u8]) -> uint {
    match response[0] {
        SOCKS5_ADDR_TYPE_IPV4 => {
            1 + 6
        },
        SOCKS5_ADDR_TYPE_DOMAIN_NAME => {
            1 + response[1] as uint + 2
        },
        SOCKS5_ADDR_TYPE_IPV6 => {
            1 + 16 + 2
        }
        _ => {
            panic!("Invalid addr type {}", response[0]);
        }
    }
}

fn do_udp(matches: &Matches, svr_addr: &AddressType, proxy_addr: &SocketAddr) {
    let udp_proxy_addr = {
        let mut stream = BufferedStream::new(
                            TcpStream::connect(proxy_addr.ip.to_string().as_slice(), proxy_addr.port).unwrap());
        let mut buf = [0u8, ..0xffff];

        stream.write([SOCKS5_VERSION, 0x01, 0x00]).unwrap();
        stream.flush().unwrap();
        stream.read(buf).unwrap();
        if buf[1] != SOCKS5_AUTH_METHOD_NONE {
            panic!("Proxy server needs authentication");
        }

        stream.write([SOCKS5_VERSION, SOCKS5_CMD_UDP_ASSOCIATE, 0x00]).unwrap();
        write_addr(svr_addr, &mut stream).unwrap();
        stream.flush().unwrap();

        let three_flag = stream.read_exact(3).unwrap();
        if three_flag[1] != SOCKS5_REPLY_SUCCEEDED {
            panic!("Failed with error code: {}", three_flag[1]);
        }

        let (_, paddr) = parse_request_header(&mut stream).unwrap();

        match paddr {
            SocketAddress(addr) => addr,
            DomainNameAddress(dmname) => {
                let addrs = get_host_addresses(dmname.domain_name.as_slice()).unwrap();
                SocketAddr {
                    ip: addrs.head().unwrap().clone(),
                    port: dmname.port,
                }
            }
        }
    };

    let local_addr = SocketAddr {
        ip: from_str(matches.opt_str("b").expect("Require local address").as_slice()).unwrap(),
        port: from_str(matches.opt_str("l").expect("Require local port").as_slice()).unwrap(),
    };

    let mut socket = UdpSocket::bind(local_addr).unwrap();

    let mut writer = MemWriter::new();
    writer.write([0x00u8, 0x00u8, 0x00u8]).unwrap();
    write_addr(svr_addr, &mut writer).unwrap();

    let inputs = stdin().read_to_end().unwrap();
    writer.write(inputs.as_slice()).unwrap();

    socket.send_to(writer.unwrap().as_slice(), udp_proxy_addr.clone()).unwrap();

    let mut buf = [0u8, ..0xffff];
    let (len, _) = socket.recv_from(buf).unwrap();

    let response = buf.slice_to(len).slice_from(3);

    let addr_len = get_addr_len(response);

    stdout().write(response.slice_from(addr_len)).unwrap();
}

fn do_tcp(_: &Matches, svr_addr: &AddressType, proxy_addr: &SocketAddr) {
    let mut stream = TcpStream::connect(proxy_addr.ip.to_string().as_slice(), proxy_addr.port).unwrap();
    let mut buf = [0u8, ..0xffff];

    stream.write([SOCKS5_VERSION, 0x01, 0x00]).unwrap();
    stream.read(buf).unwrap();
    if buf[1] != SOCKS5_AUTH_METHOD_NONE {
        panic!("Proxy server needs authentication");
    }

    stream.write([SOCKS5_VERSION, SOCKS5_CMD_TCP_CONNECT, 0x00]).unwrap();
    write_addr(svr_addr, &mut stream).unwrap();

    let inputs = stdin().read_to_end().unwrap();
    stream.write(inputs.as_slice()).unwrap();

    let mut output = stdout();

    {
        let buf_len = stream.read_at_least(1, buf).unwrap();
        let header = buf.slice_to(buf_len);

        if header[1] != SOCKS5_REPLY_SUCCEEDED {
            panic!("Failed with error: {}", header[1]);
        }

        let addr_len = get_addr_len(header.slice_from(3));
        output.write(header.slice_from(3 + addr_len)).unwrap();
    }

    loop {
        match stream.read(buf) {
            Ok(len) => {
                output.write(buf.slice_to(len)).unwrap();
            },
            Err(..) => {
                break;
            }
        }
    }
    println!("");
}

fn main() {

    let opts = [
        optflag("h", "help", "Print help message"),
        optopt("s", "server-addr", "Server address", ""),
        optopt("p", "server-port", "Server port", ""),
        optopt("b", "local-addr", "Local address for binding", ""),
        optopt("l", "local-port", "Local port for binding", ""),
        optopt("x", "proxy-addr", "Proxy address", ""),
        optopt("o", "proxy-port", "Proxy port", ""),
        optopt("t", "protocol", "Protocol to use", "tcp"),
    ];

    let matches = getopts(os::args().tail(), opts).unwrap();

    if matches.opt_present("h") {
        println!("{}", usage(format!("Usage: {} [Options]", os::args()[0]).as_slice(),
                            opts));
        return;
    }

    let is_tcp = match matches.opt_str("t").expect("Required to specify protocol").as_slice() {
        "tcp" => true,
        "udp" => false,
        _ => panic!("Unsupported protocol")
    };

    let proxy_addr = SocketAddr {
        ip: from_str(matches.opt_str("x").expect("Require proxy address").as_slice()).unwrap(),
        port: from_str(matches.opt_str("o").expect("Require proxy port").as_slice()).unwrap(),
    };

    let svr_port: Port = from_str(matches.opt_str("p").expect("Require server port").as_slice()).unwrap();
    let svr_addr = match from_str::<IpAddr>(matches.opt_str("s").expect("Require server address").as_slice()) {
        Some(ip) => SocketAddress(SocketAddr {ip: ip, port: svr_port}),
        None => DomainNameAddress(DomainNameAddr { domain_name: matches.opt_str("s").unwrap(), port: svr_port}),
    };

    if is_tcp {
        do_tcp(&matches, &svr_addr, &proxy_addr);
    } else {
        do_udp(&matches, &svr_addr, &proxy_addr);
    }
}
<commit_msg>improved client<commit_after>#![feature(phase)]

extern crate getopts;
#[phase(plugin, link)]
extern crate log;

use getopts::{optopt, optflag, getopts, usage, Matches};

use std::os;

use std::io::net::udp::UdpSocket;
use std::io::net::tcp::TcpStream;
use std::io::net::ip::SocketAddr;
use std::io::net::ip::{IpAddr, Ipv4Addr, Ipv6Addr, Port};
use std::io::net::addrinfo::get_host_addresses;
use std::io::{MemWriter, BufferedStream};
use std::io::stdio::{stdin, stdout};
use std::io::{IoResult, Reader, Writer};

pub const SOCKS5_VERSION : u8 = 0x05;

pub const SOCKS5_AUTH_METHOD_NONE            : u8 = 0x00;
pub const SOCKS5_AUTH_METHOD_GSSAPI          : u8 = 0x01;
pub const SOCKS5_AUTH_METHOD_PASSWORD        : u8 = 0x02;
pub const SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE  : u8 = 0xff;

pub const SOCKS5_CMD_TCP_CONNECT   : u8 = 0x01;
pub const SOCKS5_CMD_TCP_BIND      : u8 = 0x02;
pub const SOCKS5_CMD_UDP_ASSOCIATE : u8 = 0x03;

pub const SOCKS5_ADDR_TYPE_IPV4        : u8 = 0x01;
pub const SOCKS5_ADDR_TYPE_DOMAIN_NAME : u8 = 0x03;
pub const SOCKS5_ADDR_TYPE_IPV6        : u8 = 0x04;

pub const SOCKS5_REPLY_SUCCEEDED                     : u8 = 0x00;
pub const SOCKS5_REPLY_GENERAL_FAILURE               : u8 = 0x01;
pub const SOCKS5_REPLY_CONNECTION_NOT_ALLOWED        : u8 = 0x02;
pub const SOCKS5_REPLY_NETWORK_UNREACHABLE           : u8 = 0x03;
pub const SOCKS5_REPLY_HOST_UNREACHABLE              : u8 = 0x04;
pub const SOCKS5_REPLY_CONNECTION_REFUSED            : u8 = 0x05;
pub const SOCKS5_REPLY_TTL_EXPIRED                   : u8 = 0x06;
pub const SOCKS5_REPLY_COMMAND_NOT_SUPPORTED         : u8 = 0x07;
pub const SOCKS5_REPLY_ADDRESS_TYPE_NOT_SUPPORTED    : u8 = 0x08;

#[deriving(Show, Clone)]
pub struct Error {
    pub code: u8,
    pub message: String,
}

impl Error {
    pub fn new(code: u8, message: &str) -> Error {
        Error {
            code: code,
            message: message.to_string(),
        }
    }
}

#[deriving(Clone, PartialEq, Eq, Hash)]
pub struct DomainNameAddr {
    pub domain_name: String,
    pub port: Port,
}

#[deriving(Clone, PartialEq, Eq, Hash)]
pub enum AddressType {
    SocketAddress(SocketAddr),
    DomainNameAddress(DomainNameAddr),
}

fn write_addr(addr: &AddressType, writer: &mut Writer) -> IoResult<()> {
    match addr {
        &SocketAddress(ref ip) => {
            match ip.ip {
                Ipv4Addr(v1, v2, v3, v4) => {
                    try!(writer.write([SOCKS5_ADDR_TYPE_IPV4, v1, v2, v3, v4]));
                },
                Ipv6Addr(v1, v2, v3, v4, v5, v6, v7, v8) => {
                    try!(writer.write_u8(SOCKS5_ADDR_TYPE_IPV6));
                    try!(writer.write_be_u16(v1));
                    try!(writer.write_be_u16(v2));
                    try!(writer.write_be_u16(v3));
                    try!(writer.write_be_u16(v4));
                    try!(writer.write_be_u16(v5));
                    try!(writer.write_be_u16(v6));
                    try!(writer.write_be_u16(v7));
                    try!(writer.write_be_u16(v8));
                }
            }
            try!(writer.write_be_u16(ip.port));
        },
        &DomainNameAddress(ref domain_name) => {
            try!(writer.write_u8(SOCKS5_ADDR_TYPE_DOMAIN_NAME));
            try!(writer.write_u8(domain_name.domain_name.len() as u8));
            try!(writer.write(domain_name.domain_name.as_slice().as_bytes()));
            try!(writer.write_be_u16(domain_name.port));
        }
    }

    Ok(())
}

pub fn parse_request_header(stream: &mut Reader) -> Result<(uint, AddressType), Error> {
    let atyp = match stream.read_byte() {
        Ok(atyp) => atyp,
        Err(_) => return Err(Error::new(SOCKS5_REPLY_GENERAL_FAILURE, "Error while reading address type"))
    };

    match atyp {
        SOCKS5_ADDR_TYPE_IPV4 => {
            let wrapper = || {
                let v4addr = Ipv4Addr(try!(stream.read_byte()),
                                      try!(stream.read_byte()),
                                      try!(stream.read_byte()),
                                      try!(stream.read_byte()));
                let port = try!(stream.read_be_u16());
                Ok((v4addr, port))
            };
            match wrapper() {
                Ok((v4addr, port)) => Ok((7u, SocketAddress(SocketAddr{ip: v4addr, port: port}))),
                Err(_) =>
                    Err(Error::new(SOCKS5_REPLY_GENERAL_FAILURE, "Error while parsing IPv4 address"))
            }
        },
        SOCKS5_ADDR_TYPE_IPV6 => {
            let wrapper = || {
                let v6addr = Ipv6Addr(try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()),
                                      try!(stream.read_be_u16()));
                let port = try!(stream.read_be_u16());
                Ok((v6addr, port))
            };

            match wrapper() {
                Ok((v6addr, port)) =>
                    Ok((19u, SocketAddress(SocketAddr{ip: v6addr, port: port}))),
                Err(_) =>
                    Err(Error::new(SOCKS5_REPLY_GENERAL_FAILURE, "Error while parsing IPv6 address"))
            }
        },
        SOCKS5_ADDR_TYPE_DOMAIN_NAME => {
            let wrapper = || {
                let addr_len = try!(stream.read_byte()) as uint;
                let raw_addr = try!(stream.read_exact(addr_len));
                let port = try!(stream.read_be_u16());

                Ok((addr_len, raw_addr, port))
            };

            match wrapper() {
                Ok((addr_len, raw_addr, port)) =>
                    Ok((4 + addr_len, DomainNameAddress(DomainNameAddr{
                                                            domain_name: String::from_utf8(raw_addr).unwrap(),
                                                            port: port,
                                                        }))),
                Err(_) => {
                    Err(Error::new(SOCKS5_REPLY_GENERAL_FAILURE, "Error while parsing domain name"))
                }
            }
        },
        _ => {
            // Address type not supported
            Err(Error::new(SOCKS5_REPLY_ADDRESS_TYPE_NOT_SUPPORTED, "Not supported address type"))
        }
    }
}

fn get_addr_len(response: &[u8]) -> uint {
    match response[0] {
        SOCKS5_ADDR_TYPE_IPV4 => {
            1 + 6
        },
        SOCKS5_ADDR_TYPE_DOMAIN_NAME => {
            1 + response[1] as uint + 2
        },
        SOCKS5_ADDR_TYPE_IPV6 => {
            1 + 16 + 2
        }
        _ => {
            panic!("Invalid addr type {}", response[0]);
        }
    }
}

fn do_udp(matches: &Matches, svr_addr: &AddressType, proxy_addr: &SocketAddr) {
    let udp_proxy_addr = {
        let mut stream = BufferedStream::new(
                            TcpStream::connect(proxy_addr.ip.to_string().as_slice(), proxy_addr.port).unwrap());
        let mut buf = [0u8, ..0xffff];

        stream.write([SOCKS5_VERSION, 0x01, 0x00]).unwrap();
        stream.flush().unwrap();
        stream.read(buf).unwrap();
        if buf[1] != SOCKS5_AUTH_METHOD_NONE {
            panic!("Proxy server needs authentication");
        }

        stream.write([SOCKS5_VERSION, SOCKS5_CMD_UDP_ASSOCIATE, 0x00]).unwrap();
        write_addr(svr_addr, &mut stream).unwrap();
        stream.flush().unwrap();

        let three_flag = stream.read_exact(3).unwrap();
        if three_flag[1] != SOCKS5_REPLY_SUCCEEDED {
            panic!("Failed with error code: {}", three_flag[1]);
        }

        let (_, paddr) = parse_request_header(&mut stream).unwrap();

        match paddr {
            SocketAddress(addr) => addr,
            DomainNameAddress(dmname) => {
                let addrs = get_host_addresses(dmname.domain_name.as_slice()).unwrap();
                SocketAddr {
                    ip: addrs.head().unwrap().clone(),
                    port: dmname.port,
                }
            }
        }
    };

    let local_addr = SocketAddr {
        ip: from_str(matches.opt_str("b").expect("Require local address").as_slice()).unwrap(),
        port: from_str(matches.opt_str("l").expect("Require local port").as_slice()).unwrap(),
    };

    let mut socket = UdpSocket::bind(local_addr).unwrap();

    let mut writer = MemWriter::new();
    writer.write([0x00u8, 0x00u8, 0x00u8]).unwrap();
    write_addr(svr_addr, &mut writer).unwrap();

    let inputs = stdin().read_to_end().unwrap();
    writer.write(inputs.as_slice()).unwrap();

    socket.send_to(writer.unwrap().as_slice(), udp_proxy_addr.clone()).unwrap();

    let mut buf = [0u8, ..0xffff];
    let (len, _) = socket.recv_from(buf).unwrap();

    let response = buf.slice_to(len).slice_from(3);

    let addr_len = get_addr_len(response);

    stdout().write(response.slice_from(addr_len)).unwrap();
}

fn do_tcp(_: &Matches, svr_addr: &AddressType, proxy_addr: &SocketAddr) {
    let mut stream = BufferedStream::new(
                        TcpStream::connect(proxy_addr.ip.to_string().as_slice(), proxy_addr.port).unwrap());
    let mut buf = [0u8, ..0xffff];

    stream.write([SOCKS5_VERSION, 0x01, 0x00]).unwrap();
    stream.flush().unwrap();
    stream.read(buf).unwrap();
    if buf[1] != SOCKS5_AUTH_METHOD_NONE {
        panic!("Proxy server needs authentication");
    }

    stream.write([SOCKS5_VERSION, SOCKS5_CMD_TCP_CONNECT, 0x00]).unwrap();
    write_addr(svr_addr, &mut stream).unwrap();

    let inputs = stdin().read_to_end().unwrap();
    stream.write(inputs.as_slice()).unwrap();
    stream.flush().unwrap();

    let mut output = stdout();

    {
        let buf_len = stream.read_at_least(1, buf).unwrap();
        let header = buf.slice_to(buf_len);

        if header[1] != SOCKS5_REPLY_SUCCEEDED {
            panic!("Failed with error: {}", header[1]);
        }

        let addr_len = get_addr_len(header.slice_from(3));
        output.write(header.slice_from(3 + addr_len)).unwrap();
    }

    loop {
        match stream.read(buf) {
            Ok(len) => {
                output.write(buf.slice_to(len)).unwrap();
            },
            Err(..) => {
                break;
            }
        }
    }
    println!("");
}

fn main() {

    let opts = [
        optflag("h", "help", "Print help message"),
        optopt("s", "server-addr", "Server address", ""),
        optopt("p", "server-port", "Server port", ""),
        optopt("b", "local-addr", "Local address for binding", ""),
        optopt("l", "local-port", "Local port for binding", ""),
        optopt("x", "proxy-addr", "Proxy address", ""),
        optopt("o", "proxy-port", "Proxy port", ""),
        optopt("t", "protocol", "Protocol to use", "tcp"),
    ];

    let matches = getopts(os::args().tail(), opts).unwrap();

    if matches.opt_present("h") {
        println!("{}", usage(format!("Usage: {} [Options]", os::args()[0]).as_slice(),
                            opts));
        return;
    }

    let is_tcp = match matches.opt_str("t").expect("Required to specify protocol").as_slice() {
        "tcp" => true,
        "udp" => false,
        _ => panic!("Unsupported protocol")
    };

    let proxy_addr = SocketAddr {
        ip: from_str(matches.opt_str("x").expect("Require proxy address").as_slice()).unwrap(),
        port: from_str(matches.opt_str("o").expect("Require proxy port").as_slice()).unwrap(),
    };

    let svr_port: Port = from_str(matches.opt_str("p").expect("Require server port").as_slice()).unwrap();
    let svr_addr = match from_str::<IpAddr>(matches.opt_str("s").expect("Require server address").as_slice()) {
        Some(ip) => SocketAddress(SocketAddr {ip: ip, port: svr_port}),
        None => DomainNameAddress(DomainNameAddr { domain_name: matches.opt_str("s").unwrap(), port: svr_port}),
    };

    if is_tcp {
        do_tcp(&matches, &svr_addr, &proxy_addr);
    } else {
        do_udp(&matches, &svr_addr, &proxy_addr);
    }
}
<|endoftext|>


<commit_before>// Copyright 2022 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::{
    os::windows::io::{AsRawHandle, RawHandle},
    ptr,
    time::Duration,
};

use win_util::{LargeInteger, SecurityAttributes, SelfRelativeSecurityDescriptor};
use winapi::{
    shared::{minwindef::FALSE, winerror::WAIT_TIMEOUT},
    um::{
        synchapi::{CancelWaitableTimer, SetWaitableTimer, WaitForSingleObject},
        winbase::{CreateWaitableTimerA, INFINITE, WAIT_OBJECT_0},
    },
};

use super::{
    super::{errno_result, win::nt_query_timer_resolution, Result},
    Timer, WaitResult,
};
use crate::descriptor::{AsRawDescriptor, FromRawDescriptor, SafeDescriptor};

impl AsRawHandle for Timer {
    fn as_raw_handle(&self) -> RawHandle {
        self.handle.as_raw_descriptor()
    }
}

impl Timer {
    /// Creates a new timer.  The timer is initally disarmed and must be armed by calling
    /// `reset`.
    pub fn new() -> Result<Timer> {
        // Safe because this doesn't modify any memory and we check the return value.
        let handle = unsafe {
            CreateWaitableTimerA(
                // Not inheritable, duplicate before passing to child prcesses
                SecurityAttributes::new_with_security_descriptor(
                    SelfRelativeSecurityDescriptor::get_singleton(),
                    /* inherit= */ false,
                )
                .as_mut(),
                // This is a synchronization timer, not a manual-reset timer.
                FALSE,
                // TODO (colindr) b/145622516 - we may have to give this a name if we later
                // want to use names to test object equality
                ptr::null_mut(),
            )
        };

        if handle.is_null() {
            return errno_result();
        }

        // Safe because we uniquely own the file descriptor.
        Ok(Timer {
            handle: unsafe { SafeDescriptor::from_raw_descriptor(handle) },
            interval: None,
        })
    }

    /// Sets the timer to expire after `dur`. If `interval` is not `None` and non-zero
    /// it represents the period for repeated expirations after the initial expiration.
    /// Otherwise the timer will expire just once.  Cancels any existing duration and
    /// repeating interval.
    pub fn reset(&mut self, dur: Duration, mut interval: Option<Duration>) -> Result<()> {
        // If interval is 0 or None it means that this timer does not repeat. We
        // set self.interval to None in this case so it can easily be checked
        // in self.wait.
        if interval == Some(Duration::from_secs(0)) {
            interval = None;
        }
        self.interval = interval;
        // Windows timers use negative values for relative times, and positive
        // values for absolute times, so we'll use negative times.

        // Windows timers also use a 64 number of 100 nanosecond intervals,
        // which we get like so: (dur.as_secs()*1e7 + dur.subsec_nanos()/100)

        let due_time = LargeInteger::new(
            -((dur.as_secs() * 10_000_000 + (dur.subsec_nanos() as u64) / 100) as i64),
        );
        let period: i32 = match interval {
            Some(int) => int.as_millis() as i32,
            None => 0,
        };

        // Safe because this doesn't modify any memory and we check the return value.
        let ret = unsafe {
            SetWaitableTimer(
                self.as_raw_descriptor(),
                &*due_time,
                period,
                None,            // no completion routine
                ptr::null_mut(), // or routine argument
                FALSE,           // no restoring system from power conservation mode
            )
        };
        if ret == 0 {
            return errno_result();
        }

        Ok(())
    }

    /// Waits until the timer expires, returing WaitResult::Expired when it expires.
    ///
    /// If timeout is not None, block for a maximum of the given `timeout` duration.
    /// If a timeout occurs, return WaitResult::Timeout.
    pub fn wait(&mut self, timeout: Option<Duration>) -> Result<WaitResult> {
        let timeout = match timeout {
            None => INFINITE,
            Some(dur) => dur.as_millis() as u32,
        };

        // Safe because this doesn't modify any memory and we check the return value.
        let ret = unsafe { WaitForSingleObject(self.as_raw_descriptor(), timeout) };

        // Should return WAIT_OBJECT_0, otherwise it's some sort of error or
        // timeout (which shouldn't happen in this case).
        match ret {
            WAIT_OBJECT_0 => Ok(WaitResult::Expired),
            WAIT_TIMEOUT => Ok(WaitResult::Timeout),
            _ => errno_result(),
        }
    }

    /// After a timer is triggered from an EventContext, mark the timer as having been waited for.
    /// If a timer is not marked waited, it will immediately trigger the event context again. This
    /// does not need to be called after calling Timer::wait.
    ///
    /// Returns true if the timer has been adjusted since the EventContext was triggered by this
    /// timer.
    pub fn mark_waited(&mut self) -> Result<bool> {
        // We use a synchronization timer on windows, meaning waiting on the timer automatically
        // un-signals the timer. We assume this is atomic so the return value is always false.
        Ok(false)
    }

    /// Disarms the timer.
    pub fn clear(&mut self) -> Result<()> {
        // Safe because this doesn't modify any memory and we check the return value.
        let ret = unsafe { CancelWaitableTimer(self.as_raw_descriptor()) };

        if ret == 0 {
            return errno_result();
        }

        self.interval = None;
        Ok(())
    }

    /// Returns the resolution of timers on the host.
    pub fn resolution() -> Result<Duration> {
        nt_query_timer_resolution().map(|(current_res, _)| current_res)
    }
}
<commit_msg>base: document early wake potential on Timer.<commit_after>// Copyright 2022 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::{
    os::windows::io::{AsRawHandle, RawHandle},
    ptr,
    time::Duration,
};

use win_util::{LargeInteger, SecurityAttributes, SelfRelativeSecurityDescriptor};
use winapi::{
    shared::{minwindef::FALSE, winerror::WAIT_TIMEOUT},
    um::{
        synchapi::{CancelWaitableTimer, SetWaitableTimer, WaitForSingleObject},
        winbase::{CreateWaitableTimerA, INFINITE, WAIT_OBJECT_0},
    },
};

use super::{
    super::{errno_result, win::nt_query_timer_resolution, Result},
    Timer, WaitResult,
};
use crate::descriptor::{AsRawDescriptor, FromRawDescriptor, SafeDescriptor};

impl AsRawHandle for Timer {
    fn as_raw_handle(&self) -> RawHandle {
        self.handle.as_raw_descriptor()
    }
}

impl Timer {
    /// Creates a new timer.  The timer is initally disarmed and must be armed by calling
    /// `reset`. Note that this timer MAY wake/trigger early due to limitations on
    /// SetWaitableTimer (see https://github.com/rust-lang/rust/issues/43376).
    pub fn new() -> Result<Timer> {
        // Safe because this doesn't modify any memory and we check the return value.
        let handle = unsafe {
            CreateWaitableTimerA(
                // Not inheritable, duplicate before passing to child prcesses
                SecurityAttributes::new_with_security_descriptor(
                    SelfRelativeSecurityDescriptor::get_singleton(),
                    /* inherit= */ false,
                )
                .as_mut(),
                // This is a synchronization timer, not a manual-reset timer.
                FALSE,
                // TODO (colindr) b/145622516 - we may have to give this a name if we later
                // want to use names to test object equality
                ptr::null_mut(),
            )
        };

        if handle.is_null() {
            return errno_result();
        }

        // Safe because we uniquely own the file descriptor.
        Ok(Timer {
            handle: unsafe { SafeDescriptor::from_raw_descriptor(handle) },
            interval: None,
        })
    }

    /// Sets the timer to expire after `dur`. If `interval` is not `None` and non-zero
    /// it represents the period for repeated expirations after the initial expiration.
    /// Otherwise the timer will expire just once.  Cancels any existing duration and
    /// repeating interval.
    pub fn reset(&mut self, dur: Duration, mut interval: Option<Duration>) -> Result<()> {
        // If interval is 0 or None it means that this timer does not repeat. We
        // set self.interval to None in this case so it can easily be checked
        // in self.wait.
        if interval == Some(Duration::from_secs(0)) {
            interval = None;
        }
        self.interval = interval;
        // Windows timers use negative values for relative times, and positive
        // values for absolute times, so we'll use negative times.

        // Windows timers also use a 64 number of 100 nanosecond intervals,
        // which we get like so: (dur.as_secs()*1e7 + dur.subsec_nanos()/100)

        let due_time = LargeInteger::new(
            -((dur.as_secs() * 10_000_000 + (dur.subsec_nanos() as u64) / 100) as i64),
        );
        let period: i32 = match interval {
            Some(int) => int.as_millis() as i32,
            None => 0,
        };

        // Safe because this doesn't modify any memory and we check the return value.
        let ret = unsafe {
            SetWaitableTimer(
                self.as_raw_descriptor(),
                &*due_time,
                period,
                None,            // no completion routine
                ptr::null_mut(), // or routine argument
                FALSE,           // no restoring system from power conservation mode
            )
        };
        if ret == 0 {
            return errno_result();
        }

        Ok(())
    }

    /// Waits until the timer expires, returing WaitResult::Expired when it expires.
    ///
    /// If timeout is not None, block for a maximum of the given `timeout` duration.
    /// If a timeout occurs, return WaitResult::Timeout.
    pub fn wait(&mut self, timeout: Option<Duration>) -> Result<WaitResult> {
        let timeout = match timeout {
            None => INFINITE,
            Some(dur) => dur.as_millis() as u32,
        };

        // Safe because this doesn't modify any memory and we check the return value.
        let ret = unsafe { WaitForSingleObject(self.as_raw_descriptor(), timeout) };

        // Should return WAIT_OBJECT_0, otherwise it's some sort of error or
        // timeout (which shouldn't happen in this case).
        match ret {
            WAIT_OBJECT_0 => Ok(WaitResult::Expired),
            WAIT_TIMEOUT => Ok(WaitResult::Timeout),
            _ => errno_result(),
        }
    }

    /// After a timer is triggered from an EventContext, mark the timer as having been waited for.
    /// If a timer is not marked waited, it will immediately trigger the event context again. This
    /// does not need to be called after calling Timer::wait.
    ///
    /// Returns true if the timer has been adjusted since the EventContext was triggered by this
    /// timer.
    pub fn mark_waited(&mut self) -> Result<bool> {
        // We use a synchronization timer on windows, meaning waiting on the timer automatically
        // un-signals the timer. We assume this is atomic so the return value is always false.
        Ok(false)
    }

    /// Disarms the timer.
    pub fn clear(&mut self) -> Result<()> {
        // Safe because this doesn't modify any memory and we check the return value.
        let ret = unsafe { CancelWaitableTimer(self.as_raw_descriptor()) };

        if ret == 0 {
            return errno_result();
        }

        self.interval = None;
        Ok(())
    }

    /// Returns the resolution of timers on the host.
    pub fn resolution() -> Result<Duration> {
        nt_query_timer_resolution().map(|(current_res, _)| current_res)
    }
}
<|endoftext|>


<commit_before>//
// Copyright (c) Pirmin Kalberer. All rights reserved.
// Licensed under the MIT License. See LICENSE file in the project root for full license information.
//

use toml;
use rustc_serialize::json::{self, Json};

/// Convert Mapbox GL Styles from [TOML format](https://pka.github.io/mapbox-gl-style-spec/) to JSON
pub fn toml_style_to_gljson(toml: &toml::Value) -> String {
    let configjson = json::encode(toml).unwrap().replace("}{", "},{").replace("][", "],[");
    Json::from_str(&configjson).unwrap().pretty().to_string()
}

#[test]
pub fn color_stops() {
    use toml::{Value, Parser};

    let style = r##"
        [circle-color]
        property = "temperature"
        stops = [{ in = 0, out = "blue" }, { in = 100, out = "red" }]"##;

    let mut parser = Parser::new(style);
    let toml = Value::Table(parser.parse().unwrap());

    let configjson = toml_style_to_gljson(&toml);
    println!("{}", configjson);
    let expected= r##"{
  "circle-color": {
    "property": "temperature",
    "stops": [
      {
        "in": 0,
        "out": "blue"
      },
      {
        "in": 100,
        "out": "red"
      }
    ]
  }
}"##;
    assert_eq!(configjson, expected);

    let style = r##"
        [layers.paint]
        fill-color = "#f2eae2"
        fill-outline-color = "#dfdbd7"

        [layers.paint.fill-opacity]
        base = 1
        stops = [{ in = 15, out = 0 }, { in = 16, out = 1 }]

        [layers.paint.fill-translate]
        base = 1
        stops = [{ in = 15 }, { in = 16 }]"##;

    let mut parser = Parser::new(style);
    let toml = Value::Table(parser.parse().unwrap());

    let configjson = toml_style_to_gljson(&toml);
    println!("{}", configjson);
    let expected= r##"{
  "layers": {
    "paint": {
      "fill-color": "#f2eae2",
      "fill-opacity": {
        "base": 1,
        "stops": [
          {
            "in": 15,
            "out": 0
          },
          {
            "in": 16,
            "out": 1
          }
        ]
      },
      "fill-outline-color": "#dfdbd7",
      "fill-translate": {
        "base": 1,
        "stops": [
          {
            "in": 15
          },
          {
            "in": 16
          }
        ]
      }
    }
  }
}"##;
    assert_eq!(configjson, expected);
}

#[test]
pub fn layer_style_from_cfg() {
    use core::read_config;

    let config = read_config("src/test/example.cfg").unwrap();

    let style = config.lookup("tileset.0.style").unwrap();

    let configjson = toml_style_to_gljson(style);
    println!("{}", configjson);
    let expected= r##"{
  "paint": {
    "background-color": "#f8f4f0"
  },
  "type": "background"
}"##;
    assert_eq!(configjson, expected);

    let style = config.lookup("tileset.0.layer.1.style").unwrap();

    let configjson = toml_style_to_gljson(style);
    println!("{}", configjson);
    let expected= r##"{
  "interactive": true,
  "paint": {
    "fill-color": {
      "base": 1,
      "stops": [
        {
          "in": 15.5,
          "out": "#f2eae2"
        },
        {
          "in": 16,
          "out": "#dfdbd7"
        }
      ]
    }
  },
  "type": "fill"
}"##;
    assert_eq!(configjson, expected);
}
<commit_msg>Zoom function conversion<commit_after>//
// Copyright (c) Pirmin Kalberer. All rights reserved.
// Licensed under the MIT License. See LICENSE file in the project root for full license information.
//

use std::collections::BTreeMap;
use rustc_serialize::json::{ToJson, Json};
use toml;
use std;

use toml::Value::{self, String, Integer, Float, Boolean, Datetime, Array, Table};

/// Convert Mapbox GL Styles from [TOML format](https://pka.github.io/mapbox-gl-style-spec/) to JSON
fn toml_style_to_gljson(toml: &toml::Value) -> std::string::String {
    let converter = TomlConverter::new();
    let json = converter.convert_value(toml);
    json.pretty().to_string()
}

struct TomlConverter;
impl TomlConverter {
    pub fn new() -> TomlConverter {
        TomlConverter
    }

    pub fn convert_value(&self, toml: &toml::Value) -> Json {
        match *toml {
            Table(ref value) => self.convert_table(value),

            Array(ref array) => {
                let mut vec = Vec::new();
                for value in array.iter() {
                    vec.push(self.convert_value(value));
                }
                vec.to_json()
            }

            String(ref value) => value.to_json(),
            Integer(ref value) => value.to_json(),
            Float(ref value) => value.to_json(),
            Boolean(ref value) => value.to_json(),
            Datetime(ref value) => value.to_json(),
        }
    }

    pub fn convert_table(&self, table: &BTreeMap<std::string::String, Value>) -> Json {
        let mut json: BTreeMap<std::string::String, Json> = BTreeMap::new();
        for (key, value) in table.iter() {
            if key == "stops" {
                let mut stops = Vec::new();
                if let Array(ref stops_arr) = *value {
                    for stop in stops_arr.iter() {
                        if let Table(ref stops_tbl) = *stop {
                            if stops_tbl.contains_key("in") {
                                if stops_tbl.contains_key("out") {
                                    stops.push(vec![
                                        self.convert_value(&stops_tbl["in"]),
                                        self.convert_value(&stops_tbl["out"])]);
                                } else {
                                    stops.push(vec![
                                        self.convert_value(&stops_tbl["in"])]);
                                }
                            }
                        }
                    }

                }
                json.insert(key.to_string(), stops.to_json());
            } else {
                json.insert(key.to_string(), self.convert_value(value));
            }
        }
        json.to_json()
    }
}

#[test]
pub fn color_stops() {
    use toml::{Value, Parser};

    let style = r##"
        [circle-color]
        property = "temperature"
        stops = [{ in = 0, out = "blue" }, { in = 100, out = "red" }]"##;

    let mut parser = Parser::new(style);
    let toml = Value::Table(parser.parse().unwrap());

    let configjson = toml_style_to_gljson(&toml);
    println!("{}", configjson);
    let expected= r##"{
  "circle-color": {
    "property": "temperature",
    "stops": [
      [
        0,
        "blue"
      ],
      [
        100,
        "red"
      ]
    ]
  }
}"##;
    assert_eq!(configjson, expected);

    let style = r##"
        [layers.paint]
        fill-color = "#f2eae2"
        fill-outline-color = "#dfdbd7"

        [layers.paint.fill-opacity]
        base = 1
        stops = [{ in = 15, out = 0 }, { in = 16, out = 1 }]

        [layers.paint.fill-translate]
        base = 1
        stops = [{ in = 15, out = [0,0] }, { in = 16, out = [-2,-2] }]"##;

    let mut parser = Parser::new(style);
    let toml = Value::Table(parser.parse().unwrap());

    let configjson = toml_style_to_gljson(&toml);
    println!("{}", configjson);
    let expected= r##"{
  "layers": {
    "paint": {
      "fill-color": "#f2eae2",
      "fill-opacity": {
        "base": 1,
        "stops": [
          [
            15,
            0
          ],
          [
            16,
            1
          ]
        ]
      },
      "fill-outline-color": "#dfdbd7",
      "fill-translate": {
        "base": 1,
        "stops": [
          [
            15,
            [
              0,
              0
            ]
          ],
          [
            16,
            [
              -2,
              -2
            ]
          ]
        ]
      }
    }
  }
}"##;
    assert_eq!(configjson, expected);
}

#[test]
pub fn filters() {
    use toml::{Value, Parser};

    let style = r##"
        [[layers]]
        id = "tunnel_path_pedestrian"
        type = "line"
        source = "mapbox"
        source-layer = "road"
        filter = [["all"], ["==", "$type", "LineString"], [["all"], ["==", "structure", "tunnel"], ["in", "class", "path", "pedestrian"]]]
        interactive = true"##;

    let mut parser = Parser::new(style);
    let toml = Value::Table(parser.parse().unwrap());

    let configjson = toml_style_to_gljson(&toml);
    println!("{}", configjson);
    let expected= r##"{
  "layers": [
    {
      "filter": [
        [
          "all"
        ],
        [
          "==",
          "$type",
          "LineString"
        ],
        [
          [
            "all"
          ],
          [
            "==",
            "structure",
            "tunnel"
          ],
          [
            "in",
            "class",
            "path",
            "pedestrian"
          ]
        ]
      ],
      "id": "tunnel_path_pedestrian",
      "interactive": true,
      "source": "mapbox",
      "source-layer": "road",
      "type": "line"
    }
  ]
}"##;
    /* TODO: Should be
    "id": "tunnel_path_pedestrian",
    "type": "line",
    "source": "mapbox",
    "source-layer": "road",
    "filter": [
        "all",
        [
            "==",
            "$type",
            "LineString"
        ],
        [
            "all",
            [
                "==",
                "structure",
                "tunnel"
            ],
            [
                "in",
                "class",
                "path",
                "pedestrian"
            ]
        ]
    ],*/
    assert_eq!(configjson, expected);
}

#[test]
pub fn layer_style_from_cfg() {
    use core::read_config;

    let config = read_config("src/test/example.cfg").unwrap();

    let style = config.lookup("tileset.0.style").unwrap();

    let configjson = toml_style_to_gljson(style);
    println!("{}", configjson);
    let expected= r##"{
  "paint": {
    "background-color": "#f8f4f0"
  },
  "type": "background"
}"##;
    assert_eq!(configjson, expected);

    let style = config.lookup("tileset.0.layer.1.style").unwrap();

    let configjson = toml_style_to_gljson(style);
    println!("{}", configjson);
    let expected= r##"{
  "interactive": true,
  "paint": {
    "fill-color": {
      "base": 1,
      "stops": [
        [
          15.5,
          "#f2eae2"
        ],
        [
          16,
          "#dfdbd7"
        ]
      ]
    }
  },
  "type": "fill"
}"##;
    assert_eq!(configjson, expected);
}
<|endoftext|>


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use libimagrt::runtime::Runtime;

pub fn get_diary_name(rt: &Runtime) -> Option<String> {
    use libimagdiary::config::get_default_diary_name;

    get_default_diary_name(rt)
        .or(rt.cli().value_of("diaryname").map(String::from))
}

<commit_msg>Add config parsing helpers<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use libimagrt::runtime::Runtime;
use libimagdiary::error::*;

use toml::Value;
use toml_query::read::TomlValueReadExt;

pub fn get_diary_name(rt: &Runtime) -> Option<String> {
    use libimagdiary::config::get_default_diary_name;

    get_default_diary_name(rt)
        .or(rt.cli().value_of("diaryname").map(String::from))
}

pub enum Timed {
    Hourly,
    Minutely,
}

/// Returns true if the diary should always create timed entries, which is whenever
///
/// ```toml
/// diary.diaries.<diary>.timed = true
/// ```
///
/// # Returns
///
/// * Ok(Some(Timed::Hourly)) if diary should create timed entries
/// * Ok(Some(Timed::Minutely)) if diary should not create timed entries
/// * Ok(None) if config is not available
/// * Err(e) if reading the toml failed, type error or something like this
///
pub fn get_diary_timed_config(rt: &Runtime, diary_name: &str) -> Result<Option<Timed>> {
    match rt.config() {
        None      => Ok(None),
        Some(cfg) => {
            let v = cfg
                .config()
                .read(&format!("diary.diaries.{}.timed", diary_name))
                .chain_err(|| DiaryErrorKind::IOError);

            match v {
                Ok(Some(&Value::String(ref s))) => parse_timed_string(s, diary_name).map(Some),

                Ok(Some(_)) => {
                    let s = format!("Type error at 'diary.diaryies.{}.timed': should be either 'h'/'hourly' or 'm'/'minutely'", diary_name);
                    Err(s).map_err(From::from)
                },

                Ok(None) => Ok(None),
                Err(e) => Err(e),
            }
        }
    }
}

pub fn parse_timed_string(s: &str, diary_name: &str) -> Result<Timed> {
    if s == "h" || s == "hourly" {
        Ok(Timed::Hourly)
    } else if s == "m" || s == "minutely" {
        Ok(Timed::Minutely)
    } else {
        let s = format!("Cannot parse config: 'diary.diaries.{}.timed = {}'",
                        diary_name, s);
        Err(s).map_err(From::from)
    }
}
<|endoftext|>


<commit_before>mod terminal;
mod ansi;

use std::io::{self, Write};
use self::terminal::{Terminal, Event, Key};
use self::ansi::{clear, color, cursor, style};
use rayon::prelude::*;
use rff::choice::Choice;

#[derive(Debug)]
pub enum Error {
    Exit,
    Write(io::Error),
    Reset(terminal::Error)
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Write(err)
    }
}

impl From<terminal::Error> for Error {
    fn from(err: terminal::Error) -> Error {
        Error::Reset(err)
    }
}

pub struct Options {
    pub choices: Vec<String>,
    pub initial: String,
}

pub struct Interface {
    choices: Vec<String>,
    matching: Vec<Choice>,
    selected: usize,
    search: String,
    terminal: Terminal
}

impl Interface {
    /// Creates a new Interface from the provided options.
    pub fn from_opts(opts: Options) -> Interface {
        let mut term = Terminal::from("/dev/tty").unwrap();
        term.set_raw_mode().unwrap();

        Interface {
            choices: opts.choices,
            matching: Vec::new(),
            selected: 0,
            search: opts.initial,
            terminal: term
        }
    }

    // Starts the interface
    pub fn run(&mut self) -> Result<&str, Error> {
        self.filter_choices();
        self.render()?;

        for event in self.terminal.events()? {
            if let Event::Key(key) = event? {
                match key {
                    Key::Ctrl('c') | Key::Ctrl('d') => {
                        self.clear()?;
                        return Err(Error::Exit);
                    }

                    Key::Char('\n') => {
                        break;
                    },

                    Key::Ctrl('n') => {
                        self.selected += 1;
                        self.clamp_selected();
                        self.render()?;
                    },

                    Key::Ctrl('p') => {
                        self.selected = self.selected.
                            checked_sub(1).
                            unwrap_or(0);

                        self.clamp_selected();
                        self.render()?;
                    },

                    Key::Char(ch) => {
                        self.search.push(ch);
                        self.filter_existing_matches();
                        self.render()?;
                    },

                    Key::Backspace | Key::Ctrl('h') => {
                        self.search.pop();
                        self.filter_choices();
                        self.render()?;
                    }

                    Key::Ctrl('u') => {
                        self.search.clear();
                        self.filter_choices();
                        self.render()?;
                    }

                    _ => {}
                }
            };
        }

        self.clear()?;
        Ok(self.result())
    }

    fn filter_choices(&mut self) {
        let mut matches = self.choices.
            par_iter().
            cloned().
            filter_map(|choice| Choice::with_positions(&self.search, choice)).
            collect::<Vec<_>>();

        matches.sort_by(|a, b| b.partial_cmp(a).unwrap());

        self.matching = matches;
    }

    fn filter_existing_matches(&mut self) {
        let mut matches = self.matching.
            par_iter().
            map(|c| c.text().to_string()).
            filter_map(|choice| Choice::with_positions(&self.search, choice)).
            collect::<Vec<_>>();

        matches.sort_by(|a, b| b.partial_cmp(a).unwrap());

        self.matching = matches;
    }

    fn render(&mut self) -> io::Result<()> {
        write!(self.terminal, "{}{}", cursor::Column(1), clear::Screen)?;
        write!(self.terminal, "> {}", self.search)?;

        let n = self.render_choices()?;

        if n > 0 {
            let column = format!("> {}", self.search).len() as u16;
            write!(self.terminal, "{}{}", cursor::Up(n), cursor::Column(column + 1))?;
        }

        Ok(())
    }

    fn render_choices(&mut self) -> io::Result<u16> {
        let max_width = self.terminal.max_width;

        let matches = self.matching.iter().take(10);

        let number_of_choices = matches.len() as u16;

        for (i, choice) in matches.enumerate() {
            let selected = i == self.selected;

            write!(self.terminal, "\r\n")?;

            if selected {
                write!(self.terminal, "{}", style::Invert)?;
            }

            if let Some(positions) = choice.positions() {
                let chars = choice.text().chars().take(max_width);

                for (i, ch) in chars.enumerate() {
                    let match_position = positions.iter().any(|p| *p == i);

                    if match_position {
                        let color = color::Fg(color::Colors::Magenta);
                        let reset = color::Fg(color::Reset);
                        write!(self.terminal, "{}{}{}", color, ch, reset)?;
                    } else {
                        write!(self.terminal, "{}", ch)?;
                    }
                }

            } else {
                let text = choice.text()
                    .chars()
                    .take(max_width)
                    .collect::<String>();

                write!(self.terminal, "{}", text)?;
            }

            if selected {
                write!(self.terminal, "{}", style::NoInvert)?;
            }
        }

        Ok(number_of_choices)
    }

    #[inline(always)]
    fn clamp_selected(&mut self) {
        let mut max = self.matching.len();
        if max > 10 { max = 10; }

        if self.selected >= max {
            self.selected = max - 1;
        }
    }

    fn clear(&mut self) -> Result<(), Error> {
        write!(self.terminal, "{}{}", cursor::Column(1), clear::Screen)?;
        self.terminal.reset()?;
        Ok(())
    }

    fn result(&mut self) -> &str {
        self.matching.iter().
            map(|c| c.text()).
            nth(self.selected).
            unwrap_or(&self.search)
    }
}
<commit_msg>Move fetching number of matches out of render loop<commit_after>mod terminal;
mod ansi;

use std::io::{self, Write};
use self::terminal::{Terminal, Event, Key};
use self::ansi::{clear, color, cursor, style};
use rayon::prelude::*;
use rff::choice::Choice;

#[derive(Debug)]
pub enum Error {
    Exit,
    Write(io::Error),
    Reset(terminal::Error)
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Write(err)
    }
}

impl From<terminal::Error> for Error {
    fn from(err: terminal::Error) -> Error {
        Error::Reset(err)
    }
}

pub struct Options {
    pub choices: Vec<String>,
    pub initial: String,
}

pub struct Interface {
    choices: Vec<String>,
    matching: Vec<Choice>,
    selected: usize,
    search: String,
    terminal: Terminal
}

impl Interface {
    /// Creates a new Interface from the provided options.
    pub fn from_opts(opts: Options) -> Interface {
        let mut term = Terminal::from("/dev/tty").unwrap();
        term.set_raw_mode().unwrap();

        Interface {
            choices: opts.choices,
            matching: Vec::new(),
            selected: 0,
            search: opts.initial,
            terminal: term
        }
    }

    // Starts the interface
    pub fn run(&mut self) -> Result<&str, Error> {
        self.filter_choices();
        self.render()?;

        for event in self.terminal.events()? {
            if let Event::Key(key) = event? {
                match key {
                    Key::Ctrl('c') | Key::Ctrl('d') => {
                        self.clear()?;
                        return Err(Error::Exit);
                    }

                    Key::Char('\n') => {
                        break;
                    },

                    Key::Ctrl('n') => {
                        self.selected += 1;
                        self.clamp_selected();
                        self.render()?;
                    },

                    Key::Ctrl('p') => {
                        self.selected = self.selected.
                            checked_sub(1).
                            unwrap_or(0);

                        self.clamp_selected();
                        self.render()?;
                    },

                    Key::Char(ch) => {
                        self.search.push(ch);
                        self.filter_existing_matches();
                        self.render()?;
                    },

                    Key::Backspace | Key::Ctrl('h') => {
                        self.search.pop();
                        self.filter_choices();
                        self.render()?;
                    }

                    Key::Ctrl('u') => {
                        self.search.clear();
                        self.filter_choices();
                        self.render()?;
                    }

                    _ => {}
                }
            };
        }

        self.clear()?;
        Ok(self.result())
    }

    fn filter_choices(&mut self) {
        let mut matches = self.choices.
            par_iter().
            cloned().
            filter_map(|choice| Choice::with_positions(&self.search, choice)).
            collect::<Vec<_>>();

        matches.sort_by(|a, b| b.partial_cmp(a).unwrap());

        self.matching = matches;
    }

    fn filter_existing_matches(&mut self) {
        let mut matches = self.matching.
            par_iter().
            map(|c| c.text().to_string()).
            filter_map(|choice| Choice::with_positions(&self.search, choice)).
            collect::<Vec<_>>();

        matches.sort_by(|a, b| b.partial_cmp(a).unwrap());

        self.matching = matches;
    }

    fn render(&mut self) -> io::Result<()> {
        write!(self.terminal, "{}{}", cursor::Column(1), clear::Screen)?;
        write!(self.terminal, "> {}", self.search)?;

        let n = self.matching.iter().take(10).len() as u16;

        self.render_matches()?;

        if n > 0 {
            let column = format!("> {}", self.search).len() as u16;
            write!(self.terminal, "{}{}", cursor::Up(n), cursor::Column(column + 1))?;
        }

        Ok(())
    }

    fn render_matches(&mut self) -> io::Result<()> {
        let max_width = self.terminal.max_width;
        let matches = self.matching.iter().take(10);

        for (i, choice) in matches.enumerate() {
            let selected = i == self.selected;

            write!(self.terminal, "\r\n")?;

            if selected {
                write!(self.terminal, "{}", style::Invert)?;
            }

            if let Some(positions) = choice.positions() {
                let chars = choice.text().chars().take(max_width);

                for (i, ch) in chars.enumerate() {
                    let match_position = positions.iter().any(|p| *p == i);

                    if match_position {
                        let color = color::Fg(color::Colors::Magenta);
                        let reset = color::Fg(color::Reset);
                        write!(self.terminal, "{}{}{}", color, ch, reset)?;
                    } else {
                        write!(self.terminal, "{}", ch)?;
                    }
                }

            } else {
                let text = choice.text()
                    .chars()
                    .take(max_width)
                    .collect::<String>();

                write!(self.terminal, "{}", text)?;
            }

            if selected {
                write!(self.terminal, "{}", style::NoInvert)?;
            }
        }

        Ok(())
    }

    #[inline(always)]
    fn clamp_selected(&mut self) {
        let mut max = self.matching.len();
        if max > 10 { max = 10; }

        if self.selected >= max {
            self.selected = max - 1;
        }
    }

    fn clear(&mut self) -> Result<(), Error> {
        write!(self.terminal, "{}{}", cursor::Column(1), clear::Screen)?;
        self.terminal.reset()?;
        Ok(())
    }

    fn result(&mut self) -> &str {
        self.matching.iter().
            map(|c| c.text()).
            nth(self.selected).
            unwrap_or(&self.search)
    }
}
<|endoftext|>