<commit_before>extern crate yaml_rust;
extern crate regex;

use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::net::TcpListener;
use std::time::Duration;
use std::sync::Arc;
use std::thread;
use std::process::exit;
use yaml_rust::YamlLoader;
use regex::bytes::RegexSetBuilder;

#[derive(Copy, Clone)]
struct App {
  print_ascii: bool,
  print_binary: bool
}

fn to_hex(bytes: &[u8]) -> String {
  let mut count = 0;
  let mut result = String::with_capacity(67);
  let mut dotline = String::with_capacity(18);

  for byte in bytes.iter() {
    if count != 0 && count%8 == 0 {
      result.push(' ');
      if count%16 == 0 {
        dotline.push('\n');
        result.push_str(&dotline);
        result.reserve(67);
        dotline.truncate(0);
      }
      else { dotline.push(' '); }
    }
    result.push_str(&format!("{:02X} ", byte));
    if *byte > 31 && *byte < 127 { dotline.push(*byte as char) }
    else if *byte == 0 { dotline.push('-'); }
    else { dotline.push('.'); }
    count += 1;
  }
  while count%16 != 0 {
    if count%8 == 0 { result.push(' '); }
    result.push_str("   ");
    count += 1;
  }
  if dotline.len() != 0 { result.push_str(&dotline); }
  result
}

fn main() {
  let mut app = App { print_ascii: false, print_binary: false };
  let binary_matches = [
    ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
    ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
    ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
    ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
    ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
    ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
    ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
    ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
    ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
    ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
    ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
    ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
    ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
    ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
    ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
    ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
    ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
    ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]"),
    ("MS-TDS PRELOGIN Request", r"^\x12\x01\x00.\x00\x00"),
    ("MS-TDS LOGIN Requst", r"^\x10\x01\x00.\x00\x00")
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut file = File::open("config.yml").unwrap();
  let mut config_str = String::new();
  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ascii");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  let app = app; // Revert to immutable

  let mut patterns = Vec::with_capacity(binary_matches.len());
  for &(_, pattern) in binary_matches.into_iter() {
    patterns.push(pattern);
  }
  let regexset = RegexSetBuilder::new(patterns)
    .unicode(false)
    .dot_matches_new_line(false)
    .build().unwrap();

  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).unwrap();
        for res in server.incoming() {
          let mut stream = match res {
            Ok(stream) => stream,
            Err(e) => { println!("ACCEPT ERROR TCP {}: {}", portno, e.to_string()); continue; }
          };
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);
          let regexset = regexset.clone();
          let banner = banner.clone();
          thread::spawn(move || {
            if banner.len() > 0 {
              match stream.write((*banner).as_bytes()) {
                Ok(_) => {
                  for line in (*banner).lines() {
                    println!("> {}", line);
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("WRITE TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("WRITE ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  return;
                }
              }
            }
            let mut buf: [u8; 2048] = [0; 2048];
            loop {
              match stream.read(&mut buf) {
                Ok(c) => {
                  if c == 0 {
                    println!("CLOSE TCP {} from {}", portno, addr);
                    break;
                  }
                  let mut printables = Vec::new();
                  let mut found = false;
                  let mut mbfound = false;
                  let mut mbstring = String::new();
                  let mut start = 0;
                  for i in 0..c {
                    if (buf[i] > 31 && buf[i] < 127) || buf[i] == 10 || buf[i] == 13 {
                      if !found {
                        start = i;
                        found = true;
                      }
                    }
                    else {
                      if found {
                        if (i-start == 1) {
                          if (start > 0) && (buf[start-1] == 0) {
                            mbstring.push(buf[i-1] as char);
                            if !mbfound { mbfound = true; }
                          }
                        }
                        else { printables.push(&buf[start..i]); }
                        found = false;
                      }
                      else if mbfound {
                        mbstring.push('\n');
                        mbfound = false;
                      }
                    }
                  }
                  if found { printables.push(&buf[start..c]); }
                  if printables.len() == 1 && printables[0].len() == c {
                    if app.print_ascii {
                      let data = String::from_utf8_lossy(printables[0]);
                      for line in data.lines() {
                        println!("| {}", line);
                      }
                    }
                    else { println!("! Read {} bytes of printable ascii", c); }
                  }
                  else {
                    println!("! Read {} bytes of binary", c);
                    for id in regexset.matches(&buf[..c]).into_iter() {
                      println!("^ Matches pattern {}", binary_matches[id].0);
                    }
                    for printable in printables {
                      if printable.len() > 3 {
                        let data = String::from_utf8_lossy(printable);
                        for line in data.lines() {
                          println!("$ {}", line);
                        }
                      }
                    }
                    for line in mbstring.lines() {
                      if line.len() > 3 { println!("% {}", line); }
                    }
                    if app.print_binary {
                      let hex = to_hex(&buf[..c]);
                      for line in hex.lines() { println!(". {}", line); }
                    }
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("READ TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("READ ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  break;
                }
              }
              match stream.take_error() {
                Ok(opt) => {
                  if opt.is_some() {
                    println!("ERROR TCP {} from {}: {}", portno, addr, opt.unwrap().to_string());
                    break;
                  }
                }
                Err(_) => {
                  println!("This shouldn't happen...");
                  break;
                }
              }
            }
          });
        }
      });
    }
    else if !port["udp"].is_badvalue() {
      println!("UDP port {}", port["udp"].as_i64().unwrap());
    }
    else {
      println!("Invalid port specification in configuration file");
    }
  }

  loop { thread::sleep(Duration::new(60, 0)); }
}
<commit_msg>Update main.rs<commit_after>extern crate yaml_rust;
extern crate regex;

use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::net::TcpListener;
use std::time::Duration;
use std::sync::Arc;
use std::thread;
use std::process::exit;
use yaml_rust::YamlLoader;
use regex::bytes::RegexSetBuilder;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

#[derive(Copy, Clone)]
struct App {
  print_ascii: bool,
  print_binary: bool
}

fn to_hex(bytes: &[u8]) -> String {
  let mut count = 0;
  let mut result = String::with_capacity(67);
  let mut dotline = String::with_capacity(18);

  for byte in bytes.iter() {
    if count != 0 && count%8 == 0 {
      result.push(' ');
      if count%16 == 0 {
        dotline.push('\n');
        result.push_str(&dotline);
        result.reserve(67);
        dotline.truncate(0);
      }
      else { dotline.push(' '); }
    }
    result.push_str(&format!("{:02X} ", byte));
    if *byte > 31 && *byte < 127 { dotline.push(*byte as char) }
    else if *byte == 0 { dotline.push('-'); }
    else { dotline.push('.'); }
    count += 1;
  }
  while count%16 != 0 {
    if count%8 == 0 { result.push(' '); }
    result.push_str("   ");
    count += 1;
  }
  if dotline.len() != 0 { result.push_str(&dotline); }
  result
}

fn main() {
  println!("PortLurker v{}", VERSION);
  println!("Bart Noordervlet (http://www.mmvl.nl/)");

  let mut app = App { print_ascii: false, print_binary: false };
  let binary_matches = [
    ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
    ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
    ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
    ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
    ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
    ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
    ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
    ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
    ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
    ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
    ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
    ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
    ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
    ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
    ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
    ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
    ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
    ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]"),
    ("MS-TDS PRELOGIN Request", r"^\x12\x01\x00.\x00\x00"),
    ("MS-TDS LOGIN Requst", r"^\x10\x01\x00.\x00\x00")
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut file = File::open("config.yml").unwrap();
  let mut config_str = String::new();
  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ascii");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  let app = app; // Revert to immutable

  let mut patterns = Vec::with_capacity(binary_matches.len());
  for &(_, pattern) in binary_matches.into_iter() {
    patterns.push(pattern);
  }
  let regexset = RegexSetBuilder::new(patterns)
    .unicode(false)
    .dot_matches_new_line(false)
    .build().unwrap();

  println!("Starting listeners on the following ports:");

  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).unwrap();
        for res in server.incoming() {
          let mut stream = match res {
            Ok(stream) => stream,
            Err(e) => { println!("ACCEPT ERROR TCP {}: {}", portno, e.to_string()); continue; }
          };
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);
          let regexset = regexset.clone();
          let banner = banner.clone();
          thread::spawn(move || {
            if banner.len() > 0 {
              match stream.write((*banner).as_bytes()) {
                Ok(_) => {
                  for line in (*banner).lines() {
                    println!("> {}", line);
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("WRITE TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("WRITE ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  return;
                }
              }
            }
            let mut buf: [u8; 2048] = [0; 2048];
            loop {
              match stream.read(&mut buf) {
                Ok(c) => {
                  if c == 0 {
                    println!("CLOSE TCP {} from {}", portno, addr);
                    break;
                  }
                  let mut printables = Vec::new();
                  let mut found = false;
                  let mut mbfound = false;
                  let mut mbstring = String::new();
                  let mut start = 0;
                  for i in 0..c {
                    if (buf[i] > 31 && buf[i] < 127) || buf[i] == 10 || buf[i] == 13 {
                      if !found {
                        start = i;
                        found = true;
                      }
                    }
                    else {
                      if found {
                        if (i-start == 1) {
                          if (start > 0) && (buf[start-1] == 0) {
                            mbstring.push(buf[i-1] as char);
                            if !mbfound { mbfound = true; }
                          }
                        }
                        else { printables.push(&buf[start..i]); }
                        found = false;
                      }
                      else if mbfound {
                        mbstring.push('\n');
                        mbfound = false;
                      }
                    }
                  }
                  if found { printables.push(&buf[start..c]); }
                  if printables.len() == 1 && printables[0].len() == c {
                    if app.print_ascii {
                      let data = String::from_utf8_lossy(printables[0]);
                      for line in data.lines() {
                        println!("| {}", line);
                      }
                    }
                    else { println!("! Read {} bytes of printable ascii", c); }
                  }
                  else {
                    println!("! Read {} bytes of binary", c);
                    for id in regexset.matches(&buf[..c]).into_iter() {
                      println!("^ Matches pattern {}", binary_matches[id].0);
                    }
                    for printable in printables {
                      if printable.len() > 3 {
                        let data = String::from_utf8_lossy(printable);
                        for line in data.lines() {
                          println!("$ {}", line);
                        }
                      }
                    }
                    for line in mbstring.lines() {
                      if line.len() > 3 { println!("% {}", line); }
                    }
                    if app.print_binary {
                      let hex = to_hex(&buf[..c]);
                      for line in hex.lines() { println!(". {}", line); }
                    }
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("READ TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("READ ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  break;
                }
              }
              match stream.take_error() {
                Ok(opt) => {
                  if opt.is_some() {
                    println!("ERROR TCP {} from {}: {}", portno, addr, opt.unwrap().to_string());
                    break;
                  }
                }
                Err(_) => {
                  println!("This shouldn't happen...");
                  break;
                }
              }
            }
          });
        }
      });
    }
    else if !port["udp"].is_badvalue() {
      println!("UDP port {}", port["udp"].as_i64().unwrap());
    }
    else {
      println!("Invalid port specification in configuration file");
    }
  }

  loop { thread::sleep(Duration::new(60, 0)); }
}
<|endoftext|>


<commit_before>//! # Triple buffering
//!
//! In this crate, we attempt to implement a triple buffering mechanism,
//! suitable for emulating shared memory cells in a thread-safe and
//! non-blocking fashion between one single writer and one single reader.

use std::cell::UnsafeCell;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Duration;

/// ## Indexing
///
/// These types are used to index into triple buffers
/// TODO: Switch to i8 / AtomicI8 once the later is stable
type TripleBufferIndex = usize;
type AtomicTripleBufferIndex = AtomicUsize;

/// ## Shared state
///
/// In a triple buffering communication protocol, the reader and writer share
/// the following storage:
///
/// - Three memory buffers suitable for storing the data at hand
/// - One index pointing to the buffer which no one is currently using
/// - One index pointing to the most recently updated buffer
///
#[derive(Debug)]
struct TripleBufferSharedState<T: Clone + PartialEq> {
    /// Three buffers of data
    buffers: [UnsafeCell<T>; 3],

    /// Index of the back-buffer (buffer which no one is currently accessing)
    back_idx: AtomicTripleBufferIndex,

    /// Index of the most recently updated buffer
    last_idx: AtomicTripleBufferIndex,
}
//
impl<T: Clone + PartialEq> Clone for TripleBufferSharedState<T> {
    fn clone(&self) -> Self {
        // The use of UnsafeCell makes buffers somwhat cumbersome to clone...
        let clone_buffer = | i: TripleBufferIndex | {
            UnsafeCell::new(
                unsafe { (*self.buffers[i].get()).clone() }
            )
        };

        // ...so better define a shortcut before we move on to actual cloning
        TripleBufferSharedState {
            buffers: [
                clone_buffer(0),
                clone_buffer(1),
                clone_buffer(2),
            ],
            back_idx: AtomicTripleBufferIndex::new(
                self.back_idx.load(Ordering::Relaxed)
            ),
            last_idx: AtomicTripleBufferIndex::new(
                self.last_idx.load(Ordering::Relaxed)
            ),
        }
    }
}
//
impl<T: Clone + PartialEq> PartialEq for TripleBufferSharedState<T> {
    fn eq(&self, other: &Self) -> bool {
        // The use of UnsafeCell makes buffers somwhat cumbersome to compare...
        let buffer_eq = | i: TripleBufferIndex | -> bool {
            unsafe { *self.buffers[i].get() == *other.buffers[i].get() }
        };

        // ...and atomics aren't much better...
        let atomics_eq = | x: &AtomicTripleBufferIndex,
                           y: &AtomicTripleBufferIndex | -> bool {
            x.load(Ordering::Relaxed) == y.load(Ordering::Relaxed)
        };

        // ...so better define some shortcuts before getting started:
        buffer_eq(0) && buffer_eq(1) && buffer_eq(2) &&
        atomics_eq(&self.back_idx, &other.back_idx) &&
        atomics_eq(&self.last_idx, &other.last_idx)
    }
}
//
unsafe impl<T: Clone + PartialEq> Sync for TripleBufferSharedState<T> {}

/// ## Producing data ("input")
///
/// The producer of data has exclusive write access to the so-called "write
/// buffer". To send an update, it writes it into the back buffer, then swaps
/// the write buffer with the back buffer, making the data readable.
///
#[derive(Debug)]
struct TripleBufferInput<T: Clone + PartialEq> {
    shared: Arc<TripleBufferSharedState<T>>,
    write_idx: TripleBufferIndex,
}
//
impl<T: Clone + PartialEq> TripleBufferInput<T> {
    /// Write a new value into the triple buffer
    fn write(&mut self, value: T) {
        // Access the shared state
        let ref shared_state = *self.shared;

        // Determine which shared buffer is the write buffer
        let active_idx = self.write_idx;

        // Move the input value into the write buffer
        let write_ptr = shared_state.buffers[active_idx].get();
        unsafe { *write_ptr = value; }

        // Swap the back buffer and the write buffer. This makes the new data
        // block available to the reader and gives us a new buffer to write to.
        self.write_idx = shared_state.back_idx.swap(
            active_idx,
            Ordering::Release  // Send previous updates to the reader
        );

        // Notify the reader that we submitted an update
        shared_state.last_idx.store(active_idx, Ordering::Relaxed);
    }
}

/// ## Consuming data ("output")
///
/// The consumer of data has exclusive read access to the so-called "read
/// buffer". To access the latest data, it swaps the read buffer with the
/// back buffer if the later contains more recent data, then reads from
/// the (possibly new) read buffer.
///
#[derive(Debug)]
struct TripleBufferOutput<T: Clone + PartialEq> {
    shared: Arc<TripleBufferSharedState<T>>,
    read_idx: TripleBufferIndex,
}
//
impl<T: Clone + PartialEq> TripleBufferOutput<T> {
    /// Access the latest fully written value from the triple buffer
    fn read(&mut self) -> &T {
        // Access the shared state
        let ref shared_state = *self.shared;

        // Check if the writer has submitted an update
        let last_idx = shared_state.last_idx.load(Ordering::Relaxed);

        // If an update is pending in the back buffer, make it our read buffer
        if self.read_idx != last_idx {
            // Swap the back buffer and the read buffer. We get exclusive read
            // access to the data, and the writer gets a new back buffer.
            self.read_idx = shared_state.back_idx.swap(
                self.read_idx,
                Ordering::Acquire  // Synchronize with the writer's updates
            );
        }

        // Access data from the current read buffer
        let read_ptr = shared_state.buffers[self.read_idx].get();
        unsafe { &*read_ptr }
    }
}

/// ## Triple buffer
///
/// A triple buffer is created as the combination of an input and an output.
///
#[derive(Debug)]
struct TripleBuffer<T: Clone + PartialEq> {
    input: TripleBufferInput<T>,
    output: TripleBufferOutput<T>,
}
//
impl<T: Clone + PartialEq> TripleBuffer<T> {
    /// Construct a triple buffer with a certain initial value
    fn new(initial: T) -> Self {
        // Start with the shared state...
        let shared_state = Arc::new(
            TripleBufferSharedState {
                buffers: [
                    UnsafeCell::new(initial.clone()),
                    UnsafeCell::new(initial.clone()),
                    UnsafeCell::new(initial)
                ],
                back_idx: AtomicTripleBufferIndex::new(0),
                last_idx: AtomicTripleBufferIndex::new(2),
            }
        );

        // ...then construct the input and output structs
        TripleBuffer {
            input: TripleBufferInput {
                shared: shared_state.clone(),
                write_idx: 1,
            },
            output: TripleBufferOutput {
                shared: shared_state,
                read_idx: 2,
            },
        }
    }
}
//
impl<T: Clone + PartialEq> Clone for TripleBuffer<T> {
    fn clone(&self) -> Self {
        // Clone the shared state...
        let shared_state = Arc::new(
            (*self.input.shared).clone()
        );

        // ...then the input and output structs
        TripleBuffer {
            input: TripleBufferInput {
                shared: shared_state.clone(),
                write_idx: self.input.write_idx,
            },
            output: TripleBufferOutput {
                shared: shared_state,
                read_idx: self.output.read_idx,
            },
        }
    }
}
//
impl<T: Clone + PartialEq> PartialEq for TripleBuffer<T> {
    fn eq(&self, other: &Self) -> bool {
        (*self.input.shared == *other.input.shared) &&
        (self.input.write_idx == other.input.write_idx) &&
        (self.output.read_idx == other.output.read_idx)
    }
}

/// ## Testing
///
/// Unit tests and benchmarks for triple buffers
///
#[cfg(test)]
mod tests {
    /// Test that triple buffers are properly initialized
    #[test]
    fn test_init() {
        // Let's create a triple buffer
        let buf = ::TripleBuffer::new(42);

        // Access the shared state
        let ref buf_shared = *buf.input.shared;
        let back_idx = buf_shared.back_idx.load(::Ordering::Relaxed);
        let last_idx = buf_shared.last_idx.load(::Ordering::Relaxed);

        // Write-/read-/back-buffer indexes must be in range
        assert!(index_in_range(buf.input.write_idx));
        assert!(index_in_range(buf.output.read_idx));
        assert!(index_in_range(back_idx));

        // Write-/read-/back-buffer indexes must be distinct
        assert!(buf.input.write_idx != buf.output.read_idx);
        assert!(buf.input.write_idx != back_idx);
        assert!(buf.output.read_idx != back_idx);

        // Read buffer must be properly initialized
        let read_ptr = buf_shared.buffers[buf.output.read_idx].get();
        assert!(unsafe { *read_ptr } == 42);

        // Last-written index must initially point to the read buffer
        assert_eq!(last_idx, buf.output.read_idx);
    }

    /// Test that (sequentially) writing to a triple buffer works
    #[test]
    fn test_seq_write() {
        // Let's create a triple buffer
        let mut buf = ::TripleBuffer::new(false);

        // Back up the initial buffer state
        let old_buf = buf.clone();
        let ref old_shared = old_buf.input.shared;

        // Perform a write
        buf.input.write(true);

        // Check new implementation state
        {
            // Start from the old buffer state
            let mut expected_buf = old_buf.clone();
            let ref expected_shared = expected_buf.input.shared;

            // We expect the former write buffer to have received the new value
            let old_write_idx = old_buf.input.write_idx;
            let write_buf_ptr = expected_shared.buffers[old_write_idx].get();
            unsafe { *write_buf_ptr = true; }

            // We expect the former write buffer to be the new back buffer
            expected_shared.back_idx.store(
                old_write_idx,
                ::Ordering::Relaxed
            );

            // We expect the last updated index to point to former write buffer
            expected_shared.last_idx.store(
                old_write_idx,
                ::Ordering::Relaxed
            );

            // We expect the old back buffer to become the new write buffer
            let old_back_idx = old_shared.back_idx.load(::Ordering::Relaxed);
            expected_buf.input.write_idx = old_back_idx;

            // Nothing else should have changed
            assert_eq!(buf, expected_buf);
        }
    }

    /// Check that (sequentially) reading from a triple buffer works
    #[test]
    fn test_seq_read() {
        // Let's create a triple buffer and write into it
        let mut buf = ::TripleBuffer::new(1.0);
        buf.input.write(4.2);

        // Test readout from dirty (freshly written) triple buffer
        {
            // Back up the initial buffer state
            let old_buf = buf.clone();
            let ref old_shared = old_buf.input.shared;

            // Read from the buffer
            let result = *buf.output.read();

            // Output value should be correct
            assert_eq!(result, 4.2);

            // Start from the old buffer state
            let mut expected_buf = old_buf.clone();
            let ref expected_shared = expected_buf.input.shared;

            // We expect the new back index to point to the former read buffer
            expected_shared.back_idx.store(
                old_buf.output.read_idx,
                ::Ordering::Relaxed
            );

            // We expect the new read index to point to the former back buffer
            let old_back_idx = old_shared.back_idx.load(::Ordering::Relaxed);
            expected_buf.output.read_idx = old_back_idx;

            // Nothing else should have changed
            assert_eq!(buf, expected_buf);
        }

        // Test readout from clean (unchanged) triple buffer
        {
            // Back up the initial buffer state
            let old_buf = buf.clone();

            // Read from the buffer
            let result = *buf.output.read();

            // Output value should be correct
            assert_eq!(result, 4.2);

            // Buffer state should be unchanged
            assert_eq!(buf, old_buf);
        }
    }

    // Check that concurrent reads and writes work
    #[test]
    fn test_conc_access() {
        // We will stress the infrastructure by performing this many writes
        // as a reader continuously reads the latest value
        let test_write_count = 1000;

        // This is the buffer that our reader and writer will share
        let buf = ::TripleBuffer::new(0u64);

        // Let's split it into an input and an output
        let (mut buf_input, mut buf_output) = (buf.input, buf.output);

        // The reader continuously checks the buffered value, and should see
        // every update without any incoherent value
        let reader = ::thread::spawn(move || {
            let mut last_value = 0u64;
            while last_value < test_write_count {
                let new_value = *buf_output.read();
                assert!(new_value-last_value <= 1);
                last_value = new_value;
            }
        });

        // The writer continuously increments the buffered value, with some
        // rate limiting to ensure that the reader can see the updates
        let writer = ::thread::spawn(move || {
            ::thread::sleep(::Duration::from_millis(1));
            for value in 1 .. test_write_count+1 {
                buf_input.write(value);
                ::thread::yield_now();
                ::thread::sleep(::Duration::from_millis(1));
            }
        });

        // Wait for the reader and writer to finish
        reader.join().unwrap();
        writer.join().unwrap();
    }

    // TODO: Benchmark the performance of concurrent reads and writes

    /// Range check for triple buffer indexes
    #[allow(unused_comparisons)]
    fn index_in_range(idx: ::TripleBufferIndex) -> bool {
        (idx >= 0) & (idx <= 2)
    }
}

fn main() {
    println!("Hello, world!");

    let mut buf = TripleBuffer::new("Hello!");
    println!("Initial value is {}", buf.output.read());

    buf.input.write("World!");
    println!("New value is {}", buf.output.read())
}
<commit_msg>Minor code cleanup<commit_after>//! # Triple buffering
//!
//! In this crate, we attempt to implement a triple buffering mechanism,
//! suitable for emulating shared memory cells in a thread-safe and
//! non-blocking fashion between one single writer and one single reader.

use std::cell::UnsafeCell;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Duration;

/// ## Indexing
///
/// These types are used to index into triple buffers
/// TODO: Switch to i8 / AtomicI8 once the later is stable
type TripleBufferIndex = usize;
type AtomicTripleBufferIndex = AtomicUsize;

/// ## Shared state
///
/// In a triple buffering communication protocol, the reader and writer share
/// the following storage:
///
/// - Three memory buffers suitable for storing the data at hand
/// - One index pointing to the buffer which no one is currently using
/// - One index pointing to the most recently updated buffer
///
#[derive(Debug)]
struct TripleBufferSharedState<T: Clone + PartialEq> {
    /// Three buffers of data
    buffers: [UnsafeCell<T>; 3],

    /// Index of the back-buffer (buffer which no one is currently accessing)
    back_idx: AtomicTripleBufferIndex,

    /// Index of the most recently updated buffer
    last_idx: AtomicTripleBufferIndex,
}
//
impl<T: Clone + PartialEq> Clone for TripleBufferSharedState<T> {
    fn clone(&self) -> Self {
        // The use of UnsafeCell makes buffers somwhat cumbersome to clone...
        let clone_buffer = | i: TripleBufferIndex | {
            UnsafeCell::new(
                unsafe { (*self.buffers[i].get()).clone() }
            )
        };

        // ...so better define a shortcut before we move on to actual cloning
        TripleBufferSharedState {
            buffers: [
                clone_buffer(0),
                clone_buffer(1),
                clone_buffer(2),
            ],
            back_idx: AtomicTripleBufferIndex::new(
                self.back_idx.load(Ordering::Relaxed)
            ),
            last_idx: AtomicTripleBufferIndex::new(
                self.last_idx.load(Ordering::Relaxed)
            ),
        }
    }
}
//
impl<T: Clone + PartialEq> PartialEq for TripleBufferSharedState<T> {
    fn eq(&self, other: &Self) -> bool {
        // The use of UnsafeCell makes buffers somwhat cumbersome to compare...
        let buffer_eq = | i: TripleBufferIndex | -> bool {
            unsafe { *self.buffers[i].get() == *other.buffers[i].get() }
        };

        // ...and atomics aren't much better...
        let atomics_eq = | x: &AtomicTripleBufferIndex,
                           y: &AtomicTripleBufferIndex | -> bool {
            x.load(Ordering::Relaxed) == y.load(Ordering::Relaxed)
        };

        // ...so better define some shortcuts before getting started:
        buffer_eq(0) && buffer_eq(1) && buffer_eq(2) &&
        atomics_eq(&self.back_idx, &other.back_idx) &&
        atomics_eq(&self.last_idx, &other.last_idx)
    }
}
//
unsafe impl<T: Clone + PartialEq> Sync for TripleBufferSharedState<T> {}

/// ## Producing data ("input")
///
/// The producer of data has exclusive write access to the so-called "write
/// buffer". To send an update, it writes it into the back buffer, then swaps
/// the write buffer with the back buffer, making the data readable.
///
#[derive(Debug)]
struct TripleBufferInput<T: Clone + PartialEq> {
    shared: Arc<TripleBufferSharedState<T>>,
    write_idx: TripleBufferIndex,
}
//
impl<T: Clone + PartialEq> TripleBufferInput<T> {
    /// Write a new value into the triple buffer
    fn write(&mut self, value: T) {
        // Access the shared state
        let ref shared_state = *self.shared;

        // Determine which shared buffer is the write buffer
        let active_idx = self.write_idx;

        // Move the input value into the write buffer
        let write_ptr = shared_state.buffers[active_idx].get();
        unsafe { *write_ptr = value; }

        // Swap the back buffer and the write buffer. This makes the new data
        // block available to the reader and gives us a new buffer to write to.
        self.write_idx = shared_state.back_idx.swap(
            active_idx,
            Ordering::Release  // Send previous updates to the reader
        );

        // Notify the reader that we submitted an update
        shared_state.last_idx.store(active_idx, Ordering::Relaxed);
    }
}

/// ## Consuming data ("output")
///
/// The consumer of data has exclusive read access to the so-called "read
/// buffer". To access the latest data, it swaps the read buffer with the
/// back buffer if the later contains more recent data, then reads from
/// the (possibly new) read buffer.
///
#[derive(Debug)]
struct TripleBufferOutput<T: Clone + PartialEq> {
    shared: Arc<TripleBufferSharedState<T>>,
    read_idx: TripleBufferIndex,
}
//
impl<T: Clone + PartialEq> TripleBufferOutput<T> {
    /// Access the latest fully written value from the triple buffer
    fn read(&mut self) -> &T {
        // Access the shared state
        let ref shared_state = *self.shared;

        // Check if the writer has submitted an update
        let last_idx = shared_state.last_idx.load(Ordering::Relaxed);

        // If an update is pending in the back buffer, make it our read buffer
        if self.read_idx != last_idx {
            // Swap the back buffer and the read buffer. We get exclusive read
            // access to the data, and the writer gets a new back buffer.
            self.read_idx = shared_state.back_idx.swap(
                self.read_idx,
                Ordering::Acquire  // Synchronize with the writer's updates
            );
        }

        // Access data from the current read buffer
        let read_ptr = shared_state.buffers[self.read_idx].get();
        unsafe { &*read_ptr }
    }
}

/// ## Triple buffer
///
/// A triple buffer is created as the combination of an input and an output.
///
#[derive(Debug)]
struct TripleBuffer<T: Clone + PartialEq> {
    input: TripleBufferInput<T>,
    output: TripleBufferOutput<T>,
}
//
impl<T: Clone + PartialEq> TripleBuffer<T> {
    /// Construct a triple buffer with a certain initial value
    fn new(initial: T) -> Self {
        // Start with the shared state...
        let shared_state = Arc::new(
            TripleBufferSharedState {
                buffers: [
                    UnsafeCell::new(initial.clone()),
                    UnsafeCell::new(initial.clone()),
                    UnsafeCell::new(initial)
                ],
                back_idx: AtomicTripleBufferIndex::new(0),
                last_idx: AtomicTripleBufferIndex::new(2),
            }
        );

        // ...then construct the input and output structs
        TripleBuffer {
            input: TripleBufferInput {
                shared: shared_state.clone(),
                write_idx: 1,
            },
            output: TripleBufferOutput {
                shared: shared_state,
                read_idx: 2,
            },
        }
    }
}
//
impl<T: Clone + PartialEq> Clone for TripleBuffer<T> {
    fn clone(&self) -> Self {
        // Clone the shared state...
        let shared_state = Arc::new(
            (*self.input.shared).clone()
        );

        // ...then the input and output structs
        TripleBuffer {
            input: TripleBufferInput {
                shared: shared_state.clone(),
                write_idx: self.input.write_idx,
            },
            output: TripleBufferOutput {
                shared: shared_state,
                read_idx: self.output.read_idx,
            },
        }
    }
}
//
impl<T: Clone + PartialEq> PartialEq for TripleBuffer<T> {
    fn eq(&self, other: &Self) -> bool {
        (*self.input.shared == *other.input.shared) &&
        (self.input.write_idx == other.input.write_idx) &&
        (self.output.read_idx == other.output.read_idx)
    }
}

/// ## Testing
///
/// Unit tests and benchmarks for triple buffers
///
#[cfg(test)]
mod tests {
    /// Test that triple buffers are properly initialized
    #[test]
    fn test_init() {
        // Let's create a triple buffer
        let buf = ::TripleBuffer::new(42);

        // Access the shared state
        let ref buf_shared = *buf.input.shared;
        let back_idx = buf_shared.back_idx.load(::Ordering::Relaxed);
        let last_idx = buf_shared.last_idx.load(::Ordering::Relaxed);

        // Write-/read-/back-buffer indexes must be in range
        assert!(index_in_range(buf.input.write_idx));
        assert!(index_in_range(buf.output.read_idx));
        assert!(index_in_range(back_idx));

        // Write-/read-/back-buffer indexes must be distinct
        assert!(buf.input.write_idx != buf.output.read_idx);
        assert!(buf.input.write_idx != back_idx);
        assert!(buf.output.read_idx != back_idx);

        // Read buffer must be properly initialized
        let read_ptr = buf_shared.buffers[buf.output.read_idx].get();
        assert!(unsafe { *read_ptr } == 42);

        // Last-written index must initially point to the read buffer
        assert_eq!(last_idx, buf.output.read_idx);
    }

    /// Test that (sequentially) writing to a triple buffer works
    #[test]
    fn test_seq_write() {
        // Let's create a triple buffer
        let mut buf = ::TripleBuffer::new(false);

        // Back up the initial buffer state
        let old_buf = buf.clone();
        let ref old_shared = old_buf.input.shared;

        // Perform a write
        buf.input.write(true);

        // Check new implementation state
        {
            // Start from the old buffer state
            let mut expected_buf = old_buf.clone();
            let ref expected_shared = expected_buf.input.shared;

            // We expect the former write buffer to have received the new value
            let old_write_idx = old_buf.input.write_idx;
            let write_buf_ptr = expected_shared.buffers[old_write_idx].get();
            unsafe { *write_buf_ptr = true; }

            // We expect the former write buffer to be the new back buffer
            expected_shared.back_idx.store(
                old_write_idx,
                ::Ordering::Relaxed
            );

            // We expect the last updated index to point to former write buffer
            expected_shared.last_idx.store(
                old_write_idx,
                ::Ordering::Relaxed
            );

            // We expect the old back buffer to become the new write buffer
            let old_back_idx = old_shared.back_idx.load(::Ordering::Relaxed);
            expected_buf.input.write_idx = old_back_idx;

            // Nothing else should have changed
            assert_eq!(buf, expected_buf);
        }
    }

    /// Check that (sequentially) reading from a triple buffer works
    #[test]
    fn test_seq_read() {
        // Let's create a triple buffer and write into it
        let mut buf = ::TripleBuffer::new(1.0);
        buf.input.write(4.2);

        // Test readout from dirty (freshly written) triple buffer
        {
            // Back up the initial buffer state
            let old_buf = buf.clone();
            let ref old_shared = old_buf.input.shared;

            // Read from the buffer
            let result = *buf.output.read();

            // Output value should be correct
            assert_eq!(result, 4.2);

            // Start from the old buffer state
            let mut expected_buf = old_buf.clone();
            let ref expected_shared = expected_buf.input.shared;

            // We expect the new back index to point to the former read buffer
            expected_shared.back_idx.store(
                old_buf.output.read_idx,
                ::Ordering::Relaxed
            );

            // We expect the new read index to point to the former back buffer
            let old_back_idx = old_shared.back_idx.load(::Ordering::Relaxed);
            expected_buf.output.read_idx = old_back_idx;

            // Nothing else should have changed
            assert_eq!(buf, expected_buf);
        }

        // Test readout from clean (unchanged) triple buffer
        {
            // Back up the initial buffer state
            let old_buf = buf.clone();

            // Read from the buffer
            let result = *buf.output.read();

            // Output value should be correct
            assert_eq!(result, 4.2);

            // Buffer state should be unchanged
            assert_eq!(buf, old_buf);
        }
    }

    // Check that concurrent reads and writes work
    #[test]
    fn test_conc_access() {
        // We will stress the infrastructure by performing this many writes
        // as a reader continuously reads the latest value
        let test_write_count = 1000;

        // This is the buffer that our reader and writer will share
        let buf = ::TripleBuffer::new(0u64);

        // Let's split it into an input and an output
        let (mut buf_input, mut buf_output) = (buf.input, buf.output);

        // The reader continuously checks the buffered value, and should see
        // every update without any incoherent value
        let reader = ::thread::spawn(move || {
            let mut last_value = 0u64;
            while last_value < test_write_count {
                let new_value = *buf_output.read();
                assert!(
                    (new_value >= last_value) && (new_value-last_value <= 1)
                );
                last_value = new_value;
            }
        });

        // The writer continuously increments the buffered value, with some
        // rate limiting to ensure that the reader can see the updates
        let writer = ::thread::spawn(move || {
            // TODO: Use a barrier here instead
            ::thread::sleep(::Duration::from_millis(1));
            for value in 1 .. test_write_count+1 {
                buf_input.write(value);
                ::thread::yield_now();
                ::thread::sleep(::Duration::from_millis(1));
            }
        });

        // Wait for the reader and writer to finish
        reader.join().unwrap();
        writer.join().unwrap();
    }

    // TODO: Benchmark the performance of concurrent reads and writes

    /// Range check for triple buffer indexes
    #[allow(unused_comparisons)]
    fn index_in_range(idx: ::TripleBufferIndex) -> bool {
        (idx >= 0) & (idx <= 2)
    }
}

fn main() {
    println!("Hello, world!");

    let mut buf = TripleBuffer::new("Hello!");
    println!("Initial value is {}", buf.output.read());

    buf.input.write("World!");
    println!("New value is {}", buf.output.read())
}
<|endoftext|>


<commit_before>#![feature(slice_patterns)]
extern crate num;
#[macro_use]
extern crate glium;
extern crate nalgebra;
use nalgebra::*;
use glium::{DisplayBuild, Surface};
use num::traits::Zero;
use glium::glutin::Event::*;
use glium::glutin::ElementState::*;
use glium::glutin::MouseScrollDelta::PixelDelta;
use glium::glutin::MouseButton::*;

fn main() {
  let display = glium::glutin::WindowBuilder::new()
                  .with_vsync()
                  .build_glium().unwrap();

  #[derive(Copy, Clone)]
  struct Vertex {position: [f32; 2]}
  implement_vertex!(Vertex, position);

  let vertex_shader_src = r#"
      #version 140
      uniform mat2 projection;
      uniform vec2 translation;
      in vec2 position;
      void main() {
          gl_Position = vec4(projection * (position + translation), 0.0, 1.0);
      }
  "#;

  let fragment_shader_src = r#"
      #version 140

      out vec4 color;

      void main() {
          color = vec4(0.0, 0.0, 0.0, 1.0);
      }
  "#;

  let indices = glium::index::NoIndices(glium::index::PrimitiveType::LineStrip);

  let params = glium::DrawParameters {
    line_width: Some(2.0),
    .. Default::default()
  };

  let program = glium::Program::from_source(
    &display,
    vertex_shader_src,
    fragment_shader_src,
    None).unwrap();

  let mut current_line : Vec<Vertex> = Vec::new();
  let mut commited_lines : Vec<glium::VertexBuffer<Vertex>> = Vec::new();

  #[derive(Copy, Clone)]
  enum State {
    None,
    Drawing,
    Dragging(Vec2<f32>)
  }

  let mut state = State::None;

  let mut scale = 1.0;
  let mut translation = Vec2::zero();
  let mut mouse_last = Vec2::zero();

  loop {

    let (frame_dx, frame_dy) = display.get_framebuffer_dimensions();
    let (frame_dx, frame_dy) = (frame_dx as f32, frame_dy as f32);

    let projection
      = Mat2::new(2.0/frame_dx, 0.0,
                  0.          , 2.0/frame_dy);

    let uniforms = uniform!(
      projection : projection * scale,
      translation: translation,
    );

    let mut target = display.draw();
    target.clear_color(1.0, 1.0, 1.0, 1.0);

    for line in &commited_lines[..] {
      target.draw(
        line,
        &indices,
        &program,
        &uniforms,
        &params).unwrap();
    }

    if let State::Drawing = state {
        target.draw(
          &glium::VertexBuffer::new(&display, &current_line).unwrap(),
          &indices,
          &program,
          &uniforms,
          &Default::default()).unwrap();
    }

    target.finish().unwrap();

    for event in display.poll_events() {
      match (event, state) {

        (MouseWheel(PixelDelta(_, dy)), _)  => {
          if dy > 0.0 {
            scale *= 1.1;
          } else {
            scale *= 0.9;
          }
        },

        (MouseInput(Pressed,Left), State::None) => {
          state = State::Drawing;
        },

        (MouseInput(Pressed,Middle), State::None) => {
            state = State::Dragging(mouse_last);
        },

        (MouseInput(Released,Left), State::Drawing) => {
            commited_lines.push(
              glium::VertexBuffer::new(&display, &current_line).unwrap());
            current_line.clear();
            state = State::None;
        },

        (MouseInput(Released,Middle), State::Dragging(_)) => {
            state = State::None;
        },

        (MouseMoved((mouse_x, mouse_y)), state) => {
          let screen
            = Vec2::new(
                2.0*(mouse_x as f32)/frame_dx - 1.0,
               -2.0*(mouse_y as f32)/frame_dy + 1.0);

          let world  = projection.inv().unwrap()
                        * (1.0/scale)
                        * screen
                        - translation;

          match state {
            State::Drawing => {
              current_line.push(Vertex{position: [world.x, world.y]});
            },
            State::Dragging(drag_start) => {
              translation.x +=  world.x - drag_start.x;
              translation.y +=  world.y - drag_start.y;
            },
            _ => {}
          }

          mouse_last = world;
        },

        (Closed, _) => return,
        _ => ()
      }
    }
  }
}
<commit_msg>minor changes<commit_after>#![feature(slice_patterns)]
extern crate num;
#[macro_use]
extern crate glium;
extern crate nalgebra;
use nalgebra::*;
use glium::{DisplayBuild, Surface};
use num::traits::Zero;
use glium::glutin::Event::*;
use glium::glutin::ElementState::*;
use glium::glutin::MouseScrollDelta::PixelDelta;
use glium::glutin::MouseButton::*;

fn main() {
  let display = glium::glutin::WindowBuilder::new()
                  .with_vsync()
                  .build_glium().unwrap();

  #[derive(Copy, Clone)]
  struct Vertex {position: [f32; 2]}
  implement_vertex!(Vertex, position);

  let vertex_shader_src = r#"
      #version 140
      uniform mat2 projection;
      uniform vec2 translation;
      in vec2 position;
      void main() {
          gl_Position = vec4(projection * (position + translation), 0.0, 1.0);
      }
  "#;

  let fragment_shader_src = r#"
      #version 140

      out vec4 color;

      void main() {
          color = vec4(0.0, 0.0, 0.0, 1.0);
      }
  "#;

  let indices = glium::index::NoIndices(glium::index::PrimitiveType::LineStrip);

  let params = glium::DrawParameters {
    line_width: Some(2.0),
    .. Default::default()
  };

  let program = glium::Program::from_source(
    &display,
    vertex_shader_src,
    fragment_shader_src,
    None).unwrap();

  let mut current_line : Vec<Vertex> = Vec::new();
  let mut commited_lines : Vec<glium::VertexBuffer<Vertex>> = Vec::new();

  #[derive(Copy, Clone)]
  enum State {
    None,
    Drawing,
    Dragging(Vec2<f32>)
  }

  let mut state = State::None;

  let mut scale = 1.0;
  let mut translation = Vec2::zero();
  let mut mouse_last = Vec2::zero();

  loop {

    let (frame_dx, frame_dy) = display.get_framebuffer_dimensions();
    let (frame_dx, frame_dy) = (frame_dx as f32, frame_dy as f32);

    let projection
      = Mat2::new(2.0/frame_dx, 0.0,
                  0.          , 2.0/frame_dy);

    for event in display.poll_events() {
      match (event, state) {

        (MouseWheel(PixelDelta(_, dy)), _)  => {
          if dy > 0.0 {
            scale *= 1.1;
          } else {
            scale *= 0.9;
          }
        },

        (MouseInput(Pressed,Left), State::None) => {
          state = State::Drawing;
        },

        (MouseInput(Pressed,Middle), State::None) => {
            state = State::Dragging(mouse_last);
        },

        (MouseInput(Released,Left), State::Drawing) => {
            commited_lines.push(
              glium::VertexBuffer::immutable(&display, &current_line).unwrap());
            current_line.clear();
            state = State::None;
        },

        (MouseInput(Released,Middle), State::Dragging(_)) => {
            state = State::None;
        },

        (MouseMoved((mouse_x, mouse_y)), state) => {
          let screen
            = Vec2::new(
                2.0*(mouse_x as f32)/frame_dx - 1.0,
               -2.0*(mouse_y as f32)/frame_dy + 1.0);

          let world  = projection.inv().unwrap()
                        * (1.0/scale)
                        * screen
                        - translation;

          match state {
            State::Drawing => {
              current_line.push(Vertex{position: [world.x, world.y]});
            },
            State::Dragging(drag_start) => {
              translation.x +=  world.x - drag_start.x;
              translation.y +=  world.y - drag_start.y;
            },
            _ => {}
          }

          mouse_last = world;
        },

        (Closed, _) => return,
        _ => ()
      }
    }

    let uniforms = uniform!(
      projection : projection * scale,
      translation: translation,
    );

    let mut target = display.draw();
    target.clear_color(1.0, 1.0, 1.0, 1.0);

    for line in &commited_lines[..] {
      target.draw(
        line,
        &indices,
        &program,
        &uniforms,
        &params).unwrap();
    }

    if let State::Drawing = state {
        target.draw(
          &glium::VertexBuffer::new(&display, &current_line).unwrap(),
          &indices,
          &program,
          &uniforms,
          &Default::default()).unwrap();
    }

    target.finish().unwrap();
  }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate nom;

use nom::{IResult, space};

named!(statement, take_until_either_and_consume!("."));
named!(percent<(char, char)>, do_parse!(
    tag!("%") >>
    a: hex    >>
    b: hex    >>
    (a, b)
));
named!(hex<char>, one_of!("0123456789ABCDEFabcdef"));
named!(pn_local_esc, escaped!(
    call!(space),
    '\\',
    one_of!("_~.-!$&\'()*+,;=/?#@%")
));

fn main() {
    assert_eq!(statement(&b"a b c. d e f."[..]), IResult::Done(&b" d e f."[..], &b"a b c"[..]));
    assert_eq!(percent(&b"%2Abc"[..]),           IResult::Done(&b"bc"[..]     , ('2', 'A')   ));
    assert_eq!(hex(&b"Abc"[..])                , IResult::Done(&b"bc"[..]     , 'A'          ));
    assert_eq!(pn_local_esc(&b"\\.a b c"[..])  , IResult::Done(&b"a b c"[..]  , &b"\\."[..]  ));
}
<commit_msg>Implement PN_CHARS_BASE<commit_after>#[macro_use]
extern crate nom;

use nom::{IResult, space, hex_digit};

named!(statement, take_until_either_and_consume!("."));

fn is_pn_chars_base(chr: &u8) -> bool {
    (*chr >=    0x41 && *chr <=    0x5A) ||
    (*chr >=    0x61 && *chr <=    0x7A) ||
    (*chr >=    0xC0 && *chr <=    0xD6) ||
    (*chr >=    0xD8 && *chr <=    0xF6) ||
    (*chr >=    0xF8 && *chr <=   0x2FF) ||
    (*chr >=   0x370 && *chr <=   0x37D) ||
    (*chr >=   0x37F && *chr <=  0x1FFF) ||
    (*chr >=  0x200C && *chr <=  0x200D) ||
    (*chr >=  0x2C00 && *chr <=  0x2FEF) ||
    (*chr >=  0x3001 && *chr <=  0xD7FF) ||
    (*chr >=  0xF900 && *chr <=  0xFDCF) ||
    (*chr >=  0xFDF0 && *chr <=  0xFFFD) ||
    (*chr >= 0x10000 && *chr <= 0xEFFFF)
}
fn _pn_chars_base(input:&[u8]) -> IResult<&[u8],&[u8]> {
    for c in input {
        if ! is_pn_chars_base(c) {
            return IResult::Error(nom::ErrorKind::Custom(20))
        }
    }
    IResult::Done(&input[1..], &input[0..1])
}
/* [163s] PN_CHARS_BASE */
named!(pn_chars_base, flat_map!(
    take!(1),
    _pn_chars_base
));

/* [169s] PLX */
named!(plx, alt!(percent | pn_local_esc));

/* [170s] PERCENT */
named!(percent, do_parse!(
    tag!("%") >>
    a: flat_map!(
        take!(2),
        hex_digit
    ) >>
    (a)
));

/* [172s]  PN_LOCAL_ESC */
named!(pn_local_esc, escaped!(
    call!(space),
    '\\',
    one_of!("_~.-!$&'()*+,;=/?#@%")
));

fn main() {
    assert_eq!(statement(&b"a b c. d e f."[..]), IResult::Done(&b" d e f."[..], &b"a b c"[..]));
    assert_eq!(pn_chars_base(&b"a%2Ab"[..]),     IResult::Done(&b"%2Ab"[..]   , &b"a"[..]    ));
    assert_eq!(plx(&b"%2Abc"[..]),               IResult::Done(&b"bc"[..]     , &b"2A"[..]   ));
    assert_eq!(plx(&b"\\.%2A"[..]),              IResult::Done(&b"%2A"[..]    , &b"\\."[..]  ));
    assert_eq!(percent(&b"%2Abc"[..]),           IResult::Done(&b"bc"[..]     , &b"2A"[..]   ));
    assert_eq!(pn_local_esc(&b"\\.a b c"[..])  , IResult::Done(&b"a b c"[..]  , &b"\\."[..]  ));
}
<|endoftext|>


<commit_before>extern crate rustc_serialize;
extern crate docopt;
use docopt::Docopt;

#[macro_use]
extern crate log;
use log::{LogRecord, LogLevel, LogLevelFilter, LogMetadata};

use std::collections::{HashSet, HashMap};
use std::fs::File;
use std::io::{BufRead, BufReader, Seek, SeekFrom};
use std::path::Path;

struct SimpleLogger;

impl log::Log for SimpleLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= LogLevel::Info
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{} - {}", record.level(), record.args());
        }
    }
}

struct CsvDesc<'a> {
    file_path: &'a  Path,
    delimiter:      char,
    quote:          Option<char>,
}

impl<'a> std::fmt::Display for CsvDesc<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{} {} {:?}", self.file_path.display(), self.delimiter, self.quote)
    }
}

fn parse_args<'a>(path_arg:         &'a String,
                  delimiter_arg:    &'a String,
                  quote_arg:        &'a String) -> Result<CsvDesc<'a>, &'static str>
{

    let csv_file_path = Path::new(path_arg);

    let csv_delimiter = match delimiter_arg.chars().next() {
                            Some(result) => result,
                            None         => return Err("incorrect delimiter"),
                        };

    let csv_quote     = quote_arg.chars().next();

    Ok(CsvDesc {file_path: &csv_file_path, delimiter: csv_delimiter, quote: csv_quote})
}

fn get_csv_cols(csv_desc: &CsvDesc) -> Result<Vec<String>, String> {

    let csv_file = match File::open(csv_desc.file_path) {
        Err(why) => panic!("couldn't open csv @ {}: {}", csv_desc.file_path.display(), why),
        Ok(file) => file,
    };

    let csv_reader = BufReader::new(csv_file);

    let mut csv_line_iter = csv_reader.lines();

    let csv_header: String = match csv_line_iter.next() {
        Some(result) => match result {
                            Err(why) => return Err(format!("error getting csv header: {}", why)),
                            Ok(header) => header,
                        },
        None         => return Err("csv header reading failed".to_string()),
    };

    let csv_cols: Vec<String> = {
        let cols_iter = csv_header.split(csv_desc.delimiter);
        match csv_desc.quote {
            Some(q) => cols_iter.map(|s| {s.trim_matches(q).to_string()}).collect(),
            None    => cols_iter.map(|s| {s.to_string()}).collect(),
        }
    };

    Ok(csv_cols)
}

fn build_index(csv_desc: &CsvDesc) -> Result<HashMap<String, u64>, String> {
    // TODO it would probably be better to keep a File in the csv descriptor instead of a Path
    let mut csv_index = HashMap::new();
    let csv_file = match File::open(csv_desc.file_path) {
        Err(why) => panic!("couldn't open csv @ {}: {}", csv_desc.file_path.display(), why),
        Ok(file) => file,
    };

    let csv_reader = BufReader::new(csv_file);
    let mut csv_line_iter = csv_reader.lines();

    let mut offset_in_file: u64 = 0;
    let mut prev_col_num = -1;
    let mut row_idx = 0;
    loop {

        let csv_row: String = match csv_line_iter.next() {
            Some(result) => match result {
                                Err(why) => return Err(format!("error getting csv row: {}", why)),
                                Ok(header) => header,
                            },
            None         => break,
        };

        let csv_cols: Vec<String> = {
            let cols_iter = csv_row.split(csv_desc.delimiter);
            match csv_desc.quote {
                Some(q) => cols_iter.map(|s| {s.trim_matches(q).to_string()}).collect(),
                None    => cols_iter.map(|s| {s.to_string()}).collect(),
            }
        };

        let curr_col_num = csv_cols.len();
        if prev_col_num != -1 && prev_col_num != curr_col_num {
            return Err(format!("{} columns in row #{}, {} columns in previous row", curr_col_num,
                                                                                row_idx,
                                                                                prev_col_num));
        }
        prev_col_num = curr_col_num;
        row_idx += 1;

        let key = format!("{}{}", csv_cols[0], csv_cols[1]);

        csv_index.insert(key, offset_in_file);
        offset_in_file += (csv_row.len() + 1) as u64;
    }

    Ok(csv_index)
}

fn get_csv_row(csv_desc: &CsvDesc, line_offset: u64) -> Result<Vec<String>, String> {

    let mut csv_file = match File::open(csv_desc.file_path) {
        Err(why) => panic!("couldn't open csv @ {}: {}", csv_desc.file_path.display(), why),
        Ok(file) => file,
    };

    csv_file.seek(SeekFrom::Start(line_offset)).unwrap(); // TODO: handle me

    let mut csv_reader = BufReader::new(csv_file);
    let mut row_buff = String::new();

    match csv_reader.read_line(&mut row_buff) {
        Ok(_n)  => {
            if row_buff.ends_with("\n") {
                row_buff.pop();
            }
        }
        Err(e) => return Err(format!("error gettig csv row: {}", e)),
    };

    let result: Vec<String> = {
        let cols_iter = row_buff.split(csv_desc.delimiter);
        match csv_desc.quote {
            Some(q) => cols_iter.map(|s| {s.trim_matches(q).to_string()}).collect(),
            None    => cols_iter.map(|s| {s.to_string()}).collect(),
        }
    };

    Ok(result)
}

static USAGE: &'static str = "
Usage: rusty-csv-diff <csv1> <delim1> <quote1> <csv2> <delim2> <quote2>
       rusty-csv-diff --help

Options:
    -h, --help        Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_csv1: String,
    arg_delim1: String,
    arg_quote1: String,
    arg_csv2: String,
    arg_delim2: String,
    arg_quote2: String,
}

fn main() {

/*

1. Parse arguments
2. Open CSV files
3. Get columns (cols_N)
4. Get intersection of those two sets of columns(cols_to_compare)
5. Create {column name : column index in cols_N} dicts
6. Create {CSV_col_value : CSV row index in file} dicts, where CSV_col_value is a unique key made of the value of several CSV columns.
   For example, {Alex38 : 76}. In this example a name and age form a unique key for the 76th CSV row.
7. Get intersection of key sets of dicts from step 6 (row_keys_to_compare)
8. Loop through row_keys_to_compare, use dicts from step 6 to get line numbers for CSV files
    8.1 Loop through cols_to_compare, use dicts from step 5 to extract column values from CSV rows(step 8)
    8.2 Compare values

Input parameters: CSV paths, delimiters, quotes

For example, ./main file_1.csv "," "'" file_2.csv " " ""

*/

    /*** 0 ***/
    log::set_logger(|max_log_level| {
        max_log_level.set(LogLevelFilter::Error);
        Box::new(SimpleLogger)
    }).unwrap();

    /*** 1 ***/
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    let csv_desc_1: CsvDesc = match parse_args(&args.arg_csv1, &args.arg_delim1, &args.arg_quote1) {
        Err(why)    => panic!("error parsing arguments for CSV #1: {}", why),
        Ok(result)  => result,
    };

    let csv_desc_2: CsvDesc = match parse_args(&args.arg_csv2, &args.arg_delim2, &args.arg_quote2) {
        Err(why)    => panic!("error parsing arguments for CSV #2: {}", why),
        Ok(result)  => result,
    };

    /*** 2&3 ***/
    let csv_cols_1: Vec<String> = match get_csv_cols(&csv_desc_1) {
        Err(why) => panic!("couldn't get columns: {}", why),
        Ok(cols) => cols,
    };

    let csv_cols_2: Vec<String> = match get_csv_cols(&csv_desc_2) {
        Err(why) => panic!("couldn't get columns: {}", why),
        Ok(cols) => cols,
    };

    /*** 5 ***/
    let mut csv_col_index_1 = HashMap::new();
    for i in 0..csv_cols_1.len() {
        let key = csv_cols_1[i].clone();
        if csv_col_index_1.contains_key(&key) {
            panic!("duplicate column found in CSV #1: {}", key);
        };
        csv_col_index_1.insert(key, i);
    }
    info!("{:?}", csv_col_index_1);

    let mut csv_col_index_2 = HashMap::new();
    for i in 0..csv_cols_2.len() {
        let key = csv_cols_2[i].clone();
        if csv_col_index_2.contains_key(&key) {
            panic!("duplicate column found in CSV #1: {}", key);
        };
        csv_col_index_2.insert(key, i);
    }
    info!("{:?}", csv_col_index_2);

    /*** 4 ***/
    let mut cols_to_compare = HashSet::new();
    for col_1 in csv_col_index_1.keys() {
        if csv_col_index_2.contains_key(col_1) {
            cols_to_compare.insert(col_1);
        };
    }
    info!("{:?}", cols_to_compare);

    /*** 6 ***/
    // let's assume that the unique key is (col_0 + col_1)
    let csv_index_1 = match build_index(&csv_desc_1) {
        Err(why)  => panic!("failed building index #1: {}", why),
        Ok(index) => index,
    };

    let csv_index_2 = match build_index(&csv_desc_2) {
        Err(why)  => panic!("failed building index #2: {}", why),
        Ok(index) => index,
    };

    /*** 7 ***/
    let mut row_keys_to_compare = HashSet::new();
    for key_1 in csv_index_1.keys() {
        if csv_index_2.contains_key(key_1) {
            row_keys_to_compare.insert(key_1);
        };
    }
    info!("{:?}", row_keys_to_compare);

    /*** 8 ***/
    for row_key in row_keys_to_compare {

        let index_1 = *csv_index_1.get(row_key).unwrap(); // TODO: handle me
        let index_2 = *csv_index_2.get(row_key).unwrap();

        let row_1 = get_csv_row(&csv_desc_1, index_1).unwrap(); // TODO: handle me
        let row_2 = get_csv_row(&csv_desc_2, index_2).unwrap();

        info!("comparing {}:", row_key);
        info!("line #1: {:?}", row_1);
        info!("line #2: {:?}", row_2);

        for col in &cols_to_compare {

            let col_index_1 = *csv_col_index_1.get(*col).unwrap(); // TODO: handle me
            let col_index_2 = *csv_col_index_2.get(*col).unwrap();

            info!("column {}, index_1={}, index_2={}", col, col_index_1, col_index_2);

            if row_1[col_index_1] != row_2[col_index_2] {
                println!("found a difference for {}, {}: {} / {}", row_key, col, row_1[col_index_1], row_2[col_index_2]);
            }
        }

    }

}
<commit_msg>Error handling TODOs<commit_after>extern crate rustc_serialize;
extern crate docopt;
use docopt::Docopt;

#[macro_use]
extern crate log;
use log::{LogRecord, LogLevel, LogLevelFilter, LogMetadata};

use std::collections::{HashSet, HashMap};
use std::fs::File;
use std::io::{BufRead, BufReader, Seek, SeekFrom};
use std::path::Path;

struct SimpleLogger;

impl log::Log for SimpleLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= LogLevel::Info
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{} - {}", record.level(), record.args());
        }
    }
}

struct CsvDesc<'a> {
    file_path: &'a  Path,
    delimiter:      char,
    quote:          Option<char>,
}

impl<'a> std::fmt::Display for CsvDesc<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{} {} {:?}", self.file_path.display(), self.delimiter, self.quote)
    }
}

fn parse_args<'a>(path_arg:         &'a String,
                  delimiter_arg:    &'a String,
                  quote_arg:        &'a String) -> Result<CsvDesc<'a>, &'static str>
{

    let csv_file_path = Path::new(path_arg);

    let csv_delimiter = match delimiter_arg.chars().next() {
                            Some(result) => result,
                            None         => return Err("incorrect delimiter"),
                        };

    let csv_quote     = quote_arg.chars().next();

    Ok(CsvDesc {file_path: &csv_file_path, delimiter: csv_delimiter, quote: csv_quote})
}

fn get_csv_cols(csv_desc: &CsvDesc) -> Result<Vec<String>, String> {

    let csv_file = match File::open(csv_desc.file_path) {
        Err(why) => panic!("couldn't open csv @ {}: {}", csv_desc.file_path.display(), why),
        Ok(file) => file,
    };

    let csv_reader = BufReader::new(csv_file);

    let mut csv_line_iter = csv_reader.lines();

    let csv_header: String = match csv_line_iter.next() {
        Some(result) => match result {
                            Err(why) => return Err(format!("error getting csv header: {}", why)),
                            Ok(header) => header,
                        },
        None         => return Err("csv header reading failed".to_string()),
    };

    let csv_cols: Vec<String> = {
        let cols_iter = csv_header.split(csv_desc.delimiter);
        match csv_desc.quote {
            Some(q) => cols_iter.map(|s| {s.trim_matches(q).to_string()}).collect(),
            None    => cols_iter.map(|s| {s.to_string()}).collect(),
        }
    };

    Ok(csv_cols)
}

fn build_index(csv_desc: &CsvDesc) -> Result<HashMap<String, u64>, String> {
    // TODO it would probably be better to keep a File in the csv descriptor instead of a Path
    let mut csv_index = HashMap::new();
    let csv_file = match File::open(csv_desc.file_path) {
        Err(why) => panic!("couldn't open csv @ {}: {}", csv_desc.file_path.display(), why),
        Ok(file) => file,
    };

    let csv_reader = BufReader::new(csv_file);
    let mut csv_line_iter = csv_reader.lines();

    let mut offset_in_file: u64 = 0;
    let mut prev_col_num = -1;
    let mut row_idx = 0;
    loop {

        let csv_row: String = match csv_line_iter.next() {
            Some(result) => match result {
                                Err(why) => return Err(format!("error getting csv row: {}", why)),
                                Ok(header) => header,
                            },
            None         => break,
        };

        let csv_cols: Vec<String> = {
            let cols_iter = csv_row.split(csv_desc.delimiter);
            match csv_desc.quote {
                Some(q) => cols_iter.map(|s| {s.trim_matches(q).to_string()}).collect(),
                None    => cols_iter.map(|s| {s.to_string()}).collect(),
            }
        };

        let curr_col_num = csv_cols.len();
        if prev_col_num != -1 && prev_col_num != curr_col_num {
            return Err(format!("{} columns in row #{}, {} columns in previous row", curr_col_num,
                                                                                row_idx,
                                                                                prev_col_num));
        }
        prev_col_num = curr_col_num;
        row_idx += 1;

        let key = format!("{}{}", csv_cols[0], csv_cols[1]);

        csv_index.insert(key, offset_in_file);
        offset_in_file += (csv_row.len() + 1) as u64;
    }

    Ok(csv_index)
}

fn get_csv_row(csv_desc: &CsvDesc, line_offset: u64) -> Result<Vec<String>, String> {

    let mut csv_file = match File::open(csv_desc.file_path) {
        Err(why) => panic!("couldn't open csv @ {}: {}", csv_desc.file_path.display(), why),
        Ok(file) => file,
    };

    if let Err(e) = csv_file.seek(SeekFrom::Start(line_offset)) {
        return Err(format!("error seeking offset: {}", e));
    }

    let mut csv_reader = BufReader::new(csv_file);
    let mut row_buff = String::new();

    match csv_reader.read_line(&mut row_buff) {
        Ok(_n)  => {
            if row_buff.ends_with("\n") {
                row_buff.pop();
            }
        }
        Err(e) => return Err(format!("error gettig csv row: {}", e)),
    };

    let result: Vec<String> = {
        let cols_iter = row_buff.split(csv_desc.delimiter);
        match csv_desc.quote {
            Some(q) => cols_iter.map(|s| {s.trim_matches(q).to_string()}).collect(),
            None    => cols_iter.map(|s| {s.to_string()}).collect(),
        }
    };

    Ok(result)
}

static USAGE: &'static str = "
Usage: rusty-csv-diff <csv1> <delim1> <quote1> <csv2> <delim2> <quote2>
       rusty-csv-diff --help

Options:
    -h, --help        Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_csv1: String,
    arg_delim1: String,
    arg_quote1: String,
    arg_csv2: String,
    arg_delim2: String,
    arg_quote2: String,
}

fn main() {

/*

1. Parse arguments
2. Open CSV files
3. Get columns (cols_N)
4. Get intersection of those two sets of columns(cols_to_compare)
5. Create {column name : column index in cols_N} dicts
6. Create {CSV_col_value : CSV row index in file} dicts, where CSV_col_value is a unique key made of the value of several CSV columns.
   For example, {Alex38 : 76}. In this example a name and age form a unique key for the 76th CSV row.
7. Get intersection of key sets of dicts from step 6 (row_keys_to_compare)
8. Loop through row_keys_to_compare, use dicts from step 6 to get line numbers for CSV files
    8.1 Loop through cols_to_compare, use dicts from step 5 to extract column values from CSV rows(step 8)
    8.2 Compare values

Input parameters: CSV paths, delimiters, quotes

For example, ./main file_1.csv "," "'" file_2.csv " " ""

*/

    /*** 0 ***/
    log::set_logger(|max_log_level| {
        max_log_level.set(LogLevelFilter::Error);
        Box::new(SimpleLogger)
    }).unwrap();

    /*** 1 ***/
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    let csv_desc_1: CsvDesc = match parse_args(&args.arg_csv1, &args.arg_delim1, &args.arg_quote1) {
        Err(why)    => panic!("error parsing arguments for CSV #1: {}", why),
        Ok(result)  => result,
    };

    let csv_desc_2: CsvDesc = match parse_args(&args.arg_csv2, &args.arg_delim2, &args.arg_quote2) {
        Err(why)    => panic!("error parsing arguments for CSV #2: {}", why),
        Ok(result)  => result,
    };

    /*** 2&3 ***/
    let csv_cols_1: Vec<String> = match get_csv_cols(&csv_desc_1) {
        Err(why) => panic!("couldn't get columns: {}", why),
        Ok(cols) => cols,
    };

    let csv_cols_2: Vec<String> = match get_csv_cols(&csv_desc_2) {
        Err(why) => panic!("couldn't get columns: {}", why),
        Ok(cols) => cols,
    };

    /*** 5 ***/
    let mut csv_col_index_1 = HashMap::new();
    for i in 0..csv_cols_1.len() {
        let key = csv_cols_1[i].clone();
        if csv_col_index_1.contains_key(&key) {
            panic!("duplicate column found in CSV #1: {}", key);
        };
        csv_col_index_1.insert(key, i);
    }
    info!("{:?}", csv_col_index_1);

    let mut csv_col_index_2 = HashMap::new();
    for i in 0..csv_cols_2.len() {
        let key = csv_cols_2[i].clone();
        if csv_col_index_2.contains_key(&key) {
            panic!("duplicate column found in CSV #1: {}", key);
        };
        csv_col_index_2.insert(key, i);
    }
    info!("{:?}", csv_col_index_2);

    /*** 4 ***/
    let mut cols_to_compare = HashSet::new();
    for col_1 in csv_col_index_1.keys() {
        if csv_col_index_2.contains_key(col_1) {
            cols_to_compare.insert(col_1);
        };
    }
    info!("{:?}", cols_to_compare);

    /*** 6 ***/
    // let's assume that the unique key is (col_0 + col_1)
    let csv_index_1 = match build_index(&csv_desc_1) {
        Err(why)  => panic!("failed building index #1: {}", why),
        Ok(index) => index,
    };

    let csv_index_2 = match build_index(&csv_desc_2) {
        Err(why)  => panic!("failed building index #2: {}", why),
        Ok(index) => index,
    };

    /*** 7 ***/
    let mut row_keys_to_compare = HashSet::new();
    for key_1 in csv_index_1.keys() {
        if csv_index_2.contains_key(key_1) {
            row_keys_to_compare.insert(key_1);
        };
    }
    info!("{:?}", row_keys_to_compare);

    /*** 8 ***/
    for row_key in row_keys_to_compare {

        let index_1 = *csv_index_1.get(row_key).unwrap();
        let index_2 = *csv_index_2.get(row_key).unwrap();

        let row_1 = match get_csv_row(&csv_desc_1, index_1) {
            Ok(row) => row,
            Err(e)  => panic!("failed getting csv row #1: {}", e),
        };

        let row_2 = match get_csv_row(&csv_desc_2, index_2) {
            Ok(row) => row,
            Err(e)  => panic!("failed getting csv row #2: {}", e),
        };

        info!("comparing {}:", row_key);
        info!("line #1: {:?}", row_1);
        info!("line #2: {:?}", row_2);

        for col in &cols_to_compare {

            let col_index_1 = *csv_col_index_1.get(*col).unwrap();
            let col_index_2 = *csv_col_index_2.get(*col).unwrap();

            info!("column {}, index_1={}, index_2={}", col, col_index_1, col_index_2);

            if row_1[col_index_1] != row_2[col_index_2] {
                println!("found a difference for {}, {}: {} / {}", row_key, col, row_1[col_index_1], row_2[col_index_2]);
            }
        }

    }

}
<|endoftext|>


<commit_before>extern crate ansi_term;
extern crate argparse;
extern crate hyper;
extern crate libsnatch;

use ansi_term::Colour::{Green, Yellow, Red};
use argparse::{ArgumentParser, Store, StoreTrue};
use hyper::client::Client;
use libsnatch::Chunks;
use libsnatch::client::GetResponse;
use libsnatch::contentlength::GetContentLength;
use libsnatch::download::download_chunks;
use libsnatch::http_version::ValidateHttpVersion;
use libsnatch::write::write_file;
use std::fs::File;
use std::io;
use std::sync::{Arc, Mutex};
use std::path::Path;
use std::process::exit;

fn main() {

    let mut file = String::from("");
    let mut threads : usize = 4;
    let mut url = String::from("");
    let mut verbose = false;

    {
        let mut argparse = ArgumentParser::new();
        argparse.set_description("Snatch, a simple, fast and interruptable download accelerator, written in Rust.");
        argparse.refer(&mut file)
                .add_option(&["-f", "--file"], Store,
                "The local file to save the remote content file")
                .required();
        argparse.refer(&mut threads)
                .add_option(&["-t", "--threads"], Store,
                "Number of threads available to download");
        argparse.refer(&mut url)
                .add_option(&["-u", "--url"], Store,
                "Remote content URL to download")
                .required();
        argparse.refer(&mut verbose)
                .add_option(&["-v", "--verbose"], StoreTrue,
                "Verbose mode");
        argparse.parse_args_or_exit();
    }

    let hyper_client = Client::new();

    // Get the first response from the server
    let client_response = hyper_client.get_http_response(&url).unwrap();

    print!("# Waiting a response from the remote server... ");

    if !client_response.version.greater_than_http_11() {
        println!("{}",
                 Yellow.bold()
                     .paint("[WARNING] The version of HTTP requests is <= HTTP1.0 - it can be \
                             difficult to split efficiently the remote content!"));
    } else {
        println!("{}", Green.bold().paint("OK !"));
    }

    let local_path = Path::new(&file);

    if local_path.exists() {
        if local_path.is_dir() {
            panic!(Red.bold().paint("[ERROR] The local path to store the remote content is already exists, and is a directory!"));
        }
        println!("{}", Red.bold().paint("[WARNING] The path to store the file already exists! Do you want to override it? [y/N]"));
        {
            let mut user_input = String::new();
            io::stdin().read_line(&mut user_input)
                .ok()
                .expect("[ERROR] Couldn't read line!");
            user_input = String::from(user_input.trim());
            if ! (user_input == "y" || user_input == "Y") {
                exit(0);
            }
        }
    }

    let remote_content_length = match client_response.headers.get_content_length(){
        Some(remote_content_length) => remote_content_length,
        None => {
            println!("{}", Red.bold().paint("[ERROR] Canno't get the content length of the remote content, from the server."));
            exit(1);
        }
    };

    println!("# Remote content length: {:?} bytes", remote_content_length);

    let mut core_chunks = Chunks::with_capacity(threads);

    for _ in 0..threads {
        core_chunks.push(Vec::new());
    }

    let mut shared_chunks = Arc::new(Mutex::new(core_chunks));

    download_chunks(remote_content_length, &mut shared_chunks, threads as u64, &url);

    let mut local_file = File::create(local_path).expect("[ERROR] Canno't create a file !");

    match write_file(&mut local_file, &shared_chunks) {
        Ok(()) => println!("{}", Green.bold().paint("Chunks have been successfuly saved!")),
        Err(error) => println!("[ERROR] {}", error),
    }

}<commit_msg>Format the code<commit_after>extern crate ansi_term;
extern crate argparse;
extern crate hyper;
extern crate libsnatch;

use ansi_term::Colour::{Green, Yellow, Red};
use argparse::{ArgumentParser, Store, StoreTrue};
use hyper::client::Client;
use libsnatch::Chunks;
use libsnatch::client::GetResponse;
use libsnatch::contentlength::GetContentLength;
use libsnatch::download::download_chunks;
use libsnatch::http_version::ValidateHttpVersion;
use libsnatch::write::write_file;
use std::fs::File;
use std::io;
use std::sync::{Arc, Mutex};
use std::path::Path;
use std::process::exit;

fn main() {

    let mut file = String::from("");
    let mut threads: usize = 4;
    let mut url = String::from("");
    let mut verbose = false;

    {
        let mut argparse = ArgumentParser::new();
        argparse.set_description("Snatch, a simple, fast and interruptable download accelerator, \
                                  written in Rust.");
        argparse.refer(&mut file)
            .add_option(&["-f", "--file"],
                        Store,
                        "The local file to save the remote content file")
            .required();
        argparse.refer(&mut threads)
            .add_option(&["-t", "--threads"],
                        Store,
                        "Number of threads available to download");
        argparse.refer(&mut url)
            .add_option(&["-u", "--url"], Store, "Remote content URL to download")
            .required();
        argparse.refer(&mut verbose)
            .add_option(&["-v", "--verbose"], StoreTrue, "Verbose mode");
        argparse.parse_args_or_exit();
    }

    let hyper_client = Client::new();

    // Get the first response from the server
    let client_response = hyper_client.get_http_response(&url).unwrap();

    print!("# Waiting a response from the remote server... ");

    if !client_response.version.greater_than_http_11() {
        println!("{}",
                 Yellow.bold()
                     .paint("[WARNING] The version of HTTP requests is <= HTTP1.0 - it can be \
                             difficult to split efficiently the remote content!"));
    } else {
        println!("{}", Green.bold().paint("OK !"));
    }

    let local_path = Path::new(&file);

    if local_path.exists() {
        if local_path.is_dir() {
            panic!(Red.bold()
                .paint("[ERROR] The local path to store the remote content is already exists, \
                        and is a directory!"));
        }
        println!("{}",
                 Red.bold()
                     .paint("[WARNING] The path to store the file already exists! Do you want \
                             to override it? [y/N]"));
        {
            let mut user_input = String::new();
            io::stdin()
                .read_line(&mut user_input)
                .ok()
                .expect("[ERROR] Couldn't read line!");
            user_input = String::from(user_input.trim());
            if !(user_input == "y" || user_input == "Y") {
                exit(0);
            }
        }
    }

    let remote_content_length = match client_response.headers.get_content_length() {
        Some(remote_content_length) => remote_content_length,
        None => {
            println!("{}",
                     Red.bold()
                         .paint("[ERROR] Canno't get the content length of the remote content, \
                                 from the server."));
            exit(1);
        }
    };

    println!("# Remote content length: {:?} bytes", remote_content_length);

    let mut core_chunks = Chunks::with_capacity(threads);

    for _ in 0..threads {
        core_chunks.push(Vec::new());
    }

    let mut shared_chunks = Arc::new(Mutex::new(core_chunks));

    download_chunks(remote_content_length,
                    &mut shared_chunks,
                    threads as u64,
                    &url);

    let mut local_file = File::create(local_path).expect("[ERROR] Canno't create a file !");

    match write_file(&mut local_file, &shared_chunks) {
        Ok(()) => println!("{}", Green.bold().paint("Chunks have been successfuly saved!")),
        Err(error) => println!("[ERROR] {}", error),
    }

}
<|endoftext|>


<commit_before>#![deny(warnings)]
#![feature(phase)]

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

use std::io::TempDir;
use std::io::fs::{PathExtensions, mod};
use std::io::process::{Command, ExitSignal, ExitStatus, InheritFd, ProcessOutput};
use std::os;

fn main() {
    let args = os::args();
    let args = args.slice_from(1);

    // If `--run` is not in the arguments: pass all the arguments to `rustc`
    if args.iter().all(|arg| arg.as_slice() != "--run") {
        let mut cmd = Command::new("rustc");
        cmd.args(args);
        cmd.stdout(InheritFd(1));
        cmd.stderr(InheritFd(2));

        info!("cwd: . | cmd: `{}`", cmd);
        match cmd.output() {
            Err(e) => panic!("`{}` failed: {}", cmd, e),
            Ok(ProcessOutput { status: exit, .. }) =>  {
                let exit_code = match exit {
                    ExitSignal(code) => code,
                    ExitStatus(code) => code,
                };

                os::set_exit_status(exit_code);

                return;
            },
        }
    }

    // Before `--`: arguments for the compiler
    // After `--`: arguments for the executable
    let mut splitted_args = args.split(|arg| arg.as_slice() == "--");
    let compiler_args = splitted_args.next().unwrap();
    let executable_args = splitted_args.next();

    // TODO `--crate-type=lib` should also be forbidden
    // XXX What if the crate has a `crate_type=*lib` attribute?
    if compiler_args.iter().any(|arg| arg.as_slice() == "--out-dir") {
        panic!("Can't use both `--out-dir` and `--run` flags at the same time");
    }

    // Build the rustc command
    let mut cmd = Command::new("rustc");
    // Make all paths absolute, filter out the `--run` flag
    let current_dir = os::getcwd().ok().expect("Couldn't fetch the current directory");
    for arg in compiler_args.iter().map(|arg| arg.as_slice()).filter(|&arg| arg != "--run") {
        let path = Path::new(arg);
        if path.exists() && path.is_relative() {
            cmd.arg(current_dir.join(path));
        } else {
            cmd.arg(arg);
        }
    }
    cmd.stdout(InheritFd(1));
    cmd.stderr(InheritFd(2));

    // Create temporary directory
    let temp_dir = TempDir::new("rust").unwrap();
    let temp_dir_path = temp_dir.path();
    let temp_dir_display = temp_dir_path.display();

    // Compile
    info!("cwd: {} | cmd: `{}`", temp_dir_display, cmd);
    match cmd.cwd(temp_dir_path).output() {
        Err(e) => panic!("`{}` failed: {}", cmd, e),
        Ok(ProcessOutput { status: exit, .. }) => if !exit.success() {
            let exit_code = match exit {
                ExitSignal(code) => code,
                ExitStatus(code) => code,
            };

            os::set_exit_status(exit_code);
            return;
        },
    }

    // Look for the produced binary
    let mut cmd = match fs::readdir(temp_dir_path) {
        Err(e) => panic!("`ls {}` failed: {}", temp_dir_display, e),
        Ok(paths) => match paths.as_slice().get(0) {
            Some(path) => Command::new(path),
            None => panic!("no binary found in {}", temp_dir_display),
        }
    };

    // Build the executable command
    match executable_args {
        None => {},
        Some(args) => {
            cmd.args(args);
        },
    }
    cmd.stdin(InheritFd(0));
    cmd.stdout(InheritFd(1));
    cmd.stderr(InheritFd(2));

    // Execute
    info!("cwd: . | cmd: `{}`", cmd);
    match cmd.output() {
        Err(e) => panic!("`{}` failed: {}", cmd, e),
        Ok(ProcessOutput { status: exit, .. }) => if !exit.success() {
            let exit_code = match exit {
                ExitSignal(code) => code,
                ExitStatus(code) => code,
            };

            os::set_exit_status(exit_code);
            return;
        },
    }
}
<commit_msg>fix: update to rust nightly<commit_after>#![deny(warnings)]

#[macro_use]
extern crate log;

use std::io::TempDir;
use std::io::fs::{PathExtensions, self};
use std::io::process::{Command, ExitSignal, ExitStatus, InheritFd, ProcessOutput};
use std::os;

fn main() {
    let args = os::args();
    let args = args.slice_from(1);

    // If `--run` is not in the arguments: pass all the arguments to `rustc`
    if args.iter().all(|arg| arg.as_slice() != "--run") {
        let mut cmd = Command::new("rustc");
        cmd.args(args);
        cmd.stdout(InheritFd(1));
        cmd.stderr(InheritFd(2));

        info!("cwd: . | cmd: `{}`", cmd);
        match cmd.output() {
            Err(e) => panic!("`{}` failed: {}", cmd, e),
            Ok(ProcessOutput { status: exit, .. }) =>  {
                let exit_code = match exit {
                    ExitSignal(code) => code,
                    ExitStatus(code) => code,
                };

                os::set_exit_status(exit_code);

                return;
            },
        }
    }

    // Before `--`: arguments for the compiler
    // After `--`: arguments for the executable
    let mut splitted_args = args.split(|arg| arg.as_slice() == "--");
    let compiler_args = splitted_args.next().unwrap();
    let executable_args = splitted_args.next();

    // TODO `--crate-type=lib` should also be forbidden
    // XXX What if the crate has a `crate_type=*lib` attribute?
    if compiler_args.iter().any(|arg| arg.as_slice() == "--out-dir") {
        panic!("Can't use both `--out-dir` and `--run` flags at the same time");
    }

    // Build the rustc command
    let mut cmd = Command::new("rustc");
    // Make all paths absolute, filter out the `--run` flag
    let current_dir = os::getcwd().ok().expect("Couldn't fetch the current directory");
    for arg in compiler_args.iter().map(|arg| arg.as_slice()).filter(|&arg| arg != "--run") {
        let path = Path::new(arg);
        if path.exists() && path.is_relative() {
            cmd.arg(current_dir.join(path));
        } else {
            cmd.arg(arg);
        }
    }
    cmd.stdout(InheritFd(1));
    cmd.stderr(InheritFd(2));

    // Create temporary directory
    let temp_dir = TempDir::new("rust").unwrap();
    let temp_dir_path = temp_dir.path();
    let temp_dir_display = temp_dir_path.display();

    // Compile
    info!("cwd: {} | cmd: `{}`", temp_dir_display, cmd);
    match cmd.cwd(temp_dir_path).output() {
        Err(e) => panic!("`{}` failed: {}", cmd, e),
        Ok(ProcessOutput { status: exit, .. }) => if !exit.success() {
            let exit_code = match exit {
                ExitSignal(code) => code,
                ExitStatus(code) => code,
            };

            os::set_exit_status(exit_code);
            return;
        },
    }

    // Look for the produced binary
    let mut cmd = match fs::readdir(temp_dir_path) {
        Err(e) => panic!("`ls {}` failed: {}", temp_dir_display, e),
        Ok(paths) => match paths.as_slice().get(0) {
            Some(path) => Command::new(path),
            None => panic!("no binary found in {}", temp_dir_display),
        }
    };

    // Build the executable command
    match executable_args {
        None => {},
        Some(args) => {
            cmd.args(args);
        },
    }
    cmd.stdin(InheritFd(0));
    cmd.stdout(InheritFd(1));
    cmd.stderr(InheritFd(2));

    // Execute
    info!("cwd: . | cmd: `{}`", cmd);
    match cmd.output() {
        Err(e) => panic!("`{}` failed: {}", cmd, e),
        Ok(ProcessOutput { status: exit, .. }) => if !exit.success() {
            let exit_code = match exit {
                ExitSignal(code) => code,
                ExitStatus(code) => code,
            };

            os::set_exit_status(exit_code);
            return;
        },
    }
}
<|endoftext|>