<commit_before>#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate hyper;
extern crate hubcaps;
extern crate afterparty;
extern crate serde_json;
extern crate serde;
extern crate rustc_serialize;
extern crate toml;
#[macro_use]
extern crate log;
extern crate env_logger;

use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::env;
use std::io::Read;

use hyper::{Client, Server, Url};
use hyper::header::{Headers, Authorization, Bearer, UserAgent};
use hubcaps::{Credentials, Github, IssueOptions};
use afterparty::*;
use serde_json::Value;
use regex::Regex;
use toml::decode_str;

#[derive(RustcEncodable, RustcDecodable)]
struct RepoData {
    pub qa_user: String,
    pub qa_flags: Vec<String>,
}

#[derive(RustcEncodable, RustcDecodable)]
struct ConfigData {
    pub api_root: String,
    pub repos: HashMap<String, RepoData>,
}

macro_rules! panic_log {
    ($($arg:tt)*) => ( {
        error!($($arg)*);
        panic!()
    } );
}

trait ExpectLog {
    type Item;
    fn expect_log(self, &str) -> Self::Item;
}

impl<T> ExpectLog for Option<T> {
    type Item = T;
    fn expect_log(self, msg: &str) -> T {
        match self {
            None => panic_log!("{}", msg),
            Some(s) => s,
        }
    }
}

impl<T, E> ExpectLog for Result<T, E>
    where E: std::fmt::Display
{
    type Item = T;
    fn expect_log(self, msg: &str) -> T {
        match self {
            Err(err) => panic_log!("{} {}", msg, err),
            Ok(s) => s,
        }
    }
}

fn read_config_file(file_name: &str) -> ConfigData {
    let mut config_file = File::open(file_name)
        .expect_log(&format!("Couldn't open config file {}", file_name));
    let mut str = String::new();
    config_file.read_to_string(&mut str)
        .expect_log(&format!("Couldn't read config file {}", file_name));
    let toml: Option<ConfigData> = decode_str(&str);
    toml.expect_log(&format!("Couldn't decode config file {}", file_name))
}

#[test]
fn it_reads_toml() {
    env_logger::init().unwrap();
    let config_data = read_config_file("test-data/test.toml");
    assert_eq!(config_data.api_root, "github.ibm.com/api/v3");
    let steve = config_data.repos.get("lkgrele/steve").unwrap();
    assert_eq!(steve.qa_user, "leif");
    assert_eq!(steve.qa_flags, ["qa", "test"]);
    let raul = config_data.repos.get("raul/robot").unwrap();
    assert_eq!(raul.qa_user, "raul");
    assert_eq!(raul.qa_flags, ["qa"]);
}

fn main() {
    env_logger::init().unwrap();

    let mut hub = Hub::new();
    hub.handle("pull_request", |delivery: &Delivery| {
        match delivery.payload {
            Event::PullRequest { ref pull_request, ref repository, .. } => {
                if pull_request.merged {
                    handle_pr(&pull_request.commits_url, &repository.full_name)
                }
            }
            _ => info!("Recived a request that wasn't a pull-request webhook"),
        }
    });

    match Server::http("0.0.0.0:0") {
        Err(err) => error!("ERROR: failed to start server, the error was: {}", err),
        Ok(server) => {
            match server.handle(hub) {
                Err(err) => error!("ERROR starting handler, the error was: {}", err),
                Ok(_) => info!("Successfully started server"),
            }
        }
    }
}

fn handle_pr(commits_url: &str, repository: &str) {
    let config_data = read_config_file(".steve");
    let auth_token = env::var_os("STEVE_GITHUB_TOKEN")
        .expect_log("Missing github token, STEVE_GITHUB_TOKEN is not defined!")
        .into_string()
        .ok()
        .expect_log("Github token contains illegal characters!");
    match config_data.repos.get(&repository.to_owned()) {
        Some(config) => {
            update_issues(&commits_url,
                          &auth_token,
                          &config,
                          &config_data.api_root,
                          &repository)
        }
        None => info!("No config found for repo {}, skipping", repository),
    }
}

fn update_issues(commits_url: &str,
                 auth_token: &str,
                 config: &RepoData,
                 api_root: &str,
                 repository: &str) {
    let client = Client::new();
    let mut headers = Headers::new();
    headers.set(Authorization(Bearer { token: auth_token.to_owned() }));
    headers.set(UserAgent("steve".to_owned()));
    let commits = client.get(commits_url).headers(headers).send().unwrap();
    assert_eq!(commits.status, hyper::Ok); //replace with actual error handling

    match serde_json::from_reader(commits) {
        Ok(json) => {
            match parse_commit_data(&json) {
                Ok(messages) => {
                    let mut issues = HashSet::new();
                    for message in messages {
                        get_issues(&mut issues, &message)
                    }

                    if !issues.is_empty() {
                        let github = Github::host(api_root,
                                                  "steve",
                                                  &client,
                                                  Credentials::Token(auth_token.to_owned()));
                        let (owner, repo_name) = match repository.split("/").collect::<Vec<_>>() {
                            strs => (strs[0], strs[1]),
                        };
                        let repo = github.repo(repo_name, owner);
                        match repo.issues().list(&Default::default()) {
                            Err(err) => error!("{:?}", err),
                            Ok(issue_data) => {
                                for issue in issues {
                                    if issue_data.len() as u64 <= issue {
                                        info!("Wanted to update issue #{}, but {} doesn't have an \
                                               issue with that number",
                                              issue,
                                              repository)
                                    } else {
                                        let ref current_flags = issue_data[issue as usize].labels;
                                        let new_flags: Vec<String> = current_flags.into_iter()
                                            .map(|l| l.name.clone())
                                            .chain(config.qa_flags.clone().into_iter())
                                            .collect();
                                        let issue_options =
                                            IssueOptions::new::<String,
                                                                String,
                                                                String,
                                                                String>(None,
                                                                        None,
                                                                        Some(config.qa_user
                                                                            .clone()),
                                                                        None,
                                                                        Some(new_flags));
                                        if let Err(err) = repo.issue(issue).edit(&issue_options) {
                                            error!("{:?}", err)
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
                Err(err) => error!("{:?}", err),
            }
        }
        Err(err) => error!("{:?}", err),
    }
}

#[test]
fn it_parses_commit_data() {
    let file = File::open("test-data/test-messages.json").expect("Couldn't open message file");
    let json = serde_json::from_reader(file);
    assert!(json.is_ok());
    let json = json.unwrap();
    let messages = parse_commit_data(&json);
    assert!(messages.is_ok());
    let messages = messages.unwrap();
    assert_eq!(messages[0], "message1");
    assert_eq!(messages[1], "swordfish");
}

fn parse_commit_data<'a>(json: &'a Value) -> Result<Vec<&'a str>, &str> {
    let arr = try!(json.as_array().ok_or("expected array as top level object"));
    arr.iter()
        .map(|obj| {
            obj.pointer("/commit/message")
                .ok_or("missing message field")
                .and_then(|m| m.as_string().ok_or("message field has wrong type"))
        })
        .collect()
}

fn update_issue(client: &Client,
                auth_token: &str,
                issue_num: u64,
                config: &RepoData,
                repository: &str,
                api_root: &str) {
    let github = Github::host(api_root,
                              "steve",
                              client,
                              Credentials::Token(auth_token.to_owned()));
    let (owner, repo_name) = match repository.split("/").collect::<Vec<_>>() {
        strs => (strs[0], strs[1]),
    };
    let _ = github.repo(repo_name, owner)
        .issue(issue_num)
        .labels()
        .add(config.qa_flags.iter().map(String::as_str).collect());
    let mut headers = Headers::new();
    headers.set(Authorization(Bearer { token: auth_token.to_owned() }));
    headers.set(UserAgent("steve".to_owned()));
    let update_json = format!("{{\"assignee\":\"{}\"}}", config.qa_user);
    let url = Url::parse(&format!("{}/repos/{}/issues/{}", api_root, repository, issue_num))
        .expect("bad url");
    let res = client.patch(url)
        .body(&update_json)
        .headers(headers.to_owned())
        .send()
        .unwrap();
    assert_eq!(res.status, hyper::Ok);
}

#[test]
fn it_finds_issue_numbers() {
    let message = r"This is a commit message
        needs qa: #72 should get caught
        needs QA #333 should also get caught
        but #44 won't, and needs QA: #33 #45
        only gets #33. Duplicates are discarded:
        needs qa #72";
    let mut expected: HashSet<u64> = HashSet::new();
    expected.insert(72);
    expected.insert(333);
    expected.insert(33);
    let mut numbers: HashSet<u64> = HashSet::new();
    get_issues(&mut numbers, message);
    assert_eq!(numbers, expected);
}

fn get_issues(numbers: &mut HashSet<u64>, message: &str) {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"(?i)needs\sqa:?\s#(\d+)").expect_log("{}");
    }
    for num in RE.captures_iter(message).filter_map(|cap| cap.at(1)) {
        numbers.insert(num.parse::<u64>().expect_log("{}"));
    }
}
<commit_msg>Specify ip address and port as env variables.<commit_after>#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate hyper;
extern crate hubcaps;
extern crate afterparty;
extern crate serde_json;
extern crate serde;
extern crate rustc_serialize;
extern crate toml;
#[macro_use]
extern crate log;
extern crate env_logger;

use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::env;
use std::io::Read;

use hyper::{Client, Server, Url};
use hyper::header::{Headers, Authorization, Bearer, UserAgent};
use hubcaps::{Credentials, Github, IssueOptions};
use afterparty::*;
use serde_json::Value;
use regex::Regex;
use toml::decode_str;

#[derive(RustcEncodable, RustcDecodable)]
struct RepoData {
    pub qa_user: String,
    pub qa_flags: Vec<String>,
}

#[derive(RustcEncodable, RustcDecodable)]
struct ConfigData {
    pub api_root: String,
    pub repos: HashMap<String, RepoData>,
}

macro_rules! panic_log {
    ($($arg:tt)*) => ( {
        error!($($arg)*);
        panic!()
    } );
}

trait ExpectLog {
    type Item;
    fn expect_log(self, &str) -> Self::Item;
}

impl<T> ExpectLog for Option<T> {
    type Item = T;
    fn expect_log(self, msg: &str) -> T {
        match self {
            None => panic_log!("{}", msg),
            Some(s) => s,
        }
    }
}

impl<T, E> ExpectLog for Result<T, E>
    where E: std::fmt::Display
{
    type Item = T;
    fn expect_log(self, msg: &str) -> T {
        match self {
            Err(err) => panic_log!("{} {}", msg, err),
            Ok(s) => s,
        }
    }
}

fn read_config_file(file_name: &str) -> ConfigData {
    let mut config_file = File::open(file_name)
        .expect_log(&format!("Couldn't open config file {}", file_name));
    let mut str = String::new();
    config_file.read_to_string(&mut str)
        .expect_log(&format!("Couldn't read config file {}", file_name));
    let toml: Option<ConfigData> = decode_str(&str);
    toml.expect_log(&format!("Couldn't decode config file {}", file_name))
}

#[test]
fn it_reads_toml() {
    env_logger::init().unwrap();
    let config_data = read_config_file("test-data/test.toml");
    assert_eq!(config_data.api_root, "github.ibm.com/api/v3");
    let steve = config_data.repos.get("lkgrele/steve").unwrap();
    assert_eq!(steve.qa_user, "leif");
    assert_eq!(steve.qa_flags, ["qa", "test"]);
    let raul = config_data.repos.get("raul/robot").unwrap();
    assert_eq!(raul.qa_user, "raul");
    assert_eq!(raul.qa_flags, ["qa"]);
}

fn main() {
    env_logger::init().unwrap();

    let mut hub = Hub::new();
    hub.handle("pull_request", |delivery: &Delivery| {
        match delivery.payload {
            Event::PullRequest { ref pull_request, ref repository, .. } => {
                if pull_request.merged {
                    handle_pr(&pull_request.commits_url, &repository.full_name)
                }
            }
            _ => info!("Recived a request that wasn't a pull-request webhook"),
        }
    });

    let ip_and_port = get_ip_and_port();

    match Server::http(ip_and_port.as_str()) {
        Err(err) => error!("ERROR: failed to start server, the error was: {}", err),
        Ok(server) => {
            match server.handle(hub) {
                Err(err) => error!("ERROR starting handler, the error was: {}", err),
                Ok(_) => info!("Successfully started server"),
            }
        }
    }
}

fn env_or(name: &str, def: &str) -> String {
    env::var_os(name).and_then(|v| v.into_string().ok()).unwrap_or(def.to_owned())
}

fn get_ip_and_port() -> String {
    let ip = env_or("STEVE_IP", "0.0.0.0");
    let port = env_or("STEVE_PORT", "80");
    format!("{}:{}", ip, port)
}

fn handle_pr(commits_url: &str, repository: &str) {
    let config_data = read_config_file(".steve");
    let auth_token = env::var_os("STEVE_GITHUB_TOKEN")
        .expect_log("Missing github token, STEVE_GITHUB_TOKEN is not defined!")
        .into_string()
        .ok()
        .expect_log("Github token contains illegal characters!");
    match config_data.repos.get(&repository.to_owned()) {
        Some(config) => {
            update_issues(&commits_url,
                          &auth_token,
                          &config,
                          &config_data.api_root,
                          &repository)
        }
        None => info!("No config found for repo {}, skipping", repository),
    }
}

fn update_issues(commits_url: &str,
                 auth_token: &str,
                 config: &RepoData,
                 api_root: &str,
                 repository: &str) {
    let client = Client::new();
    let mut headers = Headers::new();
    headers.set(Authorization(Bearer { token: auth_token.to_owned() }));
    headers.set(UserAgent("steve".to_owned()));
    let commits = client.get(commits_url).headers(headers).send().unwrap();
    assert_eq!(commits.status, hyper::Ok); //replace with actual error handling

    match serde_json::from_reader(commits) {
        Ok(json) => {
            match parse_commit_data(&json) {
                Ok(messages) => {
                    let mut issues = HashSet::new();
                    for message in messages {
                        get_issues(&mut issues, &message)
                    }

                    if !issues.is_empty() {
                        let github = Github::host(api_root,
                                                  "steve",
                                                  &client,
                                                  Credentials::Token(auth_token.to_owned()));
                        let (owner, repo_name) = match repository.split("/").collect::<Vec<_>>() {
                            strs => (strs[0], strs[1]),
                        };
                        let repo = github.repo(repo_name, owner);
                        match repo.issues().list(&Default::default()) {
                            Err(err) => error!("{:?}", err),
                            Ok(issue_data) => {
                                for issue in issues {
                                    if issue_data.len() as u64 <= issue {
                                        info!("Wanted to update issue #{}, but {} doesn't have an \
                                               issue with that number",
                                              issue,
                                              repository)
                                    } else {
                                        let ref current_flags = issue_data[issue as usize].labels;
                                        let new_flags: Vec<String> = current_flags.into_iter()
                                            .map(|l| l.name.clone())
                                            .chain(config.qa_flags.clone().into_iter())
                                            .collect();
                                        let issue_options =
                                            IssueOptions::new::<String,
                                                                String,
                                                                String,
                                                                String>(None,
                                                                        None,
                                                                        Some(config.qa_user
                                                                            .clone()),
                                                                        None,
                                                                        Some(new_flags));
                                        if let Err(err) = repo.issue(issue).edit(&issue_options) {
                                            error!("{:?}", err)
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
                Err(err) => error!("{:?}", err),
            }
        }
        Err(err) => error!("{:?}", err),
    }
}

#[test]
fn it_parses_commit_data() {
    let file = File::open("test-data/test-messages.json").expect("Couldn't open message file");
    let json = serde_json::from_reader(file);
    assert!(json.is_ok());
    let json = json.unwrap();
    let messages = parse_commit_data(&json);
    assert!(messages.is_ok());
    let messages = messages.unwrap();
    assert_eq!(messages[0], "message1");
    assert_eq!(messages[1], "swordfish");
}

fn parse_commit_data<'a>(json: &'a Value) -> Result<Vec<&'a str>, &str> {
    let arr = try!(json.as_array().ok_or("expected array as top level object"));
    arr.iter()
        .map(|obj| {
            obj.pointer("/commit/message")
                .ok_or("missing message field")
                .and_then(|m| m.as_string().ok_or("message field has wrong type"))
        })
        .collect()
}

fn update_issue(client: &Client,
                auth_token: &str,
                issue_num: u64,
                config: &RepoData,
                repository: &str,
                api_root: &str) {
    let github = Github::host(api_root,
                              "steve",
                              client,
                              Credentials::Token(auth_token.to_owned()));
    let (owner, repo_name) = match repository.split("/").collect::<Vec<_>>() {
        strs => (strs[0], strs[1]),
    };
    let _ = github.repo(repo_name, owner)
        .issue(issue_num)
        .labels()
        .add(config.qa_flags.iter().map(String::as_str).collect());
    let mut headers = Headers::new();
    headers.set(Authorization(Bearer { token: auth_token.to_owned() }));
    headers.set(UserAgent("steve".to_owned()));
    let update_json = format!("{{\"assignee\":\"{}\"}}", config.qa_user);
    let url = Url::parse(&format!("{}/repos/{}/issues/{}", api_root, repository, issue_num))
        .expect("bad url");
    let res = client.patch(url)
        .body(&update_json)
        .headers(headers.to_owned())
        .send()
        .unwrap();
    assert_eq!(res.status, hyper::Ok);
}

#[test]
fn it_finds_issue_numbers() {
    let message = r"This is a commit message
        needs qa: #72 should get caught
        needs QA #333 should also get caught
        but #44 won't, and needs QA: #33 #45
        only gets #33. Duplicates are discarded:
        needs qa #72";
    let mut expected: HashSet<u64> = HashSet::new();
    expected.insert(72);
    expected.insert(333);
    expected.insert(33);
    let mut numbers: HashSet<u64> = HashSet::new();
    get_issues(&mut numbers, message);
    assert_eq!(numbers, expected);
}

fn get_issues(numbers: &mut HashSet<u64>, message: &str) {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"(?i)needs\sqa:?\s#(\d+)").expect_log("{}");
    }
    for num in RE.captures_iter(message).filter_map(|cap| cap.at(1)) {
        numbers.insert(num.parse::<u64>().expect_log("{}"));
    }
}
<|endoftext|>


<commit_before>use std::f64::consts;
use std::vec::Vec;
use std::fs::File;
use std::path::Path;
use std::io::{Read, Write};

const LENGTH: f64 = 2.0*consts::PI; //length of PIC box
const NG: i32 = 64; //number of zones for mesh
const MASS: i32 = 1; //mass of each particle
const CHARGE: i32 = -1; //charge of each particle
const V_STREAM: f64 = 4.0; //initial velocity of each stream, this number will turned into positive and negative!

fn main() {
	let particle_loading: i32 = 4; //particle loading
	let np: i32 = particle_loading*NG; //number of particles
	let k: i32 = 1; //used for the peturbation of the sin wave
	let dt: f64 = LENGTH/(NG as f64)/(4.0); //timestep
	let delta_x: f64 = LENGTH/(np as f64); //L/np, or the delta x of the lattice
	let t_max: f64 = 16 as f64*consts::PI/((2 as f64*(np as f64)/LENGTH).sqrt()); //maximum time
	let mut velocity: Vec<f64> = Vec::new(); //velocities of each particle, index is
	let mut x_position: Vec<f64> = Vec::new(); //x-position of each particle

	//seed the initial positions and velocities
	for i in 0..(np*2) {
		if i < np {
			x_position.push(i as f64*LENGTH/np as f64);
			velocity.push(-V_STREAM);
		}
		else {
			x_position.push((i-np) as f64*LENGTH/np as f64);
			velocity.push(V_STREAM);
		}
		//rust only accepts usize for indexes, sanity check
		// println!("x: {}, v: {}", x_position[i as usize], velocity[i as usize]);
	}

	//peturb the x positions
	perturb(&mut x_position, k);

	let mut rho: Vec<f64> = vec![0.0; NG as usize]; //mesh densities
	let mut phi: Vec<f64> = vec![0.0; NG as usize]; //vector of phi's to solve for the poisson
	let mut electric_mesh: Vec<f64> = vec![0.0; NG as usize]; //vector of electric field from potential

	let solnSize = np*2;
	let mut solution: Vec<f64> = vec![0.0; solnSize as usize];
	catch(&x_position, &velocity, &mut solution, np);
	let timestep = (t_max/delta_x).ceil() as u32;

	for k in 0..timestep {
		//unload from the solution
		release(&mut x_position, &mut velocity, &solution, np);
		write_out(&x_position, &velocity, &rho, k);
		//find the particle number density
		density(&mut rho, &mut x_position, np);
		//solve the poisson equation
		poisson1d(&mut phi, &mut rho);
		//find the mesh
		electric_field(&mut phi, &mut electric_mesh);
		for i in 0..x_position.len() {
			let idx = i as usize;

		}

	}
}

//evaluates the particle number density using CIC weight.
//Each mesh loops through all the particles to add in their contribution, calculates the electron density for that mesh block alone.
fn density(rho: &mut Vec<f64>, x_position: & Vec<f64>, np: i32) {
	for rho_i in 0..rho.len() {
		let mut sum: f64 = 0.0;
		for x_p in 0..np*2 {
			sum += (MASS as f64)*cic_weight(rho[rho_i as usize] - x_position[x_p as usize], np);
		}
		// println!("rho_i is {}", rho_i);
		rho[rho_i as usize] = 1.0/(LENGTH/NG as f64)*sum;
	}
}

//Cloud in Cell weighting
//			  { 1 + x/deltaX for -deltaX < x
// W_cic(X) = | 1 - x/deltaX for x < deltaX
//			  } 0 otherwise
fn cic_weight(x: f64, np: i32) -> f64{
	let delta_x: f64 = LENGTH/(np as f64);
	if x > -delta_x {
		return 1.0 + x/delta_x;
	}
	else if x < delta_x {
		return 1.0 - x/delta_x;
	}
	else {
		return 0.0;
	}
}

//solves the 1D poisson equation using the Gauss-Seidel relaxation
fn poisson1d<'a>(phi: &'a mut Vec<f64>, rho: &Vec<f64>) {
	let dx_g: f64 = LENGTH as f64/NG as f64;
	let error_tolerance: f64 = 0.000001; //10^-6
	let source = source_l2_norm(&rho, dx_g);
	let mut residual = residual_l2_norm(&rho, &phi, dx_g);
	// println!("R: {}, eS: {}", residual, error_tolerance*source);
	while residual > error_tolerance*source {
		// println!("R: {}, eS: {}", residual, error_tolerance*source);
		//print_vec(&phi);
		println!("Before new phi R: {}", residual);
		for i in 0..NG-1 {
			// rust uses C-like functions, so % is remainder (don't go negative!)
			let im: i32 = ((i - 1 + NG) % NG);
			let ip: i32 = ((i + 1) % NG);

			//println!("Old Phi: {}", phi[i as usize]);
			//println!("New Phi: {}", 0.5*(phi[im as usize] + phi[ip as usize]) + dx_g.powi(2)/2.0*rho[i as usize]);
			// println!("i: {}, ng: {}, im: {}, ip: {}", i, NG, im, ip);
			phi[i as usize] = 0.5*(phi[im as usize] + phi[ip as usize]) + dx_g.powi(2)/2.0*rho[i as usize];
			//println!("New Phi: {}", phi[i as usize]);
		}
		//print_vec(&phi);
		residual = residual_l2_norm(&rho, &phi, dx_g);
		println!("After new phi R: {}", residual);
	}
}

//calculates the L2Norm for the source
//probably don't want to call this outside of the poisson solver?
fn source_l2_norm(rho: &Vec<f64>, dx_g: f64) -> f64 {
	let mut source: f64 = 0.0;
	//sum the entire array
	for i in 0..rho.len() {
		source += rho[i as usize].powi(2);
	}
	//divide it by 1/Ng
	source /= NG as f64;
	//then square root it
	source = source.sqrt();
	return source;
}

//calculates the L2Norm for the residual
fn residual_l2_norm(rho: &Vec<f64>, phi: &Vec<f64>, dx_g: f64) -> f64 {
	let mut residual: f64 = 0.0;
	for i in 0..rho.len() {
		residual += rho[i as usize];
		if i == 0 {
			residual += (phi[i+1 as usize] - 2.0*phi[i as usize] + phi[(rho.len()-1) as usize])/dx_g.powi(2);
		}
		else if i+1 == rho.len() as usize {
			residual += (phi[0] - 2.0*phi[i as usize] + phi[(i-1) as usize])/dx_g.powi(2);
		}
		else {
			residual += (phi[(i+1 as usize)] - 2.0*phi[i as usize] + phi[(i-1) as usize])/dx_g.powi(2);
		}
	}
	residual /= NG as f64;
	residual = residual.sqrt();

	return residual;
}

//calculates the electric field from potential, where
//E is the potential
//uses the second order finite difference
fn electric_field (phi: &mut Vec<f64>, electric_mesh: & mut Vec<f64>) {
	//get dx_g
	let dx: f64 = LENGTH/(NG as f64);
	//set the iteration length for easier reference
	let iter_length = phi.len();
	println!("Phi Length: {}", phi.len());
	println!("Electric Mesh Length: {}", electric_mesh.len());
	for i in 1..iter_length-1 {
		// println!("Hit {}", i);
		//finite difference
		let phi_next = phi[i+1 as usize];
		let phi_prev = phi[i-1 as usize];
		electric_mesh[i] = (phi_next - 2.0*phi[i as usize] + phi_prev)/(dx.powi(2));
	}
	//particles at the edge need to wrap around, so wrap them around
	electric_mesh[0] = (phi[1] - 2.0*phi[0] + phi[iter_length-1])/(dx.powi(2));
	electric_mesh[iter_length-1] = (phi[0] - 2.0*phi[iter_length-1] + phi[iter_length-2])/(dx.powi(2));
}

//gets called whenever a timestep is taken, extrapolates x and v for a particle for a single timestep
fn leapfrog(electric_mesh: &Vec<f64>, x_position: &mut Vec<f64>, velocity: &mut Vec<f64>, dt: f64, delta_x: f64) {
	for i in 0..x_position.len() {
		let (x_new, v_new) = leap(x_position[i as usize], delta_x, velocity[i as usize], dt, electric_mesh);
		x_position[i as usize] = x_new;
		velocity[i as usize] = v_new;
	}
}

//calculates the drift/kick/drift for a specific value
//returns the new X and new V
//THIS SHOULD *ONLY* BE CALLED FROM THE LEAPFROG METHOD
fn leap(x: f64, delta_x: f64, v: f64, dt: f64, electric_mesh: &Vec<f64>) -> (f64, f64) {
	let j: i32 = (x/delta_x).floor() as i32;
	let y: f64 = x/delta_x - j as f64;
	//drift
	let x_half: f64 = x + 0.5*v*dt;
	//kick
	let v_new: f64 = v + accel(j, y, electric_mesh)*dt;
	//drift
	let x_new: f64 = x_half + 0.5*v_new*dt;
	return (x_new, v_new)
}

//calculates the accleration
fn accel(j: i32, y: f64, electric_mesh: &Vec<f64>) -> f64{
	let mut dphi_dx: f64 = 0.0;
	if j+1 == NG {
		dphi_dx = electric_mesh[j as usize]*(1.0-y) + electric_mesh[0]*y;
	}
	else {
		dphi_dx = electric_mesh[j as usize]*(1.0-y) + electric_mesh[(j+1) as usize]*y;
	}
	let a: f64 = -CHARGE as f64/MASS as f64*dphi_dx;
	return a;
}

//petrubs the function by a sine wave with waveLENGTH = L/k and amplitude of .01K
fn perturb(x_position: &mut Vec<f64>, k: i32) {
	let wavelength: f64 = LENGTH/(k as f64);
	let amplitude: f64 = 0.01 * LENGTH;
	for i in 0..x_position.len() {
		x_position[i as usize] = x_position[i as usize] + amplitude*(wavelength*x_position[i as usize]).sin();
	}
}

// catches particles into a temp array
fn catch(x_position: &Vec<f64>, velocity: &Vec<f64>, solution: &mut Vec<f64>, np: i32) {
	for i in 0..np {
		let pos = i as usize;
		solution[pos] = x_position[pos];
		solution[pos + np as usize] = velocity[pos];
	}
}

//dumps the temp array back into the particle vectors
fn release(x_position: &mut Vec<f64>, velocity: &mut Vec<f64>, solution: &Vec<f64>, np: i32) {
	for i in 0..np {
		let pos = i as usize;
		x_position[pos] = solution[pos];
		velocity[pos] = solution[pos + np as usize];
	}
}

//prints the list of X positions. THIS IS A COPY, NOT A DIRECT REFERENCE
fn print_vec(vector: &Vec<f64>) {
	for i in 0..vector.len() {
		println!("Item {}: {}", i, vector[i as usize])
	}
}

//writes the output to a file
fn write_out(x_position: &Vec<f64>, velocity: &Vec<f64>, rho: &Vec<f64>, timestep: u32) {
	let path = format!("{}-pvv", timestep);
	let pos_vs_velocity = Path::new(&path);
	//let pos_vs_rho = Path::new(format!("out/{}-pvr", timestep));
	let mut outStr = "x,v \n".to_owned();
	for i in 0..x_position.len() {
		let idx = i as usize;
		let tmp = (format!("{}, {}\n", x_position[idx], velocity[idx])).to_owned();
		outStr.push_str(&tmp);
	}
	println!("{}", path);
	let mut f = File::create(&pos_vs_velocity).unwrap();
	f.write_all(&outStr.as_bytes());
}<commit_msg>Fix the acceleration and finite difference functions<commit_after>use std::f64::consts;
use std::vec::Vec;
use std::fs::File;
use std::path::Path;
use std::io::{Read, Write};

const LENGTH: f64 = 2.0*consts::PI; //length of PIC box
const NG: i32 = 64; //number of zones for mesh
const MASS: i32 = 1; //mass of each particle
const CHARGE: i32 = -1; //charge of each particle
const V_STREAM: f64 = 4.0; //initial velocity of each stream, this number will turned into positive and negative!

fn main() {
	let particle_loading: i32 = 4; //particle loading
	let np: i32 = particle_loading*NG; //number of particles
	let k: i32 = 1; //used for the peturbation of the sin wave
	let dt: f64 = LENGTH/(NG as f64)/(4.0); //timestep
	let delta_x: f64 = LENGTH/(np as f64); //L/np, or the delta x of the lattice
	let t_max: f64 = 16 as f64*consts::PI/((2 as f64*(np as f64)/LENGTH).sqrt()); //maximum time
	let mut velocity: Vec<f64> = Vec::new(); //velocities of each particle, index is
	let mut x_position: Vec<f64> = Vec::new(); //x-position of each particle

	//seed the initial positions and velocities
	for i in 0..(np*2) {
		if i < np {
			x_position.push(i as f64*LENGTH/np as f64);
			velocity.push(-V_STREAM);
		}
		else {
			x_position.push((i-np) as f64*LENGTH/np as f64);
			velocity.push(V_STREAM);
		}
		//rust only accepts usize for indexes, sanity check
		// println!("x: {}, v: {}", x_position[i as usize], velocity[i as usize]);
	}

	//peturb the x positions
	perturb(&mut x_position, k);

	let mut rho: Vec<f64> = vec![0.0; NG as usize]; //mesh densities
	let mut phi: Vec<f64> = vec![0.0; NG as usize]; //vector of phi's to solve for the poisson
	let mut electric_mesh: Vec<f64> = vec![0.0; NG as usize]; //vector of electric field from potential

	let solnSize = np*2;
	let mut solution: Vec<f64> = vec![0.0; solnSize as usize];
	catch(&x_position, &velocity, &mut solution, np);
	let timestep = (t_max/delta_x).ceil() as u32;

	for k in 0..timestep {
		//unload from the solution
		release(&mut x_position, &mut velocity, &solution, np);
		write_out(&x_position, &velocity, &rho, k);
		//find the particle number density
		density(&mut rho, &mut x_position, np);
		//solve the poisson equation
		poisson1d(&mut phi, &mut rho);
		//find the mesh
		electric_field(&mut phi, &mut electric_mesh);
		for i in 0..x_position.len() {
			let idx = i as usize;

		}

	}
}

//evaluates the particle number density using CIC weight.
//Each mesh loops through all the particles to add in their contribution, calculates the electron density for that mesh block alone.
fn density(rho: &mut Vec<f64>, x_position: & Vec<f64>, np: i32) {
	for rho_i in 0..rho.len() {
		let mut sum: f64 = 0.0;
		for x_p in 0..np*2 {
			sum += (MASS as f64)*cic_weight(rho[rho_i as usize] - x_position[x_p as usize], np);
		}
		// println!("rho_i is {}", rho_i);
		rho[rho_i as usize] = 1.0/(LENGTH/NG as f64)*sum;
	}
}

//Cloud in Cell weighting
//			  { 1 + x/deltaX for -deltaX < x
// W_cic(X) = | 1 - x/deltaX for x < deltaX
//			  } 0 otherwise
fn cic_weight(x: f64, np: i32) -> f64{
	let delta_x: f64 = LENGTH/(np as f64);
	if x > -delta_x {
		return 1.0 + x/delta_x;
	}
	else if x < delta_x {
		return 1.0 - x/delta_x;
	}
	else {
		return 0.0;
	}
}

//solves the 1D poisson equation using the Gauss-Seidel relaxation
fn poisson1d<'a>(phi: &'a mut Vec<f64>, rho: &Vec<f64>) {
	let dx_g: f64 = LENGTH as f64/NG as f64;
	let error_tolerance: f64 = 0.000001; //10^-6
	let source = source_l2_norm(&rho, dx_g);
	let mut residual = residual_l2_norm(&rho, &phi, dx_g);
	// println!("R: {}, eS: {}", residual, error_tolerance*source);
	while residual > error_tolerance*source {
		// println!("R: {}, eS: {}", residual, error_tolerance*source);
		//print_vec(&phi);
		println!("Before new phi R: {}", residual);
		for i in 0..NG-1 {
			// rust uses C-like functions, so % is remainder (don't go negative!)
			let im: i32 = ((i - 1 + NG) % NG);
			let ip: i32 = ((i + 1) % NG);

			//println!("Old Phi: {}", phi[i as usize]);
			//println!("New Phi: {}", 0.5*(phi[im as usize] + phi[ip as usize]) + dx_g.powi(2)/2.0*rho[i as usize]);
			// println!("i: {}, ng: {}, im: {}, ip: {}", i, NG, im, ip);
			phi[i as usize] = 0.5*(phi[im as usize] + phi[ip as usize]) + dx_g.powi(2)/2.0*rho[i as usize];
			//println!("New Phi: {}", phi[i as usize]);
		}
		//print_vec(&phi);
		residual = residual_l2_norm(&rho, &phi, dx_g);
		println!("After new phi R: {}", residual);
	}
}

//calculates the L2Norm for the source
//probably don't want to call this outside of the poisson solver?
fn source_l2_norm(rho: &Vec<f64>, dx_g: f64) -> f64 {
	let mut source: f64 = 0.0;
	//sum the entire array
	for i in 0..rho.len() {
		source += rho[i as usize].powi(2);
	}
	//divide it by 1/Ng
	source /= NG as f64;
	//then square root it
	source = source.sqrt();
	return source;
}

//calculates the L2Norm for the residual
fn residual_l2_norm(rho: &Vec<f64>, phi: &Vec<f64>, dx_g: f64) -> f64 {
	let mut residual: f64 = 0.0;
	for i in 0..rho.len() {
		residual += rho[i as usize];
		if i == 0 {
			residual += (phi[i+1 as usize] - 2.0*phi[i as usize] + phi[(rho.len()-1) as usize])/dx_g.powi(2);
		}
		else if i+1 == rho.len() as usize {
			residual += (phi[0] - 2.0*phi[i as usize] + phi[(i-1) as usize])/dx_g.powi(2);
		}
		else {
			residual += (phi[(i+1 as usize)] - 2.0*phi[i as usize] + phi[(i-1) as usize])/dx_g.powi(2);
		}
	}
	residual /= NG as f64;
	residual = residual.sqrt();

	return residual;
}

//calculates the electric field from potential, where
//E is the potential
//uses the second order finite difference
fn electric_field (phi: &mut Vec<f64>, electric_mesh: & mut Vec<f64>) {
	//get dx_g
	let dx: f64 = LENGTH/(NG as f64);
	//set the iteration length for easier reference
	let iter_length = phi.len();
	println!("Phi Length: {}", phi.len());
	println!("Electric Mesh Length: {}", electric_mesh.len());
	for i in 1..iter_length-1 {
		// println!("Hit {}", i);
		//finite difference
		let phi_next = phi[i+1 as usize];
		let phi_prev = phi[i-1 as usize];
		electric_mesh[i] = (phi_next - phi_prev)(2.0*dx)
	}
	//particles at the edge need to wrap around, so wrap them around
	electric_mesh[0] = (phi[1] - phi[iter_length-1])/(2.0*dx);
	electric_mesh[iter_length-1] = (phi[0] - phi[iter_length-2])/(2.0*dx);
}

//gets called whenever a timestep is taken, extrapolates x and v for a particle for a single timestep
fn leapfrog(electric_mesh: &Vec<f64>, x_position: &mut Vec<f64>, velocity: &mut Vec<f64>, dt: f64, delta_x: f64, k:i32) {
	for i in 0..x_position.len() {
		let (x_new, v_new) = leap(x_position[i as usize], delta_x, velocity[i as usize], dt, k, electric_mesh);
		x_position[i as usize] = x_new;
		velocity[i as usize] = v_new;
	}
}

//calculates the drift/kick/drift for a specific value
//returns the new X and new V
//THIS SHOULD *ONLY* BE CALLED FROM THE LEAPFROG METHOD
fn leap(x: f64, delta_x: f64, v: f64, dt: f64, k: i32, electric_mesh: &Vec<f64>) -> (f64, f64) {
	let j: i32 = (x/delta_x).floor() as i32;
	let y: f64 = x/delta_x - j as f64;
	//drift
	let x_half: f64 = x + 0.5*v*dt;
	//kick
	let v_new: f64 = v + accel(j, k, electric_mesh)*dt;
	//drift
	let x_new: f64 = x_half + 0.5*v_new*dt;
	return (x_new, v_new)
}

//calculates the accleration
fn accel(j: i32, k: f64, electric_mesh: &Vec<f64>) -> f64{
	let mut dphi_dx: f64 = 0.0;
	dphi_dx = electric_mesh[j]*(k*j as f64).cos()
	let a: f64 = -CHARGE as f64/MASS as f64*dphi_dx;
	return a;
}

//petrubs the function by a sine wave with waveLENGTH = L/k and amplitude of .01K
fn perturb(x_position: &mut Vec<f64>, k: i32) {
	let wavelength: f64 = LENGTH/(k as f64);
	let amplitude: f64 = 0.01 * LENGTH;
	for i in 0..x_position.len() {
		x_position[i as usize] = x_position[i as usize] + amplitude*(wavelength*x_position[i as usize]).sin();
	}
}

// catches particles into a temp array
fn catch(x_position: &Vec<f64>, velocity: &Vec<f64>, solution: &mut Vec<f64>, np: i32) {
	for i in 0..np {
		let pos = i as usize;
		solution[pos] = x_position[pos];
		solution[pos + np as usize] = velocity[pos];
	}
}

//dumps the temp array back into the particle vectors
fn release(x_position: &mut Vec<f64>, velocity: &mut Vec<f64>, solution: &Vec<f64>, np: i32) {
	for i in 0..np {
		let pos = i as usize;
		x_position[pos] = solution[pos];
		velocity[pos] = solution[pos + np as usize];
	}
}

//prints the list of X positions. THIS IS A COPY, NOT A DIRECT REFERENCE
fn print_vec(vector: &Vec<f64>) {
	for i in 0..vector.len() {
		println!("Item {}: {}", i, vector[i as usize])
	}
}

//writes the output to a file
fn write_out(x_position: &Vec<f64>, velocity: &Vec<f64>, rho: &Vec<f64>, timestep: u32) {
	let path = format!("{}-pvv", timestep);
	let pos_vs_velocity = Path::new(&path);
	//let pos_vs_rho = Path::new(format!("out/{}-pvr", timestep));
	let mut outStr = "x,v \n".to_owned();
	for i in 0..x_position.len() {
		let idx = i as usize;
		let tmp = (format!("{}, {}\n", x_position[idx], velocity[idx])).to_owned();
		outStr.push_str(&tmp);
	}
	println!("{}", path);
	let mut f = File::create(&pos_vs_velocity).unwrap();
	f.write_all(&outStr.as_bytes());
}<|endoftext|>


<commit_before>extern crate clap;
extern crate pcap;
extern crate pnet;
extern crate serde_json;
extern crate dns_parser;

use clap::{Arg, App};
use pcap::Capture;
use pnet::packet::{Packet, ethernet, ip, ipv4, tcp, udp};
use serde_json::{Value, Map, Number};
use dns_parser::Packet as DnsPacket;

struct EthernetHeader {
	source: String,
	destination: String,
	ethertype: String,
}

impl EthernetHeader {
	fn new(p: &ethernet::EthernetPacket) -> EthernetHeader {
		EthernetHeader {
			source: p.get_source().to_string(),
			destination: p.get_destination().to_string(),
			ethertype: p.get_ethertype().to_string().to_lowercase(),
    	}
	}

	fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

		header.insert("source".to_string(), Value::String(self.source.clone()));
		header.insert("destination".to_string(), Value::String(self.destination.clone()));
		header.insert("type".to_string(), Value::String(self.ethertype.clone()));

		header
	}
}

struct Ipv4Header {
	version: Number,
	ihl: Number,
	dscp: Number,
	ecn: Number,
	total_length: Number,
	identification: Number,
	flags: Number,
	fragment_offset: Number,
	ttl: Number,
	next_level_protocol: String,
	checksum: Number,
	source: String,
	destination: String,
}

impl Ipv4Header {
	fn new(p: &ipv4::Ipv4Packet) -> Ipv4Header {
		Ipv4Header {
			version: Number::from(p.get_version()),
			ihl: Number::from(p.get_header_length()),
			dscp: Number::from(p.get_dscp()),
			ecn: Number::from(p.get_ecn()),
			total_length: Number::from(p.get_total_length()),
			identification: Number::from(p.get_identification()),
			flags: Number::from(p.get_flags()),
			fragment_offset: Number::from(p.get_fragment_offset()),
			ttl: Number::from(p.get_ttl()),
			next_level_protocol: p.get_next_level_protocol().to_string().to_lowercase(),
			checksum: Number::from(p.get_checksum()),
			source: p.get_source().to_string(),
			destination: p.get_destination().to_string(),
		}
	}

	fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

    	header.insert("version".to_string(), Value::Number(self.version.clone()));
    	header.insert("ihl".to_string(), Value::Number(self.ihl.clone()));
    	header.insert("dscp".to_string(), Value::Number(self.dscp.clone()));
    	header.insert("ecn".to_string(), Value::Number(self.ecn.clone()));
    	header.insert("total_length".to_string(), Value::Number(self.total_length.clone()));
    	header.insert("identification".to_string(), Value::Number(self.identification.clone()));
    	header.insert("flags.".to_string(), Value::Number(self.flags.clone()));
    	header.insert("fragment_offset".to_string(), Value::Number(self.fragment_offset.clone()));
    	header.insert("ttl".to_string(), Value::Number(self.ttl.clone()));
    	header.insert("protocol".to_string(), Value::String(self.next_level_protocol.clone()));
    	header.insert("checksum".to_string(), Value::Number(self.checksum.clone()));
    	header.insert("source_address".to_string(), Value::String(self.source.clone()));
    	header.insert("destination_address".to_string(), Value::String(self.destination.clone()));

		header
	}
}

struct TcpHeader {
	source_port: Number,
	destination_port: Number,
	sequence_number: Number,
	ack_number: Number,
	data_offset: Number,
	flags: Number,
	window_size: Number,
	checksum: Number,
	urgent_pointer: Number,
}

impl TcpHeader {
	fn new(p: &tcp::TcpPacket) -> TcpHeader {
		TcpHeader {
			source_port: Number::from(p.get_source()),
			destination_port: Number::from(p.get_destination()),
			sequence_number: Number::from(p.get_sequence()),
			ack_number: Number::from(p.get_acknowledgement()),
			data_offset: Number::from(p.get_data_offset()),
			flags: Number::from(p.get_flags()),
			window_size: Number::from(p.get_window()),
			checksum: Number::from(p.get_checksum()),
			urgent_pointer: Number::from(p.get_urgent_ptr()),
		}
	}

	fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

		header.insert("source_port".to_string(), Value::Number(self.source_port.clone()));
		header.insert("destination_port".to_string(), Value::Number(self.destination_port.clone()));
		header.insert("sequence_number".to_string(), Value::Number(self.sequence_number.clone()));
		header.insert("ack_number".to_string(), Value::Number(self.ack_number.clone()));
		header.insert("data_offset".to_string(), Value::Number(self.data_offset.clone()));
		header.insert("flags".to_string(), Value::Number(self.flags.clone()));
		header.insert("window_size".to_string(), Value::Number(self.window_size.clone()));
		header.insert("checksum".to_string(), Value::Number(self.checksum.clone()));
		header.insert("urgent_pointer".to_string(), Value::Number(self.urgent_pointer.clone()));

		header
	}
}

struct UdpHeader {
	source_port: Number,
	destination_port: Number,
	length: Number,
	checksum: Number,
}

impl UdpHeader {
	fn new(p: &udp::UdpPacket) -> UdpHeader {
		UdpHeader {
			source_port: Number::from(p.get_source()),
			destination_port: Number::from(p.get_destination()),
			length: Number::from(p.get_length()),
			checksum: Number::from(p.get_checksum()),
		}
	}

	fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

		header.insert("source_port".to_string(), Value::Number(self.source_port.clone()));
		header.insert("destination_port".to_string(), Value::Number(self.destination_port.clone()));
		header.insert("length".to_string(), Value::Number(self.length.clone()));
		header.insert("checksum".to_string(), Value::Number(self.checksum.clone()));

		header
	}
}

struct DnsHeader {
	id: Number,
	total_questions: Number,
	total_answer_rrs: Number,
	total_authority_rrs: Number,
	total_additional_rrs: Number,
}

impl DnsHeader {
	fn new(p: &DnsPacket) -> DnsHeader {
		DnsHeader {
			id: Number::from(p.header.id),
			total_questions: Number::from(p.header.questions),
			total_answer_rrs: Number::from(p.header.answers),
			total_authority_rrs: Number::from(p.header.nameservers),
			total_additional_rrs: Number::from(p.header.additional),
		}
	}

	fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

		header.insert("id".to_string(), Value::Number(self.id.clone()));
		header.insert("total_questions".to_string(), Value::Number(self.total_questions.clone()));
		header.insert("total_answer_rrs".to_string(), Value::Number(self.total_answer_rrs.clone()));
		header.insert("total_authority_rrs".to_string(), Value::Number(self.total_authority_rrs.clone()));
		header.insert("total_additional_rrs".to_string(), Value::Number(self.total_additional_rrs.clone()));

		header
	}
}

fn parse_headers(p: pcap::Packet) -> Map<String, Value> {
	let mut headers = Map::new();

	let timestamp = p.header.ts.tv_sec;

	let ethernet_packet = ethernet::EthernetPacket::new(p.data).unwrap();
			let ethernet_header = EthernetHeader::new(&ethernet_packet).to_json_map();

	        match ethernet_packet.get_ethertype() {
	        	ethernet::EtherTypes::Ipv4 => {
	        		let ipv4_packet = ipv4::Ipv4Packet::new(ethernet_packet.payload()).unwrap();
					let ipv4_header = Ipv4Header::new(&ipv4_packet).to_json_map();

	        		match ipv4_packet.get_next_level_protocol() {
	        			ip::IpNextHeaderProtocols::Tcp => {
	        				let tcp_packet = tcp::TcpPacket::new(ipv4_packet.payload()).unwrap();
	        				let tcp_header = TcpHeader::new(&tcp_packet).to_json_map();
	        				headers.insert("tcp".to_string(), Value::Object(tcp_header));

	        				match DnsPacket::parse(tcp_packet.payload()) {
	        					Ok(dns_packet) => {
	        						let dns_header = DnsHeader::new(&dns_packet).to_json_map();
	        					    headers.insert("dns".to_string(), Value::Object(dns_header));
	        					},
	        					Err(_) => (),
	        				}
	        			},
	        			ip::IpNextHeaderProtocols::Udp => {
	        				let udp_packet = udp::UdpPacket::new(ipv4_packet.payload()).unwrap();
	        				let udp_header = UdpHeader::new(&udp_packet).to_json_map();
	        				headers.insert("udp".to_string(), Value::Object(udp_header));

	        				match DnsPacket::parse(udp_packet.payload()) {
	        					Ok(dns_packet) => {
	        						let dns_header = DnsHeader::new(&dns_packet).to_json_map();
	        						headers.insert("dns".to_string(), Value::Object(dns_header));
	        					},
	        					Err(_) => (),
	        				}
	        			},
	        			_ => (),
	        		}
	        		headers.insert("ipv4".to_string(), Value::Object(ipv4_header));
	        	},
	        	_ => (),
	        }
	        headers.insert("ethernet".to_string(), Value::Object(ethernet_header));
	        headers.insert("timestamp".to_string(), Value::Number(Number::from(timestamp)));

	headers
}

fn sniff(interface: &str, promiscuous: bool, snaplen: i32, timeout: i32, filter: &str) {
    let mut capture = Capture::from_device(interface).unwrap()
                      	.promisc(promiscuous)
                        .snaplen(snaplen)
                        .timeout(timeout)
                        .open()
                        .unwrap();

    if !filter.is_empty() {
    	capture.filter(filter).unwrap();
    }

    while let Ok(packet) = capture.next() {
    	let packet_headers = parse_headers(packet);
		println!("{}", serde_json::to_string(&packet_headers).unwrap());
    }
    println!();
}

fn main() {
	const VERSION: &str = "0.1.0";
	const AUTHOR: &str = "Kevin Brebanov <kevin.brebanov@gmail.com>";
	const ABOUT: &str = "DNS packet sniffer that outputs to JSON";

    let arguments = App::new("snout")
                            .version(VERSION)
                            .author(AUTHOR)
                            .about(ABOUT)
                            .arg(Arg::with_name("interface")
                                .short("i")
                                .long("interface")
                                .value_name("INTERFACE")
                                .required(true)
                                .help("Network interface to sniff packets on")
                                .takes_value(true))
                            .arg(Arg::with_name("promiscuous")
                                .short("p")
                                .long("promiscuous")
                                .required(false)
                                .help("Enable promiscuous mode"))
                            .arg(Arg::with_name("snaplen")
                                .short("s")
                                .long("snaplen")
                                .value_name("SNAPLEN")
                                .required(false)
                                .help("Snapshot length")
                                .takes_value(true))
                            .arg(Arg::with_name("timeout")
                                .short("t")
                                .long("timeout")
                                .value_name("TIMEOUT")
                                .required(false)
                                .help("Read timeout in milliseconds")
                                .takes_value(true))
                            .arg(Arg::with_name("filter")
                                 .short("f")
                                 .long("filter")
                                 .value_name("FILTER")
                                 .required(false)
                                 .help("Berkeley Packet Filter (BPF)")
                                 .takes_value(true))
                            .get_matches();

    let interface = arguments.value_of("interface").unwrap();
    let mut promiscuous = false;
    if arguments.is_present("promiscuous") {
        promiscuous = true;
    }
    let snaplen: i32 = arguments.value_of("snaplen").unwrap_or("65535").parse().unwrap();
    let timeout: i32 = arguments.value_of("timeout").unwrap_or("0").parse().unwrap();
    let filter: &str = arguments.value_of("filter").unwrap_or("");

    sniff(interface, promiscuous, snaplen, timeout, filter);
}
<commit_msg>use if let syntax for DNS packet parsing<commit_after>extern crate clap;
extern crate pcap;
extern crate pnet;
extern crate serde_json;
extern crate dns_parser;

use clap::{Arg, App};
use pcap::Capture;
use pnet::packet::{Packet, ethernet, ip, ipv4, tcp, udp};
use serde_json::{Value, Map, Number};
use dns_parser::Packet as DnsPacket;

struct EthernetHeader {
	source: String,
	destination: String,
	ethertype: String,
}

impl EthernetHeader {
	fn new(p: &ethernet::EthernetPacket) -> EthernetHeader {
		EthernetHeader {
			source: p.get_source().to_string(),
			destination: p.get_destination().to_string(),
			ethertype: p.get_ethertype().to_string().to_lowercase(),
    	}
	}

	fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

		header.insert("source".to_string(), Value::String(self.source.clone()));
		header.insert("destination".to_string(), Value::String(self.destination.clone()));
		header.insert("type".to_string(), Value::String(self.ethertype.clone()));

		header
	}
}

struct Ipv4Header {
	version: Number,
	ihl: Number,
	dscp: Number,
	ecn: Number,
	total_length: Number,
	identification: Number,
	flags: Number,
	fragment_offset: Number,
	ttl: Number,
	next_level_protocol: String,
	checksum: Number,
	source: String,
	destination: String,
}

impl Ipv4Header {
	fn new(p: &ipv4::Ipv4Packet) -> Ipv4Header {
		Ipv4Header {
			version: Number::from(p.get_version()),
			ihl: Number::from(p.get_header_length()),
			dscp: Number::from(p.get_dscp()),
			ecn: Number::from(p.get_ecn()),
			total_length: Number::from(p.get_total_length()),
			identification: Number::from(p.get_identification()),
			flags: Number::from(p.get_flags()),
			fragment_offset: Number::from(p.get_fragment_offset()),
			ttl: Number::from(p.get_ttl()),
			next_level_protocol: p.get_next_level_protocol().to_string().to_lowercase(),
			checksum: Number::from(p.get_checksum()),
			source: p.get_source().to_string(),
			destination: p.get_destination().to_string(),
		}
	}

	fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

    	header.insert("version".to_string(), Value::Number(self.version.clone()));
    	header.insert("ihl".to_string(), Value::Number(self.ihl.clone()));
    	header.insert("dscp".to_string(), Value::Number(self.dscp.clone()));
    	header.insert("ecn".to_string(), Value::Number(self.ecn.clone()));
    	header.insert("total_length".to_string(), Value::Number(self.total_length.clone()));
    	header.insert("identification".to_string(), Value::Number(self.identification.clone()));
    	header.insert("flags.".to_string(), Value::Number(self.flags.clone()));
    	header.insert("fragment_offset".to_string(), Value::Number(self.fragment_offset.clone()));
    	header.insert("ttl".to_string(), Value::Number(self.ttl.clone()));
    	header.insert("protocol".to_string(), Value::String(self.next_level_protocol.clone()));
    	header.insert("checksum".to_string(), Value::Number(self.checksum.clone()));
    	header.insert("source_address".to_string(), Value::String(self.source.clone()));
    	header.insert("destination_address".to_string(), Value::String(self.destination.clone()));

		header
	}
}

struct TcpHeader {
	source_port: Number,
	destination_port: Number,
	sequence_number: Number,
	ack_number: Number,
	data_offset: Number,
	flags: Number,
	window_size: Number,
	checksum: Number,
	urgent_pointer: Number,
}

impl TcpHeader {
	fn new(p: &tcp::TcpPacket) -> TcpHeader {
		TcpHeader {
			source_port: Number::from(p.get_source()),
			destination_port: Number::from(p.get_destination()),
			sequence_number: Number::from(p.get_sequence()),
			ack_number: Number::from(p.get_acknowledgement()),
			data_offset: Number::from(p.get_data_offset()),
			flags: Number::from(p.get_flags()),
			window_size: Number::from(p.get_window()),
			checksum: Number::from(p.get_checksum()),
			urgent_pointer: Number::from(p.get_urgent_ptr()),
		}
	}

	fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

		header.insert("source_port".to_string(), Value::Number(self.source_port.clone()));
		header.insert("destination_port".to_string(), Value::Number(self.destination_port.clone()));
		header.insert("sequence_number".to_string(), Value::Number(self.sequence_number.clone()));
		header.insert("ack_number".to_string(), Value::Number(self.ack_number.clone()));
		header.insert("data_offset".to_string(), Value::Number(self.data_offset.clone()));
		header.insert("flags".to_string(), Value::Number(self.flags.clone()));
		header.insert("window_size".to_string(), Value::Number(self.window_size.clone()));
		header.insert("checksum".to_string(), Value::Number(self.checksum.clone()));
		header.insert("urgent_pointer".to_string(), Value::Number(self.urgent_pointer.clone()));

		header
	}
}

struct UdpHeader {
	source_port: Number,
	destination_port: Number,
	length: Number,
	checksum: Number,
}

impl UdpHeader {
	fn new(p: &udp::UdpPacket) -> UdpHeader {
		UdpHeader {
			source_port: Number::from(p.get_source()),
			destination_port: Number::from(p.get_destination()),
			length: Number::from(p.get_length()),
			checksum: Number::from(p.get_checksum()),
		}
	}

	fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

		header.insert("source_port".to_string(), Value::Number(self.source_port.clone()));
		header.insert("destination_port".to_string(), Value::Number(self.destination_port.clone()));
		header.insert("length".to_string(), Value::Number(self.length.clone()));
		header.insert("checksum".to_string(), Value::Number(self.checksum.clone()));

		header
	}
}

struct DnsHeader {
	id: Number,
	total_questions: Number,
	total_answer_rrs: Number,
	total_authority_rrs: Number,
	total_additional_rrs: Number,
}

impl DnsHeader {
	fn new(p: &DnsPacket) -> DnsHeader {
		DnsHeader {
			id: Number::from(p.header.id),
			total_questions: Number::from(p.header.questions),
			total_answer_rrs: Number::from(p.header.answers),
			total_authority_rrs: Number::from(p.header.nameservers),
			total_additional_rrs: Number::from(p.header.additional),
		}
	}

	fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

		header.insert("id".to_string(), Value::Number(self.id.clone()));
		header.insert("total_questions".to_string(), Value::Number(self.total_questions.clone()));
		header.insert("total_answer_rrs".to_string(), Value::Number(self.total_answer_rrs.clone()));
		header.insert("total_authority_rrs".to_string(), Value::Number(self.total_authority_rrs.clone()));
		header.insert("total_additional_rrs".to_string(), Value::Number(self.total_additional_rrs.clone()));

		header
	}
}

fn parse_headers(p: pcap::Packet) -> Map<String, Value> {
	let mut headers = Map::new();

	let timestamp = p.header.ts.tv_sec;

	let ethernet_packet = ethernet::EthernetPacket::new(p.data).unwrap();
			let ethernet_header = EthernetHeader::new(&ethernet_packet).to_json_map();

	        match ethernet_packet.get_ethertype() {
	        	ethernet::EtherTypes::Ipv4 => {
	        		let ipv4_packet = ipv4::Ipv4Packet::new(ethernet_packet.payload()).unwrap();
					let ipv4_header = Ipv4Header::new(&ipv4_packet).to_json_map();

	        		match ipv4_packet.get_next_level_protocol() {
	        			ip::IpNextHeaderProtocols::Tcp => {
	        				let tcp_packet = tcp::TcpPacket::new(ipv4_packet.payload()).unwrap();
	        				let tcp_header = TcpHeader::new(&tcp_packet).to_json_map();
	        				headers.insert("tcp".to_string(), Value::Object(tcp_header));

							if let Ok(dns_packet) = DnsPacket::parse(tcp_packet.payload()) {
								let dns_header = DnsHeader::new(&dns_packet).to_json_map();
	        					headers.insert("dns".to_string(), Value::Object(dns_header));
							}
	        			},
	        			ip::IpNextHeaderProtocols::Udp => {
	        				let udp_packet = udp::UdpPacket::new(ipv4_packet.payload()).unwrap();
	        				let udp_header = UdpHeader::new(&udp_packet).to_json_map();
	        				headers.insert("udp".to_string(), Value::Object(udp_header));

							if let Ok(dns_packet) = DnsPacket::parse(udp_packet.payload()) {
								let dns_header = DnsHeader::new(&dns_packet).to_json_map();
	        					headers.insert("dns".to_string(), Value::Object(dns_header));
							}
	        			},
	        			_ => (),
	        		}
	        		headers.insert("ipv4".to_string(), Value::Object(ipv4_header));
	        	},
	        	_ => (),
	        }
	        headers.insert("ethernet".to_string(), Value::Object(ethernet_header));
	        headers.insert("timestamp".to_string(), Value::Number(Number::from(timestamp)));

	headers
}

fn sniff(interface: &str, promiscuous: bool, snaplen: i32, timeout: i32, filter: &str) {
    let mut capture = Capture::from_device(interface).unwrap()
                      	.promisc(promiscuous)
                        .snaplen(snaplen)
                        .timeout(timeout)
                        .open()
                        .unwrap();

    if !filter.is_empty() {
    	capture.filter(filter).unwrap();
    }

    while let Ok(packet) = capture.next() {
    	let packet_headers = parse_headers(packet);
		println!("{}", serde_json::to_string(&packet_headers).unwrap());
    }
    println!();
}

fn main() {
	const VERSION: &str = "0.1.0";
	const AUTHOR: &str = "Kevin Brebanov <kevin.brebanov@gmail.com>";
	const ABOUT: &str = "DNS packet sniffer that outputs to JSON";

    let arguments = App::new("snout")
                            .version(VERSION)
                            .author(AUTHOR)
                            .about(ABOUT)
                            .arg(Arg::with_name("interface")
                                .short("i")
                                .long("interface")
                                .value_name("INTERFACE")
                                .required(true)
                                .help("Network interface to sniff packets on")
                                .takes_value(true))
                            .arg(Arg::with_name("promiscuous")
                                .short("p")
                                .long("promiscuous")
                                .required(false)
                                .help("Enable promiscuous mode"))
                            .arg(Arg::with_name("snaplen")
                                .short("s")
                                .long("snaplen")
                                .value_name("SNAPLEN")
                                .required(false)
                                .help("Snapshot length")
                                .takes_value(true))
                            .arg(Arg::with_name("timeout")
                                .short("t")
                                .long("timeout")
                                .value_name("TIMEOUT")
                                .required(false)
                                .help("Read timeout in milliseconds")
                                .takes_value(true))
                            .arg(Arg::with_name("filter")
                                 .short("f")
                                 .long("filter")
                                 .value_name("FILTER")
                                 .required(false)
                                 .help("Berkeley Packet Filter (BPF)")
                                 .takes_value(true))
                            .get_matches();

    let interface = arguments.value_of("interface").unwrap();
    let mut promiscuous = false;
    if arguments.is_present("promiscuous") {
        promiscuous = true;
    }
    let snaplen: i32 = arguments.value_of("snaplen").unwrap_or("65535").parse().unwrap();
    let timeout: i32 = arguments.value_of("timeout").unwrap_or("0").parse().unwrap();
    let filter: &str = arguments.value_of("filter").unwrap_or("");

    sniff(interface, promiscuous, snaplen, timeout, filter);
}
<|endoftext|>


<commit_before>#![allow(unstable, dead_code)]

extern crate libc;

mod args;
mod ansi;
mod matching;
mod screen;

use std::os;
use std::io;
use std::cmp::min;
use screen::Screen;
use screen::Key::*;

fn main() {
  let args = match args::parse_args() {
    Some(args) => args,
    None => {
      os::set_exit_status(1);
      return;
    },
  };

  if args.help { return; }

  let choices = read_choices();
  let mut search = args.initial_search.clone();
  if args.use_first {
    let matches = matching::compute_matches(&choices, search.as_slice());
    println!("{}", matches[0]);
    return;
  }

  let mut screen = Screen::open_screen();
  let mut index = 0;

  let visible_choices = min(20, screen.height - 1);

  let start_line = screen.height - visible_choices - 1;
  loop {
    screen.hide_cursor();
    screen.blank_screen(start_line);
    screen.move_cursor(start_line, 0);
    let matches = matching::compute_matches(&choices, search.as_slice());
    let mut i = 1;
    screen.write(format!("> {} ({} choices)\n", search.as_slice(), choices.len()).as_slice());
    for choice in matches.iter() {
      if i == index + 1 {
        screen.write_inverted(choice.as_slice());
      } else {
        screen.write(choice.as_slice());
      }
      if i >= visible_choices as usize {
        break;
      } else {
        screen.write("\n");
      }
      i += 1;
    }

    screen.move_cursor(start_line, 2 + search.len() as u16);
    screen.show_cursor();

    match screen.tty.getchar() {
      Char(x) => {
        search.push(x);
        index = 0;
      }
      Backspace => { search.pop(); }
      Control('h') => { search.pop(); }
      Control('u') => { search.clear(); }
      Control('c') => { return; }
      Control('n') => { index = min(index + 1, min(visible_choices as usize - 1, matches.len() - 1)); }
      Control('p') => { index = if index == 0 { 0 } else { index - 1 }; }
      Enter => {
        screen.move_cursor(start_line + visible_choices, 0);
        screen.write("\n");
        println!("{}", matches[index]);
        break;
      }
      _ => panic!("Unexpected input"),
    }
  }
}

fn read_choices() -> Vec<String> {
  let mut stdin = io::stdio::stdin();
  let mut lines = Vec::new();

  while let Ok(mut s) = stdin.read_line() {
    trim(&mut s);
    lines.push(s);
  }

  lines
}

fn trim(s: &mut String) {
  while let Some(x) = s.pop() {
    if x != '\n' {
      s.push(x);
      return;
    }
  }
}
<commit_msg>Blank screen after calling compute_matches<commit_after>#![allow(unstable, dead_code)]

extern crate libc;

mod args;
mod ansi;
mod matching;
mod screen;

use std::os;
use std::io;
use std::cmp::min;
use screen::Screen;
use screen::Key::*;

fn main() {
  let args = match args::parse_args() {
    Some(args) => args,
    None => {
      os::set_exit_status(1);
      return;
    },
  };

  if args.help { return; }

  let choices = read_choices();
  let mut search = args.initial_search.clone();
  if args.use_first {
    let matches = matching::compute_matches(&choices, search.as_slice());
    println!("{}", matches[0]);
    return;
  }

  let mut screen = Screen::open_screen();
  let mut index = 0;

  let visible_choices = min(20, screen.height - 1);

  let start_line = screen.height - visible_choices - 1;
  loop {
    let matches = matching::compute_matches(&choices, search.as_slice());
    screen.hide_cursor();
    screen.blank_screen(start_line);
    screen.move_cursor(start_line, 0);
    let mut i = 1;
    screen.write(format!("> {} ({} choices)\n", search.as_slice(), choices.len()).as_slice());
    for choice in matches.iter() {
      if i == index + 1 {
        screen.write_inverted(choice.as_slice());
      } else {
        screen.write(choice.as_slice());
      }
      if i >= visible_choices as usize {
        break;
      } else {
        screen.write("\n");
      }
      i += 1;
    }

    screen.move_cursor(start_line, 2 + search.len() as u16);
    screen.show_cursor();

    match screen.tty.getchar() {
      Char(x) => {
        search.push(x);
        index = 0;
      }
      Backspace => { search.pop(); }
      Control('h') => { search.pop(); }
      Control('u') => { search.clear(); }
      Control('c') => { return; }
      Control('n') => { index = min(index + 1, min(visible_choices as usize - 1, matches.len() - 1)); }
      Control('p') => { index = if index == 0 { 0 } else { index - 1 }; }
      Enter => {
        screen.move_cursor(start_line + visible_choices, 0);
        screen.write("\n");
        println!("{}", matches[index]);
        break;
      }
      _ => panic!("Unexpected input"),
    }
  }
}

fn read_choices() -> Vec<String> {
  let mut stdin = io::stdio::stdin();
  let mut lines = Vec::new();

  while let Ok(mut s) = stdin.read_line() {
    trim(&mut s);
    lines.push(s);
  }

  lines
}

fn trim(s: &mut String) {
  while let Some(x) = s.pop() {
    if x != '\n' {
      s.push(x);
      return;
    }
  }
}
<|endoftext|>


<commit_before>//use std::sync::Weak;
extern crate rand;
use rand::{thread_rng, Rng};

#[derive(PartialEq, Eq, Debug, Copy, Clone)]
enum PlayerColor {
    White,
    Black
}

#[allow(dead_code)] // Dead code is allowed, because I want x, y, z even if I don't use it.
#[derive(Debug)]
struct Point {
    x : i8,
    y : i8,
    z : i8,
    flat_coordinate : i8,
    lines : Vec<i8> // This vector stores the IDs of all lines through it.
}

impl Point {
    // constructor, by convention
    pub fn new(x:i8, y:i8, z:i8) -> Point {
        Point {x:x, y:y, z:z, flat_coordinate:flatten(x, y, z), lines:Vec::new()} //, lines : Vec::new()}
    }
}

struct Line {
    points : [Point; 4]
}

impl Line {
    // constructor, by convention
    pub fn new(x:i8, y:i8, z:i8, dx:i8, dy:i8, dz:i8) -> Line {
        let point1 = Point::new(x, y, z);
        let point2 = Point::new(x+1*dx, y+1*dy, z+1*dz);
        let point3 = Point::new(x+2*dx, y+2*dy, z+2*dz);
        let point4 = Point::new(x+3*dx, y+3*dy, z+3*dz);
        Line { points : [point1, point2, point3, point4]}
    }
}

#[derive(PartialEq, Eq, Debug, Copy, Clone)]
enum PointState {
    Piece(PlayerColor),
    Empty
}

#[derive(PartialEq, Eq, Debug, Clone, Copy)]
enum LineState {
    Empty,
    Pure { color: PlayerColor, count: i8 },
    Mixed,
    Win(PlayerColor),
}

struct GameState {
    points : [PointState; 64],
    lines  : [LineState; 76]  // something something mutable?
}

impl GameState {
    pub fn new() -> GameState {
        GameState { points : [PointState::Empty; 64], lines : [LineState::Empty; 76]}
    }
}

fn flatten(x:i8, y:i8, z:i8) -> i8 {
    return x + 4*y + 16*z
}

// fn expand(index:i8) -> (i8, i8, i8) {
//     return (index % 4, (index / 4) % 4, index / 16)
// }

fn add_ball(line_state : LineState, new_color : PlayerColor) -> LineState {
    match line_state {
        LineState::Empty => LineState::Pure { color : new_color, count : 1},
        LineState::Pure { color : current_color, count : old_count} =>
            if current_color != new_color { LineState::Mixed } else {
                if old_count == 3 {LineState::Win(current_color)}
                else {LineState::Pure {color : current_color, count : old_count+1}}
            },
        LineState::Mixed => LineState::Mixed,
        LineState::Win(_) => panic!("A filled line can't accept any more balls.")
    }
}

fn z_value(game_state : &GameState, x : i8, y : i8) -> Option<i8> {
    for z in 0..4 {
        if game_state.points[flatten(x, y, z) as usize] == PointState::Empty {
            return Some(z)
        }
    }
    return None
}

fn legal_moves(game_state : &GameState) -> Vec<(i8, i8)> {
    let mut result = Vec::new();
    for x in 0..4 {
        for y in 0..4 {
            let height = z_value(game_state, x, y);
            match height {
                Some(_) => result.push((x, y)),
                None => ()
            }
        }
    }
    return result;
}

fn main() {
    // Initialize a vector of all Points.
    // This has the property, that point_box[i].flat_coordinate = i.
    let mut point_box = Vec::new();
    for z in 0..4 {
        for y in 0..4 {
            for x in 0..4 {
                point_box.push(Point::new(x, y, z));
            }
        }
    }
    // Check if this really produced the right amount of points.
    assert_eq!(point_box.len(), 64);
    let mut i = 0;
    for p in &point_box {
        assert_eq!(p.flat_coordinate, i);
        i += 1;
    }

    // Initialize a vector of all Lines.
    let mut line_box = Vec::new();
    for a in 0..4 {
        for b in 0..4 {
            line_box.push(Line::new(a, b, 0, 0, 0, 1));
            line_box.push(Line::new(0, a, b, 1, 0, 0));
            line_box.push(Line::new(b, 0, a, 0, 1, 0))
        }
        // Diagonals in two spacial directions
        line_box.push(Line::new(a, 0, 0, 0, 1, 1));
        line_box.push(Line::new(0, a, 0, 1, 0, 1));
        line_box.push(Line::new(0, 0, a, 1, 1, 0));
        line_box.push(Line::new(a, 3, 0, 0,-1, 1));
        line_box.push(Line::new(0, a, 3, 1, 0,-1));
        line_box.push(Line::new(3, 0, a,-1, 1, 0));
    }
    // Diagonals in all three directions at once
    line_box.push(Line::new(0, 0, 0, 1, 1, 1));
    line_box.push(Line::new(3, 0, 0,-1, 1, 1));
    line_box.push(Line::new(3, 3, 0,-1,-1, 1));
    line_box.push(Line::new(0, 3, 0, 1,-1, 1));

    // Verify if the line_box has the right length.
    assert_eq!(line_box.len(), 76);

    // Refenence the line ID in the points.
    let mut i = 0;
    for line in line_box {
        for point in line.points.iter() {
            point_box[point.flat_coordinate as usize].lines.push(i);
        }
        i += 1;
    }

    // Do some testing with the LineState Code
    let mut test_state = LineState::Empty;
    test_state = add_ball(test_state, PlayerColor::White);
    println!("{:?}", test_state);
    test_state = add_ball(test_state, PlayerColor::White);
    println!("{:?}", test_state);
    test_state = add_ball(test_state, PlayerColor::White);
    println!("{:?}", test_state);
    test_state = add_ball(test_state, PlayerColor::Black);
    println!("{:?}", test_state);

    // for point in point_box {
    //     println!("{}", point.lines.len());
    // }

    println!("{:?}", z_value(&GameState::new(), 0, 2));

    fn play_at(state : &mut GameState, x:i8, y:i8, color : PlayerColor) {//-> GameState {
        let z = z_value(&state, x, y);
        let flat_coordinate = match z {
            Some(z) => flatten(x, y, z),
            None => panic!("Added a ball on a forbidden row")
        };
        state.points[flat_coordinate as usize] = PointState::Piece(color);
        //return state; // Here I only return it to hand back ownership to the calling code.
    }

    let mut mystate = GameState::new();
    println!("{:?}", z_value(&mystate, 0, 2));
    {
        let moves = legal_moves(&mystate);
        let position = thread_rng().choose(&moves);
        match position {
            Some(&(x, y)) => play_at(&mut mystate, x, y, PlayerColor::White),
            None => panic!("can't play on that board")
        }
    }
    println!("{:?}", z_value(&mystate, 0, 2));

    let choices = [1, 2, 4, 8, 16, 32];
    let mut rng = thread_rng();
    println!("{:?}", rng.choose(&point_box));
    println!("{:?}", rng.choose(&point_box));
    println!("{:?}", rng.choose(&point_box));
    println!("{:?}", rng.choose(&point_box));
    println!("{:?}", rng.choose(&point_box));
    println!("{:?}", rng.choose(&point_box));
    assert_eq!(rng.choose(&choices[..0]), None);
}
<commit_msg>encapsulate the game structure and enable randomly playing balls.<commit_after>//use std::sync::Weak;
extern crate rand;
use rand::{thread_rng, Rng};

#[derive(PartialEq, Eq, Debug, Copy, Clone)]
enum PlayerColor {
    White,
    Black
}

fn flip_color(c : PlayerColor) -> PlayerColor {
    match c {
        PlayerColor::White => PlayerColor::Black,
        PlayerColor::Black => PlayerColor::White
    }
}

#[allow(dead_code)] // Dead code is allowed, because I want x, y, z even if I don't use it.
#[derive(Debug)]
struct Point {
    x : i8,
    y : i8,
    z : i8,
    flat_coordinate : i8,
    lines : Vec<i8> // This vector stores the IDs of all lines through it.
}

impl Point {
    // constructor, by convention
    pub fn new(x:i8, y:i8, z:i8) -> Point {
        Point {x:x, y:y, z:z, flat_coordinate:flatten(x, y, z), lines:Vec::new()} //, lines : Vec::new()}
    }
}

struct Line {
    points : [Point; 4]
}

impl Line {
    // constructor, by convention
    pub fn new(x:i8, y:i8, z:i8, dx:i8, dy:i8, dz:i8) -> Line {
        let point1 = Point::new(x, y, z);
        let point2 = Point::new(x+1*dx, y+1*dy, z+1*dz);
        let point3 = Point::new(x+2*dx, y+2*dy, z+2*dz);
        let point4 = Point::new(x+3*dx, y+3*dy, z+3*dz);
        Line { points : [point1, point2, point3, point4]}
    }
}

#[derive(PartialEq, Eq, Debug, Copy, Clone)]
enum PointState {
    Piece(PlayerColor),
    Empty
}

#[derive(PartialEq, Eq, Debug, Clone, Copy)]
enum LineState {
    Empty,
    Pure { color: PlayerColor, count: i8 },
    Mixed,
    Win(PlayerColor),
}

struct GameState {
    points : [PointState; 64],
    lines  : [LineState; 76],  // something something mutable?
    current_color : PlayerColor
}

impl GameState {
    pub fn new() -> GameState {
        GameState {
            points : [PointState::Empty; 64],
            lines : [LineState::Empty; 76],
            current_color : PlayerColor::White
        }
    }
}

struct GameStructure {
    points : Vec<Point>,
    lines  : Vec<Line>
}

impl GameStructure {
    pub fn new() -> GameStructure {
        // Initialize a vector of all Points.
        let mut point_box = Vec::new();
        for z in 0..4 {
            for y in 0..4 {
                for x in 0..4 {
                    point_box.push(Point::new(x, y, z));
                }
            }
        }
        // Check if this really produced the right amount of points.
        // Assert the property, that point_box[i].flat_coordinate = i.
        assert_eq!(point_box.len(), 64);
        let mut i = 0;
        for p in &point_box {
            assert_eq!(p.flat_coordinate, i);
            i += 1;
        }

        // Initialize a vector of all Lines.
        let mut line_box = Vec::new();
        for a in 0..4 {
            for b in 0..4 {
                line_box.push(Line::new(a, b, 0, 0, 0, 1));
                line_box.push(Line::new(0, a, b, 1, 0, 0));
                line_box.push(Line::new(b, 0, a, 0, 1, 0))
            }
            // Diagonals in two spacial directions
            line_box.push(Line::new(a, 0, 0, 0, 1, 1));
            line_box.push(Line::new(0, a, 0, 1, 0, 1));
            line_box.push(Line::new(0, 0, a, 1, 1, 0));
            line_box.push(Line::new(a, 3, 0, 0,-1, 1));
            line_box.push(Line::new(0, a, 3, 1, 0,-1));
            line_box.push(Line::new(3, 0, a,-1, 1, 0));
        }
        // Diagonals in all three directions at once
        line_box.push(Line::new(0, 0, 0, 1, 1, 1));
        line_box.push(Line::new(3, 0, 0,-1, 1, 1));
        line_box.push(Line::new(3, 3, 0,-1,-1, 1));
        line_box.push(Line::new(0, 3, 0, 1,-1, 1));

        // Verify if the line_box has the right length.
        assert_eq!(line_box.len(), 76);

        // Refenence the line ID in the points.
        let mut i = 0;
        for line in &line_box {
            for point in line.points.iter() {
                point_box[point.flat_coordinate as usize].lines.push(i);
            }
            i += 1;
        }
        GameStructure { points : point_box, lines : line_box }
    }
}

fn flatten(x:i8, y:i8, z:i8) -> i8 {
    return x + 4*y + 16*z
}

// fn expand(index:i8) -> (i8, i8, i8) {
//     return (index % 4, (index / 4) % 4, index / 16)
// }

fn add_ball(line_state : LineState, new_color : PlayerColor) -> LineState {
    match line_state {
        LineState::Empty => LineState::Pure { color : new_color, count : 1},
        LineState::Pure { color : current_color, count : old_count} =>
            if current_color != new_color { LineState::Mixed } else {
                if old_count == 3 {LineState::Win(current_color)}
                else {LineState::Pure {color : current_color, count : old_count+1}}
            },
        LineState::Mixed => LineState::Mixed,
        LineState::Win(_) => panic!("A filled line can't accept any more balls.")
    }
}

fn z_value(game_state : &GameState, x : i8, y : i8) -> Option<i8> {
    for z in 0..4 {
        if game_state.points[flatten(x, y, z) as usize] == PointState::Empty {
            return Some(z)
        }
    }
    return None
}

fn legal_moves(game_state : &GameState) -> Vec<(i8, i8)> {
    let mut result = Vec::new();
    for x in 0..4 {
        for y in 0..4 {
            let height = z_value(game_state, x, y);
            match height {
                Some(_) => result.push((x, y)),
                None => ()
            }
        }
    }
    return result;
}

fn play_at(structure : &GameStructure, state : &mut GameState, x:i8, y:i8) {
    let z = z_value(&state, x, y);
    let flat_coordinate = match z {
        Some(z) => flatten(x, y, z),
        None => panic!("Added a ball on a forbidden row")
    };
    state.points[flat_coordinate as usize] = PointState::Piece(state.current_color);
    for line in structure.points[flat_coordinate as usize].lines.clone() {
        state.lines[line as usize] = add_ball(state.lines[line as usize], state.current_color);
    }
    state.current_color = flip_color(state.current_color);
}

fn play_at_random(structure : &GameStructure, state : &mut GameState) {
    let moves = legal_moves(&state);
    let position = thread_rng().choose(&moves);
    match position {
        Some(&(x, y)) => play_at(structure, state, x, y),
        None => panic!("can't play on that board")
    }
}

fn main() {
    let game_structure = GameStructure::new();

    let mut mystate = GameState::new();
    println!("{:?}", z_value(&mystate, 0, 2));
    for _ in 0..40 {
        play_at_random(&game_structure, &mut mystate);
    }
    println!("{:?}", z_value(&mystate, 0, 2));
    println!("{:?}", mystate.lines[4]);
}
<|endoftext|>