<commit_before>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_qldb::{Client as QLDBClient, Error, Region, PKG_VERSION};
use structopt::StructOpt;
use tokio_stream::StreamExt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// The AWS Region.
    #[structopt(short, long)]
    region: Option<String>,

    /// Whether to display additional information.
    #[structopt(short, long)]
    verbose: bool,
}

// List ledgers.
// snippet-start:[qldb.rust.list-ledgers]
async fn show_ledgers(client: &QLDBClient) -> Result<(), Error> {
    let mut pages = client.list_ledgers().into_paginator().page_size(2).send();

    while let Some(page) = pages.next().await {
        println!("* {:?}", page); //prints an entire page of ledgers
        for ledger in page.unwrap().ledgers().unwrap() {
            println!("* {:?}", ledger); //prints the LedgerSummary of a single ledger
        }
    }

    Ok(())
}
// snippet-end:[qldb.rust.list-ledgers]

/// Lists your Amazon Quantum Ledger Database (Amazon QLDB) ledgers.
/// # Arguments
///
/// * `[-r REGION]` - The Region in which the client is created.
///    If not supplied, uses the value of the **AWS_REGION** environment variable.
///    If the environment variable is not set, defaults to **us-west-2**.
/// * `[-v]` - Whether to display additional information.
#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt::init();

    let Opt { region, verbose } = Opt::from_args();

    let region_provider = RegionProviderChain::first_try(region.map(Region::new))
        .or_default_provider()
        .or_else(Region::new("us-west-2"));

    println!();

    if verbose {
        println!("OLDB client version: {}", PKG_VERSION);
        println!(
            "Region:              {}",
            region_provider.region().await.unwrap().as_ref()
        );
        println!();
    }

    let shared_config = aws_config::from_env().region(region_provider).load().await;
    let client = QLDBClient::new(&shared_config);

    show_ledgers(&client).await
}
<commit_msg>Fixed a couple comments.<commit_after>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_qldb::{Client as QLDBClient, Error, Region, PKG_VERSION};
use structopt::StructOpt;
use tokio_stream::StreamExt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// The AWS Region.
    #[structopt(short, long)]
    region: Option<String>,

    /// Whether to display additional information.
    #[structopt(short, long)]
    verbose: bool,
}

// List ledgers.
// snippet-start:[qldb.rust.list-ledgers]
async fn show_ledgers(client: &QLDBClient) -> Result<(), Error> {
    let mut pages = client.list_ledgers().into_paginator().page_size(2).send();

    while let Some(page) = pages.next().await {
        println!("* {:?}", page); //Prints an entire page of ledgers.
        for ledger in page.unwrap().ledgers().unwrap() {
            println!("* {:?}", ledger); //Prints the LedgerSummary of a single ledger.
        }
    }

    Ok(())
}
// snippet-end:[qldb.rust.list-ledgers]

/// Lists your Amazon Quantum Ledger Database (Amazon QLDB) ledgers.
/// # Arguments
///
/// * `[-r REGION]` - The Region in which the client is created.
///    If not supplied, uses the value of the **AWS_REGION** environment variable.
///    If the environment variable is not set, defaults to **us-west-2**.
/// * `[-v]` - Whether to display additional information.
#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt::init();

    let Opt { region, verbose } = Opt::from_args();

    let region_provider = RegionProviderChain::first_try(region.map(Region::new))
        .or_default_provider()
        .or_else(Region::new("us-west-2"));

    println!();

    if verbose {
        println!("OLDB client version: {}", PKG_VERSION);
        println!(
            "Region:              {}",
            region_provider.region().await.unwrap().as_ref()
        );
        println!();
    }

    let shared_config = aws_config::from_env().region(region_provider).load().await;
    let client = QLDBClient::new(&shared_config);

    show_ledgers(&client).await
}
<|endoftext|>


<commit_before>use std::f32::consts::PI;
use std::sync::{Arc, Mutex,
                mpsc::{Sender, Receiver, TryRecvError},
                atomic::{AtomicUsize, Ordering},
};

use cgmath::{Point3, Vector3, Vector4, prelude::*};

use glium::Rect;

use rand::{self, prelude::*};

use crate::bvh::BVHNode;
use crate::camera::Camera;
use crate::color::Color;
use crate::material::Material;
use crate::pt_renderer::{Intersect, Ray, RenderCoordinator};
use crate::scene::Scene;
use crate::triangle::Hit;

// TODO: tune EPSILON since crytek-sponza has shadow acne
const EPSILON: f32 = 1e-5;
const _EB: f32 = 5.0; // Desired expectation value of bounces
// The matching survival probability from negative binomial distribution
const RR_PROB: f32 = _EB / (_EB + 1.0);

pub struct RenderWorker {
    scene: Arc<Scene>,
    camera: Camera,
    coordinator: Arc<Mutex<RenderCoordinator>>,
    message_rx: Receiver<()>,
    result_tx: Sender<(Rect, Vec<f32>)>,
    ray_count: Arc<AtomicUsize>,
}

impl RenderWorker {
    pub fn new(scene: Arc<Scene>, camera: Camera, coordinator: Arc<Mutex<RenderCoordinator>>,
               message_rx: Receiver<()>, result_tx: Sender<(Rect, Vec<f32>)>,
               ray_count: Arc<AtomicUsize>) -> RenderWorker {
        RenderWorker {
            scene, camera, coordinator,
            message_rx, result_tx,
            ray_count,
        }
    }

    pub fn run(&self) {
        let (width, height) = {
            let coordinator = self.coordinator.lock().unwrap();
            (coordinator.width, coordinator.height)
        };
        let clip_to_world = self.camera.get_world_to_clip().invert().unwrap();
        let mut node_stack: Vec<(&BVHNode, f32)> = Vec::new();
        loop {
            match self.message_rx.try_recv() {
                Err(TryRecvError::Empty) => (),
                Ok(_) => return,
                Err(TryRecvError::Disconnected) => {
                    println!("Threads were not properly stopped before disconnecting channel!");
                    return;
                }
            }
            let block = {
                let mut coordinator = self.coordinator.lock().unwrap();
                coordinator.next_block()
            };
            if let Some(rect) = block {
                let mut block = vec![0.0f32; (3 * rect.width * rect.height) as usize];
                for h in 0..rect.height {
                    for w in 0..rect.width {
                        let samples_per_dir = 2u32;
                        let mut c = Color::black();
                        for j in 0..samples_per_dir {
                            for i in 0..samples_per_dir {
                                let dx = (i as f32 + rand::random::<f32>()) / samples_per_dir as f32;
                                let dy = (j as f32 + rand::random::<f32>()) / samples_per_dir as f32;
                                let clip_x = 2.0 * ((rect.left + w) as f32 + dx) / width as f32 - 1.0;
                                let clip_y = 2.0 * ((rect.bottom + h) as f32 + dy) / height as f32 - 1.0;
                                let clip_p = Vector4::new(clip_x, clip_y, 1.0, 1.0);
                                let world_p = clip_to_world * clip_p;
                                let dir = ((world_p / world_p.w).truncate() - self.camera.pos.to_vec())
                                    .normalize();
                                let ray = Ray::new(self.camera.pos, dir, 10.0 * self.camera.scale);
                                c += self.trace_ray(&ray, &mut node_stack, 0);
                            }
                        }
                        c /= samples_per_dir.pow(2) as f32;
                        let pixel_i = 3 * (h * rect.width + w) as usize;
                        block[pixel_i]     = c.r;
                        block[pixel_i + 1] = c.g;
                        block[pixel_i + 2] = c.b;
                    }
                }
                self.result_tx.send((rect, block)).expect("Receiver closed!");
            } else {
                return;
            }
        }
    }

    fn trace_ray<'a>(&'a self, ray: &Ray, node_stack: &mut Vec<(&'a BVHNode, f32)>, bounce: u32) -> Color {
        let mut c = Color::black();
        if let Some(hit) = self.find_hit(&ray, node_stack) {
            let material = &self.scene.materials[hit.tri.material_i];
            let mut normal = hit.normal();
            // Flip the normal if its pointing to the opposite side from the hit
            if normal.dot(ray.dir) > 0.0 {
                normal *= -1.0;
            }
            if bounce == 0 {
                if let Some(emissive) = material.emissive {
                    c += normal.dot(-ray.dir).max(0.0) * emissive;
                }
            }
            let (emissive, light_pos, light_normal, light_pdf) = self.sample_light();
            let bump_pos = hit.pos() + EPSILON * normal;
            let hit_to_light = light_pos - bump_pos;
            let light_dir = hit_to_light.normalize();
            let shadow_ray = Ray::new(bump_pos, light_dir, hit_to_light.magnitude() - EPSILON);
            if self.find_hit(&shadow_ray, node_stack).is_none() {
                let cos_l = light_normal.dot(-light_dir).max(0.0);
                let cos_t = normal.dot(light_dir).max(0.0);
                c += emissive * self.brdf(&hit, &ray, &shadow_ray, material)
                    * cos_l * cos_t / (hit_to_light.magnitude2() * light_pdf);
            }
            let rr: f32 = rand::random();
            if rr < RR_PROB {
                let (new_dir, mut pdf) = self.sample_dir(normal);
                pdf *= RR_PROB;
                let new_ray = Ray::new(bump_pos, new_dir, 10.0 * self.camera.scale);
                c += normal.dot(new_dir) * self.brdf(&hit, &ray, &new_ray, material)
                    * self.trace_ray(&new_ray, node_stack, bounce + 1) / pdf;
            }
        }
        c
    }

    fn brdf(&self, hit: &Hit, _in_ray: &Ray, _out_ray: &Ray, material: &Material) -> Color {
        let tex_coords = hit.tex_coords();
        material.diffuse(tex_coords) / PI
    }

    fn sample_dir(&self, normal: Vector3<f32>) -> (Vector3<f32>, f32) {
        let dir = 2.0 * PI * rand::random::<f32>();
        let length: f32 = rand::random();
        let x = length * dir.cos();
        let y = length * dir.sin();
        let z = (1.0 - length.powi(2)).sqrt();
        let nx = if normal.x.abs() > normal.y.abs() {
            Vector3::new(normal.z, 0.0, -normal.x).normalize()
        } else {
            Vector3::new(0.0, -normal.z, normal.y).normalize()
        };
        let ny = normal.cross(nx);
        (x * nx + y * ny + z * normal, z / PI)
    }

    pub fn sample_light(&self) -> (Color, Point3<f32>, Vector3<f32>, f32) {
        if self.scene.lights.is_empty() {
            let light_intensity = 10.0 * self.camera.scale;
            (light_intensity * Color::white(), self.camera.pos, self.camera.dir, 1.0)
        } else {
            let i = rand::thread_rng().gen_range(0, self.scene.lights.len());
            let light = &self.scene.lights[i];
            let pdf = 1.0 / (self.scene.lights.len() as f32 * light.area());
            let (point, normal) = light.random_point();
            let light_material = &self.scene.materials[light.material_i];
            let emissive = light_material.emissive.expect("Light wasn't emissive");
            (emissive, point, normal, pdf)
        }
    }

    fn find_hit<'a>(&'a self, ray: &Ray, node_stack: &mut Vec<(&'a BVHNode, f32)>) -> Option<Hit> {
        self.ray_count.fetch_add(1, Ordering::Relaxed);
        let bvh = &self.scene.bvh;
        node_stack.push((bvh.root(), 0.0f32));
        let mut closest_hit: Option<Hit> = None;
        while let Some((node, t)) = node_stack.pop() {
            // We've already found closer hit
            if closest_hit.as_ref().map_or(false, |closest| closest.t <= t) { continue }
            if node.is_leaf() {
                for tri in &self.scene.triangles[node.start_i..node.end_i] {
                    if let Some(hit) = tri.intersect(&ray) {
                        if let Some(closest) = closest_hit.take() {
                            if hit.t < closest.t {
                                closest_hit = Some(hit);
                            } else {
                                closest_hit = Some(closest);
                            }
                        } else {
                            closest_hit = Some(hit);
                        }
                    }
                }
            } else {
                let (left, right) = bvh.get_children(node)
                    .expect("Non leaf node had no child nodes!");
                // TODO: benchmark this properly on the larger scenes
                // TODO: Could this work without pushing the next node to the stack
                if false { // not sorted
                    if let Some(t_left) = left.intersect(&ray) {
                        node_stack.push((left, t_left));
                    }
                    if let Some(t_right) = right.intersect(&ray) {
                        node_stack.push((right, t_right));
                    }
                } else { // sort nodes
                    let left_intersect = left.intersect(&ray);
                    let right_intersect = right.intersect(&ray);
                    if let Some(t_left) = left_intersect {
                        if let Some(t_right) = right_intersect {
                            // Put the closer hit on top
                            if t_left >= t_right {
                                node_stack.push((left, t_left));
                                node_stack.push((right, t_right));
                            } else {
                                node_stack.push((right, t_right));
                                node_stack.push((left, t_left));
                            }
                        } else {
                            node_stack.push((left, t_left));
                        }
                    } else if let Some(t_right) = right_intersect {
                        node_stack.push((right, t_right));
                    }
                }
            }
        }
        closest_hit
    }
}
<commit_msg>Update ray length based on closest hits<commit_after>use std::f32::consts::PI;
use std::sync::{Arc, Mutex,
                mpsc::{Sender, Receiver, TryRecvError},
                atomic::{AtomicUsize, Ordering},
};

use cgmath::{Point3, Vector3, Vector4, prelude::*};

use glium::Rect;

use rand::{self, prelude::*};

use crate::bvh::BVHNode;
use crate::camera::Camera;
use crate::color::Color;
use crate::material::Material;
use crate::pt_renderer::{Intersect, Ray, RenderCoordinator};
use crate::scene::Scene;
use crate::triangle::Hit;

// TODO: tune EPSILON since crytek-sponza has shadow acne
const EPSILON: f32 = 1e-5;
const _EB: f32 = 5.0; // Desired expectation value of bounces
// The matching survival probability from negative binomial distribution
const RR_PROB: f32 = _EB / (_EB + 1.0);

pub struct RenderWorker {
    scene: Arc<Scene>,
    camera: Camera,
    coordinator: Arc<Mutex<RenderCoordinator>>,
    message_rx: Receiver<()>,
    result_tx: Sender<(Rect, Vec<f32>)>,
    ray_count: Arc<AtomicUsize>,
}

impl RenderWorker {
    pub fn new(scene: Arc<Scene>, camera: Camera, coordinator: Arc<Mutex<RenderCoordinator>>,
               message_rx: Receiver<()>, result_tx: Sender<(Rect, Vec<f32>)>,
               ray_count: Arc<AtomicUsize>) -> RenderWorker {
        RenderWorker {
            scene, camera, coordinator,
            message_rx, result_tx,
            ray_count,
        }
    }

    pub fn run(&self) {
        let (width, height) = {
            let coordinator = self.coordinator.lock().unwrap();
            (coordinator.width, coordinator.height)
        };
        let clip_to_world = self.camera.get_world_to_clip().invert().unwrap();
        let mut node_stack: Vec<(&BVHNode, f32)> = Vec::new();
        loop {
            match self.message_rx.try_recv() {
                Err(TryRecvError::Empty) => (),
                Ok(_) => return,
                Err(TryRecvError::Disconnected) => {
                    println!("Threads were not properly stopped before disconnecting channel!");
                    return;
                }
            }
            let block = {
                let mut coordinator = self.coordinator.lock().unwrap();
                coordinator.next_block()
            };
            if let Some(rect) = block {
                let mut block = vec![0.0f32; (3 * rect.width * rect.height) as usize];
                for h in 0..rect.height {
                    for w in 0..rect.width {
                        let samples_per_dir = 2u32;
                        let mut c = Color::black();
                        for j in 0..samples_per_dir {
                            for i in 0..samples_per_dir {
                                let dx = (i as f32 + rand::random::<f32>()) / samples_per_dir as f32;
                                let dy = (j as f32 + rand::random::<f32>()) / samples_per_dir as f32;
                                let clip_x = 2.0 * ((rect.left + w) as f32 + dx) / width as f32 - 1.0;
                                let clip_y = 2.0 * ((rect.bottom + h) as f32 + dy) / height as f32 - 1.0;
                                let clip_p = Vector4::new(clip_x, clip_y, 1.0, 1.0);
                                let world_p = clip_to_world * clip_p;
                                let dir = ((world_p / world_p.w).truncate() - self.camera.pos.to_vec())
                                    .normalize();
                                let ray = Ray::new(self.camera.pos, dir, std::f32::INFINITY);
                                c += self.trace_ray(ray, &mut node_stack, 0);
                            }
                        }
                        c /= samples_per_dir.pow(2) as f32;
                        let pixel_i = 3 * (h * rect.width + w) as usize;
                        block[pixel_i]     = c.r;
                        block[pixel_i + 1] = c.g;
                        block[pixel_i + 2] = c.b;
                    }
                }
                self.result_tx.send((rect, block)).expect("Receiver closed!");
            } else {
                return;
            }
        }
    }

    fn trace_ray<'a>(&'a self, mut ray: Ray, node_stack: &mut Vec<(&'a BVHNode, f32)>, bounce: u32) -> Color {
        let mut c = Color::black();
        if let Some(hit) = self.find_hit(&mut ray, node_stack) {
            let material = &self.scene.materials[hit.tri.material_i];
            let mut normal = hit.normal();
            // Flip the normal if its pointing to the opposite side from the hit
            if normal.dot(ray.dir) > 0.0 {
                normal *= -1.0;
            }
            if bounce == 0 {
                if let Some(emissive) = material.emissive {
                    c += normal.dot(-ray.dir).max(0.0) * emissive;
                }
            }
            let (emissive, light_pos, light_normal, light_pdf) = self.sample_light();
            let bump_pos = hit.pos() + EPSILON * normal;
            let hit_to_light = light_pos - bump_pos;
            let light_dir = hit_to_light.normalize();
            let mut shadow_ray = Ray::new(bump_pos, light_dir, hit_to_light.magnitude() - EPSILON);
            if self.find_hit(&mut shadow_ray, node_stack).is_none() {
                let cos_l = light_normal.dot(-light_dir).max(0.0);
                let cos_t = normal.dot(light_dir).max(0.0);
                c += emissive * self.brdf(&hit, &ray, &shadow_ray, material)
                    * cos_l * cos_t / (hit_to_light.magnitude2() * light_pdf);
            }
            let rr: f32 = rand::random();
            if rr < RR_PROB {
                let (new_dir, mut pdf) = self.sample_dir(normal);
                pdf *= RR_PROB;
                let new_ray = Ray::new(bump_pos, new_dir, 10.0 * self.camera.scale);
                c += normal.dot(new_dir) * self.brdf(&hit, &ray, &new_ray, material)
                    * self.trace_ray(new_ray, node_stack, bounce + 1) / pdf;
            }
        }
        c
    }

    fn brdf(&self, hit: &Hit, _in_ray: &Ray, _out_ray: &Ray, material: &Material) -> Color {
        let tex_coords = hit.tex_coords();
        material.diffuse(tex_coords) / PI
    }

    fn sample_dir(&self, normal: Vector3<f32>) -> (Vector3<f32>, f32) {
        let dir = 2.0 * PI * rand::random::<f32>();
        let length: f32 = rand::random();
        let x = length * dir.cos();
        let y = length * dir.sin();
        let z = (1.0 - length.powi(2)).sqrt();
        let nx = if normal.x.abs() > normal.y.abs() {
            Vector3::new(normal.z, 0.0, -normal.x).normalize()
        } else {
            Vector3::new(0.0, -normal.z, normal.y).normalize()
        };
        let ny = normal.cross(nx);
        (x * nx + y * ny + z * normal, z / PI)
    }

    pub fn sample_light(&self) -> (Color, Point3<f32>, Vector3<f32>, f32) {
        if self.scene.lights.is_empty() {
            let light_intensity = 10.0 * self.camera.scale;
            (light_intensity * Color::white(), self.camera.pos, self.camera.dir, 1.0)
        } else {
            let i = rand::thread_rng().gen_range(0, self.scene.lights.len());
            let light = &self.scene.lights[i];
            let pdf = 1.0 / (self.scene.lights.len() as f32 * light.area());
            let (point, normal) = light.random_point();
            let light_material = &self.scene.materials[light.material_i];
            let emissive = light_material.emissive.expect("Light wasn't emissive");
            (emissive, point, normal, pdf)
        }
    }

    fn find_hit<'a>(&'a self, ray: &mut Ray, node_stack: &mut Vec<(&'a BVHNode, f32)>) -> Option<Hit> {
        self.ray_count.fetch_add(1, Ordering::Relaxed);
        let bvh = &self.scene.bvh;
        node_stack.push((bvh.root(), 0.0f32));
        let mut closest_hit: Option<Hit> = None;
        while let Some((node, t)) = node_stack.pop() {
            // We've already found closer hit
            if closest_hit.as_ref().map_or(false, |closest| closest.t <= t) { continue }
            if node.is_leaf() {
                for tri in &self.scene.triangles[node.start_i..node.end_i] {
                    if let Some(hit) = tri.intersect(&ray) {
                        ray.length = hit.t;
                        closest_hit = Some(hit);
                    }
                }
            } else {
                let (left, right) = bvh.get_children(node)
                    .expect("Non leaf node had no child nodes!");
                // TODO: benchmark this properly on the larger scenes
                // TODO: Could this work without pushing the next node to the stack
                if false { // not sorted
                    if let Some(t_left) = left.intersect(&ray) {
                        node_stack.push((left, t_left));
                    }
                    if let Some(t_right) = right.intersect(&ray) {
                        node_stack.push((right, t_right));
                    }
                } else { // sort nodes
                    let left_intersect = left.intersect(&ray);
                    let right_intersect = right.intersect(&ray);
                    if let Some(t_left) = left_intersect {
                        if let Some(t_right) = right_intersect {
                            // Put the closer hit on top
                            if t_left >= t_right {
                                node_stack.push((left, t_left));
                                node_stack.push((right, t_right));
                            } else {
                                node_stack.push((right, t_right));
                                node_stack.push((left, t_left));
                            }
                        } else {
                            node_stack.push((left, t_left));
                        }
                    } else if let Some(t_right) = right_intersect {
                        node_stack.push((right, t_right));
                    }
                }
            }
        }
        closest_hit
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::collections::hash_map;
use std::rc::Rc;

use render::Renderer;
use routing::{DagHandle, Edge, Effect, GraphWatcher, NodeData, NodeHandle};

#[derive(Default)]
pub struct RefRenderer {
    nodes: HashMap<NodeHandle, Node>,
}

struct Node {
    data: MyNodeData,
    inbound: HashSet<Edge>
}

enum MyNodeData {
    /// This node is a non-primitive effect.
    UserNode(Rc<Effect>),
    /// This node is an instance of another DAG.
    Graph(DagHandle),
    /// Primitive Delay(samples) effect
    Delay(u64),
    /// Primitive Constant(value) effect,
    Constant(f32),
    /// This node is a DAG definition. i.e. it holds the output edges of a DAG.
    DagIO,
}

impl Renderer for RefRenderer {
    fn get_sample(&mut self, time: u64, ch: u8) -> f32 {
        // Try to find the edge that goes to -> (Null, slot=0, ch=ch)
        let root_handle = NodeHandle::toplevel();
        match self.nodes.get(&root_handle) {
            // for an empty graph, yield silence.
            None => 0f32,
            Some(node) => {
                // find all edges to ([Null], slot=0, ch=ch)
                let matches = node.inbound.iter().filter(|edge| {
                    edge.to_slot() == 0 && edge.to_ch() == ch
                });
                matches.map(|edge| {
                    self.get_value(edge, time, &vec!())
                }).sum()
            }
        }
    }
}
impl RefRenderer {
    /// Get the value on an edge at a particular time
    /// When backtracking from the output, we push each Node onto the context if we enter inside of
    ///   it (i.e. if it's a nested DAG) & pop when exiting.
    fn get_value(&self, edge: &Edge, time: u64, context: &Vec<NodeHandle>) -> f32 {
        let from = edge.from_full();
        if *from.node_handle() == None {
            // Reading from one of the inputs to the top of `context`
            let mut new_context = context.clone();
            let head = new_context.pop().unwrap();
            // Sum the inputs to the matchin slot/ch
            let in_edges = self.nodes[&head].inbound.iter().filter(|in_edge| {
                in_edge.to_slot() == edge.from_slot() && in_edge.to_ch() == edge.from_ch()
            });
            let in_values = in_edges.map(|in_edge| {
                self.get_value(in_edge, time, &new_context)
            });
            in_values.sum()
        } else {
            // Reading from another node within the DAG
            let node = &self.nodes[&from];
            match node.data {
                MyNodeData::UserNode(ref _effect) => unimplemented!(),
                // Output = sum of all edges to Null of the same slot & ch, within the given DAG.
                MyNodeData::Graph(ref dag_handle) => {
                    let mut new_context = context.clone();
                    new_context.push(from);
                    let subdag = &self.nodes[&NodeHandle::new(dag_handle.clone(), None)];
                    let out_edges = subdag.inbound.iter().filter(|out_edge| {
                        out_edge.to_slot() == edge.from_slot() && out_edge.to_ch() == edge.from_ch()
                    });
                    let out_values = out_edges.map(|out_edge| {
                        self.get_value(out_edge, time, &new_context)
                    });
                    out_values.sum()
                }
                // Output = sum of all inputs to slot 0 of the same ch.
                MyNodeData::Delay(ref frames) => {
                    assert!(edge.from_slot() == 1);
                    let in_edges = node.inbound.iter().filter(|in_edge| {
                        in_edge.to_slot() == 0 && in_edge.to_ch() == edge.from_ch()
                    });
                    let in_values = in_edges.map(|in_edge| {
                        // TODO: handle time underflow
                        self.get_value(in_edge, time-frames, context)
                    });
                    in_values.sum()
                },
                MyNodeData::Constant(ref value) => value.clone(),
                _ => panic!("Internal RefRenderer error: illegal node type"),
            }
        }
    }
}

impl GraphWatcher for RefRenderer {
    fn on_add_node(&mut self, handle: &NodeHandle, data: &NodeData) {
        let my_node_data = match *data {
            NodeData::Graph(ref handle) => MyNodeData::Graph(handle.clone()),
            NodeData::Effect(ref effect) => {
                match effect.meta().get_primitive_url() {
                    Some(ref url) => {
                        let mut params: HashMap<_, _> = url.query_pairs().collect();
                        match url.path() {
                            "/Delay" => {
                                let frames: u64 = match params.entry(Cow::from("frames")) {
                                    hash_map::Entry::Occupied(e) => e.remove().parse().unwrap(),
                                    hash_map::Entry::Vacant(_) => 0u64,
                                };
                                // Make sure we consumed all arguments.
                                assert!(params.is_empty());
                                MyNodeData::Delay(frames)
                            },
                            "/Constant" => {
                                let value: f32 = match params.entry(Cow::from("value")) {
                                    hash_map::Entry::Occupied(e) => e.remove().parse().unwrap(),
                                    hash_map::Entry::Vacant(_) => 0f32,
                                };
                                // Make sure we consumed all arguments.
                                assert!(params.is_empty());
                                MyNodeData::Constant(value)
                            },
                            _ => panic!("Unrecognized primitive effect: {} (full url: {})", url.path(), url),
                        }
                    }
                    None => MyNodeData::UserNode(effect.clone())
                }
            }
        };
        self.nodes.insert(handle.clone(), Node::new(my_node_data));
        // If the node is part of a new DAG, allocate data so that future edges
        // to null within the DAG can be held.
        self.nodes.entry(NodeHandle::new_dag(handle.dag_handle().clone())).or_insert_with(|| {
            Node::new(MyNodeData::DagIO)
        });
    }
    fn on_del_node(&mut self, handle: &NodeHandle) {
        self.nodes.remove(handle);
    }
    fn on_add_edge(&mut self, edge: &Edge) {
        self.nodes.get_mut(&edge.to_full()).unwrap().inbound.insert(edge.clone());
    }
    fn on_del_edge(&mut self, edge: &Edge) {
        self.nodes.get_mut(&edge.to_full()).unwrap().inbound.remove(edge);
    }
}

impl Node {
    fn new(data: MyNodeData) -> Self {
        Node {
            data: data,
            inbound: HashSet::new(),
        }
    }
}
<commit_msg>Implemented primitive Multiply node<commit_after>use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::collections::hash_map;
use std::rc::Rc;

use render::Renderer;
use routing::{DagHandle, Edge, Effect, GraphWatcher, NodeData, NodeHandle};

#[derive(Default)]
pub struct RefRenderer {
    nodes: HashMap<NodeHandle, Node>,
}

struct Node {
    data: MyNodeData,
    inbound: HashSet<Edge>
}

enum MyNodeData {
    /// This node is a non-primitive effect.
    UserNode(Rc<Effect>),
    /// This node is an instance of another DAG.
    Graph(DagHandle),
    /// Primitive Delay(samples) effect
    Delay(u64),
    /// Primitive Constant(value) effect.
    /// Also serves as a unit step;
    /// Returns the f32 value for t >= 0, else 0.
    Constant(f32),
    /// Primitive effect to multiply TWO input streams sample-wise.
    Multiply,
    /// This node is a DAG definition. i.e. it holds the output edges of a DAG.
    DagIO,
}

impl Renderer for RefRenderer {
    fn get_sample(&mut self, time: u64, ch: u8) -> f32 {
        // Try to find the edge that goes to -> (Null, slot=0, ch=ch)
        let root_handle = NodeHandle::toplevel();
        match self.nodes.get(&root_handle) {
            // for an empty graph, yield silence.
            None => 0f32,
            Some(node) => {
                // find all edges to ([Null], slot=0, ch=ch)
                let matches = node.inbound.iter().filter(|edge| {
                    edge.to_slot() == 0 && edge.to_ch() == ch
                });
                matches.map(|edge| {
                    self.get_value(edge, time, &vec!())
                }).sum()
            }
        }
    }
}
impl RefRenderer {
    /// Get the value on an edge at a particular time
    /// When backtracking from the output, we push each Node onto the context if we enter inside of
    ///   it (i.e. if it's a nested DAG) & pop when exiting.
    fn get_value(&self, edge: &Edge, time: u64, context: &Vec<NodeHandle>) -> f32 {
        let from = edge.from_full();
        if *from.node_handle() == None {
            // Reading from one of the inputs to the top of `context`
            let mut new_context = context.clone();
            let head = new_context.pop().unwrap();
            // Sum the inputs to the matching slot/ch
            let in_edges = self.nodes[&head].inbound.iter().filter(|in_edge| {
                in_edge.to_slot() == edge.from_slot() && in_edge.to_ch() == edge.from_ch()
            });
            let in_values = in_edges.map(|in_edge| {
                self.get_value(in_edge, time, &new_context)
            });
            in_values.sum()
        } else {
            // Reading from another node within the DAG
            let node = &self.nodes[&from];
            match node.data {
                MyNodeData::UserNode(ref _effect) => unimplemented!(),
                // Output = sum of all edges to Null of the same slot & ch, within the given DAG.
                MyNodeData::Graph(ref dag_handle) => {
                    let mut new_context = context.clone();
                    new_context.push(from);
                    let subdag = &self.nodes[&NodeHandle::new(dag_handle.clone(), None)];
                    let out_edges = subdag.inbound.iter().filter(|out_edge| {
                        out_edge.to_slot() == edge.from_slot() && out_edge.to_ch() == edge.from_ch()
                    });
                    let out_values = out_edges.map(|out_edge| {
                        self.get_value(out_edge, time, &new_context)
                    });
                    out_values.sum()
                }
                // Output = sum of all inputs to slot 0 of the same ch.
                MyNodeData::Delay(ref frames) => {
                    // The only nonzero output is slot=1.
                    if edge.from_slot() != 1 {
                        println!("Warning: attempt to read from Delay slot != 1");
                        0f32
                    } else {
                        let in_edges = node.inbound.iter().filter(|in_edge| {
                            in_edge.to_slot() == 0 && in_edge.to_ch() == edge.from_ch()
                        });
                        let in_values = in_edges.map(|in_edge| {
                            // TODO: handle time underflow
                            self.get_value(in_edge, time-frames, context)
                        });
                        in_values.sum()
                    }
                },
                MyNodeData::Constant(ref value) => value.clone(),
                MyNodeData::Multiply => {
                    // The only nonzero output is slot=1.
                    if edge.from_slot() != 1 {
                        println!("Warning: attempt to read from Multiply slot != 1");
                        0f32
                    } else {
                        // Sum all inputs from slot=0 and slot=2 into two separate
                        // variables, then multiply them.
                        let edges_a = node.inbound.iter().filter(|in_edge| {
                            in_edge.to_slot() == 0 && in_edge.to_ch() == edge.from_ch()
                        });
                        let edges_b = node.inbound.iter().filter(|in_edge| {
                            in_edge.to_slot() == 2 && in_edge.to_ch() == edge.from_ch()
                        });
                        let val_a: f32 = edges_a.map(|edge| self.get_value(edge, time, context)).sum();
                        let val_b: f32 = edges_b.map(|edge| self.get_value(edge, time, context)).sum();
                        val_a * val_b
                    }
                },
                _ => panic!("Internal RefRenderer error: illegal node type"),
            }
        }
    }
}

impl GraphWatcher for RefRenderer {
    fn on_add_node(&mut self, handle: &NodeHandle, data: &NodeData) {
        let my_node_data = match *data {
            NodeData::Graph(ref handle) => MyNodeData::Graph(handle.clone()),
            NodeData::Effect(ref effect) => {
                match effect.meta().get_primitive_url() {
                    Some(ref url) => {
                        let mut params: HashMap<_, _> = url.query_pairs().collect();
                        match url.path() {
                            "/Delay" => {
                                let frames: u64 = match params.entry(Cow::from("frames")) {
                                    hash_map::Entry::Occupied(e) => e.remove().parse().unwrap(),
                                    hash_map::Entry::Vacant(_) => 0u64,
                                };
                                // Make sure we consumed all arguments.
                                assert!(params.is_empty());
                                MyNodeData::Delay(frames)
                            },
                            "/Constant" => {
                                let value: f32 = match params.entry(Cow::from("value")) {
                                    hash_map::Entry::Occupied(e) => e.remove().parse().unwrap(),
                                    hash_map::Entry::Vacant(_) => 0f32,
                                };
                                // Make sure we consumed all arguments.
                                assert!(params.is_empty());
                                MyNodeData::Constant(value)
                            },
                            _ => panic!("Unrecognized primitive effect: {} (full url: {})", url.path(), url),
                        }
                    }
                    None => MyNodeData::UserNode(effect.clone())
                }
            }
        };
        self.nodes.insert(handle.clone(), Node::new(my_node_data));
        // If the node is part of a new DAG, allocate data so that future edges
        // to null within the DAG can be held.
        self.nodes.entry(NodeHandle::new_dag(handle.dag_handle().clone())).or_insert_with(|| {
            Node::new(MyNodeData::DagIO)
        });
    }
    fn on_del_node(&mut self, handle: &NodeHandle) {
        self.nodes.remove(handle);
    }
    fn on_add_edge(&mut self, edge: &Edge) {
        self.nodes.get_mut(&edge.to_full()).unwrap().inbound.insert(edge.clone());
    }
    fn on_del_edge(&mut self, edge: &Edge) {
        self.nodes.get_mut(&edge.to_full()).unwrap().inbound.remove(edge);
    }
}

impl Node {
    fn new(data: MyNodeData) -> Self {
        Node {
            data: data,
            inbound: HashSet::new(),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

use std::collections::{HashMap, HashSet};
use std::io;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use futures::future::{self, Future};
use futures::sync::oneshot;

use boxfuture::{BoxFuture, Boxable};
use context::{Context, Core};
use core::{Failure, Key, Params, TypeConstraint, TypeId, Value};
use fs::{self, GlobMatching, PosixFS};
use graph::{EntryId, Graph, Node, NodeContext};
use nodes::{NodeKey, Select, Tracer, TryInto, Visualizer};
use parking_lot::Mutex;
use rule_graph;
use selectors;

///
/// A Session represents a related series of requests (generally: one run of the pants CLI) on an
/// underlying Scheduler, and is a useful scope for metrics.
///
/// Both Scheduler and Session are exposed to python and expected to be used by multiple threads, so
/// they use internal mutability in order to avoid exposing locks to callers.
///
pub struct Session {
  // The total size of the graph at Session-creation time.
  preceding_graph_size: usize,
  // The set of roots that have been requested within this session.
  roots: Mutex<HashSet<Root>>,
}

impl Session {
  pub fn new(scheduler: &Scheduler) -> Session {
    Session {
      preceding_graph_size: scheduler.core.graph.len(),
      roots: Mutex::new(HashSet::new()),
    }
  }

  fn extend(&self, new_roots: &[Root]) {
    let mut roots = self.roots.lock();
    roots.extend(new_roots.iter().cloned());
  }

  fn root_nodes(&self) -> Vec<NodeKey> {
    let roots = self.roots.lock();
    roots.iter().map(|r| r.clone().into()).collect()
  }
}

pub struct ExecutionRequest {
  // Set of roots for an execution, in the order they were declared.
  pub roots: Vec<Root>,
}

impl ExecutionRequest {
  pub fn new() -> ExecutionRequest {
    ExecutionRequest { roots: Vec::new() }
  }

  ///
  /// Roots are limited to `Select`, which is known to produce a Value. This method
  /// exists to satisfy Graph APIs which need instances of the NodeKey enum.
  ///
  fn root_nodes(&self) -> Vec<NodeKey> {
    self.roots.iter().map(|r| r.clone().into()).collect()
  }
}

///
/// Represents the state of an execution of a Graph.
///
pub struct Scheduler {
  pub core: Arc<Core>,
}

impl Scheduler {
  pub fn new(core: Core) -> Scheduler {
    Scheduler {
      core: Arc::new(core),
    }
  }

  pub fn visualize(&self, session: &Session, path: &Path) -> io::Result<()> {
    self
      .core
      .graph
      .visualize(Visualizer::default(), &session.root_nodes(), path)
  }

  pub fn trace(&self, request: &ExecutionRequest, path: &Path) -> Result<(), String> {
    self
      .core
      .graph
      .trace::<Tracer>(&request.root_nodes(), path)?;
    Ok(())
  }

  pub fn add_root_select(
    &self,
    request: &mut ExecutionRequest,
    subject: Key,
    product: TypeConstraint,
  ) -> Result<(), String> {
    let edges = self.find_root_edges_or_update_rule_graph(
      subject.type_id().clone(),
      &selectors::Select::new(product),
    )?;
    request
      .roots
      .push(Select::new(product, Params::new_single(subject), &edges));
    Ok(())
  }

  fn find_root_edges_or_update_rule_graph(
    &self,
    subject_type: TypeId,
    select: &selectors::Select,
  ) -> Result<rule_graph::RuleEdges, String> {
    self
      .core
      .rule_graph
      .find_root_edges(subject_type, select.clone())
      .ok_or_else(|| {
        format!(
          "No installed rules can satisfy {} for a root subject of type {}.",
          rule_graph::select_str(&select),
          rule_graph::type_str(subject_type)
        )
      })
  }

  ///
  /// Invalidate the invalidation roots represented by the given Paths.
  ///
  pub fn invalidate(&self, paths: &HashSet<PathBuf>) -> usize {
    let invalidation_result = self.core.graph.invalidate_from_roots(move |node| {
      if let Some(fs_subject) = node.fs_subject() {
        paths.contains(fs_subject)
      } else {
        false
      }
    });
    // TODO: Expose.
    invalidation_result.cleared + invalidation_result.dirtied
  }

  ///
  /// Invalidate all filesystem dependencies in the graph.
  ///
  pub fn invalidate_all_paths(&self) -> usize {
    let invalidation_result = self
      .core
      .graph
      .invalidate_from_roots(|node| node.fs_subject().is_some());
    // TODO: Expose.
    invalidation_result.cleared + invalidation_result.dirtied
  }

  ///
  /// Return Scheduler and per-Session metrics.
  ///
  pub fn metrics(&self, session: &Session) -> HashMap<&str, i64> {
    let mut m = HashMap::new();
    m.insert(
      "affected_file_count",
      self
        .core
        .graph
        .reachable_digest_count(&session.root_nodes()) as i64,
    );
    m.insert("preceding_graph_size", session.preceding_graph_size as i64);
    m.insert("resulting_graph_size", self.core.graph.len() as i64);
    m
  }

  ///
  /// Attempts to complete all of the given roots, retrying the entire set (up to `count`
  /// times) if any of them fail with `Failure::Invalidated`.
  ///
  /// In common usage, graph entries won't be repeatedly invalidated, but in a case where they
  /// were (say by an automated process changing files under pants), we'd want to eventually
  /// give up.
  ///
  fn execute_helper(
    context: RootContext,
    roots: Vec<Root>,
    count: usize,
  ) -> BoxFuture<Vec<Result<Value, Failure>>, ()> {
    let executor = context.core.runtime.get().executor();
    // Attempt all roots in parallel, failing fast to retry for `Invalidated`.
    let roots_res = future::join_all(
      roots
        .clone()
        .into_iter()
        .map(|root| {
          context
            .core
            .graph
            .create(root.clone().into(), &context)
            .then::<_, Result<Result<Value, Failure>, Failure>>(move |r| {
              match r {
                Err(Failure::Invalidated) if count > 0 => {
                  // A node was invalidated: fail quickly so that all roots can be retried.
                  Err(Failure::Invalidated)
                }
                other => {
                  // Otherwise (if it is a success, some other type of Failure, or if we've run
                  // out of retries) recover to complete the join, which will cause the results to
                  // propagate to the user.
                  debug!(
                    "Root {} completed.",
                    NodeKey::Select(Box::new(root)).format()
                  );
                  Ok(other.map(|res| {
                    res
                      .try_into()
                      .unwrap_or_else(|_| panic!("A Node implementation was ambiguous."))
                  }))
                }
              }
            })
        }).collect::<Vec<_>>(),
    );

    // If the join failed (due to `Invalidated`, since that is the only error we propagate), retry
    // the entire set of roots.
    oneshot::spawn(
      roots_res.or_else(move |_| Scheduler::execute_helper(context, roots, count - 1)),
      &executor,
    ).to_boxed()
  }

  ///
  /// Compute the results for roots in the given request.
  ///
  pub fn execute<'e>(
    &self,
    request: &'e ExecutionRequest,
    session: &Session,
  ) -> Vec<(&'e Key, &'e TypeConstraint, RootResult)> {
    // Bootstrap tasks for the roots, and then wait for all of them.
    debug!("Launching {} roots.", request.roots.len());

    session.extend(&request.roots);

    // Wait for all roots to complete. Failure here should be impossible, because each
    // individual Future in the join was (eventually) mapped into success.
    let context = RootContext {
      core: self.core.clone(),
    };
    let results = Scheduler::execute_helper(context, request.roots.clone(), 8)
      .wait()
      .expect("Execution failed.");

    request
      .roots
      .iter()
      .zip(results.into_iter())
      .map(|(s, r)| (s.params.expect_single(), &s.selector.product, r))
      .collect()
  }

  pub fn capture_snapshot_from_arbitrary_root<P: AsRef<Path>>(
    &self,
    root_path: P,
    path_globs: fs::PathGlobs,
  ) -> BoxFuture<fs::Snapshot, String> {
    // Note that we don't use a Graph here, and don't cache any intermediate steps, we just place
    // the resultant Snapshot into the store and return it. This is important, because we're reading
    // things from arbitrary filepaths which we don't want to cache in the graph, as we don't watch
    // them for changes.
    // We assume that this Snapshot is of an immutable piece of the filesystem.

    let posix_fs = Arc::new(try_future!(PosixFS::new(
      root_path,
      self.core.fs_pool.clone(),
      &[]
    )));
    let store = self.core.store.clone();

    posix_fs
      .expand(path_globs)
      .map_err(|err| format!("Error expanding globs: {:?}", err))
      .and_then(|path_stats| {
        fs::Snapshot::from_path_stats(
          store.clone(),
          &fs::OneOffStoreFileByDigest::new(store, posix_fs),
          path_stats,
        )
      }).to_boxed()
  }
}

impl Drop for Scheduler {
  fn drop(&mut self) {
    // Because Nodes may hold references to the Core in their closure, this is intended to
    // break cycles between Nodes and the Core.
    self.core.graph.clear();
  }
}

///
/// Root requests are limited to Selectors that produce (python) Values.
///
type Root = Select;

pub type RootResult = Result<Value, Failure>;

///
/// NB: This basic wrapper exists to allow us to implement the `NodeContext` trait (which lives
/// outside of this crate) for the `Arc` struct (which also lives outside our crate), which is not
/// possible without the wrapper due to "trait coherence".
///
#[derive(Clone)]
struct RootContext {
  core: Arc<Core>,
}

impl NodeContext for RootContext {
  type Node = NodeKey;

  fn clone_for(&self, entry_id: EntryId) -> Context {
    Context::new(entry_id, self.core.clone())
  }

  fn graph(&self) -> &Graph<NodeKey> {
    &self.core.graph
  }

  fn spawn<F>(&self, future: F)
  where
    F: Future<Item = (), Error = ()> + Send + 'static,
  {
    self.core.runtime.get().executor().spawn(future);
  }
}
<commit_msg>Move root executions off the main thread to allow for polling (#6570)<commit_after>// Copyright 2017 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

use std::collections::{HashMap, HashSet};
use std::io;
use std::path::{Path, PathBuf};
use std::sync::{mpsc, Arc};
use std::time::Duration;

use futures::future::{self, Future};

use boxfuture::{BoxFuture, Boxable};
use context::{Context, Core};
use core::{Failure, Key, Params, TypeConstraint, TypeId, Value};
use fs::{self, GlobMatching, PosixFS};
use graph::{EntryId, Graph, Node, NodeContext};
use nodes::{NodeKey, Select, Tracer, TryInto, Visualizer};
use parking_lot::Mutex;
use rule_graph;
use selectors;

///
/// A Session represents a related series of requests (generally: one run of the pants CLI) on an
/// underlying Scheduler, and is a useful scope for metrics.
///
/// Both Scheduler and Session are exposed to python and expected to be used by multiple threads, so
/// they use internal mutability in order to avoid exposing locks to callers.
///
pub struct Session {
  // The total size of the graph at Session-creation time.
  preceding_graph_size: usize,
  // The set of roots that have been requested within this session.
  roots: Mutex<HashSet<Root>>,
}

impl Session {
  pub fn new(scheduler: &Scheduler) -> Session {
    Session {
      preceding_graph_size: scheduler.core.graph.len(),
      roots: Mutex::new(HashSet::new()),
    }
  }

  fn extend(&self, new_roots: &[Root]) {
    let mut roots = self.roots.lock();
    roots.extend(new_roots.iter().cloned());
  }

  fn root_nodes(&self) -> Vec<NodeKey> {
    let roots = self.roots.lock();
    roots.iter().map(|r| r.clone().into()).collect()
  }
}

pub struct ExecutionRequest {
  // Set of roots for an execution, in the order they were declared.
  pub roots: Vec<Root>,
}

impl ExecutionRequest {
  pub fn new() -> ExecutionRequest {
    ExecutionRequest { roots: Vec::new() }
  }

  ///
  /// Roots are limited to `Select`, which is known to produce a Value. This method
  /// exists to satisfy Graph APIs which need instances of the NodeKey enum.
  ///
  fn root_nodes(&self) -> Vec<NodeKey> {
    self.roots.iter().map(|r| r.clone().into()).collect()
  }
}

///
/// Represents the state of an execution of a Graph.
///
pub struct Scheduler {
  pub core: Arc<Core>,
}

impl Scheduler {
  pub fn new(core: Core) -> Scheduler {
    Scheduler {
      core: Arc::new(core),
    }
  }

  pub fn visualize(&self, session: &Session, path: &Path) -> io::Result<()> {
    self
      .core
      .graph
      .visualize(Visualizer::default(), &session.root_nodes(), path)
  }

  pub fn trace(&self, request: &ExecutionRequest, path: &Path) -> Result<(), String> {
    self
      .core
      .graph
      .trace::<Tracer>(&request.root_nodes(), path)?;
    Ok(())
  }

  pub fn add_root_select(
    &self,
    request: &mut ExecutionRequest,
    subject: Key,
    product: TypeConstraint,
  ) -> Result<(), String> {
    let edges = self.find_root_edges_or_update_rule_graph(
      subject.type_id().clone(),
      &selectors::Select::new(product),
    )?;
    request
      .roots
      .push(Select::new(product, Params::new_single(subject), &edges));
    Ok(())
  }

  fn find_root_edges_or_update_rule_graph(
    &self,
    subject_type: TypeId,
    select: &selectors::Select,
  ) -> Result<rule_graph::RuleEdges, String> {
    self
      .core
      .rule_graph
      .find_root_edges(subject_type, select.clone())
      .ok_or_else(|| {
        format!(
          "No installed rules can satisfy {} for a root subject of type {}.",
          rule_graph::select_str(&select),
          rule_graph::type_str(subject_type)
        )
      })
  }

  ///
  /// Invalidate the invalidation roots represented by the given Paths.
  ///
  pub fn invalidate(&self, paths: &HashSet<PathBuf>) -> usize {
    let invalidation_result = self.core.graph.invalidate_from_roots(move |node| {
      if let Some(fs_subject) = node.fs_subject() {
        paths.contains(fs_subject)
      } else {
        false
      }
    });
    // TODO: Expose.
    invalidation_result.cleared + invalidation_result.dirtied
  }

  ///
  /// Invalidate all filesystem dependencies in the graph.
  ///
  pub fn invalidate_all_paths(&self) -> usize {
    let invalidation_result = self
      .core
      .graph
      .invalidate_from_roots(|node| node.fs_subject().is_some());
    // TODO: Expose.
    invalidation_result.cleared + invalidation_result.dirtied
  }

  ///
  /// Return Scheduler and per-Session metrics.
  ///
  pub fn metrics(&self, session: &Session) -> HashMap<&str, i64> {
    let mut m = HashMap::new();
    m.insert(
      "affected_file_count",
      self
        .core
        .graph
        .reachable_digest_count(&session.root_nodes()) as i64,
    );
    m.insert("preceding_graph_size", session.preceding_graph_size as i64);
    m.insert("resulting_graph_size", self.core.graph.len() as i64);
    m
  }

  ///
  /// Attempts to complete all of the given roots, retrying the entire set (up to `count`
  /// times) if any of them fail with `Failure::Invalidated`. Sends the result on the given
  /// mpsc Sender, which allows the caller to poll a channel for the result without blocking
  /// uninterruptibly on a Future.
  ///
  /// In common usage, graph entries won't be repeatedly invalidated, but in a case where they
  /// were (say by an automated process changing files under pants), we'd want to eventually
  /// give up.
  ///
  fn execute_helper(
    context: RootContext,
    sender: mpsc::Sender<Vec<Result<Value, Failure>>>,
    roots: Vec<Root>,
    count: usize,
  ) {
    let executor = context.core.runtime.get().executor();
    // Attempt all roots in parallel, failing fast to retry for `Invalidated`.
    let roots_res = future::join_all(
      roots
        .clone()
        .into_iter()
        .map(|root| {
          context
            .core
            .graph
            .create(root.clone().into(), &context)
            .then::<_, Result<Result<Value, Failure>, Failure>>(move |r| {
              match r {
                Err(Failure::Invalidated) if count > 0 => {
                  // A node was invalidated: fail quickly so that all roots can be retried.
                  Err(Failure::Invalidated)
                }
                other => {
                  // Otherwise (if it is a success, some other type of Failure, or if we've run
                  // out of retries) recover to complete the join, which will cause the results to
                  // propagate to the user.
                  debug!(
                    "Root {} completed.",
                    NodeKey::Select(Box::new(root)).format()
                  );
                  Ok(other.map(|res| {
                    res
                      .try_into()
                      .unwrap_or_else(|_| panic!("A Node implementation was ambiguous."))
                  }))
                }
              }
            })
        }).collect::<Vec<_>>(),
    );

    // If the join failed (due to `Invalidated`, since that is the only error we propagate), retry
    // the entire set of roots.
    executor.spawn(roots_res.then(move |res| {
      if let Ok(res) = res {
        sender.send(res).map_err(|_| ())
      } else {
        Scheduler::execute_helper(context, sender, roots, count - 1);
        Ok(())
      }
    }));
  }

  ///
  /// Compute the results for roots in the given request.
  ///
  pub fn execute<'e>(
    &self,
    request: &'e ExecutionRequest,
    session: &Session,
  ) -> Vec<(&'e Key, &'e TypeConstraint, RootResult)> {
    // Bootstrap tasks for the roots, and then wait for all of them.
    debug!("Launching {} roots.", request.roots.len());

    session.extend(&request.roots);

    // Wait for all roots to complete. Failure here should be impossible, because each
    // individual Future in the join was (eventually) mapped into success.
    let context = RootContext {
      core: self.core.clone(),
    };
    let (sender, receiver) = mpsc::channel();
    Scheduler::execute_helper(context, sender, request.roots.clone(), 8);
    let results = loop {
      if let Ok(res) = receiver.recv_timeout(Duration::from_millis(100)) {
        break res;
      }
    };

    request
      .roots
      .iter()
      .zip(results.into_iter())
      .map(|(s, r)| (s.params.expect_single(), &s.selector.product, r))
      .collect()
  }

  pub fn capture_snapshot_from_arbitrary_root<P: AsRef<Path>>(
    &self,
    root_path: P,
    path_globs: fs::PathGlobs,
  ) -> BoxFuture<fs::Snapshot, String> {
    // Note that we don't use a Graph here, and don't cache any intermediate steps, we just place
    // the resultant Snapshot into the store and return it. This is important, because we're reading
    // things from arbitrary filepaths which we don't want to cache in the graph, as we don't watch
    // them for changes.
    // We assume that this Snapshot is of an immutable piece of the filesystem.

    let posix_fs = Arc::new(try_future!(PosixFS::new(
      root_path,
      self.core.fs_pool.clone(),
      &[]
    )));
    let store = self.core.store.clone();

    posix_fs
      .expand(path_globs)
      .map_err(|err| format!("Error expanding globs: {:?}", err))
      .and_then(|path_stats| {
        fs::Snapshot::from_path_stats(
          store.clone(),
          &fs::OneOffStoreFileByDigest::new(store, posix_fs),
          path_stats,
        )
      }).to_boxed()
  }
}

impl Drop for Scheduler {
  fn drop(&mut self) {
    // Because Nodes may hold references to the Core in their closure, this is intended to
    // break cycles between Nodes and the Core.
    self.core.graph.clear();
  }
}

///
/// Root requests are limited to Selectors that produce (python) Values.
///
type Root = Select;

pub type RootResult = Result<Value, Failure>;

///
/// NB: This basic wrapper exists to allow us to implement the `NodeContext` trait (which lives
/// outside of this crate) for the `Arc` struct (which also lives outside our crate), which is not
/// possible without the wrapper due to "trait coherence".
///
#[derive(Clone)]
struct RootContext {
  core: Arc<Core>,
}

impl NodeContext for RootContext {
  type Node = NodeKey;

  fn clone_for(&self, entry_id: EntryId) -> Context {
    Context::new(entry_id, self.core.clone())
  }

  fn graph(&self) -> &Graph<NodeKey> {
    &self.core.graph
  }

  fn spawn<F>(&self, future: F)
  where
    F: Future<Item = (), Error = ()> + Send + 'static,
  {
    self.core.runtime.get().executor().spawn(future);
  }
}
<|endoftext|>


<commit_before>// This file is part of the uutils coreutils package.
//
// (c) Jian Zeng <anonymousknight96@gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

// spell-checker:ignore (ToDO) COMFOLLOW Passwd RFILE RFILE's derefer dgid duid groupname

use uucore::display::Quotable;
pub use uucore::entries::{self, Group, Locate, Passwd};
use uucore::format_usage;
use uucore::perms::{chown_base, options, IfFrom};

use uucore::error::{FromIo, UResult, USimpleError};

use clap::{crate_version, Arg, ArgMatches, Command};

use std::fs;
use std::os::unix::fs::MetadataExt;

static ABOUT: &str = "change file owner and group";

const USAGE: &str = "\
    {} [OPTION]... [OWNER][:[GROUP]] FILE...
    {} [OPTION]... --reference=RFILE FILE...";

fn parse_gid_uid_and_filter(matches: &ArgMatches) -> UResult<(Option<u32>, Option<u32>, IfFrom)> {
    let filter = if let Some(spec) = matches.value_of(options::FROM) {
        match parse_spec(spec, ':')? {
            (Some(uid), None) => IfFrom::User(uid),
            (None, Some(gid)) => IfFrom::Group(gid),
            (Some(uid), Some(gid)) => IfFrom::UserGroup(uid, gid),
            (None, None) => IfFrom::All,
        }
    } else {
        IfFrom::All
    };

    let dest_uid: Option<u32>;
    let dest_gid: Option<u32>;
    if let Some(file) = matches.value_of(options::REFERENCE) {
        let meta = fs::metadata(&file)
            .map_err_context(|| format!("failed to get attributes of {}", file.quote()))?;
        dest_gid = Some(meta.gid());
        dest_uid = Some(meta.uid());
    } else {
        let (u, g) = parse_spec(matches.value_of(options::ARG_OWNER).unwrap(), ':')?;
        dest_uid = u;
        dest_gid = g;
    }
    Ok((dest_gid, dest_uid, filter))
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    chown_base(
        uu_app(),
        args,
        options::ARG_OWNER,
        parse_gid_uid_and_filter,
        false,
    )
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .infer_long_args(true)
        .arg(
            Arg::new(options::HELP)
                .long(options::HELP)
                .help("Print help information."),
        )
        .arg(
            Arg::new(options::verbosity::CHANGES)
                .short('c')
                .long(options::verbosity::CHANGES)
                .help("like verbose but report only when a change is made"),
        )
        .arg(
            Arg::new(options::dereference::DEREFERENCE)
                .long(options::dereference::DEREFERENCE)
                .help(
                    "affect the referent of each symbolic link (this is the default), \
                    rather than the symbolic link itself",
                ),
        )
        .arg(
            Arg::new(options::dereference::NO_DEREFERENCE)
                .short('h')
                .long(options::dereference::NO_DEREFERENCE)
                .help(
                    "affect symbolic links instead of any referenced file \
                    (useful only on systems that can change the ownership of a symlink)",
                ),
        )
        .arg(
            Arg::new(options::FROM)
                .long(options::FROM)
                .help(
                    "change the owner and/or group of each file only if its \
                    current owner and/or group match those specified here. \
                    Either may be omitted, in which case a match is not required \
                    for the omitted attribute",
                )
                .value_name("CURRENT_OWNER:CURRENT_GROUP"),
        )
        .arg(
            Arg::new(options::preserve_root::PRESERVE)
                .long(options::preserve_root::PRESERVE)
                .help("fail to operate recursively on '/'"),
        )
        .arg(
            Arg::new(options::preserve_root::NO_PRESERVE)
                .long(options::preserve_root::NO_PRESERVE)
                .help("do not treat '/' specially (the default)"),
        )
        .arg(
            Arg::new(options::verbosity::QUIET)
                .long(options::verbosity::QUIET)
                .help("suppress most error messages"),
        )
        .arg(
            Arg::new(options::RECURSIVE)
                .short('R')
                .long(options::RECURSIVE)
                .help("operate on files and directories recursively"),
        )
        .arg(
            Arg::new(options::REFERENCE)
                .long(options::REFERENCE)
                .help("use RFILE's owner and group rather than specifying OWNER:GROUP values")
                .value_name("RFILE")
                .min_values(1),
        )
        .arg(
            Arg::new(options::verbosity::SILENT)
                .short('f')
                .long(options::verbosity::SILENT),
        )
        .arg(
            Arg::new(options::traverse::TRAVERSE)
                .short(options::traverse::TRAVERSE.chars().next().unwrap())
                .help("if a command line argument is a symbolic link to a directory, traverse it")
                .overrides_with_all(&[options::traverse::EVERY, options::traverse::NO_TRAVERSE]),
        )
        .arg(
            Arg::new(options::traverse::EVERY)
                .short(options::traverse::EVERY.chars().next().unwrap())
                .help("traverse every symbolic link to a directory encountered")
                .overrides_with_all(&[options::traverse::TRAVERSE, options::traverse::NO_TRAVERSE]),
        )
        .arg(
            Arg::new(options::traverse::NO_TRAVERSE)
                .short(options::traverse::NO_TRAVERSE.chars().next().unwrap())
                .help("do not traverse any symbolic links (default)")
                .overrides_with_all(&[options::traverse::TRAVERSE, options::traverse::EVERY]),
        )
        .arg(
            Arg::new(options::verbosity::VERBOSE)
                .long(options::verbosity::VERBOSE)
                .short('v')
                .help("output a diagnostic for every file processed"),
        )
}

/// Parse the owner/group specifier string into a user ID and a group ID.
///
/// The `spec` can be of the form:
///
/// * `"owner:group"`,
/// * `"owner"`,
/// * `":group"`,
///
/// and the owner or group can be specified either as an ID or a
/// name. The `sep` argument specifies which character to use as a
/// separator between the owner and group; calling code should set
/// this to `':'`.
fn parse_spec(spec: &str, sep: char) -> UResult<(Option<u32>, Option<u32>)> {
    assert!(['.', ':'].contains(&sep));
    let mut args = spec.splitn(2, sep);
    let user = args.next().unwrap_or("");
    let group = args.next().unwrap_or("");

    let uid = if !user.is_empty() {
        Some(match Passwd::locate(user) {
            Ok(u) => u.uid, // We have been able to get the uid
            Err(_) =>
            // we have NOT been able to find the uid
            // but we could be in the case where we have user.group
            {
                if spec.contains('.') && !spec.contains(':') && sep == ':' {
                    // but the input contains a '.' but not a ':'
                    // we might have something like username.groupname
                    // So, try to parse it this way
                    return parse_spec(spec, '.');
                } else {
                    // It's possible that the `user` string contains a
                    // numeric user ID, in which case, we respect that.
                    match user.parse() {
                        Ok(uid) => uid,
                        Err(_) => {
                            return Err(USimpleError::new(
                                1,
                                format!("invalid user: {}", spec.quote()),
                            ))
                        }
                    }
                }
            }
        })
    } else {
        None
    };
    let gid = if !group.is_empty() {
        Some(match Group::locate(group) {
            Ok(g) => g.gid,
            Err(_) => match group.parse() {
                Ok(gid) => gid,
                Err(_) => {
                    return Err(USimpleError::new(
                        1,
                        format!("invalid group: {}", spec.quote()),
                    ));
                }
            },
        })
    } else {
        None
    };
    Ok((uid, gid))
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_parse_spec() {
        assert!(matches!(parse_spec(":", ':'), Ok((None, None))));
        assert!(matches!(parse_spec(".", ':'), Ok((None, None))));
        assert!(matches!(parse_spec(".", '.'), Ok((None, None))));
        assert!(format!("{}", parse_spec("::", ':').err().unwrap()).starts_with("invalid group: "));
        assert!(format!("{}", parse_spec("..", ':').err().unwrap()).starts_with("invalid group: "));
    }

    /// Test for parsing IDs that don't correspond to a named user or group.
    #[test]
    fn test_parse_spec_nameless_ids() {
        // This assumes that there is no named user with ID 12345.
        assert!(matches!(parse_spec("12345", ':'), Ok((Some(12345), None))));
        // This assumes that there is no named group with ID 54321.
        assert!(matches!(parse_spec(":54321", ':'), Ok((None, Some(54321)))));
        assert!(matches!(
            parse_spec("12345:54321", ':'),
            Ok((Some(12345), Some(54321)))
        ));
    }
}
<commit_msg>WIP Trying to diagnose 'invalid group: 1001:121' error in CI environment<commit_after>// This file is part of the uutils coreutils package.
//
// (c) Jian Zeng <anonymousknight96@gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

// spell-checker:ignore (ToDO) COMFOLLOW Passwd RFILE RFILE's derefer dgid duid groupname

use uucore::display::Quotable;
pub use uucore::entries::{self, Group, Locate, Passwd};
use uucore::format_usage;
use uucore::perms::{chown_base, options, IfFrom};

use uucore::error::{FromIo, UResult, USimpleError};

use clap::{crate_version, Arg, ArgMatches, Command};

use std::fs;
use std::os::unix::fs::MetadataExt;

static ABOUT: &str = "change file owner and group";

const USAGE: &str = "\
    {} [OPTION]... [OWNER][:[GROUP]] FILE...
    {} [OPTION]... --reference=RFILE FILE...";

fn parse_gid_uid_and_filter(matches: &ArgMatches) -> UResult<(Option<u32>, Option<u32>, IfFrom)> {
    let filter = if let Some(spec) = matches.value_of(options::FROM) {
        match parse_spec(spec, ':')? {
            (Some(uid), None) => IfFrom::User(uid),
            (None, Some(gid)) => IfFrom::Group(gid),
            (Some(uid), Some(gid)) => IfFrom::UserGroup(uid, gid),
            (None, None) => IfFrom::All,
        }
    } else {
        IfFrom::All
    };

    let dest_uid: Option<u32>;
    let dest_gid: Option<u32>;
    if let Some(file) = matches.value_of(options::REFERENCE) {
        let meta = fs::metadata(&file)
            .map_err_context(|| format!("failed to get attributes of {}", file.quote()))?;
        dest_gid = Some(meta.gid());
        dest_uid = Some(meta.uid());
    } else {
        let (u, g) = parse_spec(matches.value_of(options::ARG_OWNER).unwrap(), ':')?;
        dest_uid = u;
        dest_gid = g;
    }
    Ok((dest_gid, dest_uid, filter))
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    chown_base(
        uu_app(),
        args,
        options::ARG_OWNER,
        parse_gid_uid_and_filter,
        false,
    )
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .infer_long_args(true)
        .arg(
            Arg::new(options::HELP)
                .long(options::HELP)
                .help("Print help information."),
        )
        .arg(
            Arg::new(options::verbosity::CHANGES)
                .short('c')
                .long(options::verbosity::CHANGES)
                .help("like verbose but report only when a change is made"),
        )
        .arg(
            Arg::new(options::dereference::DEREFERENCE)
                .long(options::dereference::DEREFERENCE)
                .help(
                    "affect the referent of each symbolic link (this is the default), \
                    rather than the symbolic link itself",
                ),
        )
        .arg(
            Arg::new(options::dereference::NO_DEREFERENCE)
                .short('h')
                .long(options::dereference::NO_DEREFERENCE)
                .help(
                    "affect symbolic links instead of any referenced file \
                    (useful only on systems that can change the ownership of a symlink)",
                ),
        )
        .arg(
            Arg::new(options::FROM)
                .long(options::FROM)
                .help(
                    "change the owner and/or group of each file only if its \
                    current owner and/or group match those specified here. \
                    Either may be omitted, in which case a match is not required \
                    for the omitted attribute",
                )
                .value_name("CURRENT_OWNER:CURRENT_GROUP"),
        )
        .arg(
            Arg::new(options::preserve_root::PRESERVE)
                .long(options::preserve_root::PRESERVE)
                .help("fail to operate recursively on '/'"),
        )
        .arg(
            Arg::new(options::preserve_root::NO_PRESERVE)
                .long(options::preserve_root::NO_PRESERVE)
                .help("do not treat '/' specially (the default)"),
        )
        .arg(
            Arg::new(options::verbosity::QUIET)
                .long(options::verbosity::QUIET)
                .help("suppress most error messages"),
        )
        .arg(
            Arg::new(options::RECURSIVE)
                .short('R')
                .long(options::RECURSIVE)
                .help("operate on files and directories recursively"),
        )
        .arg(
            Arg::new(options::REFERENCE)
                .long(options::REFERENCE)
                .help("use RFILE's owner and group rather than specifying OWNER:GROUP values")
                .value_name("RFILE")
                .min_values(1),
        )
        .arg(
            Arg::new(options::verbosity::SILENT)
                .short('f')
                .long(options::verbosity::SILENT),
        )
        .arg(
            Arg::new(options::traverse::TRAVERSE)
                .short(options::traverse::TRAVERSE.chars().next().unwrap())
                .help("if a command line argument is a symbolic link to a directory, traverse it")
                .overrides_with_all(&[options::traverse::EVERY, options::traverse::NO_TRAVERSE]),
        )
        .arg(
            Arg::new(options::traverse::EVERY)
                .short(options::traverse::EVERY.chars().next().unwrap())
                .help("traverse every symbolic link to a directory encountered")
                .overrides_with_all(&[options::traverse::TRAVERSE, options::traverse::NO_TRAVERSE]),
        )
        .arg(
            Arg::new(options::traverse::NO_TRAVERSE)
                .short(options::traverse::NO_TRAVERSE.chars().next().unwrap())
                .help("do not traverse any symbolic links (default)")
                .overrides_with_all(&[options::traverse::TRAVERSE, options::traverse::EVERY]),
        )
        .arg(
            Arg::new(options::verbosity::VERBOSE)
                .long(options::verbosity::VERBOSE)
                .short('v')
                .help("output a diagnostic for every file processed"),
        )
}

/// Parse the owner/group specifier string into a user ID and a group ID.
///
/// The `spec` can be of the form:
///
/// * `"owner:group"`,
/// * `"owner"`,
/// * `":group"`,
///
/// and the owner or group can be specified either as an ID or a
/// name. The `sep` argument specifies which character to use as a
/// separator between the owner and group; calling code should set
/// this to `':'`.
fn parse_spec(spec: &str, sep: char) -> UResult<(Option<u32>, Option<u32>)> {
    assert!(['.', ':'].contains(&sep));
    let mut args = spec.splitn(2, sep);
    let user = args.next().unwrap_or("");
    let group = args.next().unwrap_or("");

    let uid = if !user.is_empty() {
        Some(match Passwd::locate(user) {
            Ok(u) => u.uid, // We have been able to get the uid
            Err(_) =>
            // we have NOT been able to find the uid
            // but we could be in the case where we have user.group
            {
                if spec.contains('.') && !spec.contains(':') && sep == ':' {
                    // but the input contains a '.' but not a ':'
                    // we might have something like username.groupname
                    // So, try to parse it this way
                    return parse_spec(spec, '.');
                } else {
                    // It's possible that the `user` string contains a
                    // numeric user ID, in which case, we respect that.
                    match user.parse() {
                        Ok(uid) => uid,
                        Err(_) => {
                            return Err(USimpleError::new(
                                1,
                                format!("invalid user: {}", spec.quote()),
                            ))
                        }
                    }
                }
            }
        })
    } else {
        None
    };
    let gid = if !group.is_empty() {
        Some(match Group::locate(group) {
            Ok(g) => g.gid,
            Err(_) => match group.parse() {
                Ok(gid) => gid,
                Err(e) => {
                    eprintln!("error {}, user {}, group {}, uid {:?}", e, user, group, uid);
                    return Err(USimpleError::new(
                        1,
                        format!("invalid group: {}", spec.quote()),
                    ));
                }
            },
        })
    } else {
        None
    };
    Ok((uid, gid))
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_parse_spec() {
        assert!(matches!(parse_spec(":", ':'), Ok((None, None))));
        assert!(matches!(parse_spec(".", ':'), Ok((None, None))));
        assert!(matches!(parse_spec(".", '.'), Ok((None, None))));
        assert!(format!("{}", parse_spec("::", ':').err().unwrap()).starts_with("invalid group: "));
        assert!(format!("{}", parse_spec("..", ':').err().unwrap()).starts_with("invalid group: "));
    }

    /// Test for parsing IDs that don't correspond to a named user or group.
    #[test]
    fn test_parse_spec_nameless_ids() {
        // This assumes that there is no named user with ID 12345.
        assert!(matches!(parse_spec("12345", ':'), Ok((Some(12345), None))));
        // This assumes that there is no named group with ID 54321.
        assert!(matches!(parse_spec(":54321", ':'), Ok((None, Some(54321)))));
        assert!(matches!(
            parse_spec("12345:54321", ':'),
            Ok((Some(12345), Some(54321)))
        ));
    }
}
<|endoftext|>