<commit_before>//! The master module provides the Master struct which instructs Workers which
//! portions of the image they should render and collects their results to combine
//! into the final image.

use std::path::PathBuf;
use std::io::prelude::*;
use std::collections::HashMap;
use std::iter;

use bincode::rustc_serialize::decode;
use image;
use mio::tcp::{TcpStream, Shutdown};
use mio::*;

use film::Image;
use exec::Config;
use exec::distrib::{worker, Instructions};
use sampler::BlockQueue;

pub static PORT: u16 = 63245;

pub struct DistributedFrame {
    pub frame: usize,
    // The number of workers who have reported results for this
    // frame so far
    pub num_reporting: usize,
    pub render: Image,
    pub completed: bool,
}

#[derive(Clone, Debug)]
struct WorkerBuffer {
    pub buf: Vec<u8>,
    pub expected_size: usize,
    pub currently_read: usize,
}

impl WorkerBuffer {
    pub fn new() -> WorkerBuffer {
        WorkerBuffer { buf: Vec::new(), expected_size: 0, currently_read: 0 }
    }
}

pub struct Master {
    // Hostnames of the workers to send work too
    workers: Vec<String>,
    connections: Vec<TcpStream>,
    // Temporary buffers to store worker results in as they're
    // read in over TCP
    worker_buffers: Vec<WorkerBuffer>,
    config: Config,
    frames: HashMap<usize, DistributedFrame>,
    img_dim: (usize, usize),
}

impl Master {
    /// Create a new master that will contact the worker nodes passed and
    /// send instructions on what parts of the scene to start rendering
    pub fn start_workers(workers: Vec<String>, config: Config, scene_file: &String,
                         img_dim: (usize, usize), frames: Option<(usize, usize)>)
                         -> (Master, EventLoop<Master>) {
        // Figure out how many blocks we have for this image and assign them to our workers
        let queue = BlockQueue::new((img_dim.0 as u32, img_dim.1 as u32), (8, 8), (0, 0));
        let blocks_per_worker = queue.len() / workers.len();
        let blocks_remainder = queue.len() % workers.len();
        println!("Have {} workers, each worker does {} blocks w/ {} remainder",
                 workers.len(), blocks_per_worker, blocks_remainder);

        let mut event_loop = EventLoop::<Master>::new().unwrap();
        let mut connections = Vec::new();

        // Connect to each worker and send instructions on what to render
        for (i, host) in workers.iter().enumerate() {
            let b_start = i * blocks_per_worker;
            let b_count =
                if i == workers.len() - 1 {
                    blocks_per_worker + blocks_remainder
                } else {
                    blocks_per_worker
                };
            let instr = Instructions::new(scene_file, PORT, frames, b_start, b_count);
            println!("Sending instructions {:?} too {}", instr, host);
            let addr = format!("{}:{}", host, worker::PORT).parse().unwrap();
            match TcpStream::connect(&addr) {
                Ok(mut stream) => {
                    println!("Connected and writing");
                    let bytes = instr.to_json().into_bytes();
                    match stream.write_all(&bytes[..]) {
                        Err(e) => println!("Failed to send instructions to {}: {:?}", host, e),
                        _ => {},
                    }
                    // We no longer need to write anything, so close the write end
                    match stream.shutdown(Shutdown::Write) {
                        Err(e) => panic!(format!("Failed to shutdown write end to worker {}: {}", host, e)),
                        Ok(_) => {},
                    }
                    println!("Adding stream to event loop");
                    // Each worker is identified in the event loop by their index in the vec
                    match event_loop.register(&stream, Token(i)) {
                        Err(e) => println!("Error registering stream from {}: {}", host, e),
                        Ok(_) => {},
                    }
                    connections.push(stream);
                },
                Err(e) => println!("Failed to contact worker {}: {:?}", host, e),
            }
        }
        let worker_buffers: Vec<_> = iter::repeat(WorkerBuffer::new()).take(workers.len()).collect();
        let master = Master { workers: workers, connections: connections,
                              worker_buffers: worker_buffers, config: config,
                              frames: HashMap::new(), img_dim: img_dim };
        (master, event_loop)
    }
    /// Read a result frame from a worker and save it into the list of frames we're collecting from
    /// all workers. Will save out the final render if all workers have reported results for this
    /// frame.
    fn save_results(&mut self, frame: worker::Frame) {
        // Find the frame this worker is referring to and add its results
        let frame_num = frame.frame as usize;
        println!("Worker reporting frame {}", frame_num);
        if let Some(df) = self.frames.get_mut(&frame_num) {
            df.render.add_pixels(&frame.pixels);
            df.num_reporting += 1;
            println!("We have parts of this frame already");
        }
        // TODO: Better way here? We get stuck because self.frames is counted as borrowed
        // from the if above, making this really awkward.
        if self.frames.get_mut(&frame_num).is_none() {
            // TODO: I need to seperate the render target from the actual image as the
            // filter here isn't used in this path and it doesn't make sense to have one
            let render = Image::new(self.img_dim);
            let mut df = DistributedFrame { frame: frame_num, num_reporting: 1, render: render, completed: false };
            df.render.add_pixels(&frame.pixels);
            self.frames.insert(df.frame, df);
            println!("This is a new frame");
        }
        if let Some(df) = self.frames.get_mut(&frame_num) {
            if df.num_reporting == self.workers.len() {
                df.completed = true;
                println!("We have completed frame {}", frame_num);
                let out_file = match self.config.out_path.extension() {
                    Some(_) => self.config.out_path.clone(),
                    None => self.config.out_path.join(PathBuf::from(format!("frame{:05}.png", df.frame))),
                };
                println!("Frame {}: rendered to '{}'\n--------------------", frame_num, out_file.display());
                let img = df.render.get_srgb8();
                let dim = df.render.dimensions();
                match image::save_buffer(&out_file.as_path(), &img[..], dim.0 as u32, dim.1 as u32, image::RGB(8)) {
                    Ok(_) => {},
                    Err(e) => println!("Error saving image, {}", e),
                };
            }
        }
    }
    fn read_worker_buffer(&mut self, worker: usize) -> bool {
        let mut buf = &mut self.worker_buffers[worker];
        // If we haven't read the size of data being sent, read that now
        if buf.currently_read < 8 {
            // First 8 bytes are a u64 specifying the number of bytes being sent
            buf.buf.extend(iter::repeat(0u8).take(8));
            while buf.currently_read != 8 {
                println!("Reading size header...");
                match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
                    Ok(n) => buf.currently_read += n,
                    Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
                }
                println!("Read {} bytes so far, must read {} in total", buf.currently_read, 8);
            }
            // How many bytes we expect to get from the worker for a frame
            buf.expected_size = decode(&buf.buf[..]).unwrap();
            // Extend the Vec so we've got enough room for the remaning bytes, minus the 8 for the
            // encoded size header
            buf.buf.extend(iter::repeat(0u8).take(buf.expected_size - 8));
        }
        println!("Reading from worker");
        match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
            Ok(n) => buf.currently_read += n,
            Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
        }
        println!("Read {} bytes so far, must read {} in total", buf.currently_read, buf.expected_size);
        buf.currently_read == buf.expected_size
    }
}

impl Handler for Master {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Master>, token: Token, event: EventSet) {
        // Some results are read for reading from a worker
        if event.is_readable() {
            let worker = token.as_usize();
            let finished = self.read_worker_buffer(worker);
            if finished {
                println!("Got all frame data, now parsing");
                let frame: worker::Frame = decode(&self.worker_buffers[worker].buf[..]).unwrap();
                println!("Saving frame results");
                self.save_results(frame);
                // Clean up our read buffer
                self.worker_buffers[worker].buf.clear();
                self.worker_buffers[worker].expected_size = 0;
                self.worker_buffers[worker].currently_read = 0;
            }
        }
        // After getting results from the worker we check if we've completed all our frames
        // and exit if so
        if !self.frames.is_empty() && self.frames.values().fold(true, |all, v| all && v.completed) {
            println!("All frames have been rendered, master exiting");
            event_loop.shutdown();
        }
    }
}

<commit_msg>Need to write in a loop, async sockets. Also check for hangups & close<commit_after>//! The master module provides the Master struct which instructs Workers which
//! portions of the image they should render and collects their results to combine
//! into the final image.

use std::path::PathBuf;
use std::io::prelude::*;
use std::collections::HashMap;
use std::net::ToSocketAddrs;
use std::iter;

use bincode::rustc_serialize::decode;
use image;
use mio::tcp::{TcpStream, Shutdown};
use mio::*;

use film::Image;
use exec::Config;
use exec::distrib::{worker, Instructions};
use sampler::BlockQueue;

pub static PORT: u16 = 63245;

pub struct DistributedFrame {
    pub frame: usize,
    // The number of workers who have reported results for this
    // frame so far
    pub num_reporting: usize,
    pub render: Image,
    pub completed: bool,
}

#[derive(Clone, Debug)]
struct WorkerBuffer {
    pub buf: Vec<u8>,
    pub expected_size: usize,
    pub currently_read: usize,
}

impl WorkerBuffer {
    pub fn new() -> WorkerBuffer {
        WorkerBuffer { buf: Vec::new(), expected_size: 8, currently_read: 0 }
    }
}

pub struct Master {
    // Hostnames of the workers to send work too
    workers: Vec<String>,
    connections: Vec<TcpStream>,
    // Temporary buffers to store worker results in as they're
    // read in over TCP
    worker_buffers: Vec<WorkerBuffer>,
    config: Config,
    frames: HashMap<usize, DistributedFrame>,
    img_dim: (usize, usize),
}

impl Master {
    /// Create a new master that will contact the worker nodes passed and
    /// send instructions on what parts of the scene to start rendering
    pub fn start_workers(workers: Vec<String>, config: Config, scene_file: &String,
                         img_dim: (usize, usize), frames: Option<(usize, usize)>)
                         -> (Master, EventLoop<Master>) {
        // Figure out how many blocks we have for this image and assign them to our workers
        let queue = BlockQueue::new((img_dim.0 as u32, img_dim.1 as u32), (8, 8), (0, 0));
        let blocks_per_worker = queue.len() / workers.len();
        let blocks_remainder = queue.len() % workers.len();
        println!("Have {} workers, each worker does {} blocks w/ {} remainder",
                 workers.len(), blocks_per_worker, blocks_remainder);

        let mut event_loop = EventLoop::<Master>::new().unwrap();
        let mut connections = Vec::new();

        // Connect to each worker and send instructions on what to render
        // I guess this should also be queue'd up and the writes actually performed asynchronously?
        // Otherwise we need to loop on write_all while waiting for the connection to open
        for (i, host) in workers.iter().enumerate() {
            let b_start = i * blocks_per_worker;
            let b_count =
                if i == workers.len() - 1 {
                    blocks_per_worker + blocks_remainder
                } else {
                    blocks_per_worker
                };
            let instr = Instructions::new(scene_file, PORT, frames, b_start, b_count);
            println!("Sending instructions {:?} too {}", instr, host);
            let addr = (&host[..], worker::PORT).to_socket_addrs().unwrap().next().unwrap();
            println!("Connecting too {:?}", addr);
            match TcpStream::connect(&addr) {
                Ok(mut stream) => {
                    println!("Connected and writing");
                    let bytes = instr.to_json().into_bytes();
                    // Loop until we successfully write the byes
                    loop {
                        match stream.write_all(&bytes[..]) {
                            Err(e) => println!("Failed to send instructions to {}: {:?}", host, e),
                            Ok(_) => break,
                        }
                    }
                    // We no longer need to write anything, so close the write end
                    match stream.shutdown(Shutdown::Write) {
                        Err(e) => panic!(format!("Failed to shutdown write end to worker {}: {}", host, e)),
                        Ok(_) => {},
                    }
                    println!("Adding stream to event loop");
                    // Each worker is identified in the event loop by their index in the vec
                    match event_loop.register(&stream, Token(i)) {
                        Err(e) => println!("Error registering stream from {}: {}", host, e),
                        Ok(_) => {},
                    }
                    connections.push(stream);
                },
                Err(e) => println!("Failed to contact worker {}: {:?}", host, e),
            }
        }
        let worker_buffers: Vec<_> = iter::repeat(WorkerBuffer::new()).take(workers.len()).collect();
        let master = Master { workers: workers, connections: connections,
                              worker_buffers: worker_buffers, config: config,
                              frames: HashMap::new(), img_dim: img_dim };
        (master, event_loop)
    }
    /// Read a result frame from a worker and save it into the list of frames we're collecting from
    /// all workers. Will save out the final render if all workers have reported results for this
    /// frame.
    fn save_results(&mut self, frame: worker::Frame) {
        // Find the frame this worker is referring to and add its results
        let frame_num = frame.frame as usize;
        println!("Worker reporting frame {}", frame_num);
        if let Some(df) = self.frames.get_mut(&frame_num) {
            df.render.add_pixels(&frame.pixels);
            df.num_reporting += 1;
            println!("We have parts of this frame already");
        }
        // TODO: Better way here? We get stuck because self.frames is counted as borrowed
        // from the if above, making this really awkward.
        if self.frames.get_mut(&frame_num).is_none() {
            // TODO: I need to seperate the render target from the actual image as the
            // filter here isn't used in this path and it doesn't make sense to have one
            let render = Image::new(self.img_dim);
            let mut df = DistributedFrame { frame: frame_num, num_reporting: 1, render: render, completed: false };
            df.render.add_pixels(&frame.pixels);
            self.frames.insert(df.frame, df);
            println!("This is a new frame");
        }
        if let Some(df) = self.frames.get_mut(&frame_num) {
            if df.num_reporting == self.workers.len() {
                df.completed = true;
                println!("We have completed frame {}", frame_num);
                let out_file = match self.config.out_path.extension() {
                    Some(_) => self.config.out_path.clone(),
                    None => self.config.out_path.join(PathBuf::from(format!("frame{:05}.png", df.frame))),
                };
                println!("Frame {}: rendered to '{}'\n--------------------", frame_num, out_file.display());
                let img = df.render.get_srgb8();
                let dim = df.render.dimensions();
                match image::save_buffer(&out_file.as_path(), &img[..], dim.0 as u32, dim.1 as u32, image::RGB(8)) {
                    Ok(_) => {},
                    Err(e) => println!("Error saving image, {}", e),
                };
            }
        }
    }
    fn read_worker_buffer(&mut self, worker: usize) -> bool {
        let mut buf = &mut self.worker_buffers[worker];
        // If we haven't read the size of data being sent, read that now
        if buf.currently_read < 8 {
            // First 8 bytes are a u64 specifying the number of bytes being sent
            buf.buf.extend(iter::repeat(0u8).take(8));
            println!("Reading size header...");
            match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
                Ok(n) => buf.currently_read += n,
                Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
            }
            println!("Read {} bytes so far, must read {} in total", buf.currently_read, buf.expected_size);
            if buf.currently_read == buf.expected_size {
                println!("Got size header");
                // How many bytes we expect to get from the worker for a frame
                buf.expected_size = decode(&buf.buf[..]).unwrap();
                println!("Expecting {} bytes", buf.expected_size);
                // Extend the Vec so we've got enough room for the remaning bytes, minus the 8 for the
                // encoded size header
                buf.buf.extend(iter::repeat(0u8).take(buf.expected_size - 8));
            }
        }
        println!("Reading from worker");
        match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
            Ok(n) => buf.currently_read += n,
            Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
        }
        println!("Read {} bytes so far, must read {} in total", buf.currently_read, buf.expected_size);
        buf.currently_read == buf.expected_size
    }
}

impl Handler for Master {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Master>, token: Token, event: EventSet) {
        let worker = token.as_usize();
        // Some results are available from a worker
        if event.is_readable() {
            println!("Readable event from worker {}", worker);
            let finished = self.read_worker_buffer(worker);
            if finished {
                println!("Got all frame data, now parsing");
                let frame: worker::Frame = decode(&self.worker_buffers[worker].buf[..]).unwrap();
                println!("Saving frame results");
                self.save_results(frame);
                // Clean up our read buffer
                self.worker_buffers[worker].buf.clear();
                self.worker_buffers[worker].expected_size = 8;
                self.worker_buffers[worker].currently_read = 0;
            }
        }
        // If the worker has terminated, shutdown the read end of the connection
        if event.is_hup() {
            self.connections[worker].shutdown(Shutdown::Read);
        }
        // After getting results from the worker we check if we've completed all our frames
        // and exit if so
        if self.frames.len() == self.config.frame_info.frames
            && self.frames.values().fold(true, |all, v| all && v.completed) {
                println!("All frames have been rendered, master exiting");
                event_loop.shutdown();
        }
    }
}

<|endoftext|>


<commit_before>use crate::context::{BlockParams, Context};
use crate::error::RenderError;
use crate::helpers::{HelperDef, HelperResult};
use crate::output::Output;
use crate::registry::Registry;
use crate::render::{Helper, RenderContext, Renderable};
use crate::value::JsonTruthy;

#[derive(Clone, Copy)]
pub struct WithHelper;

impl HelperDef for WithHelper {
    fn call<'reg: 'rc, 'rc>(
        &self,
        h: &Helper<'reg, 'rc>,
        r: &'reg Registry,
        ctx: &'rc Context,
        rc: &mut RenderContext<'reg>,
        out: &mut dyn Output,
    ) -> HelperResult {
        let param = h
            .param(0)
            .ok_or_else(|| RenderError::new("Param not found for helper \"with\""))?;

        rc.promote_local_vars();
        let local_path_root = param.path_root();

        let result = {
            let mut local_rc = rc.derive();

            let not_empty = param.value().is_truthy(false);
            let template = if not_empty { h.template() } else { h.inverse() };

            if let Some(path_root) = local_path_root {
                local_rc.push_local_path_root(path_root.to_vec());
            }

            if not_empty {
                let new_path = param.context_path();
                if let Some(new_path) = new_path {
                    local_rc.set_path(new_path.clone());
                }

                if let Some(block_param) = h.block_param() {
                    let mut params = BlockParams::new();
                    if new_path.is_some() {
                        params.add_path(block_param, local_rc.get_path().clone())?;
                    } else {
                        params.add_value(block_param, param.value().clone())?;
                    }

                    local_rc.push_block_context(params)?;
                }
            }

            match template {
                Some(t) => t.render(r, ctx, &mut local_rc, out),
                None => Ok(()),
            }
        };

        rc.demote_local_vars();
        result
    }
}

pub static WITH_HELPER: WithHelper = WithHelper;

#[cfg(test)]
mod test {
    use crate::registry::Registry;
    use crate::value::to_json;

    #[derive(Serialize)]
    struct Address {
        city: String,
        country: String,
    }

    #[derive(Serialize)]
    struct Person {
        name: String,
        age: i16,
        addr: Address,
        titles: Vec<String>,
    }

    #[test]
    fn test_with() {
        let addr = Address {
            city: "Beijing".to_string(),
            country: "China".to_string(),
        };

        let person = Person {
            name: "Ning Sun".to_string(),
            age: 27,
            addr,
            titles: vec!["programmer".to_string(), "cartographier".to_string()],
        };

        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#with addr}}{{city}}{{/with}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t1", "{{#with notfound}}hello{{else}}world{{/with}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t2", "{{#with addr/country}}{{this}}{{/with}}")
            .is_ok());

        let r0 = handlebars.render("t0", &person);
        assert_eq!(r0.ok().unwrap(), "Beijing".to_string());

        let r1 = handlebars.render("t1", &person);
        assert_eq!(r1.ok().unwrap(), "world".to_string());

        let r2 = handlebars.render("t2", &person);
        assert_eq!(r2.ok().unwrap(), "China".to_string());
    }

    #[test]
    fn test_with_block_param() {
        let addr = Address {
            city: "Beijing".to_string(),
            country: "China".to_string(),
        };

        let person = Person {
            name: "Ning Sun".to_string(),
            age: 27,
            addr,
            titles: vec!["programmer".to_string(), "cartographier".to_string()],
        };

        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#with addr as |a|}}{{a.city}}{{/with}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t1", "{{#with notfound as |c|}}hello{{else}}world{{/with}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t2", "{{#with addr/country as |t|}}{{t}}{{/with}}")
            .is_ok());

        let r0 = handlebars.render("t0", &person);
        assert_eq!(r0.ok().unwrap(), "Beijing".to_string());

        let r1 = handlebars.render("t1", &person);
        assert_eq!(r1.ok().unwrap(), "world".to_string());

        let r2 = handlebars.render("t2", &person);
        assert_eq!(r2.ok().unwrap(), "China".to_string());
    }

    #[test]
    fn test_with_in_each() {
        let addr = Address {
            city: "Beijing".to_string(),
            country: "China".to_string(),
        };

        let person = Person {
            name: "Ning Sun".to_string(),
            age: 27,
            addr,
            titles: vec!["programmer".to_string(), "cartographier".to_string()],
        };

        let addr2 = Address {
            city: "Beijing".to_string(),
            country: "China".to_string(),
        };

        let person2 = Person {
            name: "Ning Sun".to_string(),
            age: 27,
            addr: addr2,
            titles: vec!["programmer".to_string(), "cartographier".to_string()],
        };

        let people = vec![person, person2];

        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string(
                "t0",
                "{{#each this}}{{#with addr}}{{city}}{{/with}}{{/each}}"
            )
            .is_ok());
        assert!(handlebars
            .register_template_string(
                "t1",
                "{{#each this}}{{#with addr}}{{../age}}{{/with}}{{/each}}"
            )
            .is_ok());
        assert!(handlebars
            .register_template_string(
                "t2",
                "{{#each this}}{{#with addr}}{{@../index}}{{/with}}{{/each}}"
            )
            .is_ok());

        let r0 = handlebars.render("t0", &people);
        assert_eq!(r0.ok().unwrap(), "BeijingBeijing".to_string());

        let r1 = handlebars.render("t1", &people);
        assert_eq!(r1.ok().unwrap(), "2727".to_string());

        let r2 = handlebars.render("t2", &people);
        assert_eq!(r2.ok().unwrap(), "01".to_string());
    }

    #[test]
    fn test_path_up() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#with a}}{{#with b}}{{../../d}}{{/with}}{{/with}}")
            .is_ok());
        let data = btreemap! {
            "a".to_string() => to_json(&btreemap! {
                "b".to_string() => vec![btreemap!{"c".to_string() => vec![1]}]
            }),
            "d".to_string() => to_json(1)
        };

        let r0 = handlebars.render("t0", &data);
        assert_eq!(r0.ok().unwrap(), "1".to_string());
    }
}
<commit_msg>(refactor) avoid render_context clone in with helper<commit_after>use crate::context::{BlockParams, Context};
use crate::error::RenderError;
use crate::helpers::{HelperDef, HelperResult};
use crate::output::Output;
use crate::registry::Registry;
use crate::render::{Helper, RenderContext, Renderable};
use crate::value::JsonTruthy;

#[derive(Clone, Copy)]
pub struct WithHelper;

impl HelperDef for WithHelper {
    fn call<'reg: 'rc, 'rc>(
        &self,
        h: &Helper<'reg, 'rc>,
        r: &'reg Registry,
        ctx: &'rc Context,
        rc: &mut RenderContext<'reg>,
        out: &mut dyn Output,
    ) -> HelperResult {
        let param = h
            .param(0)
            .ok_or_else(|| RenderError::new("Param not found for helper \"with\""))?;

        rc.promote_local_vars();
        let local_path_root = param.path_root();

        let result = {
            let not_empty = param.value().is_truthy(false);
            let template = if not_empty { h.template() } else { h.inverse() };

            if let Some(path_root) = local_path_root {
                rc.push_local_path_root(path_root.to_vec());
            }

            if not_empty {
                let new_path = param.context_path();
                if let Some(new_path) = new_path {
                    rc.set_path(new_path.clone());
                }

                if let Some(block_param) = h.block_param() {
                    let mut params = BlockParams::new();
                    if new_path.is_some() {
                        params.add_path(block_param, rc.get_path().clone())?;
                    } else {
                        params.add_value(block_param, param.value().clone())?;
                    }

                    rc.push_block_context(params)?;
                }
            }

            let result = match template {
                Some(t) => t.render(r, ctx, rc, out),
                None => Ok(()),
            };

            if h.has_block_param() {
                rc.pop_block_context();
            }

            if local_path_root.is_some() {
                rc.pop_local_path_root();
            }

            result
        };

        rc.demote_local_vars();
        result
    }
}

pub static WITH_HELPER: WithHelper = WithHelper;

#[cfg(test)]
mod test {
    use crate::registry::Registry;
    use crate::value::to_json;

    #[derive(Serialize)]
    struct Address {
        city: String,
        country: String,
    }

    #[derive(Serialize)]
    struct Person {
        name: String,
        age: i16,
        addr: Address,
        titles: Vec<String>,
    }

    #[test]
    fn test_with() {
        let addr = Address {
            city: "Beijing".to_string(),
            country: "China".to_string(),
        };

        let person = Person {
            name: "Ning Sun".to_string(),
            age: 27,
            addr,
            titles: vec!["programmer".to_string(), "cartographier".to_string()],
        };

        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#with addr}}{{city}}{{/with}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t1", "{{#with notfound}}hello{{else}}world{{/with}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t2", "{{#with addr/country}}{{this}}{{/with}}")
            .is_ok());

        let r0 = handlebars.render("t0", &person);
        assert_eq!(r0.ok().unwrap(), "Beijing".to_string());

        let r1 = handlebars.render("t1", &person);
        assert_eq!(r1.ok().unwrap(), "world".to_string());

        let r2 = handlebars.render("t2", &person);
        assert_eq!(r2.ok().unwrap(), "China".to_string());
    }

    #[test]
    fn test_with_block_param() {
        let addr = Address {
            city: "Beijing".to_string(),
            country: "China".to_string(),
        };

        let person = Person {
            name: "Ning Sun".to_string(),
            age: 27,
            addr,
            titles: vec!["programmer".to_string(), "cartographier".to_string()],
        };

        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#with addr as |a|}}{{a.city}}{{/with}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t1", "{{#with notfound as |c|}}hello{{else}}world{{/with}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t2", "{{#with addr/country as |t|}}{{t}}{{/with}}")
            .is_ok());

        let r0 = handlebars.render("t0", &person);
        assert_eq!(r0.ok().unwrap(), "Beijing".to_string());

        let r1 = handlebars.render("t1", &person);
        assert_eq!(r1.ok().unwrap(), "world".to_string());

        let r2 = handlebars.render("t2", &person);
        assert_eq!(r2.ok().unwrap(), "China".to_string());
    }

    #[test]
    fn test_with_in_each() {
        let addr = Address {
            city: "Beijing".to_string(),
            country: "China".to_string(),
        };

        let person = Person {
            name: "Ning Sun".to_string(),
            age: 27,
            addr,
            titles: vec!["programmer".to_string(), "cartographier".to_string()],
        };

        let addr2 = Address {
            city: "Beijing".to_string(),
            country: "China".to_string(),
        };

        let person2 = Person {
            name: "Ning Sun".to_string(),
            age: 27,
            addr: addr2,
            titles: vec!["programmer".to_string(), "cartographier".to_string()],
        };

        let people = vec![person, person2];

        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string(
                "t0",
                "{{#each this}}{{#with addr}}{{city}}{{/with}}{{/each}}"
            )
            .is_ok());
        assert!(handlebars
            .register_template_string(
                "t1",
                "{{#each this}}{{#with addr}}{{../age}}{{/with}}{{/each}}"
            )
            .is_ok());
        assert!(handlebars
            .register_template_string(
                "t2",
                "{{#each this}}{{#with addr}}{{@../index}}{{/with}}{{/each}}"
            )
            .is_ok());

        let r0 = handlebars.render("t0", &people);
        assert_eq!(r0.ok().unwrap(), "BeijingBeijing".to_string());

        let r1 = handlebars.render("t1", &people);
        assert_eq!(r1.ok().unwrap(), "2727".to_string());

        let r2 = handlebars.render("t2", &people);
        assert_eq!(r2.ok().unwrap(), "01".to_string());
    }

    #[test]
    fn test_path_up() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#with a}}{{#with b}}{{../../d}}{{/with}}{{/with}}")
            .is_ok());
        let data = btreemap! {
            "a".to_string() => to_json(&btreemap! {
                "b".to_string() => vec![btreemap!{"c".to_string() => vec![1]}]
            }),
            "d".to_string() => to_json(1)
        };

        let r0 = handlebars.render("t0", &data);
        assert_eq!(r0.ok().unwrap(), "1".to_string());
    }
}
<|endoftext|>


<commit_before>/* diosix high-level hypervisor's locking primitives
 *
 * Provides a standard spin lock and a mutex
 *
 * The mutex is reentrant, which means when a physical
 * core holds a mutex and then tries to acquire it
 * again, this operation will succeed.
 *
 * this is so that, eg, if a core holds a mutex
 * and is interrupted, it can regain access
 * to the locked data, use it, and release it.
 *
 * use lock() to acquire a mutex.
 * it is unlocked when it goes out of scope.
 * the mutex also maintains accounting stats
 * and is named to aid debugging.
 *
 * (c) Chris Williams, 2021.
 *
 * See LICENSE for usage and copying.
 */

use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use core::cell::UnsafeCell;
use core::ops::{Deref, DerefMut};
use super::pcore::PhysicalCore;

/* if a lock() call spins more than DEADLOCK_THRESHOLD times
   then it's considered a deadlocked mutex */
const DEADLOCK_THRESHOLD: usize = 1000000;

/* define a snip lock primitive */
pub struct SpinLock
{
    lock: AtomicBool
}

impl SpinLock
{
    pub fn new() -> SpinLock
    {
        SpinLock { lock: AtomicBool::new(false) }
    }

    /* spin until the lock value == must_equal, and then atomically do lock value = new_value */
    fn spin(&self, must_equal: bool, new_value: bool)
    {
        loop
        {
            if self.lock.compare_exchange(must_equal, new_value, Ordering::Acquire, Ordering::Relaxed) == Ok(must_equal)
            {
                return;
            }
        }
    }

    /* acquire the lock, and block until successful */
    pub fn lock(&self)
    {
        self.spin(false, true);
    }

    /* release the lock */
    pub fn unlock(&self)
    {
        self.spin(true, false);
    }
}

pub struct Mutex<T>
{
    /* the data we're protecting */
    content: UnsafeCell<T>,

    /* owner_lock protects owned and owner.
       if the owned is false, then the mutex is considered free.
       if the owned is true, the mutex is considered held by a physical core whose ID == owner */
    owner_lock: SpinLock,
    owned: AtomicBool,
    owner: AtomicUsize,

    /* accounting */
    lock_attempts: AtomicUsize,
    lock_count: AtomicUsize,
    description: &'static str
}

/* Mutex uses the same API as std's Mutex. Create a Mutex using new() and then
   call lock() to block until mutex successfully acquired. Drop the mutex to release */
impl<T> Mutex<T>
{
    pub fn new(description: &'static str, data: T) -> Mutex<T>
    {
        Mutex
        {
            content: UnsafeCell::new(data),
            owner_lock: SpinLock::new(),
            owned: AtomicBool::new(false),
            owner: AtomicUsize::new(0),
            lock_attempts: AtomicUsize::new(0),
            lock_count: AtomicUsize::new(0),
            description
        }
    }

    /* spin until ready to return reference to protected data */
    pub fn lock(&self) -> MutexGuard<'_, T>
    {
        let mut attempts = 0;

        let this_pcore_id = PhysicalCore::get_id();
        loop
        {
            /* hold the spin lock while checking the metadata */
            self.owner_lock.lock();
            self.lock_attempts.fetch_add(1, Ordering::Relaxed);
            attempts = attempts + 1;
            if attempts == DEADLOCK_THRESHOLD
            {
                hvdebug!("BUG: {} mutex ({:p}) may be deadlocked", self.description, &self.content);
            }

            /* determine if the mutex is available, or may even
               already be held by this physical core */
            if self.owned.load(Ordering::SeqCst) == false
            {
                /* lock is available so claim it */
                self.owned.store(true, Ordering::SeqCst);
                self.owner.store(this_pcore_id, Ordering::SeqCst);
                break;
            }
            else
            {
                /* mutex is already held though this pcore may own it anyway */
                if self.owner.load(Ordering::SeqCst) == this_pcore_id
                {
                    break;
                }
            }

            /* give another core a chance to acquire the mutex */
            self.owner_lock.unlock();
        }

        /* don't forget to unlock the metadata
           before returning a reference to the content */
        self.lock_count.fetch_add(1, Ordering::Relaxed);
        self.owner_lock.unlock();
        MutexGuard { mutex: &self }
    }

    /* unlock the mutex */
    fn unlock(&self)
    {
        self.owner_lock.lock();
        self.owned.store(false, Ordering::SeqCst);
        self.owner_lock.unlock();
    }

    /* return true if the mutex is locked, or false if not */
    pub fn is_locked(&self) -> bool
    {
        self.owner_lock.lock();
        let locked = self.owned.load(Ordering::SeqCst);
        self.owner_lock.unlock();
        locked
    }
}

/* pretty print a mutex's stats */
impl<T> core::fmt::Debug for MutexGuard<'_, T>
{
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result
    {
        write!(f, "{} attempts to acquire {}, {} succeeded",
            self.mutex.lock_attempts.load(Ordering::Relaxed),
            self.mutex.description,
            self.mutex.lock_count.load(Ordering::Relaxed))
    }
}

pub struct MutexGuard<'a, T>
{
    mutex: &'a Mutex<T>,
}

impl<T> Deref for MutexGuard<'_, T>
{
    type Target = T;

    fn deref(&self) -> &Self::Target
    {
        unsafe { &*self.mutex.content.get() }
    }
}

impl<T> DerefMut for MutexGuard<'_, T>
{
    fn deref_mut(&mut self) -> &mut Self::Target
    {
        unsafe { &mut *self.mutex.content.get() }
    }
}

impl<T> Drop for MutexGuard<'_, T>
{
    fn drop(&mut self)
    {
        self.mutex.unlock()
    }
}

/* keep rustc happy */
unsafe impl<T> Send for Mutex<T> where T: Send {}
unsafe impl<T> Sync for Mutex<T> where T: Send {}
unsafe impl<T> Send for MutexGuard<'_, T> where T: Send {}
unsafe impl<T> Sync for MutexGuard<'_, T> where T: Send + Sync {}
<commit_msg>Clarify in comments what needs to be dropped to release a lock<commit_after>/* diosix high-level hypervisor's locking primitives
 *
 * Provides a standard spin lock and a mutex
 *
 * The mutex is reentrant, which means when a physical
 * core holds a mutex and then tries to acquire it
 * again, this operation will succeed.
 *
 * this is so that, eg, if a core holds a mutex
 * and is interrupted, it can regain access
 * to the locked data, use it, and release it.
 *
 * use lock() to acquire a mutex.
 * it is unlocked when it goes out of scope.
 * the mutex also maintains accounting stats
 * and is named to aid debugging.
 *
 * (c) Chris Williams, 2021.
 *
 * See LICENSE for usage and copying.
 */

use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use core::cell::UnsafeCell;
use core::ops::{Deref, DerefMut};
use super::pcore::PhysicalCore;

/* if a lock() call spins more than DEADLOCK_THRESHOLD times
   then it's considered a deadlocked mutex */
const DEADLOCK_THRESHOLD: usize = 1000000;

/* define a snip lock primitive */
pub struct SpinLock
{
    lock: AtomicBool
}

impl SpinLock
{
    pub fn new() -> SpinLock
    {
        SpinLock { lock: AtomicBool::new(false) }
    }

    /* spin until the lock value == must_equal, and then atomically do lock value = new_value */
    fn spin(&self, must_equal: bool, new_value: bool)
    {
        loop
        {
            if self.lock.compare_exchange(must_equal, new_value, Ordering::Acquire, Ordering::Relaxed) == Ok(must_equal)
            {
                return;
            }
        }
    }

    /* acquire the lock, and block until successful */
    pub fn lock(&self)
    {
        self.spin(false, true);
    }

    /* release the lock */
    pub fn unlock(&self)
    {
        self.spin(true, false);
    }
}

pub struct Mutex<T>
{
    /* the data we're protecting */
    content: UnsafeCell<T>,

    /* owner_lock protects owned and owner.
       if the owned is false, then the mutex is considered free.
       if the owned is true, the mutex is considered held by a physical core whose ID == owner */
    owner_lock: SpinLock,
    owned: AtomicBool,
    owner: AtomicUsize,

    /* accounting */
    lock_attempts: AtomicUsize,
    lock_count: AtomicUsize,
    description: &'static str
}

/* Mutex uses the same API as std's Mutex. Create a Mutex using new() and then
   call lock() to block until mutex successfully acquired. Drop the mutex guard to release */
impl<T> Mutex<T>
{
    pub fn new(description: &'static str, data: T) -> Mutex<T>
    {
        Mutex
        {
            content: UnsafeCell::new(data),
            owner_lock: SpinLock::new(),
            owned: AtomicBool::new(false),
            owner: AtomicUsize::new(0),
            lock_attempts: AtomicUsize::new(0),
            lock_count: AtomicUsize::new(0),
            description
        }
    }

    /* spin until ready to return reference to protected data */
    pub fn lock(&self) -> MutexGuard<'_, T>
    {
        let mut attempts = 0;

        let this_pcore_id = PhysicalCore::get_id();
        loop
        {
            /* hold the spin lock while checking the metadata */
            self.owner_lock.lock();
            self.lock_attempts.fetch_add(1, Ordering::Relaxed);
            attempts = attempts + 1;
            if attempts == DEADLOCK_THRESHOLD
            {
                hvdebug!("BUG: {} mutex ({:p}) may be deadlocked", self.description, &self.content);
            }

            /* determine if the mutex is available, or may even
               already be held by this physical core */
            if self.owned.load(Ordering::SeqCst) == false
            {
                /* lock is available so claim it */
                self.owned.store(true, Ordering::SeqCst);
                self.owner.store(this_pcore_id, Ordering::SeqCst);
                break;
            }
            else
            {
                /* mutex is already held though this pcore may own it anyway */
                if self.owner.load(Ordering::SeqCst) == this_pcore_id
                {
                    break;
                }
            }

            /* give another core a chance to acquire the mutex */
            self.owner_lock.unlock();
        }

        /* don't forget to unlock the metadata
           before returning a reference to the content */
        self.lock_count.fetch_add(1, Ordering::Relaxed);
        self.owner_lock.unlock();
        MutexGuard { mutex: &self }
    }

    /* unlock the mutex */
    fn unlock(&self)
    {
        self.owner_lock.lock();
        self.owned.store(false, Ordering::SeqCst);
        self.owner_lock.unlock();
    }

    /* return true if the mutex is locked, or false if not */
    pub fn is_locked(&self) -> bool
    {
        self.owner_lock.lock();
        let locked = self.owned.load(Ordering::SeqCst);
        self.owner_lock.unlock();
        locked
    }
}

/* pretty print a mutex's stats */
impl<T> core::fmt::Debug for MutexGuard<'_, T>
{
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result
    {
        write!(f, "{} attempts to acquire {}, {} succeeded",
            self.mutex.lock_attempts.load(Ordering::Relaxed),
            self.mutex.description,
            self.mutex.lock_count.load(Ordering::Relaxed))
    }
}

pub struct MutexGuard<'a, T>
{
    mutex: &'a Mutex<T>,
}

impl<T> Deref for MutexGuard<'_, T>
{
    type Target = T;

    fn deref(&self) -> &Self::Target
    {
        unsafe { &*self.mutex.content.get() }
    }
}

impl<T> DerefMut for MutexGuard<'_, T>
{
    fn deref_mut(&mut self) -> &mut Self::Target
    {
        unsafe { &mut *self.mutex.content.get() }
    }
}

impl<T> Drop for MutexGuard<'_, T>
{
    fn drop(&mut self)
    {
        self.mutex.unlock()
    }
}

/* keep rustc happy */
unsafe impl<T> Send for Mutex<T> where T: Send {}
unsafe impl<T> Sync for Mutex<T> where T: Send {}
unsafe impl<T> Send for MutexGuard<'_, T> where T: Send {}
unsafe impl<T> Sync for MutexGuard<'_, T> where T: Send + Sync {}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashSet};

use core::PackageId;
use core::registry::PackageRegistry;
use core::{Resolve, SourceId, Workspace};
use core::resolver::Method;
use ops;
use util::config::Config;
use util::CargoResult;

pub struct UpdateOptions<'a> {
    pub config: &'a Config,
    pub to_update: &'a [String],
    pub precise: Option<&'a str>,
    pub aggressive: bool,
}

pub fn generate_lockfile(ws: &Workspace) -> CargoResult<()> {
    let mut registry = PackageRegistry::new(ws.config())?;
    let resolve = ops::resolve_with_previous(&mut registry, ws,
                                             Method::Everything,
                                             None, None, &[], true)?;
    ops::write_pkg_lockfile(ws, &resolve)?;
    Ok(())
}

pub fn update_lockfile(ws: &Workspace, opts: &UpdateOptions)
                       -> CargoResult<()> {

    if opts.aggressive && opts.precise.is_some() {
        bail!("cannot specify both aggressive and precise simultaneously")
    }

    if ws.members().is_empty() {
        bail!("you can't generate a lockfile for an empty workspace.")
    }

    let previous_resolve = match ops::load_pkg_lockfile(ws)? {
        Some(resolve) => resolve,
        None => return generate_lockfile(ws),
    };
    let mut registry = PackageRegistry::new(opts.config)?;
    let mut to_avoid = HashSet::new();

    if opts.to_update.is_empty() {
        to_avoid.extend(previous_resolve.iter());
    } else {
        let mut sources = Vec::new();
        for name in opts.to_update {
            let dep = previous_resolve.query(name)?;
            if opts.aggressive {
                fill_with_deps(&previous_resolve, dep, &mut to_avoid,
                               &mut HashSet::new());
            } else {
                to_avoid.insert(dep);
                sources.push(match opts.precise {
                    Some(precise) => {
                        // TODO: see comment in `resolve.rs` as well, but this
                        //       seems like a pretty hokey reason to single out
                        //       the registry as well.
                        let precise = if dep.source_id().is_registry() {
                            format!("{}={}", dep.name(), precise)
                        } else {
                            precise.to_string()
                        };
                        dep.source_id().clone().with_precise(Some(precise))
                    }
                    None => {
                        dep.source_id().clone().with_precise(None)
                    }
                });
            }
        }
        registry.add_sources(&sources)?;
    }

    let resolve = ops::resolve_with_previous(&mut registry,
                                                  ws,
                                                  Method::Everything,
                                                  Some(&previous_resolve),
                                                  Some(&to_avoid),
                                                  &[],
                                                  true)?;

    // Summarize what is changing for the user.
    let print_change = |status: &str, msg: String| {
        opts.config.shell().status(status, msg)
    };
    for (removed, added) in compare_dependency_graphs(&previous_resolve, &resolve) {
        if removed.len() == 1 && added.len() == 1 {
            let msg = if removed[0].source_id().is_git() {
                format!("{} -> #{}", removed[0],
                        &added[0].source_id().precise().unwrap()[..8])
            } else {
                format!("{} -> v{}", removed[0], added[0].version())
            };
            print_change("Updating", msg)?;
        } else {
            for package in removed.iter() {
                print_change("Removing", format!("{}", package))?;
            }
            for package in added.iter() {
                print_change("Adding", format!("{}", package))?;
            }
        }
    }

    ops::write_pkg_lockfile(ws, &resolve)?;
    return Ok(());

    fn fill_with_deps<'a>(resolve: &'a Resolve, dep: &'a PackageId,
                          set: &mut HashSet<&'a PackageId>,
                          visited: &mut HashSet<&'a PackageId>) {
        if !visited.insert(dep) {
            return
        }
        set.insert(dep);
        for dep in resolve.deps(dep) {
            fill_with_deps(resolve, dep, set, visited);
        }
    }

    fn compare_dependency_graphs<'a>(previous_resolve: &'a Resolve,
                                     resolve: &'a Resolve) ->
                                     Vec<(Vec<&'a PackageId>, Vec<&'a PackageId>)> {
        fn key(dep: &PackageId) -> (&str, &SourceId) {
            (dep.name(), dep.source_id())
        }

        // Removes all package ids in `b` from `a`. Note that this is somewhat
        // more complicated because the equality for source ids does not take
        // precise versions into account (e.g. git shas), but we want to take
        // that into account here.
        fn vec_subtract<'a>(a: &[&'a PackageId],
                            b: &[&'a PackageId]) -> Vec<&'a PackageId> {
            a.iter().filter(|a| {
                // If this package id is not found in `b`, then it's definitely
                // in the subtracted set
                let i = match b.binary_search(a) {
                    Ok(i) => i,
                    Err(..) => return true,
                };

                // If we've found `a` in `b`, then we iterate over all instances
                // (we know `b` is sorted) and see if they all have different
                // precise versions. If so, then `a` isn't actually in `b` so
                // we'll let it through.
                //
                // Note that we only check this for non-registry sources,
                // however, as registries contain enough version information in
                // the package id to disambiguate
                if a.source_id().is_registry() {
                    return false
                }
                b[i..].iter().take_while(|b| a == b).all(|b| {
                    a.source_id().precise() != b.source_id().precise()
                })
            }).cloned().collect()
        }

        // Map (package name, package source) to (removed versions, added versions).
        let mut changes = BTreeMap::new();
        let empty = (Vec::new(), Vec::new());
        for dep in previous_resolve.iter() {
            changes.entry(key(dep)).or_insert(empty.clone()).0.push(dep);
        }
        for dep in resolve.iter() {
            changes.entry(key(dep)).or_insert(empty.clone()).1.push(dep);
        }

        for v in changes.values_mut() {
            let (ref mut old, ref mut new) = *v;
            old.sort();
            new.sort();
            let removed = vec_subtract(old, new);
            let added = vec_subtract(new, old);
            *old = removed;
            *new = added;
        }
        debug!("{:#?}", changes);

        changes.into_iter().map(|(_, v)| v).collect()
    }
}
<commit_msg>cargo_generate_lockfile: use color to also indicate the change<commit_after>use std::collections::{BTreeMap, HashSet};

use termcolor::Color::{self, Cyan, Green, Red};

use core::PackageId;
use core::registry::PackageRegistry;
use core::{Resolve, SourceId, Workspace};
use core::resolver::Method;
use ops;
use util::config::Config;
use util::CargoResult;

pub struct UpdateOptions<'a> {
    pub config: &'a Config,
    pub to_update: &'a [String],
    pub precise: Option<&'a str>,
    pub aggressive: bool,
}

pub fn generate_lockfile(ws: &Workspace) -> CargoResult<()> {
    let mut registry = PackageRegistry::new(ws.config())?;
    let resolve = ops::resolve_with_previous(&mut registry, ws,
                                             Method::Everything,
                                             None, None, &[], true)?;
    ops::write_pkg_lockfile(ws, &resolve)?;
    Ok(())
}

pub fn update_lockfile(ws: &Workspace, opts: &UpdateOptions)
                       -> CargoResult<()> {

    if opts.aggressive && opts.precise.is_some() {
        bail!("cannot specify both aggressive and precise simultaneously")
    }

    if ws.members().is_empty() {
        bail!("you can't generate a lockfile for an empty workspace.")
    }

    let previous_resolve = match ops::load_pkg_lockfile(ws)? {
        Some(resolve) => resolve,
        None => return generate_lockfile(ws),
    };
    let mut registry = PackageRegistry::new(opts.config)?;
    let mut to_avoid = HashSet::new();

    if opts.to_update.is_empty() {
        to_avoid.extend(previous_resolve.iter());
    } else {
        let mut sources = Vec::new();
        for name in opts.to_update {
            let dep = previous_resolve.query(name)?;
            if opts.aggressive {
                fill_with_deps(&previous_resolve, dep, &mut to_avoid,
                               &mut HashSet::new());
            } else {
                to_avoid.insert(dep);
                sources.push(match opts.precise {
                    Some(precise) => {
                        // TODO: see comment in `resolve.rs` as well, but this
                        //       seems like a pretty hokey reason to single out
                        //       the registry as well.
                        let precise = if dep.source_id().is_registry() {
                            format!("{}={}", dep.name(), precise)
                        } else {
                            precise.to_string()
                        };
                        dep.source_id().clone().with_precise(Some(precise))
                    }
                    None => {
                        dep.source_id().clone().with_precise(None)
                    }
                });
            }
        }
        registry.add_sources(&sources)?;
    }

    let resolve = ops::resolve_with_previous(&mut registry,
                                                  ws,
                                                  Method::Everything,
                                                  Some(&previous_resolve),
                                                  Some(&to_avoid),
                                                  &[],
                                                  true)?;

    // Summarize what is changing for the user.
    let print_change = |status: &str, msg: String, color: Color| {
        opts.config.shell().status_with_color(status, msg, color)
    };
    for (removed, added) in compare_dependency_graphs(&previous_resolve, &resolve) {
        if removed.len() == 1 && added.len() == 1 {
            let msg = if removed[0].source_id().is_git() {
                format!("{} -> #{}", removed[0],
                        &added[0].source_id().precise().unwrap()[..8])
            } else {
                format!("{} -> v{}", removed[0], added[0].version())
            };
            print_change("Updating", msg, Cyan)?;
        } else {
            for package in removed.iter() {
                print_change("Removing", format!("{}", package), Red)?;
            }
            for package in added.iter() {
                print_change("Adding", format!("{}", package), Green)?;
            }
        }
    }

    ops::write_pkg_lockfile(ws, &resolve)?;
    return Ok(());

    fn fill_with_deps<'a>(resolve: &'a Resolve, dep: &'a PackageId,
                          set: &mut HashSet<&'a PackageId>,
                          visited: &mut HashSet<&'a PackageId>) {
        if !visited.insert(dep) {
            return
        }
        set.insert(dep);
        for dep in resolve.deps(dep) {
            fill_with_deps(resolve, dep, set, visited);
        }
    }

    fn compare_dependency_graphs<'a>(previous_resolve: &'a Resolve,
                                     resolve: &'a Resolve) ->
                                     Vec<(Vec<&'a PackageId>, Vec<&'a PackageId>)> {
        fn key(dep: &PackageId) -> (&str, &SourceId) {
            (dep.name(), dep.source_id())
        }

        // Removes all package ids in `b` from `a`. Note that this is somewhat
        // more complicated because the equality for source ids does not take
        // precise versions into account (e.g. git shas), but we want to take
        // that into account here.
        fn vec_subtract<'a>(a: &[&'a PackageId],
                            b: &[&'a PackageId]) -> Vec<&'a PackageId> {
            a.iter().filter(|a| {
                // If this package id is not found in `b`, then it's definitely
                // in the subtracted set
                let i = match b.binary_search(a) {
                    Ok(i) => i,
                    Err(..) => return true,
                };

                // If we've found `a` in `b`, then we iterate over all instances
                // (we know `b` is sorted) and see if they all have different
                // precise versions. If so, then `a` isn't actually in `b` so
                // we'll let it through.
                //
                // Note that we only check this for non-registry sources,
                // however, as registries contain enough version information in
                // the package id to disambiguate
                if a.source_id().is_registry() {
                    return false
                }
                b[i..].iter().take_while(|b| a == b).all(|b| {
                    a.source_id().precise() != b.source_id().precise()
                })
            }).cloned().collect()
        }

        // Map (package name, package source) to (removed versions, added versions).
        let mut changes = BTreeMap::new();
        let empty = (Vec::new(), Vec::new());
        for dep in previous_resolve.iter() {
            changes.entry(key(dep)).or_insert(empty.clone()).0.push(dep);
        }
        for dep in resolve.iter() {
            changes.entry(key(dep)).or_insert(empty.clone()).1.push(dep);
        }

        for v in changes.values_mut() {
            let (ref mut old, ref mut new) = *v;
            old.sort();
            new.sort();
            let removed = vec_subtract(old, new);
            let added = vec_subtract(new, old);
            *old = removed;
            *new = added;
        }
        debug!("{:#?}", changes);

        changes.into_iter().map(|(_, v)| v).collect()
    }
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashSet};

use log::debug;
use termcolor::Color::{self, Cyan, Green, Red};

use crate::core::registry::PackageRegistry;
use crate::core::resolver::Method;
use crate::core::PackageId;
use crate::core::{Resolve, SourceId, Workspace};
use crate::ops;
use crate::util::config::Config;
use crate::util::CargoResult;

pub struct UpdateOptions<'a> {
    pub config: &'a Config,
    pub to_update: Vec<String>,
    pub precise: Option<&'a str>,
    pub aggressive: bool,
    pub dry_run: bool,
}

pub fn generate_lockfile(ws: &Workspace<'_>) -> CargoResult<()> {
    let mut registry = PackageRegistry::new(ws.config())?;
    let resolve =
        ops::resolve_with_previous(&mut registry, ws, Method::Everything, None, None, &[], true)?;
    ops::write_pkg_lockfile(ws, &resolve)?;
    Ok(())
}

pub fn update_lockfile(ws: &Workspace<'_>, opts: &UpdateOptions<'_>) -> CargoResult<()> {
    if opts.aggressive && opts.precise.is_some() {
        failure::bail!("cannot specify both aggressive and precise simultaneously")
    }

    if ws.members().count() == 0 {
        failure::bail!("you can't generate a lockfile for an empty workspace.")
    }

    if opts.config.cli_unstable().offline {
        failure::bail!("you can't update in the offline mode");
    }

    // Updates often require a lot of modifications to the registry, so ensure
    // that we're synchronized against other Cargos.
    let _lock = ws.config().acquire_package_cache_lock()?;

    let previous_resolve = match ops::load_pkg_lockfile(ws)? {
        Some(resolve) => resolve,
        None => {
            generate_lockfile(ws)?;
            match opts.precise {
                None => return Ok(()),

                // Precise option specified, so re-run to do the rest of the update.
                Some(_) => return update_lockfile(ws, opts)
            }
        }
    };
    let mut registry = PackageRegistry::new(opts.config)?;
    let mut to_avoid = HashSet::new();

    if opts.to_update.is_empty() {
        to_avoid.extend(previous_resolve.iter());
    } else {
        let mut sources = Vec::new();
        for name in opts.to_update.iter() {
            let dep = previous_resolve.query(name)?;
            if opts.aggressive {
                fill_with_deps(&previous_resolve, dep, &mut to_avoid, &mut HashSet::new());
            } else {
                to_avoid.insert(dep);
                sources.push(match opts.precise {
                    Some(precise) => {
                        // TODO: see comment in `resolve.rs` as well, but this
                        //       seems like a pretty hokey reason to single out
                        //       the registry as well.
                        let precise = if dep.source_id().is_registry() {
                            format!("{}={}->{}", dep.name(), dep.version(), precise)
                        } else {
                            precise.to_string()
                        };
                        dep.source_id().with_precise(Some(precise))
                    }
                    None => dep.source_id().with_precise(None),
                });
            }
        }

        registry.add_sources(sources)?;
    }

    let resolve = ops::resolve_with_previous(
        &mut registry,
        ws,
        Method::Everything,
        Some(&previous_resolve),
        Some(&to_avoid),
        &[],
        true,
    )?;

    // Summarize what is changing for the user.
    let print_change = |status: &str, msg: String, color: Color| {
        opts.config.shell().status_with_color(status, msg, color)
    };
    for (removed, added) in compare_dependency_graphs(&previous_resolve, &resolve) {
        if removed.len() == 1 && added.len() == 1 {
            let msg = if removed[0].source_id().is_git() {
                format!(
                    "{} -> #{}",
                    removed[0],
                    &added[0].source_id().precise().unwrap()[..8]
                )
            } else {
                format!("{} -> v{}", removed[0], added[0].version())
            };
            print_change("Updating", msg, Green)?;
        } else {
            for package in removed.iter() {
                print_change("Removing", format!("{}", package), Red)?;
            }
            for package in added.iter() {
                print_change("Adding", format!("{}", package), Cyan)?;
            }
        }
    }
    if opts.dry_run {
        opts.config
            .shell()
            .warn("not updating lockfile due to dry run")?;
    } else {
        ops::write_pkg_lockfile(ws, &resolve)?;
    }
    return Ok(());

    fn fill_with_deps<'a>(
        resolve: &'a Resolve,
        dep: PackageId,
        set: &mut HashSet<PackageId>,
        visited: &mut HashSet<PackageId>,
    ) {
        if !visited.insert(dep) {
            return;
        }
        set.insert(dep);
        for dep in resolve.deps_not_replaced(dep) {
            fill_with_deps(resolve, dep, set, visited);
        }
    }

    fn compare_dependency_graphs(
        previous_resolve: &Resolve,
        resolve: &Resolve,
    ) -> Vec<(Vec<PackageId>, Vec<PackageId>)> {
        fn key(dep: PackageId) -> (&'static str, SourceId) {
            (dep.name().as_str(), dep.source_id())
        }

        // Removes all package IDs in `b` from `a`. Note that this is somewhat
        // more complicated because the equality for source IDs does not take
        // precise versions into account (e.g., git shas), but we want to take
        // that into account here.
        fn vec_subtract(a: &[PackageId], b: &[PackageId]) -> Vec<PackageId> {
            a.iter()
                .filter(|a| {
                    // If this package ID is not found in `b`, then it's definitely
                    // in the subtracted set.
                    let i = match b.binary_search(a) {
                        Ok(i) => i,
                        Err(..) => return true,
                    };

                    // If we've found `a` in `b`, then we iterate over all instances
                    // (we know `b` is sorted) and see if they all have different
                    // precise versions. If so, then `a` isn't actually in `b` so
                    // we'll let it through.
                    //
                    // Note that we only check this for non-registry sources,
                    // however, as registries contain enough version information in
                    // the package ID to disambiguate.
                    if a.source_id().is_registry() {
                        return false;
                    }
                    b[i..]
                        .iter()
                        .take_while(|b| a == b)
                        .all(|b| a.source_id().precise() != b.source_id().precise())
                })
                .cloned()
                .collect()
        }

        // Map `(package name, package source)` to `(removed versions, added versions)`.
        let mut changes = BTreeMap::new();
        let empty = (Vec::new(), Vec::new());
        for dep in previous_resolve.iter() {
            changes
                .entry(key(dep))
                .or_insert_with(|| empty.clone())
                .0
                .push(dep);
        }
        for dep in resolve.iter() {
            changes
                .entry(key(dep))
                .or_insert_with(|| empty.clone())
                .1
                .push(dep);
        }

        for v in changes.values_mut() {
            let (ref mut old, ref mut new) = *v;
            old.sort();
            new.sort();
            let removed = vec_subtract(old, new);
            let added = vec_subtract(new, old);
            *old = removed;
            *new = added;
        }
        debug!("{:#?}", changes);

        changes.into_iter().map(|(_, v)| v).collect()
    }
}
<commit_msg>Re-write to avoid recursion.<commit_after>use std::collections::{BTreeMap, HashSet};

use log::debug;
use termcolor::Color::{self, Cyan, Green, Red};

use crate::core::registry::PackageRegistry;
use crate::core::resolver::Method;
use crate::core::PackageId;
use crate::core::{Resolve, SourceId, Workspace};
use crate::ops;
use crate::util::config::Config;
use crate::util::CargoResult;

pub struct UpdateOptions<'a> {
    pub config: &'a Config,
    pub to_update: Vec<String>,
    pub precise: Option<&'a str>,
    pub aggressive: bool,
    pub dry_run: bool,
}

pub fn generate_lockfile(ws: &Workspace<'_>) -> CargoResult<()> {
    let mut registry = PackageRegistry::new(ws.config())?;
    let resolve =
        ops::resolve_with_previous(&mut registry, ws, Method::Everything, None, None, &[], true)?;
    ops::write_pkg_lockfile(ws, &resolve)?;
    Ok(())
}

pub fn update_lockfile(ws: &Workspace<'_>, opts: &UpdateOptions<'_>) -> CargoResult<()> {
    if opts.aggressive && opts.precise.is_some() {
        failure::bail!("cannot specify both aggressive and precise simultaneously")
    }

    if ws.members().count() == 0 {
        failure::bail!("you can't generate a lockfile for an empty workspace.")
    }

    if opts.config.cli_unstable().offline {
        failure::bail!("you can't update in the offline mode");
    }

    // Updates often require a lot of modifications to the registry, so ensure
    // that we're synchronized against other Cargos.
    let _lock = ws.config().acquire_package_cache_lock()?;

    let previous_resolve = match ops::load_pkg_lockfile(ws)? {
        Some(resolve) => resolve,
        None => {
            match opts.precise {
                None => return generate_lockfile(ws),

                // Precise option specified, so calculate a previous_resolve required
                // by precise package update later.
                Some(_) => {
                    let mut registry = PackageRegistry::new(opts.config)?;
                    ops::resolve_with_previous(&mut registry, ws, Method::Everything,
                                               None, None, &[], true)?
                }
            }
        }
    };
    let mut registry = PackageRegistry::new(opts.config)?;
    let mut to_avoid = HashSet::new();

    if opts.to_update.is_empty() {
        to_avoid.extend(previous_resolve.iter());
    } else {
        let mut sources = Vec::new();
        for name in opts.to_update.iter() {
            let dep = previous_resolve.query(name)?;
            if opts.aggressive {
                fill_with_deps(&previous_resolve, dep, &mut to_avoid, &mut HashSet::new());
            } else {
                to_avoid.insert(dep);
                sources.push(match opts.precise {
                    Some(precise) => {
                        // TODO: see comment in `resolve.rs` as well, but this
                        //       seems like a pretty hokey reason to single out
                        //       the registry as well.
                        let precise = if dep.source_id().is_registry() {
                            format!("{}={}->{}", dep.name(), dep.version(), precise)
                        } else {
                            precise.to_string()
                        };
                        dep.source_id().with_precise(Some(precise))
                    }
                    None => dep.source_id().with_precise(None),
                });
            }
        }

        registry.add_sources(sources)?;
    }

    let resolve = ops::resolve_with_previous(
        &mut registry,
        ws,
        Method::Everything,
        Some(&previous_resolve),
        Some(&to_avoid),
        &[],
        true,
    )?;

    // Summarize what is changing for the user.
    let print_change = |status: &str, msg: String, color: Color| {
        opts.config.shell().status_with_color(status, msg, color)
    };
    for (removed, added) in compare_dependency_graphs(&previous_resolve, &resolve) {
        if removed.len() == 1 && added.len() == 1 {
            let msg = if removed[0].source_id().is_git() {
                format!(
                    "{} -> #{}",
                    removed[0],
                    &added[0].source_id().precise().unwrap()[..8]
                )
            } else {
                format!("{} -> v{}", removed[0], added[0].version())
            };
            print_change("Updating", msg, Green)?;
        } else {
            for package in removed.iter() {
                print_change("Removing", format!("{}", package), Red)?;
            }
            for package in added.iter() {
                print_change("Adding", format!("{}", package), Cyan)?;
            }
        }
    }
    if opts.dry_run {
        opts.config
            .shell()
            .warn("not updating lockfile due to dry run")?;
    } else {
        ops::write_pkg_lockfile(ws, &resolve)?;
    }
    return Ok(());

    fn fill_with_deps<'a>(
        resolve: &'a Resolve,
        dep: PackageId,
        set: &mut HashSet<PackageId>,
        visited: &mut HashSet<PackageId>,
    ) {
        if !visited.insert(dep) {
            return;
        }
        set.insert(dep);
        for dep in resolve.deps_not_replaced(dep) {
            fill_with_deps(resolve, dep, set, visited);
        }
    }

    fn compare_dependency_graphs(
        previous_resolve: &Resolve,
        resolve: &Resolve,
    ) -> Vec<(Vec<PackageId>, Vec<PackageId>)> {
        fn key(dep: PackageId) -> (&'static str, SourceId) {
            (dep.name().as_str(), dep.source_id())
        }

        // Removes all package IDs in `b` from `a`. Note that this is somewhat
        // more complicated because the equality for source IDs does not take
        // precise versions into account (e.g., git shas), but we want to take
        // that into account here.
        fn vec_subtract(a: &[PackageId], b: &[PackageId]) -> Vec<PackageId> {
            a.iter()
                .filter(|a| {
                    // If this package ID is not found in `b`, then it's definitely
                    // in the subtracted set.
                    let i = match b.binary_search(a) {
                        Ok(i) => i,
                        Err(..) => return true,
                    };

                    // If we've found `a` in `b`, then we iterate over all instances
                    // (we know `b` is sorted) and see if they all have different
                    // precise versions. If so, then `a` isn't actually in `b` so
                    // we'll let it through.
                    //
                    // Note that we only check this for non-registry sources,
                    // however, as registries contain enough version information in
                    // the package ID to disambiguate.
                    if a.source_id().is_registry() {
                        return false;
                    }
                    b[i..]
                        .iter()
                        .take_while(|b| a == b)
                        .all(|b| a.source_id().precise() != b.source_id().precise())
                })
                .cloned()
                .collect()
        }

        // Map `(package name, package source)` to `(removed versions, added versions)`.
        let mut changes = BTreeMap::new();
        let empty = (Vec::new(), Vec::new());
        for dep in previous_resolve.iter() {
            changes
                .entry(key(dep))
                .or_insert_with(|| empty.clone())
                .0
                .push(dep);
        }
        for dep in resolve.iter() {
            changes
                .entry(key(dep))
                .or_insert_with(|| empty.clone())
                .1
                .push(dep);
        }

        for v in changes.values_mut() {
            let (ref mut old, ref mut new) = *v;
            old.sort();
            new.sort();
            let removed = vec_subtract(old, new);
            let added = vec_subtract(new, old);
            *old = removed;
            *new = added;
        }
        debug!("{:#?}", changes);

        changes.into_iter().map(|(_, v)| v).collect()
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::collections::range::RangeArgument;
use std::fmt::Debug;
use std::iter::{self, FromIterator};
use std::slice;
use std::marker::PhantomData;
use std::ops::{Index, IndexMut, Range};
use std::fmt;
use std::vec;
use std::u32;

use rustc_serialize as serialize;

/// Represents some newtyped `usize` wrapper.
///
/// (purpose: avoid mixing indexes for different bitvector domains.)
pub trait Idx: Copy + 'static + Eq + Debug {
    fn new(idx: usize) -> Self;
    fn index(self) -> usize;
}

impl Idx for usize {
    fn new(idx: usize) -> Self { idx }
    fn index(self) -> usize { self }
}

impl Idx for u32 {
    fn new(idx: usize) -> Self { assert!(idx <= u32::MAX as usize); idx as u32 }
    fn index(self) -> usize { self as usize }
}

#[macro_export]
macro_rules! newtype_index {
    // ---- public rules ----

    // Use default constants
    ($name:ident) => (
        newtype_index!(
            @type[$name]
            @max[::std::u32::MAX]
            @debug_name[unsafe {::std::intrinsics::type_name::<$name>() }]);
    );

    // Define any constants
    ($name:ident { $($tokens:tt)+ }) => (
        newtype_index!(
            @type[$name]
            @max[::std::u32::MAX]
            @debug_name[unsafe {::std::intrinsics::type_name::<$name>() }]
            $($tokens)+);
    );

    // ---- private rules ----

    // Base case, user-defined constants (if any) have already been defined
    (@type[$type:ident] @max[$max:expr] @debug_name[$debug_name:expr]) => (
        #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord,
            RustcEncodable, RustcDecodable)]
        pub struct $type(u32);

        impl Idx for $type {
            fn new(value: usize) -> Self {
                assert!(value < ($max) as usize);
                $type(value as u32)
            }
            fn index(self) -> usize {
                self.0 as usize
            }
        }

        impl ::std::fmt::Debug for $type {
            fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
                write!(fmt, "{}{}", $debug_name, self.0)
            }
        }
    );

    // Rewrite final without comma to one that includes comma
    (@type[$type:ident] @max[$max:expr] @debug_name[$debug_name:expr] $name:ident = $constant:expr) => (
        newtype_index!(@type[$type] @max[$max] @debug_name[$debug_name] $name = $constant,);
    );

    // Rewrite final const without comma to one that includes comma
    (@type[$type:ident] @max[$_max:expr] @debug_name[$debug_name:expr] const $name:ident = $constant:expr) => (
        newtype_index!(@type[$type] @max[$max] @debug_name[$debug_name] const $name = $constant,);
    );

    // Replace existing default for max
    (@type[$type:ident] @max[$_max:expr] @debug_name[$debug_name:expr] MAX = $max:expr, $($tokens:tt)*) => (
        newtype_index!(@type[$type] @max[$max] @debug_name[$debug_name] $(tokens)*);
    );

    // Replace existing default for debug_name
    (@type[$type:ident] @max[$max:expr] @debug_name[$_debug_name:expr] DEBUG_NAME = $debug_name:expr, $($tokens:tt)*) => (
        newtype_index!(@type[$type] @max[$max] @debug_name[$debug_name] $($tokens)*);
    );

    // Assign a user-defined constant (as final param)
    (@type[$type:ident] @max[$max:expr] @debug_name[$debug_name:expr] const $name:ident = $constant:expr, $($tokens:tt)*) => (
        pub const $name: $type = $type($constant);
        newtype_index!(@type[$type] @max[$max] @debug_name[$debug_name] $($tokens)*);
    );
}

#[derive(Clone, PartialEq, Eq)]
pub struct IndexVec<I: Idx, T> {
    pub raw: Vec<T>,
    _marker: PhantomData<Fn(&I)>
}

impl<I: Idx, T: serialize::Encodable> serialize::Encodable for IndexVec<I, T> {
    fn encode<S: serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        serialize::Encodable::encode(&self.raw, s)
    }
}

impl<I: Idx, T: serialize::Decodable> serialize::Decodable for IndexVec<I, T> {
    fn decode<D: serialize::Decoder>(d: &mut D) -> Result<Self, D::Error> {
        serialize::Decodable::decode(d).map(|v| {
            IndexVec { raw: v, _marker: PhantomData }
        })
    }
}

impl<I: Idx, T: fmt::Debug> fmt::Debug for IndexVec<I, T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.raw, fmt)
    }
}

pub type Enumerated<I, J> = iter::Map<iter::Enumerate<J>, IntoIdx<I>>;

impl<I: Idx, T> IndexVec<I, T> {
    #[inline]
    pub fn new() -> Self {
        IndexVec { raw: Vec::new(), _marker: PhantomData }
    }

    #[inline]
    pub fn with_capacity(capacity: usize) -> Self {
        IndexVec { raw: Vec::with_capacity(capacity), _marker: PhantomData }
    }

    #[inline]
    pub fn from_elem<S>(elem: T, universe: &IndexVec<I, S>) -> Self
        where T: Clone
    {
        IndexVec { raw: vec![elem; universe.len()], _marker: PhantomData }
    }

    #[inline]
    pub fn from_elem_n(elem: T, n: usize) -> Self
        where T: Clone
    {
        IndexVec { raw: vec![elem; n], _marker: PhantomData }
    }

    #[inline]
    pub fn push(&mut self, d: T) -> I {
        let idx = I::new(self.len());
        self.raw.push(d);
        idx
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.raw.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.raw.is_empty()
    }

    #[inline]
    pub fn into_iter(self) -> vec::IntoIter<T> {
        self.raw.into_iter()
    }

    #[inline]
    pub fn into_iter_enumerated(self) -> Enumerated<I, vec::IntoIter<T>>
    {
        self.raw.into_iter().enumerate().map(IntoIdx { _marker: PhantomData })
    }

    #[inline]
    pub fn iter(&self) -> slice::Iter<T> {
        self.raw.iter()
    }

    #[inline]
    pub fn iter_enumerated(&self) -> Enumerated<I, slice::Iter<T>>
    {
        self.raw.iter().enumerate().map(IntoIdx { _marker: PhantomData })
    }

    #[inline]
    pub fn indices(&self) -> iter::Map<Range<usize>, IntoIdx<I>> {
        (0..self.len()).map(IntoIdx { _marker: PhantomData })
    }

    #[inline]
    pub fn iter_mut(&mut self) -> slice::IterMut<T> {
        self.raw.iter_mut()
    }

    #[inline]
    pub fn iter_enumerated_mut(&mut self) -> Enumerated<I, slice::IterMut<T>>
    {
        self.raw.iter_mut().enumerate().map(IntoIdx { _marker: PhantomData })
    }

    #[inline]
    pub fn drain<'a, R: RangeArgument<usize>>(
        &'a mut self, range: R) -> impl Iterator<Item=T> + 'a {
        self.raw.drain(range)
    }

    #[inline]
    pub fn drain_enumerated<'a, R: RangeArgument<usize>>(
        &'a mut self, range: R) -> impl Iterator<Item=(I, T)> + 'a {
        self.raw.drain(range).enumerate().map(IntoIdx { _marker: PhantomData })
    }

    #[inline]
    pub fn last(&self) -> Option<I> {
        self.len().checked_sub(1).map(I::new)
    }

    #[inline]
    pub fn shrink_to_fit(&mut self) {
        self.raw.shrink_to_fit()
    }

    #[inline]
    pub fn swap(&mut self, a: usize, b: usize) {
        self.raw.swap(a, b)
    }

    #[inline]
    pub fn truncate(&mut self, a: usize) {
        self.raw.truncate(a)
    }

    #[inline]
    pub fn get(&self, index: I) -> Option<&T> {
        self.raw.get(index.index())
    }

    #[inline]
    pub fn get_mut(&mut self, index: I) -> Option<&mut T> {
        self.raw.get_mut(index.index())
    }
}

impl<I: Idx, T: Clone> IndexVec<I, T> {
    #[inline]
    pub fn resize(&mut self, new_len: usize, value: T) {
        self.raw.resize(new_len, value)
    }
}

impl<I: Idx, T: Ord> IndexVec<I, T> {
    #[inline]
    pub fn binary_search(&self, value: &T) -> Result<I, I> {
        match self.raw.binary_search(value) {
            Ok(i) => Ok(Idx::new(i)),
            Err(i) => Err(Idx::new(i)),
        }
    }
}

impl<I: Idx, T> Index<I> for IndexVec<I, T> {
    type Output = T;

    #[inline]
    fn index(&self, index: I) -> &T {
        &self.raw[index.index()]
    }
}

impl<I: Idx, T> IndexMut<I> for IndexVec<I, T> {
    #[inline]
    fn index_mut(&mut self, index: I) -> &mut T {
        &mut self.raw[index.index()]
    }
}

impl<I: Idx, T> Default for IndexVec<I, T> {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

impl<I: Idx, T> Extend<T> for IndexVec<I, T> {
    #[inline]
    fn extend<J: IntoIterator<Item = T>>(&mut self, iter: J) {
        self.raw.extend(iter);
    }
}

impl<I: Idx, T> FromIterator<T> for IndexVec<I, T> {
    #[inline]
    fn from_iter<J>(iter: J) -> Self where J: IntoIterator<Item=T> {
        IndexVec { raw: FromIterator::from_iter(iter), _marker: PhantomData }
    }
}

impl<I: Idx, T> IntoIterator for IndexVec<I, T> {
    type Item = T;
    type IntoIter = vec::IntoIter<T>;

    #[inline]
    fn into_iter(self) -> vec::IntoIter<T> {
        self.raw.into_iter()
    }

}

impl<'a, I: Idx, T> IntoIterator for &'a IndexVec<I, T> {
    type Item = &'a T;
    type IntoIter = slice::Iter<'a, T>;

    #[inline]
    fn into_iter(self) -> slice::Iter<'a, T> {
        self.raw.iter()
    }
}

impl<'a, I: Idx, T> IntoIterator for &'a mut IndexVec<I, T> {
    type Item = &'a mut T;
    type IntoIter = slice::IterMut<'a, T>;

    #[inline]
    fn into_iter(self) -> slice::IterMut<'a, T> {
        self.raw.iter_mut()
    }
}

pub struct IntoIdx<I: Idx> { _marker: PhantomData<fn(&I)> }
impl<I: Idx, T> FnOnce<((usize, T),)> for IntoIdx<I> {
    type Output = (I, T);

    extern "rust-call" fn call_once(self, ((n, t),): ((usize, T),)) -> Self::Output {
        (I::new(n), t)
    }
}

impl<I: Idx, T> FnMut<((usize, T),)> for IntoIdx<I> {
    extern "rust-call" fn call_mut(&mut self, ((n, t),): ((usize, T),)) -> Self::Output {
        (I::new(n), t)
    }
}

impl<I: Idx> FnOnce<(usize,)> for IntoIdx<I> {
    type Output = I;

    extern "rust-call" fn call_once(self, (n,): (usize,)) -> Self::Output {
        I::new(n)
    }
}

impl<I: Idx> FnMut<(usize,)> for IntoIdx<I> {
    extern "rust-call" fn call_mut(&mut self, (n,): (usize,)) -> Self::Output {
        I::new(n)
    }
}
<commit_msg>Split lines longer than 100 columns<commit_after>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::collections::range::RangeArgument;
use std::fmt::Debug;
use std::iter::{self, FromIterator};
use std::slice;
use std::marker::PhantomData;
use std::ops::{Index, IndexMut, Range};
use std::fmt;
use std::vec;
use std::u32;

use rustc_serialize as serialize;

/// Represents some newtyped `usize` wrapper.
///
/// (purpose: avoid mixing indexes for different bitvector domains.)
pub trait Idx: Copy + 'static + Eq + Debug {
    fn new(idx: usize) -> Self;
    fn index(self) -> usize;
}

impl Idx for usize {
    fn new(idx: usize) -> Self { idx }
    fn index(self) -> usize { self }
}

impl Idx for u32 {
    fn new(idx: usize) -> Self { assert!(idx <= u32::MAX as usize); idx as u32 }
    fn index(self) -> usize { self as usize }
}

#[macro_export]
macro_rules! newtype_index {
    // ---- public rules ----

    // Use default constants
    ($name:ident) => (
        newtype_index!(
            @type[$name]
            @max[::std::u32::MAX]
            @debug_name[unsafe {::std::intrinsics::type_name::<$name>() }]);
    );

    // Define any constants
    ($name:ident { $($tokens:tt)+ }) => (
        newtype_index!(
            @type[$name]
            @max[::std::u32::MAX]
            @debug_name[unsafe {::std::intrinsics::type_name::<$name>() }]
            $($tokens)+);
    );

    // ---- private rules ----

    // Base case, user-defined constants (if any) have already been defined
    (@type[$type:ident] @max[$max:expr] @debug_name[$debug_name:expr]) => (
        #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord,
            RustcEncodable, RustcDecodable)]
        pub struct $type(u32);

        impl Idx for $type {
            fn new(value: usize) -> Self {
                assert!(value < ($max) as usize);
                $type(value as u32)
            }
            fn index(self) -> usize {
                self.0 as usize
            }
        }

        impl ::std::fmt::Debug for $type {
            fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
                write!(fmt, "{}{}", $debug_name, self.0)
            }
        }
    );

    // Rewrite final without comma to one that includes comma
    (@type[$type:ident] @max[$max:expr] @debug_name[$debug_name:expr]
            $name:ident = $constant:expr) => (
        newtype_index!(@type[$type] @max[$max] @debug_name[$debug_name] $name = $constant,);
    );

    // Rewrite final const without comma to one that includes comma
    (@type[$type:ident] @max[$_max:expr] @debug_name[$debug_name:expr]
            const $name:ident = $constant:expr) => (
        newtype_index!(@type[$type] @max[$max] @debug_name[$debug_name] const $name = $constant,);
    );

    // Replace existing default for max
    (@type[$type:ident] @max[$_max:expr] @debug_name[$debug_name:expr]
            MAX = $max:expr, $($tokens:tt)*) => (
        newtype_index!(@type[$type] @max[$max] @debug_name[$debug_name] $(tokens)*);
    );

    // Replace existing default for debug_name
    (@type[$type:ident] @max[$max:expr] @debug_name[$_debug_name:expr]
            DEBUG_NAME = $debug_name:expr, $($tokens:tt)*) => (
        newtype_index!(@type[$type] @max[$max] @debug_name[$debug_name] $($tokens)*);
    );

    // Assign a user-defined constant (as final param)
    (@type[$type:ident] @max[$max:expr] @debug_name[$debug_name:expr]
            const $name:ident = $constant:expr, $($tokens:tt)*) => (
        pub const $name: $type = $type($constant);
        newtype_index!(@type[$type] @max[$max] @debug_name[$debug_name] $($tokens)*);
    );
}

#[derive(Clone, PartialEq, Eq)]
pub struct IndexVec<I: Idx, T> {
    pub raw: Vec<T>,
    _marker: PhantomData<Fn(&I)>
}

impl<I: Idx, T: serialize::Encodable> serialize::Encodable for IndexVec<I, T> {
    fn encode<S: serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        serialize::Encodable::encode(&self.raw, s)
    }
}

impl<I: Idx, T: serialize::Decodable> serialize::Decodable for IndexVec<I, T> {
    fn decode<D: serialize::Decoder>(d: &mut D) -> Result<Self, D::Error> {
        serialize::Decodable::decode(d).map(|v| {
            IndexVec { raw: v, _marker: PhantomData }
        })
    }
}

impl<I: Idx, T: fmt::Debug> fmt::Debug for IndexVec<I, T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.raw, fmt)
    }
}

pub type Enumerated<I, J> = iter::Map<iter::Enumerate<J>, IntoIdx<I>>;

impl<I: Idx, T> IndexVec<I, T> {
    #[inline]
    pub fn new() -> Self {
        IndexVec { raw: Vec::new(), _marker: PhantomData }
    }

    #[inline]
    pub fn with_capacity(capacity: usize) -> Self {
        IndexVec { raw: Vec::with_capacity(capacity), _marker: PhantomData }
    }

    #[inline]
    pub fn from_elem<S>(elem: T, universe: &IndexVec<I, S>) -> Self
        where T: Clone
    {
        IndexVec { raw: vec![elem; universe.len()], _marker: PhantomData }
    }

    #[inline]
    pub fn from_elem_n(elem: T, n: usize) -> Self
        where T: Clone
    {
        IndexVec { raw: vec![elem; n], _marker: PhantomData }
    }

    #[inline]
    pub fn push(&mut self, d: T) -> I {
        let idx = I::new(self.len());
        self.raw.push(d);
        idx
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.raw.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.raw.is_empty()
    }

    #[inline]
    pub fn into_iter(self) -> vec::IntoIter<T> {
        self.raw.into_iter()
    }

    #[inline]
    pub fn into_iter_enumerated(self) -> Enumerated<I, vec::IntoIter<T>>
    {
        self.raw.into_iter().enumerate().map(IntoIdx { _marker: PhantomData })
    }

    #[inline]
    pub fn iter(&self) -> slice::Iter<T> {
        self.raw.iter()
    }

    #[inline]
    pub fn iter_enumerated(&self) -> Enumerated<I, slice::Iter<T>>
    {
        self.raw.iter().enumerate().map(IntoIdx { _marker: PhantomData })
    }

    #[inline]
    pub fn indices(&self) -> iter::Map<Range<usize>, IntoIdx<I>> {
        (0..self.len()).map(IntoIdx { _marker: PhantomData })
    }

    #[inline]
    pub fn iter_mut(&mut self) -> slice::IterMut<T> {
        self.raw.iter_mut()
    }

    #[inline]
    pub fn iter_enumerated_mut(&mut self) -> Enumerated<I, slice::IterMut<T>>
    {
        self.raw.iter_mut().enumerate().map(IntoIdx { _marker: PhantomData })
    }

    #[inline]
    pub fn drain<'a, R: RangeArgument<usize>>(
        &'a mut self, range: R) -> impl Iterator<Item=T> + 'a {
        self.raw.drain(range)
    }

    #[inline]
    pub fn drain_enumerated<'a, R: RangeArgument<usize>>(
        &'a mut self, range: R) -> impl Iterator<Item=(I, T)> + 'a {
        self.raw.drain(range).enumerate().map(IntoIdx { _marker: PhantomData })
    }

    #[inline]
    pub fn last(&self) -> Option<I> {
        self.len().checked_sub(1).map(I::new)
    }

    #[inline]
    pub fn shrink_to_fit(&mut self) {
        self.raw.shrink_to_fit()
    }

    #[inline]
    pub fn swap(&mut self, a: usize, b: usize) {
        self.raw.swap(a, b)
    }

    #[inline]
    pub fn truncate(&mut self, a: usize) {
        self.raw.truncate(a)
    }

    #[inline]
    pub fn get(&self, index: I) -> Option<&T> {
        self.raw.get(index.index())
    }

    #[inline]
    pub fn get_mut(&mut self, index: I) -> Option<&mut T> {
        self.raw.get_mut(index.index())
    }
}

impl<I: Idx, T: Clone> IndexVec<I, T> {
    #[inline]
    pub fn resize(&mut self, new_len: usize, value: T) {
        self.raw.resize(new_len, value)
    }
}

impl<I: Idx, T: Ord> IndexVec<I, T> {
    #[inline]
    pub fn binary_search(&self, value: &T) -> Result<I, I> {
        match self.raw.binary_search(value) {
            Ok(i) => Ok(Idx::new(i)),
            Err(i) => Err(Idx::new(i)),
        }
    }
}

impl<I: Idx, T> Index<I> for IndexVec<I, T> {
    type Output = T;

    #[inline]
    fn index(&self, index: I) -> &T {
        &self.raw[index.index()]
    }
}

impl<I: Idx, T> IndexMut<I> for IndexVec<I, T> {
    #[inline]
    fn index_mut(&mut self, index: I) -> &mut T {
        &mut self.raw[index.index()]
    }
}

impl<I: Idx, T> Default for IndexVec<I, T> {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

impl<I: Idx, T> Extend<T> for IndexVec<I, T> {
    #[inline]
    fn extend<J: IntoIterator<Item = T>>(&mut self, iter: J) {
        self.raw.extend(iter);
    }
}

impl<I: Idx, T> FromIterator<T> for IndexVec<I, T> {
    #[inline]
    fn from_iter<J>(iter: J) -> Self where J: IntoIterator<Item=T> {
        IndexVec { raw: FromIterator::from_iter(iter), _marker: PhantomData }
    }
}

impl<I: Idx, T> IntoIterator for IndexVec<I, T> {
    type Item = T;
    type IntoIter = vec::IntoIter<T>;

    #[inline]
    fn into_iter(self) -> vec::IntoIter<T> {
        self.raw.into_iter()
    }

}

impl<'a, I: Idx, T> IntoIterator for &'a IndexVec<I, T> {
    type Item = &'a T;
    type IntoIter = slice::Iter<'a, T>;

    #[inline]
    fn into_iter(self) -> slice::Iter<'a, T> {
        self.raw.iter()
    }
}

impl<'a, I: Idx, T> IntoIterator for &'a mut IndexVec<I, T> {
    type Item = &'a mut T;
    type IntoIter = slice::IterMut<'a, T>;

    #[inline]
    fn into_iter(self) -> slice::IterMut<'a, T> {
        self.raw.iter_mut()
    }
}

pub struct IntoIdx<I: Idx> { _marker: PhantomData<fn(&I)> }
impl<I: Idx, T> FnOnce<((usize, T),)> for IntoIdx<I> {
    type Output = (I, T);

    extern "rust-call" fn call_once(self, ((n, t),): ((usize, T),)) -> Self::Output {
        (I::new(n), t)
    }
}

impl<I: Idx, T> FnMut<((usize, T),)> for IntoIdx<I> {
    extern "rust-call" fn call_mut(&mut self, ((n, t),): ((usize, T),)) -> Self::Output {
        (I::new(n), t)
    }
}

impl<I: Idx> FnOnce<(usize,)> for IntoIdx<I> {
    type Output = I;

    extern "rust-call" fn call_once(self, (n,): (usize,)) -> Self::Output {
        I::new(n)
    }
}

impl<I: Idx> FnMut<(usize,)> for IntoIdx<I> {
    extern "rust-call" fn call_mut(&mut self, (n,): (usize,)) -> Self::Output {
        I::new(n)
    }
}
<|endoftext|>