<commit_before>extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate docopt;

use docopt::Docopt;

use std::thread;

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_shared::*;
use timely::example_shared::operators::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

use columnar::Columnar;

mod typedrw;
mod graphmap;
use graphmap::GraphMMap;

mod sorting;
use sorting::{SegmentList, radix_sort_32};

static USAGE: &'static str = "
Usage: pagerank <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      file containing list of host:port for workers
";

#[derive(Copy, Clone, Debug)]
struct NodeRank { node: u32, rank: f32, }
impl Columnar for NodeRank { type Stack = Vec<NodeRank>; }

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_spawn(communicators, source);
    }
    else if workers > 1 { pagerank_spawn(ProcessCommunicator::new_vector(workers), source); }
    else { pagerank_spawn(vec![ThreadCommunicator], source); };
}

fn pagerank_spawn<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank_thread(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

// returns [src/peers] degrees, (dst, deg) pairs, and a list of [src/peers] endpoints
fn transpose(mut edges: Vec<Vec<(u32, u32)>>, peers: usize, nodes: usize) -> (Vec<u32>, Vec<(u32, u32)>, Vec<u32>)  {

    let mut deg = vec![0; (nodes as usize / peers) + 1];
    for list in &edges {
        for &(s, _) in list {
            deg[s as usize / peers] += 1;
        }
    }

    // let start = time::precise_time_s();
    radix_sort_32(&mut edges, &mut Vec::new(), &|&(_,d)| d);
    // println!("sorted in {}s", time::precise_time_s() - start);

    let mut rev = Vec::<(u32,u32)>::with_capacity(deg.len());
    let mut trn = Vec::with_capacity(edges.len() * 1024);
    for list in edges {
        for (s,d) in list {
            let len = rev.len();
            if (len == 0) || (rev[len-1].0 < d) {
                rev.push((d, 0u32));
            }

            let len = rev.len();
            rev[len-1].1 += 1;
            trn.push(s / peers as u32);
        }
    }

    return (deg, rev, trn);
}

// pagerank dataflow graph has a set of edges as input, and a binary vertex that for each epoch of
// received edges initiates an iterative subcomputation to compute the pagerank.

fn pagerank_thread<C>(communicator: C, filename: String, _workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    let start = time::precise_time_s();

    let nodes = GraphMMap::new(&filename).nodes();

    let mut segments = SegmentList::new(1024); // list of edge segments

    let mut src = vec![];   // holds ranks
    let mut deg = vec![];   // holds source degrees
    let mut rev = vec![];   // holds (dst, deg) pairs
    let mut trn = vec![];   // holds transposed sources

    let mut going = start;

    let mut input = root.subcomputation(|builder| {

        let (input, edges) = builder.new_input::<(u32, u32)>();
        let (cycle, ranks) = builder.loop_variable::<NodeRank>(RootTimestamp::new(20), Local(1));

        edges.binary_notify(&ranks,
                            Exchange::new(|x: &(u32,u32)| x.0 as u64),
                            Exchange::new(|x: &NodeRank| x.node as u64),
                            format!("pagerank"),
                            vec![RootTimestamp::new(0)],
                            move |input1, input2, output, notificator| {

            // receive incoming edges
            while let Some((_index, data)) = input1.pull() {
                segments.push(data.drain_temp());
            }

            while let Some((iter, _)) = notificator.next() {

                // if the very first iteration, prepare some stuff
                if iter.inner == 0 {
                    let (a, b, c) = transpose(segments.finalize(), peers, nodes);
                    deg = a; rev = b; trn = c;
                    src = vec![0.0f32; deg.len()];
                }

                if iter.inner == 10 && index == 0 { going = time::precise_time_s(); }
                if iter.inner == 20 && index == 0 { println!("average: {}", (time::precise_time_s() - going) / 10.0 ); }

                for s in 0..src.len() { src[s] = 0.15 + 0.85 * src[s] / deg[s] as f32; }

                let mut trn_slice = &trn[..];
                let mut rev_slice = &rev[..];
                while rev_slice.len() > 0 {
                    // TODO: session should just flush...
                    let mut session = output.session(&iter);
                    let next = std::cmp::min(200_000, rev_slice.len());
                    for &(dst, deg) in &rev_slice[..next] {
                        let mut accum = 0.0;
                        for &s in &trn_slice[..deg as usize] {
                            accum += src[s as usize];
                        }
                        trn_slice = &trn_slice[deg as usize..];
                        session.give(NodeRank { node: dst, rank: accum });
                    }
                    rev_slice = &rev_slice[next..];
                }

                for s in &mut src { *s = 0.0; }
            }

            while let Some((iter, data)) = input2.pull() {
                notificator.notify_at(&iter);
                for x in data.drain_temp() {
                    src[x.node as usize / peers] += x.rank;
                }
            }
        })
        .connect_loop(cycle);

        input
    });

    {
        let graph = GraphMMap::new(&filename);

        let mut edges = Vec::new();
        for node in 0..graph.nodes() {
            if node % peers == index {
                for dst in graph.edges(node) {
                    edges.push((node as u32, *dst as u32));
                }
                if edges.len() > 100_000 {
                    input.send_at(0, edges.drain_temp());
                    root.step();
                }
            }
        }

        input.send_at(0, edges.drain_temp());
    }
    input.close();
    while root.step() { }

    if index == 0 { println!("elapsed: {}", time::precise_time_s() - start); }
}

    // .unary_notify(
    //     Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
    //     format!("Aggregation"),
    //     vec![],
    //     move |input, output, iterator| {
    //         while let Some((iter, data)) = input.pull() {
    //             iterator.notify_at(&iter);
    //             for (node, rank) in data.drain_temp() {
    //                 acc[node as usize / workers] += rank;
    //             }
    //         }
    //
    //         while let Some((item, _)) = iterator.next() {
    //
    //             output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
    //                                      .map(|(u,f)| (((u * workers + local_index) as u32), f)));
    //
    //             for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
    //             assert!(acc.len() == (1 + (nodes/workers)));
    //         }
    //     }
    // )
<commit_msg>-n and -h allow a prefix of the hosts file<commit_after>extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate docopt;

use docopt::Docopt;

use std::fs::File;
use std::thread;
use std::io::{Read, BufRead, BufReader};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_shared::*;
use timely::example_shared::operators::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;

use timely::drain::DrainExt;

use columnar::Columnar;

mod typedrw;
mod graphmap;
use graphmap::GraphMMap;

mod sorting;
use sorting::{SegmentList, radix_sort_32};

static USAGE: &'static str = "
Usage: pagerank <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      file containing list of host:port for workers
";

#[derive(Copy, Clone, Debug)]
struct NodeRank { node: u32, rank: f32, }
impl Columnar for NodeRank { type Stack = Vec<NodeRank>; }

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let addresses: Vec<_> = if hosts != "" {
            let reader = BufReader::new(File::open(hosts).unwrap());
            reader.lines().take(processes as usize).map(|x| x.unwrap()).collect()
        }
        else {
            (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect()
        };

        if addresses.len() != processes as usize { panic!("only {} hosts for -p: {}", addresses.len(), processes); }

        let communicators = initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking");

        pagerank_spawn(communicators, source);
    }
    else if workers > 1 { pagerank_spawn(ProcessCommunicator::new_vector(workers), source); }
    else { pagerank_spawn(vec![ThreadCommunicator], source); };
}

fn pagerank_spawn<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank_thread(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

// returns [src/peers] degrees, (dst, deg) pairs, and a list of [src/peers] endpoints
fn transpose(mut edges: Vec<Vec<(u32, u32)>>, peers: usize, nodes: usize) -> (Vec<u32>, Vec<(u32, u32)>, Vec<u32>)  {

    let mut deg = vec![0; (nodes as usize / peers) + 1];
    for list in &edges {
        for &(s, _) in list {
            deg[s as usize / peers] += 1;
        }
    }

    radix_sort_32(&mut edges, &mut Vec::new(), &|&(_,d)| d);

    let mut rev = Vec::<(u32,u32)>::with_capacity(deg.len());
    let mut trn = Vec::with_capacity(edges.len() * 1024);
    for list in edges {
        for (s,d) in list {
            let len = rev.len();
            if (len == 0) || (rev[len-1].0 < d) {
                rev.push((d, 0u32));
            }

            let len = rev.len();
            rev[len-1].1 += 1;
            trn.push(s / peers as u32);
        }
    }

    return (deg, rev, trn);
}

// pagerank dataflow graph has a set of edges as input, and a binary vertex that for each epoch of
// received edges initiates an iterative subcomputation to compute the pagerank.

fn pagerank_thread<C>(communicator: C, filename: String, _workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    let start = time::precise_time_s();

    let nodes = GraphMMap::new(&filename).nodes();

    let mut segments = SegmentList::new(1024); // list of edge segments

    let mut src = vec![];   // holds ranks
    let mut deg = vec![];   // holds source degrees
    let mut rev = vec![];   // holds (dst, deg) pairs
    let mut trn = vec![];   // holds transposed sources

    let mut going = start;

    let mut input = root.subcomputation(|builder| {

        let (input, edges) = builder.new_input::<(u32, u32)>();
        let (cycle, ranks) = builder.loop_variable::<NodeRank>(RootTimestamp::new(20), Local(1));

        edges.binary_notify(&ranks,
                            Exchange::new(|x: &(u32,u32)| x.0 as u64),
                            Exchange::new(|x: &NodeRank| x.node as u64),
                            format!("pagerank"),
                            vec![RootTimestamp::new(0)],
                            move |input1, input2, output, notificator| {

            // receive incoming edges (should only be iter 0)
            while let Some((_index, data)) = input1.pull() {
                segments.push(data.drain_temp());
            }

            // all inputs received for iter, commence multiplication
            while let Some((iter, _)) = notificator.next() {

                // if the very first iteration, prepare some stuff.
                // specifically, transpose edges and sort by destination.
                if iter.inner == 0 {
                    let (a, b, c) = transpose(segments.finalize(), peers, nodes);
                    deg = a; rev = b; trn = c;
                    src = vec![0.0f32; deg.len()];
                }

                // record some timings in order to estimate per-iteration times
                if iter.inner == 10 && index == 0 { going = time::precise_time_s(); }
                if iter.inner == 20 && index == 0 { println!("average: {}", (time::precise_time_s() - going) / 10.0 ); }

                // prepare src for transmitting to destinations
                for s in 0..src.len() { src[s] = 0.15 + 0.85 * src[s] / deg[s] as f32; }

                // wander through destinations
                let mut trn_slice = &trn[..];
                let mut rev_slice = &rev[..];
                while rev_slice.len() > 0 {
                    // TODO: session should flush automatically...
                    let mut session = output.session(&iter);
                    let next = std::cmp::min(200_000, rev_slice.len());
                    for &(dst, deg) in &rev_slice[..next] {
                        let mut accum = 0.0;
                        for &s in &trn_slice[..deg as usize] {
                            accum += src[s as usize];
                        }
                        trn_slice = &trn_slice[deg as usize..];
                        session.give(NodeRank { node: dst, rank: accum });
                    }
                    rev_slice = &rev_slice[next..];
                }

                // clean out src to accumulate into
                for s in &mut src { *s = 0.0; }
            }

            // receive data from workers, accumulate in src
            while let Some((iter, data)) = input2.pull() {
                notificator.notify_at(&iter);
                for x in data.drain_temp() {
                    src[x.node as usize / peers] += x.rank;
                }
            }
        })
        .connect_loop(cycle);

        input
    });

    {
        let graph = GraphMMap::new(&filename);

        let mut edges = Vec::new();
        for node in 0..graph.nodes() {
            if node % peers == index {
                for dst in graph.edges(node) {
                    edges.push((node as u32, *dst as u32));
                }
                if edges.len() > 100_000 {
                    input.send_at(0, edges.drain_temp());
                    root.step();
                }
            }
        }

        input.send_at(0, edges.drain_temp());
    }
    input.close();
    while root.step() { }

    if index == 0 { println!("elapsed: {}", time::precise_time_s() - start); }
}

    // .unary_notify(
    //     Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
    //     format!("Aggregation"),
    //     vec![],
    //     move |input, output, iterator| {
    //         while let Some((iter, data)) = input.pull() {
    //             iterator.notify_at(&iter);
    //             for (node, rank) in data.drain_temp() {
    //                 acc[node as usize / workers] += rank;
    //             }
    //         }
    //
    //         while let Some((item, _)) = iterator.next() {
    //
    //             output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
    //                                      .map(|(u,f)| (((u * workers + local_index) as u32), f)));
    //
    //             for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
    //             assert!(acc.len() == (1 + (nodes/workers)));
    //         }
    //     }
    // )
<|endoftext|>


<commit_before>extern crate html5ever;
extern crate html5ever_dom_sink;
extern crate hyper;
extern crate iron;
extern crate rustc_serialize;
extern crate tempfile;
extern crate tendril;
extern crate url;

mod json;

use std::fs::OpenOptions;
use std::io::Read;
use std::io::Write;
use std::process::Command;
use std::process::Stdio;
use std::str::FromStr;

use html5ever::Attribute;
use html5ever::tree_builder::interface::TreeSink;
use html5ever_dom_sink::common::NodeEnum;
use html5ever_dom_sink::rcdom::Handle;
use html5ever_dom_sink::rcdom::RcDom;
use hyper::Client;
use hyper::header::Connection;
use iron::Iron;
use iron::IronResult;
use iron::Request;
use iron::Response;
use iron::headers::ContentType;
use iron::mime::Mime;
use iron::mime::SubLevel;
use iron::mime::TopLevel;
use rustc_serialize::json::Json;
use tempfile::NamedTempFile;
use url::percent_encoding;

use json::JsonPathElement::{Key, Only};

// TODO: consider doing s/en.wikipedia.org/this app's url/ on the HTML before serving it. This
// currently works fine, but might not over HTTPS.
// TODO: Do nested functions work? If so, use those where appropriate.
// TODO: there are some places where I've handrolled try!() equivalents. Fix those.
// TODO: make sure I'm returning Results everywhere, and propagating errors correctly. Remove all
// uses of unwrap() that might panic.

// TODO: return a Result
fn call_wikimedia_api(parameters: Vec<(&str, &str)>) -> String {
    let post_body =
        parameters.into_iter().map(|p| format!("{}={}", p.0, p.1))
        .collect::<Vec<_>>().join("&") + "&format=json";

    let client = Client::new();
    let mut res = client.post("https://en.wikipedia.org/w/api.php")
        .body(&post_body)
        .header(Connection::close())
        .send().unwrap();

    let mut body = String::new();
    res.read_to_string(&mut body).unwrap();
    body
}

// TODO: this name is terrible.
// Returns pairs of (revid, parentid)
fn get_revert_revision_ids(page: &str) -> Result<Vec<(u64, u64)>, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", page),
             ("rvprop", "comment|ids"), ("rvlimit", "60")]);

    let json = Json::from_str(&json_str).unwrap();
    let revisions = try!(
        json::get_json_array(&json, &[Key("query"), Key("pages"), Only, Key("revisions")]));

    // Filter for revisions that mention vandalism, and that have revid and parentid (which they all
    // should, but merits checking).
    Ok(revisions.into_iter().filter(|revision| {
        match (json::get_json_string(revision, &[Key("comment")]),
               json::get_json_number(revision, &[Key("revid")]),
               json::get_json_number(revision, &[Key("parentid")])) {
            (Ok(ref comment), Ok(_), Ok(_)) => comment.contains("vandal"),
            _ => false,
        }})
       .map(|revision| {
           (json::get_json_number(revision, &[Key("revid")]).unwrap(),
            json::get_json_number(revision, &[Key("parentid")]).unwrap())
       }).collect())
}

fn get_latest_revision_id(page: &str) -> Result<u64, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", page), ("rvprop", "ids"),
             ("rvlimit", "1")]);
    let json = Json::from_str(&json_str).unwrap();
    json::get_json_number(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("revid")])
}

fn get_revision_content(page: &str, id: u64) -> Result<String, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", page), ("rvprop", "content"),
             ("rvlimit", "1"), ("rvstartid", &id.to_string())]);
    let json = Json::from_str(&json_str).unwrap();
    match json::get_json_string(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("*")]) {
        Ok(content) => Ok(content.to_string()),
        Err(msg) => Err(msg),
    }
}

fn write_to_temp_file(contents: &str) -> NamedTempFile {
    let tempfile = NamedTempFile::new().unwrap();
    let mut file = OpenOptions::new().write(true).open(tempfile.path()).unwrap();
    file.write_all(contents.as_bytes()).unwrap();
    file.flush().unwrap();
    tempfile
}

// TODO: I'm not so sure these parameter names aren't terrible.
fn merge(old: &str, new1: &str, new2: &str) -> Option<String> {
    let old_tempfile = write_to_temp_file(old);
    let new1_tempfile = write_to_temp_file(new1);
    let new2_tempfile = write_to_temp_file(new2);

    let mut process = Command::new("diff3");
    process.arg("-m").args(&[new1_tempfile.path(), old_tempfile.path(), new2_tempfile.path()])
        .stdout(Stdio::piped()).stderr(Stdio::null());
    let output = process.output().unwrap();
    if output.status.success() {
        Some(String::from_utf8(output.stdout).unwrap())
    } else {
        None
    }
}

fn render(wikitext: &str) -> Result<String, String> {
    let encoded_wikitext =
        percent_encoding::percent_encode(wikitext.as_bytes(), percent_encoding::QUERY_ENCODE_SET);
    let html = call_wikimedia_api(
        vec![("action", "parse"), ("prop", "text"), ("disablepp", ""), ("contentmodel", "wikitext"),
             ("text", &encoded_wikitext)]);
    // TODO: check return value
    let json = Json::from_str(&html).unwrap();
    match json::get_json_string(&json, &[Key("parse"), Key("text"), Key("*")]) {
        Ok(contents) => Ok(contents.to_string()),
        Err(message) => Err(message),
    }
}

fn has_matching_id(attributes: &Vec<Attribute>, id: &str) -> bool {
    // TODO: could also do this with filter() and is_empty(). Not sure if that would be better.
    for attribute in attributes {
        // TODO: do I seriously have to construct a StrTendril here?
        // There has to be a better way.
        if attribute.name.local.as_slice() == "id" &&
            attribute.value == tendril::StrTendril::from_str(id).unwrap() {
                return true;
            }
    }
    return false;
}

fn find_node_by_id(handle: &Handle, id: &str) -> Result<Handle, String> {
    let node = handle.borrow();
    match node.node {
        NodeEnum::Element(_, ref attributes) if has_matching_id(attributes, id) => Ok(handle.clone()),
        _ => {
            for child in &node.children {
                match find_node_by_id(child, id) {
                    // TODO: this looks weird. Is this an abnormal way to do things?
                    Ok(node) => return Ok(node),
                    _ => continue,
                }
            }
            Err(format!("No node with ID {} found", id))
        },
    }
}

fn replace_node_with_placeholder(original_html: &str, div_id: &str, placeholder: &str) -> Result<String, String> {
    // TODO: check errors
    let html = tendril::StrTendril::from_str(original_html).unwrap();
    let mut dom: RcDom = html5ever::parse(html5ever::one_input(html), Default::default());

    let handle = try!(find_node_by_id(&dom.get_document(), div_id));
    let child_handles =
        (&handle.borrow().children).into_iter().map(|child| child.clone()).collect::<Vec<_>>();
    for child_handle in child_handles {
        dom.remove_from_parent(child_handle);
    }
    dom.append(handle,
               html5ever::tree_builder::interface::NodeOrText::AppendText(
                   tendril::StrTendril::from_str(placeholder).unwrap()));
    let mut serialized: Vec<u8> = vec![];
    html5ever::serialize::serialize(&mut serialized, &dom.document, Default::default());
    // TODO: error handling
    Ok(String::from_utf8(serialized).unwrap())
}

fn get_current_page(page: &str) -> String {
    let client = Client::new();
    let mut res = client.get(
        &format!("https://en.wikipedia.org/wiki/{}", page))
        .header(Connection::close())
        .send().unwrap();

    let mut body = String::new();
    res.read_to_string(&mut body).unwrap();

    body
}

// TODO: this name is terrible.
fn get_page_with_vandalism_restored(page: &str) -> Result<String, String> {
    let latest_revid = get_latest_revision_id(page).unwrap();
    let mut accumulated_contents = get_revision_content(page, latest_revid).unwrap();
    let mut revisions = vec![];
    for (revert_revid, vandalism_revid) in get_revert_revision_ids(page).ok().unwrap() {
        let reverted_contents = get_revision_content(page, revert_revid).unwrap();
        let vandalized_contents = get_revision_content(page, vandalism_revid).unwrap();
        match merge(&reverted_contents, &vandalized_contents, &accumulated_contents) {
            Some(merged_contents) => {
                // TODO: replace this with a log statement, if there's a good logging framework.
                println!(
                    "For page \"{}\", restored vandalism https://en.wikipedia.org/w/index.php?title=Zachary_Taylor&diff=prev&oldid={}",
                         page, revert_revid);
                accumulated_contents = merged_contents;
                revisions.push(revert_revid);
            }
            None => (),
        }
    }

    // TODO: replace this with a log statement, if there's a good logging framework.
    println!("For page \"{}\", restored vandalisms reverted in: {:?}", page, revisions);

    let rendered_body = render(&accumulated_contents).unwrap();
    let current_page_contents = get_current_page(page);
    // TODO: randomize the placeholder string per-request
    let page_contents_with_placeholder =
        replace_node_with_placeholder(
            &current_page_contents, "mw-content-text", "WMW_PLACEHOLDER_TEXT").unwrap();
    Ok(page_contents_with_placeholder.replace("WMW_PLACEHOLDER_TEXT", &rendered_body))
}

fn serve_request(request: &mut Request) -> IronResult<Response> {
    if request.url.path.len() == 2 && request.url.path[0] == "wiki" {
        let mut response = Response::with(
            (iron::status::Ok, get_page_with_vandalism_restored(&request.url.path[1]).unwrap()));
        response.headers.set(ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));
        Ok(response)
    } else {
        println!("Would forward request for {} to en.wikipedia.org", request.url.path.join("/"));
        Ok(Response::with((iron::status::Ok, "Would forward to en.wikipedia.org")))
    }
}

fn main() {
    Iron::new(serve_request).http("localhost:3000").unwrap();
}

#[cfg(test)]
mod tests {
    use super::{merge, replace_node_with_placeholder};

    #[test]
    fn test_merge_clean() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed.\n";
        let new2 = "First line changed.\n\nSecond line.\n";
        assert_eq!("First line changed.\n\nSecond line changed.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_merge_conflicting() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed one way.\n";
        let new2 = "First line changed.\n\nSecond line changed a different way.\n";
        assert_eq!(None, merge(old, new1, new2));
    }

    #[test]
    fn test_merge_special_characters() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line 𐅃.\n";
        let new2 = "First line さようなら.\n\nSecond line.\n";
        assert_eq!("First line さようなら.\n\nSecond line 𐅃.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_replace_html_content() {
        let original_html = "

original text

Other text

";
        let expected_html = "

replaced text

Other text

";
        let processed_html = replace_node_with_placeholder(original_html, "mw-content-text", "replaced text").unwrap();
        assert_eq!(expected_html, processed_html);
    }
}
<commit_msg>Forward non-article requests to Wikipedia.<commit_after>extern crate html5ever;
extern crate html5ever_dom_sink;
extern crate hyper;
extern crate iron;
extern crate rustc_serialize;
extern crate tempfile;
extern crate tendril;
extern crate url;

mod json;

use std::fs::OpenOptions;
use std::io::Read;
use std::io::Write;
use std::process::Command;
use std::process::Stdio;
use std::str::FromStr;

use html5ever::Attribute;
use html5ever::tree_builder::interface::TreeSink;
use html5ever_dom_sink::common::NodeEnum;
use html5ever_dom_sink::rcdom::Handle;
use html5ever_dom_sink::rcdom::RcDom;
use hyper::Client;
use hyper::header::Connection;
use iron::Iron;
use iron::IronResult;
use iron::Request;
use iron::Response;
use iron::headers::ContentType;
use iron::mime::Mime;
use iron::mime::SubLevel;
use iron::mime::TopLevel;
use rustc_serialize::json::Json;
use tempfile::NamedTempFile;
use url::percent_encoding;

use json::JsonPathElement::{Key, Only};

// TODO: consider doing s/en.wikipedia.org/this app's url/ on the HTML before serving it. This
// currently works fine, but might not over HTTPS.
// TODO: Do nested functions work? If so, use those where appropriate.
// TODO: there are some places where I've handrolled try!() equivalents. Fix those.
// TODO: make sure I'm returning Results everywhere, and propagating errors correctly. Remove all
// uses of unwrap() that might panic.

// TODO: return a Result
fn call_wikimedia_api(parameters: Vec<(&str, &str)>) -> String {
    let post_body =
        parameters.into_iter().map(|p| format!("{}={}", p.0, p.1))
        .collect::<Vec<_>>().join("&") + "&format=json";

    let client = Client::new();
    // TODO: rename to "response".
    let mut res = client.post("https://en.wikipedia.org/w/api.php")
        .body(&post_body)
        .header(Connection::close())
        .send().unwrap();

    let mut body = String::new();
    res.read_to_string(&mut body).unwrap();
    body
}

// TODO: this name is terrible.
// Returns pairs of (revid, parentid)
fn get_revert_revision_ids(page: &str) -> Result<Vec<(u64, u64)>, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", page),
             ("rvprop", "comment|ids"), ("rvlimit", "60")]);

    let json = Json::from_str(&json_str).unwrap();
    let revisions = try!(
        json::get_json_array(&json, &[Key("query"), Key("pages"), Only, Key("revisions")]));

    // Filter for revisions that mention vandalism, and that have revid and parentid (which they all
    // should, but merits checking).
    Ok(revisions.into_iter().filter(|revision| {
        match (json::get_json_string(revision, &[Key("comment")]),
               json::get_json_number(revision, &[Key("revid")]),
               json::get_json_number(revision, &[Key("parentid")])) {
            (Ok(ref comment), Ok(_), Ok(_)) => comment.contains("vandal"),
            _ => false,
        }})
       .map(|revision| {
           (json::get_json_number(revision, &[Key("revid")]).unwrap(),
            json::get_json_number(revision, &[Key("parentid")]).unwrap())
       }).collect())
}

fn get_latest_revision_id(page: &str) -> Result<u64, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", page), ("rvprop", "ids"),
             ("rvlimit", "1")]);
    let json = Json::from_str(&json_str).unwrap();
    json::get_json_number(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("revid")])
}

fn get_revision_content(page: &str, id: u64) -> Result<String, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", page), ("rvprop", "content"),
             ("rvlimit", "1"), ("rvstartid", &id.to_string())]);
    let json = Json::from_str(&json_str).unwrap();
    match json::get_json_string(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("*")]) {
        Ok(content) => Ok(content.to_string()),
        Err(msg) => Err(msg),
    }
}

fn write_to_temp_file(contents: &str) -> NamedTempFile {
    let tempfile = NamedTempFile::new().unwrap();
    let mut file = OpenOptions::new().write(true).open(tempfile.path()).unwrap();
    file.write_all(contents.as_bytes()).unwrap();
    file.flush().unwrap();
    tempfile
}

// TODO: I'm not so sure these parameter names aren't terrible.
fn merge(old: &str, new1: &str, new2: &str) -> Option<String> {
    let old_tempfile = write_to_temp_file(old);
    let new1_tempfile = write_to_temp_file(new1);
    let new2_tempfile = write_to_temp_file(new2);

    let mut process = Command::new("diff3");
    process.arg("-m").args(&[new1_tempfile.path(), old_tempfile.path(), new2_tempfile.path()])
        .stdout(Stdio::piped()).stderr(Stdio::null());
    let output = process.output().unwrap();
    if output.status.success() {
        Some(String::from_utf8(output.stdout).unwrap())
    } else {
        None
    }
}

fn render(wikitext: &str) -> Result<String, String> {
    let encoded_wikitext =
        percent_encoding::percent_encode(wikitext.as_bytes(), percent_encoding::QUERY_ENCODE_SET);
    let html = call_wikimedia_api(
        vec![("action", "parse"), ("prop", "text"), ("disablepp", ""), ("contentmodel", "wikitext"),
             ("text", &encoded_wikitext)]);
    // TODO: check return value
    let json = Json::from_str(&html).unwrap();
    match json::get_json_string(&json, &[Key("parse"), Key("text"), Key("*")]) {
        Ok(contents) => Ok(contents.to_string()),
        Err(message) => Err(message),
    }
}

fn has_matching_id(attributes: &Vec<Attribute>, id: &str) -> bool {
    // TODO: could also do this with filter() and is_empty(). Not sure if that would be better.
    for attribute in attributes {
        // TODO: do I seriously have to construct a StrTendril here?
        // There has to be a better way.
        if attribute.name.local.as_slice() == "id" &&
            attribute.value == tendril::StrTendril::from_str(id).unwrap() {
                return true;
            }
    }
    return false;
}

fn find_node_by_id(handle: &Handle, id: &str) -> Result<Handle, String> {
    let node = handle.borrow();
    match node.node {
        NodeEnum::Element(_, ref attributes) if has_matching_id(attributes, id) => Ok(handle.clone()),
        _ => {
            for child in &node.children {
                match find_node_by_id(child, id) {
                    // TODO: this looks weird. Is this an abnormal way to do things?
                    Ok(node) => return Ok(node),
                    _ => continue,
                }
            }
            Err(format!("No node with ID {} found", id))
        },
    }
}

fn replace_node_with_placeholder(original_html: &str, div_id: &str, placeholder: &str) -> Result<String, String> {
    // TODO: check errors
    let html = tendril::StrTendril::from_str(original_html).unwrap();
    let mut dom: RcDom = html5ever::parse(html5ever::one_input(html), Default::default());

    let handle = try!(find_node_by_id(&dom.get_document(), div_id));
    let child_handles =
        (&handle.borrow().children).into_iter().map(|child| child.clone()).collect::<Vec<_>>();
    for child_handle in child_handles {
        dom.remove_from_parent(child_handle);
    }
    dom.append(handle,
               html5ever::tree_builder::interface::NodeOrText::AppendText(
                   tendril::StrTendril::from_str(placeholder).unwrap()));
    let mut serialized: Vec<u8> = vec![];
    html5ever::serialize::serialize(&mut serialized, &dom.document, Default::default());
    // TODO: error handling
    Ok(String::from_utf8(serialized).unwrap())
}

fn get_current_page(page: &str) -> String {
    let client = Client::new();
    let mut res = client.get(
        &format!("https://en.wikipedia.org/wiki/{}", page))
        .header(Connection::close())
        .send().unwrap();

    let mut body = String::new();
    res.read_to_string(&mut body).unwrap();

    body
}

// TODO: this name is terrible.
fn get_page_with_vandalism_restored(page: &str) -> Result<String, String> {
    let latest_revid = get_latest_revision_id(page).unwrap();
    let mut accumulated_contents = get_revision_content(page, latest_revid).unwrap();
    let mut revisions = vec![];
    for (revert_revid, vandalism_revid) in get_revert_revision_ids(page).ok().unwrap() {
        let reverted_contents = get_revision_content(page, revert_revid).unwrap();
        let vandalized_contents = get_revision_content(page, vandalism_revid).unwrap();
        match merge(&reverted_contents, &vandalized_contents, &accumulated_contents) {
            Some(merged_contents) => {
                // TODO: replace this with a log statement, if there's a good logging framework.
                println!(
                    "For page \"{}\", restored vandalism https://en.wikipedia.org/w/index.php?title=Zachary_Taylor&diff=prev&oldid={}",
                         page, revert_revid);
                accumulated_contents = merged_contents;
                revisions.push(revert_revid);
            }
            None => (),
        }
    }

    // TODO: replace this with a log statement, if there's a good logging framework.
    println!("For page \"{}\", restored vandalisms reverted in: {:?}", page, revisions);

    let rendered_body = render(&accumulated_contents).unwrap();
    let current_page_contents = get_current_page(page);
    // TODO: randomize the placeholder string per-request
    let page_contents_with_placeholder =
        replace_node_with_placeholder(
            &current_page_contents, "mw-content-text", "WMW_PLACEHOLDER_TEXT").unwrap();
    Ok(page_contents_with_placeholder.replace("WMW_PLACEHOLDER_TEXT", &rendered_body))
}

fn serve_request(request: &mut Request) -> IronResult<Response> {
    if request.url.path.len() == 2 && request.url.path[0] == "wiki" {
        let mut response = Response::with(
            (iron::status::Ok, get_page_with_vandalism_restored(&request.url.path[1]).unwrap()));
        response.headers.set(ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));
        Ok(response)
    } else {
        let client = Client::new();
        // TODO: error handling
        let mut wikipedia_response =
            client.get(&format!("https://en.wikipedia.org/{}", request.url.path.join("/")))
            .header(Connection::close())
            .send().unwrap();
        let mut wikipedia_body: Vec<u8> = Vec::new();
        wikipedia_response.read_to_end(&mut wikipedia_body);

        let mut response = Response::with(wikipedia_body);
        response.status = Some(wikipedia_response.status);
        response.headers = wikipedia_response.headers.clone();
        //response.body = wikipedia_body;
        println!("Forwarded request for {} to en.wikipedia.org", request.url.path.join("/"));
        Ok(response)
    }
}

fn main() {
    Iron::new(serve_request).http("localhost:3000").unwrap();
}

#[cfg(test)]
mod tests {
    use super::{merge, replace_node_with_placeholder};

    #[test]
    fn test_merge_clean() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed.\n";
        let new2 = "First line changed.\n\nSecond line.\n";
        assert_eq!("First line changed.\n\nSecond line changed.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_merge_conflicting() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed one way.\n";
        let new2 = "First line changed.\n\nSecond line changed a different way.\n";
        assert_eq!(None, merge(old, new1, new2));
    }

    #[test]
    fn test_merge_special_characters() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line 𐅃.\n";
        let new2 = "First line さようなら.\n\nSecond line.\n";
        assert_eq!("First line さようなら.\n\nSecond line 𐅃.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_replace_html_content() {
        let original_html = "

original text

Other text

";
        let expected_html = "

replaced text

Other text

";
        let processed_html = replace_node_with_placeholder(original_html, "mw-content-text", "replaced text").unwrap();
        assert_eq!(expected_html, processed_html);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate time;
extern crate image;
extern crate threadpool;
extern crate rustc_serialize;

use clap::{Arg, App};
use time::PreciseTime;
use rustc_serialize::json;
use threadpool::ThreadPool;

use std::fs::File;
use std::path::Path;
use std::error::Error;
use std::io::prelude::*;
use std::sync::mpsc::channel;

#[derive(Debug, RustcDecodable, RustcEncodable)]
struct Size {
    width: f32,
    height: f32,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
struct Point {
    x: f32,
    y: f32,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
struct Rect {
    origin: Point,
    size: Size,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
struct Config {
    ppu: u32,
    limit: u32,
    color_steps: f32,
    color_components: u8,
    color_palette: Vec<Vec<f32>>,
    window: Rect,
}

fn idx2point(idx: u32, width: u32) -> Point {
    let x = idx % width;
    let y = idx / width;
    Point { x: x as f32, y: y as f32}
}

fn point2idx(p: Point, width: u32) -> u32 {
    width * (p.y as u32) + (p.x as u32)
}

fn mandelbrot(cx: f32, cy: f32, limit: u32) -> u32 {
    let mut x = cx;
    let mut y = cy;
    let mut count = 0;
    while count < limit {
        let xy = x * y;
        let xx = x * x;
        let yy = y * y;
        let sum = xx + yy;
        if sum > 4.0 {
            break
        }
        count += 1;
        x = xx - yy + cx;
        y = xy * 2.0 + cy;
    }
    count as u32
}

fn gen_mandelbrot(size: &Size, config: &Config) -> Vec<u32> {
    let window = &config.window;
    let limit = config.limit;

    let data_size = size.width as u32 * size.height as u32;
    let mut data = Vec::with_capacity(data_size as usize);

    for i in 0..data_size {
        let p = idx2point(i, size.width as u32);

        let px: f32 = p.x / size.width;
        let py: f32 = p.y / size.height;

        let cx = window.origin.x + px * window.size.width;
        let cy = (window.origin.y + window.size.height) - py * window.size.height;

        let c = mandelbrot(cx, cy, limit);

        data.push(c);
    }

    data
}

fn rbg_from_palette(palette: &Vec<Vec<f32>>, idx: usize) -> (f32, f32, f32) {
    let color = &palette[idx];
    (color[0], color[1], color[2])
}

fn color_for_val_with_config(val: u32, config: &Config) -> (u8, u8, u8) {
    let (r, g, b);

    let limit = config.limit;
    let steps = config.color_steps;
    let palette = &config.color_palette;

    if val == limit as u32 {
        r = 0;
        g = 0;
        b = 0;
    } else {
        let val = (val as f32 % steps) * (palette.len() as f32) / steps;
        let left = val as usize % palette.len();
        let right = (left + 1) % palette.len();

        let p = val - left as f32;
        let (r1, g1, b1) = rbg_from_palette(palette, left);
        let (r2, g2, b2) = rbg_from_palette(palette, right);
        r = (r1 + (r2 - r1) * p) as u8;
        g = (g1 + (g2 - g1) * p) as u8;
        b = (b1 + (b2 - b1) * p) as u8;
    }
    (r, g, b)
}

fn validate_config(conf: &Config) {

    // Check if limit is 'to large'
    if conf.limit > 10000 {
        panic!("Config Error: limit is over 10,000");
    }

    // Check if color component count is valid
    match conf.color_components {
        3 => println!("Using RBG colors"),
        _ => panic!("Unsuported color component count"),
    };

    // Check colors in palette match component count
    for v in &conf.color_palette {
        if v.len() != conf.color_components as usize {
            panic!("Config Error: Color {:?} does not match color component count", v);
        }
    }
}

fn main() {
    let start = PreciseTime::now();

    let args = App::new("Mandelbrot Generator")
        .version(&crate_version!()[..])
        .author("DJ Edmonson <djedmonson@gmail.com>")
        .about("Generates a mandelbrot image")
        .arg(Arg::with_name("CONFIG")
             .long("config")
             .help("Config JSON file to use. Output will be at <input_file_path>.png")
             .required(true)
             .takes_value(true))
        .arg(Arg::with_name("output-palette")
             .long("output-palette")
             .help("Generate image with 100px squares of the provided colors in order. Outputs to <input_file_path>-palette.png"))
        .get_matches();

    let config_file_path = Path::new(args.value_of("CONFIG").unwrap());
    println!("Getting config from {}", config_file_path.display());

    let mut config_file = match File::open(&config_file_path) {
        Err(why) => panic!("Could not open config file at {}: {}",
                           config_file_path.display(),
                           Error::description(&why)),
        Ok(f) => f,
    };

    let mut config_json = String::new();
    match config_file.read_to_string(&mut config_json) {
        Err(why) => panic!("Could not read config file at {}: {}",
                           config_file_path.display(),
                           Error::description(&why)),
        Ok(_) => println!("Read config file"),
    };

    let config: Config = match json::decode(&config_json) {
        Err(why) => panic!("Error parsing config JSON: {}",
                           why),
        Ok(conf) => conf,
    };

    validate_config(&config);

    println!("Bootstrap time:\n{}", start.to(PreciseTime::now()));

    if args.is_present("output-palette") {
        let root = match config_file_path.file_stem() {
            None => unreachable!(),
            Some(r) => r.to_str().unwrap().to_string(),
        };

        let output_path_string = root + "-palette.png";
        let output_path = Path::new(&output_path_string);
        println!("Generating color palette image at {}", output_path.display());

        let width = config.color_palette.len() * 100;
        let height = 100;

        let mut imgbuf = image::ImageBuffer::new(width as u32, height as u32);

        for (x, _, pixel) in imgbuf.enumerate_pixels_mut() {
            let color_idx = (x as f32 / width as f32) * config.color_palette.len() as f32;

            let (r, b, g) = rbg_from_palette(&config.color_palette, color_idx as usize);

            *pixel = image::Rgb([r as u8, b as u8, g as u8]);
        }

        let _ = imgbuf.save(output_path);
    } else {
        let output_path = config_file_path.with_extension("png");

        let img_width = config.ppu as f32 * config.window.size.width;
        let img_height = config.ppu as f32 * config.window.size.height;

        println!("Generating image at {} with size {}x{}", output_path.display(), img_width, img_height);

        let size = Size {width: img_width, height: img_height};

        let render_start = PreciseTime::now();
        let mut phase_start = PreciseTime::now();

        let imgdata = gen_mandelbrot(&size, &config);

        println!("Generation Duration:\n{}", phase_start.to(PreciseTime::now()));

        phase_start = PreciseTime::now();

        let mut imgbuf = image::ImageBuffer::new(size.width as u32, size.height as u32);

        for (x, y, pixel) in imgbuf.enumerate_pixels_mut() {
            let idx = point2idx(Point{ x: x as f32, y: y as f32}, size.width as u32) as usize;

            let it = imgdata[idx];

            let (r, g, b) = color_for_val_with_config(it, &config);

            *pixel = image::Rgb([r as u8, g as u8, b as u8]);
        }
        println!("Render Duration:\n{}", phase_start.to(PreciseTime::now()));
        phase_start = PreciseTime::now();

        let _ = imgbuf.save(output_path);
        println!("Save Duration:\n{}", phase_start.to(PreciseTime::now()));
        println!("Total Render Duration:\n{}", render_start.to(PreciseTime::now()));
    }

    println!("Total time:\n{}", start.to(PreciseTime::now()));
}

<commit_msg>Ouput PPM image, much faster<commit_after>#[macro_use]
extern crate clap;
extern crate time;
extern crate image;
extern crate threadpool;
extern crate rustc_serialize;

use clap::{Arg, App};
use time::PreciseTime;
use rustc_serialize::json;
use threadpool::ThreadPool;

use std::fs::File;
use std::path::Path;
use std::error::Error;
use std::io::prelude::*;
use std::sync::mpsc::channel;

#[derive(Debug, RustcDecodable, RustcEncodable)]
struct Size {
    width: f32,
    height: f32,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
struct Point {
    x: f32,
    y: f32,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
struct Rect {
    origin: Point,
    size: Size,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
struct Config {
    ppu: u32,
    limit: u32,
    color_steps: f32,
    color_components: u8,
    color_palette: Vec<Vec<f32>>,
    window: Rect,
}

fn idx2point(idx: u32, width: u32) -> Point {
    let x = idx % width;
    let y = idx / width;
    Point { x: x as f32, y: y as f32}
}

fn point2idx(p: Point, width: u32) -> u32 {
    width * (p.y as u32) + (p.x as u32)
}

fn mandelbrot(cx: f32, cy: f32, limit: u32) -> u32 {
    let mut x = cx;
    let mut y = cy;
    let mut count = 0;
    while count < limit {
        let xy = x * y;
        let xx = x * x;
        let yy = y * y;
        let sum = xx + yy;
        if sum > 4.0 {
            break
        }
        count += 1;
        x = xx - yy + cx;
        y = xy * 2.0 + cy;
    }
    count as u32
}

fn gen_mandelbrot(size: &Size, config: &Config) -> Vec<u32> {
    let window = &config.window;
    let limit = config.limit;

    let data_size = size.width as u32 * size.height as u32;
    let mut data = Vec::with_capacity(data_size as usize);

    for i in 0..data_size {
        let p = idx2point(i, size.width as u32);

        let px: f32 = p.x / size.width;
        let py: f32 = p.y / size.height;

        let cx = window.origin.x + px * window.size.width;
        let cy = (window.origin.y + window.size.height) - py * window.size.height;

        let c = mandelbrot(cx, cy, limit);

        data.push(c);
    }

    data
}

fn rbg_from_palette(palette: &Vec<Vec<f32>>, idx: usize) -> (f32, f32, f32) {
    let color = &palette[idx];
    (color[0], color[1], color[2])
}

fn color_for_val_with_config(val: u32, config: &Config) -> (u8, u8, u8) {
    let (r, g, b);

    let limit = config.limit;
    let steps = config.color_steps;
    let palette = &config.color_palette;

    if val == limit as u32 {
        r = 0;
        g = 0;
        b = 0;
    } else {
        let val = (val as f32 % steps) * (palette.len() as f32) / steps;
        let left = val as usize % palette.len();
        let right = (left + 1) % palette.len();

        let p = val - left as f32;
        let (r1, g1, b1) = rbg_from_palette(palette, left);
        let (r2, g2, b2) = rbg_from_palette(palette, right);
        r = (r1 + (r2 - r1) * p) as u8;
        g = (g1 + (g2 - g1) * p) as u8;
        b = (b1 + (b2 - b1) * p) as u8;
    }
    (r, g, b)
}

fn validate_config(conf: &Config) {

    // Check if limit is 'to large'
    if conf.limit > 10000 {
        panic!("Config Error: limit is over 10,000");
    }

    // Check if color component count is valid
    match conf.color_components {
        3 => println!("Using RBG colors"),
        _ => panic!("Unsuported color component count"),
    };

    // Check colors in palette match component count
    for v in &conf.color_palette {
        if v.len() != conf.color_components as usize {
            panic!("Config Error: Color {:?} does not match color component count", v);
        }
    }
}

fn main() {
    let start = PreciseTime::now();

    let args = App::new("Mandelbrot Generator")
        .version(&crate_version!()[..])
        .author("DJ Edmonson <djedmonson@gmail.com>")
        .about("Generates a mandelbrot image")
        .arg(Arg::with_name("CONFIG")
             .long("config")
             .help("Config JSON file to use. Output will be at <input_file_path>.png")
             .required(true)
             .takes_value(true))
        .arg(Arg::with_name("output-palette")
             .long("output-palette")
             .help("Generate image with 100px squares of the provided colors in order. Outputs to <input_file_path>-palette.png"))
        .get_matches();

    let config_file_path = Path::new(args.value_of("CONFIG").unwrap());
    println!("Getting config from {}", config_file_path.display());

    let mut config_file = match File::open(&config_file_path) {
        Err(why) => panic!("Could not open config file at {}: {}",
                           config_file_path.display(),
                           Error::description(&why)),
        Ok(f) => f,
    };

    let mut config_json = String::new();
    match config_file.read_to_string(&mut config_json) {
        Err(why) => panic!("Could not read config file at {}: {}",
                           config_file_path.display(),
                           Error::description(&why)),
        Ok(_) => println!("Read config file"),
    };

    let config: Config = match json::decode(&config_json) {
        Err(why) => panic!("Error parsing config JSON: {}",
                           why),
        Ok(conf) => conf,
    };

    validate_config(&config);

    println!("Bootstrap time:\n{}", start.to(PreciseTime::now()));

    if args.is_present("output-palette") {
        let root = match config_file_path.file_stem() {
            None => unreachable!(),
            Some(r) => r.to_str().unwrap().to_string(),
        };

        let output_path_string = root + "-palette.png";
        let output_path = Path::new(&output_path_string);
        println!("Generating color palette image at {}", output_path.display());

        let width = config.color_palette.len() * 100;
        let height = 100;

        let mut imgbuf = image::ImageBuffer::new(width as u32, height as u32);

        for (x, _, pixel) in imgbuf.enumerate_pixels_mut() {
            let color_idx = (x as f32 / width as f32) * config.color_palette.len() as f32;

            let (r, b, g) = rbg_from_palette(&config.color_palette, color_idx as usize);

            *pixel = image::Rgb([r as u8, b as u8, g as u8]);
        }

        let _ = imgbuf.save(output_path);
    } else {
        let output_path = config_file_path.with_extension("ppm");

        let mut output_file =  match File::create(&output_path) {
            Err(why) => panic!("Could not open output file at {}: {}",
                               output_path.display(),
                               Error::description(&why)),
            Ok(f) => f,
        };

        let img_width = config.ppu as f32 * config.window.size.width;
        let img_height = config.ppu as f32 * config.window.size.height;

        println!("Generating image at {} with size {}x{}", output_path.display(), img_width, img_height);

        let size = Size {width: img_width, height: img_height};

        let render_start = PreciseTime::now();
        let mut phase_start = PreciseTime::now();

        let imgdata = gen_mandelbrot(&size, &config);

        println!("Generation Duration:\n{}", phase_start.to(PreciseTime::now()));

        phase_start = PreciseTime::now();

        writeln!(output_file, "P6").unwrap();
        writeln!(output_file, "{} {}", img_width as usize, img_height as usize).unwrap();
        writeln!(output_file, "255").unwrap();

        let mut linebuf = vec![0; img_width as usize * 3];

        for y in 0 .. img_height as usize {
            for x in 0 .. img_width as usize {
                let idx = point2idx(Point{ x: x as f32, y: y as f32}, size.width as u32) as usize;

                let it = imgdata[idx];

                let (r, g, b) = color_for_val_with_config(it, &config);

                let offset = 3 * x;

                linebuf[offset]     = r;
                linebuf[offset + 1] = g;
                linebuf[offset + 2] = b
            }
            output_file.write(&linebuf).unwrap();
        }
        println!("Render Duration:\n{}", phase_start.to(PreciseTime::now()));
        println!("Total Render Duration:\n{}", render_start.to(PreciseTime::now()));
    }

    println!("Total time:\n{}", start.to(PreciseTime::now()));
}

<|endoftext|>


<commit_before>extern crate ppbert;
#[macro_use] extern crate clap;

use std::io::{self, BufReader, Read};
use std::fs::File;
use std::process::exit;
use std::time::Instant;

use clap::{Arg, App};

use ppbert::bertterm::{
    self,
    BertTerm,
    DEFAULT_INDENT_WIDTH,
    DEFAULT_MAX_TERMS_PER_LINE
};
use ppbert::error::Result;
use ppbert::parser;

const PROG_NAME: &str = "ppbert";

fn main() {
    let matches = App::new(PROG_NAME)
        .version(crate_version!())
        .author("Vincent Foley")
        .about("Pretty print structure encoded in Erlang's External Term Format")
        .arg(Arg::with_name("input_files")
             .value_name("FILES")
             .multiple(true))
        .arg(Arg::with_name("indent_width")
             .help("Indents with <num> spaces")
             .value_name("num")
             .short("i")
             .long("--indent-width")
             .takes_value(true))
        .arg(Arg::with_name("max_per_line")
             .help("Prints at most <num> basic terms per line")
             .value_name("num")
             .short("m")
             .long("--max-terms-per-line")
             .takes_value(true))
        .arg(Arg::with_name("verbose")
             .help("Enables verbose mode")
             .short("v")
             .long("--verbose"))
        .arg(Arg::with_name("parse")
             .help("Parses the input, doesn't pretty print it")
             .short("p")
             .long("--parse"))
        .arg(Arg::with_name("bert2")
             .help("Parses .bert2 files")
             .short("2")
             .long("bert2"))
        .arg(Arg::with_name("json")
             .help("Outputs in JSON")
             .short("j")
             .long("json"))
        .arg(Arg::with_name("transform-proplists")
             .help("Transforms proplists into JSON objects (only valid with --json)")
             .long("transform-proplists"))
        .get_matches();

    let files: Vec<&str> = match matches.values_of("input_files") {
        Some(files) => files.collect(),
        None => vec!["-"]
    };

    let indent_level = value_t!(matches, "indent_width", usize)
        .unwrap_or(DEFAULT_INDENT_WIDTH);
    let max_per_line = value_t!(matches, "max_per_line", usize)
        .unwrap_or(DEFAULT_MAX_TERMS_PER_LINE);
    let verbose = matches.is_present("verbose");
    let parse_only = matches.is_present("parse");

    let transform_proplists = matches.is_present("transform-proplists");
    let parse_fn = if matches.is_present("bert2") { parse_bert2 } else { parse_bert1 };
    let output_fn =
        if matches.is_present("json") {
            if transform_proplists {
                bertterm::pp_json_proplist
            } else {
                bertterm::pp_json
            }
        } else {
            if transform_proplists {
                eprintln!("{}: warning: --transform-proplists is only valid with the --json flag",
                          PROG_NAME);
            }
            bertterm::pp_bert
        };

    let mut return_code = 0;
    for file in files {
        let res = handle_file(file, parse_only, verbose,
                              indent_level, max_per_line,
                              parse_fn, output_fn);
        match res {
            Ok(()) => (),
            Err(e) => {
                return_code = 1;
                eprintln!("{}: {}: {}", PROG_NAME, file, e);
            }
        }
    }
    exit(return_code);
}

fn handle_file<T>(
    file: &str,
    parse_only: bool,
    verbose: bool,
    indent: usize,
    terms_per_line: usize,
    parse_fn: fn(Vec<u8>) -> Result<T>,
    pp_fn: fn(T, usize, usize) -> ()
) -> Result<()> {

    // Read file or stdin into buffer
    let mut buf: Vec<u8> = Vec::new();
    let now = Instant::now();
    if file == "-" {
        let stdin = io::stdin();
        let mut stdin = stdin.lock();
        stdin.read_to_end(&mut buf)?;
    } else {
        let f = File::open(file)?;
        let mut rdr = BufReader::new(f);
        rdr.read_to_end(&mut buf)?;
    }
    let dur0 = now.elapsed();
    if verbose {
        eprintln!("{}: read time: {}.{:09}",
                  PROG_NAME, dur0.as_secs(), dur0.subsec_nanos());
    }

    // Parse input
    let now = Instant::now();
    let parse_output = parse_fn(buf)?;
    let dur1 = now.elapsed();

    if verbose {
        eprintln!("{}: parse time: {}.{:09}",
                  PROG_NAME, dur1.as_secs(), dur1.subsec_nanos());
    }

    // Early exit if parse-only
    if parse_only {
        return Ok(());
    }

    // Pretty print
    let now = Instant::now();
    pp_fn(parse_output, indent, terms_per_line);
    let dur2 = now.elapsed();

    if verbose {
        eprintln!("{}: print time: {}.{:09}",
                  PROG_NAME, dur2.as_secs(), dur2.subsec_nanos());
    }

    return Ok(());
}

fn parse_bert1(buf: Vec<u8>) -> Result<Vec<BertTerm>> {
    let mut parser = parser::Parser::new(buf);
    let term = parser.parse()?;
    return Ok(vec![term]);
}

fn parse_bert2(buf: Vec<u8>) -> Result<Vec<BertTerm>> {
    let mut parser = parser::Parser::new(buf);
    return parser.parse_bert2();
}
<commit_msg>Use shadowing for durations rather than different names<commit_after>extern crate ppbert;
#[macro_use] extern crate clap;

use std::io::{self, BufReader, Read};
use std::fs::File;
use std::process::exit;
use std::time::Instant;

use clap::{Arg, App};

use ppbert::bertterm::{
    self,
    BertTerm,
    DEFAULT_INDENT_WIDTH,
    DEFAULT_MAX_TERMS_PER_LINE
};
use ppbert::error::Result;
use ppbert::parser;

const PROG_NAME: &str = "ppbert";

fn main() {
    let matches = App::new(PROG_NAME)
        .version(crate_version!())
        .author("Vincent Foley")
        .about("Pretty print structure encoded in Erlang's External Term Format")
        .arg(Arg::with_name("input_files")
             .value_name("FILES")
             .multiple(true))
        .arg(Arg::with_name("indent_width")
             .help("Indents with <num> spaces")
             .value_name("num")
             .short("i")
             .long("--indent-width")
             .takes_value(true))
        .arg(Arg::with_name("max_per_line")
             .help("Prints at most <num> basic terms per line")
             .value_name("num")
             .short("m")
             .long("--max-terms-per-line")
             .takes_value(true))
        .arg(Arg::with_name("verbose")
             .help("Enables verbose mode")
             .short("v")
             .long("--verbose"))
        .arg(Arg::with_name("parse")
             .help("Parses the input, doesn't pretty print it")
             .short("p")
             .long("--parse"))
        .arg(Arg::with_name("bert2")
             .help("Parses .bert2 files")
             .short("2")
             .long("bert2"))
        .arg(Arg::with_name("json")
             .help("Outputs in JSON")
             .short("j")
             .long("json"))
        .arg(Arg::with_name("transform-proplists")
             .help("Transforms proplists into JSON objects (only valid with --json)")
             .long("transform-proplists"))
        .get_matches();

    let files: Vec<&str> = match matches.values_of("input_files") {
        Some(files) => files.collect(),
        None => vec!["-"]
    };

    let indent_level = value_t!(matches, "indent_width", usize)
        .unwrap_or(DEFAULT_INDENT_WIDTH);
    let max_per_line = value_t!(matches, "max_per_line", usize)
        .unwrap_or(DEFAULT_MAX_TERMS_PER_LINE);
    let verbose = matches.is_present("verbose");
    let parse_only = matches.is_present("parse");

    let transform_proplists = matches.is_present("transform-proplists");
    let parse_fn = if matches.is_present("bert2") { parse_bert2 } else { parse_bert1 };
    let output_fn =
        if matches.is_present("json") {
            if transform_proplists {
                bertterm::pp_json_proplist
            } else {
                bertterm::pp_json
            }
        } else {
            if transform_proplists {
                eprintln!("{}: warning: --transform-proplists is only valid with the --json flag",
                          PROG_NAME);
            }
            bertterm::pp_bert
        };

    let mut return_code = 0;
    for file in files {
        let res = handle_file(file, parse_only, verbose,
                              indent_level, max_per_line,
                              parse_fn, output_fn);
        match res {
            Ok(()) => (),
            Err(e) => {
                return_code = 1;
                eprintln!("{}: {}: {}", PROG_NAME, file, e);
            }
        }
    }
    exit(return_code);
}

fn handle_file<T>(
    file: &str,
    parse_only: bool,
    verbose: bool,
    indent: usize,
    terms_per_line: usize,
    parse_fn: fn(Vec<u8>) -> Result<T>,
    pp_fn: fn(T, usize, usize) -> ()
) -> Result<()> {

    // Read file or stdin into buffer
    let mut buf: Vec<u8> = Vec::new();
    let now = Instant::now();
    if file == "-" {
        let stdin = io::stdin();
        let mut stdin = stdin.lock();
        stdin.read_to_end(&mut buf)?;
    } else {
        let f = File::open(file)?;
        let mut rdr = BufReader::new(f);
        rdr.read_to_end(&mut buf)?;
    }
    let dur = now.elapsed();
    if verbose {
        eprintln!("{}: read time: {}.{:09}",
                  PROG_NAME, dur.as_secs(), dur.subsec_nanos());
    }

    // Parse input
    let now = Instant::now();
    let parse_output = parse_fn(buf)?;
    let dur = now.elapsed();

    if verbose {
        eprintln!("{}: parse time: {}.{:09}",
                  PROG_NAME, dur.as_secs(), dur.subsec_nanos());
    }

    // Early exit if parse-only
    if parse_only {
        return Ok(());
    }

    // Pretty print
    let now = Instant::now();
    pp_fn(parse_output, indent, terms_per_line);
    let dur = now.elapsed();

    if verbose {
        eprintln!("{}: print time: {}.{:09}",
                  PROG_NAME, dur.as_secs(), dur.subsec_nanos());
    }

    return Ok(());
}

fn parse_bert1(buf: Vec<u8>) -> Result<Vec<BertTerm>> {
    let mut parser = parser::Parser::new(buf);
    let term = parser.parse()?;
    return Ok(vec![term]);
}

fn parse_bert2(buf: Vec<u8>) -> Result<Vec<BertTerm>> {
    let mut parser = parser::Parser::new(buf);
    return parser.parse_bert2();
}
<|endoftext|>


<commit_before>use std::io;

#[derive(Debug, PartialEq, Clone)]
enum Token {
    LParen,
    RParen,
    Symbol(String),
    Number(i32),
    EOF,
}

pub struct LexerState {
    s: String,
    pos: usize,
    tok_buf: Option<Token>,
}

#[derive(Debug, Clone, PartialEq)]
enum SExpr {
    Symbol(String),
    Number(i32),
    Bool(bool),
    List(Vec<SExpr>),

    Define(String, Box<SExpr>),
    Let(Vec<(String, SExpr)>, Box<SExpr>),
    If(Box<SExpr>, Box<SExpr>, Box<SExpr>),
    App(Box<SExpr>, Vec<SExpr>),
}

fn unread(ls: &mut LexerState, tok: Token) {
    if let Some(_) = ls.tok_buf {
        println!("error: unread buffer full");
    }
    else {
        ls.tok_buf = Some(tok)
    }
}

fn is_valid_symbol_start(c: char) -> bool {
    // TODO: avoid allocatiing this in each call
    let SymbolStartChars = vec!['+', '-', '*', '/', '#'];

    let mut ret = false;
    if c.is_alphabetic() { ret = true; }
    else {
        for s in SymbolStartChars {
            if c == s { ret = true; break; }
            else { continue; }
        }
    }

    return ret;
}

fn get_token(ls: &mut LexerState) -> Token {
    if let Some(tok) = ls.tok_buf.clone() {
        ls.tok_buf = None;
        return tok;
    }
    else {
        let mut iter = ls.s[ls.pos..].chars().peekable();
        while let Some(&c) = iter.peek() {
            if c.is_numeric() {
                let mut acc = String::new();
                let mut n = c;
                while n.is_numeric() {
                    acc.push(n);
                    iter.next();
                    ls.pos += 1;
                    n = match iter.peek() {
                        Some(&x) => x,
                        None => break,
                    };
                }
                return Token::Number(acc.parse().unwrap());
            }
            else if is_valid_symbol_start(c) {
                let mut acc = String::new();
                let mut s = c;
                while (s.is_alphanumeric() ||
                       is_valid_symbol_start(s)) {
                    acc.push(s);
                    iter.next();
                    ls.pos += 1;
                    s = match iter.peek() {
                        Some(&x) => x,
                        None => break,
                    };
                }
                return Token::Symbol(acc);
            }
            else {
                match c {
                    ' ' => {
                        iter.next();
                        ls.pos += 1;
                        continue
                    },
                    '(' => {
                        iter.next();
                        ls.pos += 1;
                        return Token::LParen
                    },
                    ')' => {
                        iter.next();
                        ls.pos += 1;
                        return Token::RParen
                    },
                    _ => panic!("unexpected char: {}", c),
                }
            }
        }
        return Token::EOF;
    }
}

fn get_list(ls: &mut LexerState) -> Vec<SExpr> {
    match get_expr(ls) {
        exp => match get_token(ls) {
            Token::RParen => return vec![exp],
            tok => {
                unread(ls, tok);
                let mut seq = get_list(ls);
                seq.insert(0, exp);
                return seq;
            },
        },
    }
}

fn get_expr(ls: &mut LexerState) -> SExpr {
    match get_token(ls) {
        Token::Symbol(s) => return SExpr::Symbol(s),
        Token::Number(n) => return SExpr::Number(n),
        Token::LParen => {
            return SExpr::List(get_list(ls));
        },
        Token::RParen => panic!("unmatched ')'"),
        _ => return SExpr::Number(32),
    }
}

fn get_ast(expr: SExpr) -> SExpr {
    match expr {
        SExpr::Symbol(sym) => {
            match &sym[..] {
                "#f" => SExpr::Bool(false),
                "#t" => SExpr::Bool(true),
                _ => SExpr::Symbol(sym),
            }
        },
        SExpr::List(elts) =>
            match elts[0].clone() {
                SExpr::Symbol(sym) => {
                    match &sym[..] {
                        "define" => {
                            // TODO: Check that elts has correct length
                            if let SExpr::Symbol(name) = elts[1].clone() {
                                let e = elts[2].clone();
                                return SExpr::Define(name, Box::new(e));
                            }
                            else {
                                panic!("expected `name` to be a symbol");
                            }
                        },
                        "if" => {
                            return SExpr::If(Box::new(get_ast(elts[1].clone())),
                                             Box::new(get_ast(elts[2].clone())),
                                             Box::new(get_ast(elts[3].clone())));
                        },
                        "let" => {
                            let bindings = match elts[1] {
                                SExpr::List(ref bs) => bs,
                                _ => panic!("let bindings"),
                            };

                            let mut astified_bindings = vec![];
                            for bind_pair in bindings {
                                let (key, val) = match bind_pair {
                                    // TODO: check length
                                    &SExpr::List(ref kv) => (kv[0].clone(), kv[1].clone()),
                                    _ => panic!("non-list in let-binding"),
                                };

                                let keyname = match key {
                                    SExpr::Symbol(k) => k,
                                    _ => panic!("let binding key is not symbol"),
                                };
                                astified_bindings.push((keyname, get_ast(val)));
                            }
                            return SExpr::Let(astified_bindings, Box::new(get_ast(elts[2].clone())));
                        },
                        _ => SExpr::App(Box::new(elts[0].clone()), elts[1..].to_vec()),
                    }
                }
                _ => panic!("NYI"),
            },
        _ => expr,
    }
}

fn read_input() -> io::Result<()> {
    let mut input = String::new();

    try!(io::stdin().read_line(&mut input));

    let mut lexer = LexerState {
        s: input,
        pos: 0,
        tok_buf: None,
    };

    println!("{:?}", get_ast(get_expr(&mut lexer)));

    Ok(())
}

#[test]
fn test_parser() {
    let mut input = String::from("(if #f (+ 42 (foo 12)) 17)");
    let mut lexer = LexerState {
        s: input,
        pos: 0,
        tok_buf: None,
    };
    assert_eq!(SExpr::If(Box::new(SExpr::Bool(false)),
                         Box::new(SExpr::App(Box::new(SExpr::Symbol("+".to_string())),
                                             vec![SExpr::Number(42),
                                                  SExpr::List(vec![SExpr::Symbol("foo".to_string()),
                                                                   SExpr::Number(12)])])),
                         Box::new(SExpr::Number(17))),
               get_ast(get_expr(&mut lexer)));
}

fn main() {
    read_input();
}
<commit_msg>`uniquify` pass<commit_after>use std::io;
use std::collections::HashMap;

#[derive(Debug, PartialEq, Clone)]
enum Token {
    LParen,
    RParen,
    Symbol(String),
    Number(i32),
    EOF,
}

pub struct LexerState {
    s: String,
    pos: usize,
    tok_buf: Option<Token>,
}

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
enum SExpr {
    Symbol(String),
    Number(i32),
    Bool(bool),
    List(Vec<SExpr>),

    // TODO: change define to be special-form for function definition
    Define(String, Box<SExpr>),
    Let(Vec<(String, SExpr)>, Box<SExpr>),
    If(Box<SExpr>, Box<SExpr>, Box<SExpr>),
    App(Box<SExpr>, Vec<SExpr>),
}

fn unread(ls: &mut LexerState, tok: Token) {
    if let Some(_) = ls.tok_buf {
        println!("error: unread buffer full");
    }
    else {
        ls.tok_buf = Some(tok)
    }
}

fn is_valid_symbol_start(c: char) -> bool {
    // TODO: avoid allocatiing this in each call
    let SymbolStartChars = vec!['+', '-', '*', '/', '#'];

    let mut ret = false;
    if c.is_alphabetic() { ret = true; }
    else {
        for s in SymbolStartChars {
            if c == s { ret = true; break; }
            else { continue; }
        }
    }

    return ret;
}

fn get_token(ls: &mut LexerState) -> Token {
    if let Some(tok) = ls.tok_buf.clone() {
        ls.tok_buf = None;
        return tok;
    }
    else {
        let mut iter = ls.s[ls.pos..].chars().peekable();
        while let Some(&c) = iter.peek() {
            if c.is_numeric() {
                let mut acc = String::new();
                let mut n = c;
                while n.is_numeric() {
                    acc.push(n);
                    iter.next();
                    ls.pos += 1;
                    n = match iter.peek() {
                        Some(&x) => x,
                        None => break,
                    };
                }
                return Token::Number(acc.parse().unwrap());
            }
            else if is_valid_symbol_start(c) {
                let mut acc = String::new();
                let mut s = c;
                while (s.is_alphanumeric() ||
                       is_valid_symbol_start(s)) {
                    acc.push(s);
                    iter.next();
                    ls.pos += 1;
                    s = match iter.peek() {
                        Some(&x) => x,
                        None => break,
                    };
                }
                return Token::Symbol(acc);
            }
            else {
                match c {
                    ' ' => {
                        iter.next();
                        ls.pos += 1;
                        continue
                    },
                    '(' => {
                        iter.next();
                        ls.pos += 1;
                        return Token::LParen
                    },
                    ')' => {
                        iter.next();
                        ls.pos += 1;
                        return Token::RParen
                    },
                    _ => panic!("unexpected char: {}", c),
                }
            }
        }
        return Token::EOF;
    }
}

fn get_list(ls: &mut LexerState) -> Vec<SExpr> {
    match get_expr(ls) {
        exp => match get_token(ls) {
            Token::RParen => return vec![exp],
            tok => {
                unread(ls, tok);
                let mut seq = get_list(ls);
                seq.insert(0, exp);
                return seq;
            },
        },
    }
}

fn get_expr(ls: &mut LexerState) -> SExpr {
    match get_token(ls) {
        Token::Symbol(s) => return SExpr::Symbol(s),
        Token::Number(n) => return SExpr::Number(n),
        Token::LParen => {
            return SExpr::List(get_list(ls));
        },
        Token::RParen => panic!("unmatched ')'"),
        _ => return SExpr::Number(32),
    }
}

fn get_ast(expr: SExpr) -> SExpr {
    match expr {
        SExpr::Symbol(sym) => {
            match &sym[..] {
                "#f" => SExpr::Bool(false),
                "#t" => SExpr::Bool(true),
                _ => SExpr::Symbol(sym),
            }
        },
        SExpr::List(elts) =>
            match elts[0].clone() {
                SExpr::Symbol(sym) => {
                    match &sym[..] {
                        "define" => {
                            // TODO: Check that elts has correct length
                            if let SExpr::Symbol(name) = elts[1].clone() {
                                let e = elts[2].clone();
                                return SExpr::Define(name, Box::new(e));
                            }
                            else {
                                panic!("expected `name` to be a symbol");
                            }
                        },
                        "if" => {
                            return SExpr::If(Box::new(get_ast(elts[1].clone())),
                                             Box::new(get_ast(elts[2].clone())),
                                             Box::new(get_ast(elts[3].clone())));
                        },
                        "let" => {
                            let bindings = match elts[1] {
                                SExpr::List(ref bs) => bs,
                                _ => panic!("let bindings"),
                            };

                            let mut astified_bindings = vec![];
                            for bind_pair in bindings {
                                let (key, val) = match bind_pair {
                                    // TODO: check length
                                    &SExpr::List(ref kv) => (kv[0].clone(), kv[1].clone()),
                                    _ => panic!("non-list in let-binding"),
                                };

                                let keyname = match key {
                                    SExpr::Symbol(k) => k,
                                    _ => panic!("let binding key is not symbol"),
                                };
                                astified_bindings.push((keyname, get_ast(val)));
                            }
                            return SExpr::Let(astified_bindings, Box::new(get_ast(elts[2].clone())));
                        },
                        _ => SExpr::App(Box::new(elts[0].clone()), elts[1..].to_vec()),
                    }
                }
                _ => panic!("NYI"),
            },
        _ => expr,
    }
}

fn read_input() -> io::Result<()> {
    let mut input = String::new();

    try!(io::stdin().read_line(&mut input));

    let mut lexer = LexerState {
        s: input,
        pos: 0,
        tok_buf: None,
    };

    let mut uniquify_mapping = HashMap::new();

    println!("{:?}", uniquify(&mut uniquify_mapping,
                              get_ast(get_expr(&mut lexer))));

    Ok(())
}

#[test]
fn test_parser() {
    let mut input = String::from("(if #f (+ 42 (foo 12)) 17)");
    let mut lexer = LexerState {
        s: input,
        pos: 0,
        tok_buf: None,
    };
    assert_eq!(SExpr::If(Box::new(SExpr::Bool(false)),
                         Box::new(SExpr::App(Box::new(SExpr::Symbol("+".to_string())),
                                             vec![SExpr::Number(42),
                                                  SExpr::List(vec![SExpr::Symbol("foo".to_string()),
                                                                   SExpr::Number(12)])])),
                         Box::new(SExpr::Number(17))),
               get_ast(get_expr(&mut lexer)));
}

static mut var_counter : i32 = 0;
fn get_unique_varname(stem: &str) -> String {
    unsafe {
        var_counter += 1;
        return stem.to_string() + &"." + &var_counter.to_string();
    }
}

fn uniquify(mapping: &mut HashMap<String, String>, expr: SExpr)
            -> SExpr {
    match expr {
        SExpr::Symbol(name) =>
            SExpr::Symbol(mapping.get(&name).unwrap().to_string()),
        SExpr::Number(_) => expr,
        SExpr::Bool(_) => expr,
        SExpr::Let(bindings, body) => {
            let mut new_bindings = vec![];
            for (k,v) in bindings {
                let uniq_k = get_unique_varname(&k);
                mapping.insert(k.clone(), uniq_k.clone());
                new_bindings.push((uniq_k, v));
            }
            return SExpr::Let(new_bindings, Box::new(uniquify(mapping, *body)));
        },
        SExpr::List(elts) => {
            let mut new_elts = vec![];
            for e in elts {
                new_elts.push(uniquify(mapping, e));
            }

            return SExpr::List(new_elts);
        }
        SExpr::Define(name, val) => {
            return SExpr::Define(name,
                                 Box::new(uniquify(mapping, *val)));
        },
        SExpr::If(cond, thn, els) => {
            return SExpr::If(Box::new(uniquify(mapping, *cond)),
                             Box::new(uniquify(mapping, *thn)),
                             Box::new(uniquify(mapping, *els)));
        },
        SExpr::App(f, args) => {
            let mut new_args = vec![];
            for a in args {
                new_args.push(uniquify(mapping, a));
            }
            return SExpr::App(f, new_args);
        },
        //_ => expr,
    }
}

fn main() {
    read_input();
}
<|endoftext|>