<commit_before>use crate::cli;
use crate::command_prelude::*;
use anyhow::{bail, format_err};
use cargo::core::dependency::DepKind;
use cargo::ops::tree::{self, EdgeKind};
use cargo::ops::Packages;
use cargo::util::print_available_packages;
use cargo::util::CargoResult;
use std::collections::HashSet;
use std::str::FromStr;

pub fn cli() -> App {
    subcommand("tree")
        .about("Display a tree visualization of a dependency graph")
        .arg(opt("quiet", "Suppress status messages").short("q"))
        .arg_manifest_path()
        .arg_package_spec_no_all(
            "Package to be used as the root of the tree",
            "Display the tree for all packages in the workspace",
            "Exclude specific workspace members",
        )
        // Deprecated, use --no-dedupe instead.
        .arg(Arg::with_name("all").long("all").short("a").hidden(true))
        // Deprecated, use --target=all instead.
        .arg(
            Arg::with_name("all-targets")
                .long("all-targets")
                .hidden(true),
        )
        .arg_features()
        .arg_target_triple(
            "Filter dependencies matching the given target-triple (default host platform). \
            Pass `all` to include all targets.",
        )
        // Deprecated, use -e=no-dev instead.
        .arg(
            Arg::with_name("no-dev-dependencies")
                .long("no-dev-dependencies")
                .hidden(true),
        )
        .arg(
            multi_opt(
                "edges",
                "KINDS",
                "The kinds of dependencies to display \
                 (features, normal, build, dev, all, \
                 no-normal, no-build, no-dev, no-proc-macro)",
            )
            .short("e"),
        )
        .arg(
            optional_multi_opt(
                "invert",
                "SPEC",
                "Invert the tree direction and focus on the given package",
            )
            .short("i"),
        )
        // Deprecated, use --prefix=none instead.
        .arg(Arg::with_name("no-indent").long("no-indent").hidden(true))
        // Deprecated, use --prefix=depth instead.
        .arg(
            Arg::with_name("prefix-depth")
                .long("prefix-depth")
                .hidden(true),
        )
        .arg(
            opt(
                "prefix",
                "Change the prefix (indentation) of how each entry is displayed",
            )
            .value_name("PREFIX")
            .possible_values(&["depth", "indent", "none"])
            .default_value("indent"),
        )
        .arg(opt(
            "no-dedupe",
            "Do not de-duplicate (repeats all shared dependencies)",
        ))
        .arg(
            opt(
                "duplicates",
                "Show only dependencies which come in multiple versions (implies -i)",
            )
            .short("d")
            .alias("duplicate"),
        )
        .arg(
            opt("charset", "Character set to use in output: utf8, ascii")
                .value_name("CHARSET")
                .possible_values(&["utf8", "ascii"])
                .default_value("utf8"),
        )
        .arg(
            opt("format", "Format string used for printing dependencies")
                .value_name("FORMAT")
                .short("f")
                .default_value("{p}"),
        )
        .arg(
            // Backwards compatibility with old cargo-tree.
            Arg::with_name("version")
                .long("version")
                .short("V")
                .hidden(true),
        )
        .after_help("Run `cargo help tree` for more detailed information.\n")
}

pub fn exec(config: &mut Config, args: &ArgMatches<'_>) -> CliResult {
    if args.is_present("version") {
        let verbose = args.occurrences_of("verbose") > 0;
        let version = cli::get_version_string(verbose);
        cargo::drop_print!(config, "{}", version);
        return Ok(());
    }
    let prefix = if args.is_present("no-indent") {
        config
            .shell()
            .warn("the --no-indent flag has been changed to --prefix=none")?;
        "none"
    } else if args.is_present("prefix-depth") {
        config
            .shell()
            .warn("the --prefix-depth flag has been changed to --prefix=depth")?;
        "depth"
    } else {
        args.value_of("prefix").unwrap()
    };
    let prefix = tree::Prefix::from_str(prefix).map_err(|e| anyhow::anyhow!("{}", e))?;

    let no_dedupe = args.is_present("no-dedupe") || args.is_present("all");
    if args.is_present("all") {
        config.shell().warn(
            "The `cargo tree` --all flag has been changed to --no-dedupe, \
             and may be removed in a future version.\n\
             If you are looking to display all workspace members, use the --workspace flag.",
        )?;
    }

    let targets = if args.is_present("all-targets") {
        config
            .shell()
            .warn("the --all-targets flag has been changed to --target=all")?;
        vec!["all".to_string()]
    } else {
        args._values_of("target")
    };
    let target = tree::Target::from_cli(targets);

    let (edge_kinds, no_proc_macro) = parse_edge_kinds(config, args)?;
    let graph_features = edge_kinds.contains(&EdgeKind::Feature);

    let packages = args.packages_from_flags()?;
    let mut invert = args
        .values_of("invert")
        .map_or_else(|| Vec::new(), |is| is.map(|s| s.to_string()).collect());
    if args.is_present_with_zero_values("invert") {
        match &packages {
            Packages::Packages(ps) => {
                // Backwards compatibility with old syntax of `cargo tree -i -p foo`.
                invert.extend(ps.clone());
            }
            _ => {
                return Err(format_err!(
                    "The `-i` flag requires a package name.\n\
\n\
The `-i` flag is used to inspect the reverse dependencies of a specific\n\
package. It will invert the tree and display the packages that depend on the\n\
given package.\n\
\n\
Note that in a workspace, by default it will only display the package's\n\
reverse dependencies inside the tree of the workspace member in the current\n\
directory. The --workspace flag can be used to extend it so that it will show\n\
the package's reverse dependencies across the entire workspace. The -p flag\n\
can be used to display the package's reverse dependencies only with the\n\
subtree of the package given to -p.\n\
"
                )
                .into());
            }
        }
    }

    let ws = args.workspace(config)?;

    if args.is_present_with_zero_values("package") {
        print_available_packages(&ws)?;
    }

    let charset = tree::Charset::from_str(args.value_of("charset").unwrap())
        .map_err(|e| anyhow::anyhow!("{}", e))?;
    let opts = tree::TreeOptions {
        cli_features: args.cli_features()?,
        packages,
        target,
        edge_kinds,
        invert,
        prefix,
        no_dedupe,
        duplicates: args.is_present("duplicates"),
        charset,
        format: args.value_of("format").unwrap().to_string(),
        graph_features,
        no_proc_macro,
    };

    if opts.graph_features && opts.duplicates {
        return Err(format_err!("the `-e features` flag does not support `--duplicates`").into());
    }

    tree::build_and_print(&ws, &opts)?;
    Ok(())
}

/// Parses `--edges` option.
///
/// Returns a tuple of `EdgeKind` map and `no_proc_marco` flag.
fn parse_edge_kinds(
    config: &Config,
    args: &ArgMatches<'_>,
) -> CargoResult<(HashSet<EdgeKind>, bool)> {
    let (kinds, no_proc_macro) = {
        let mut no_proc_macro = false;
        let mut kinds = args.values_of("edges").map_or_else(
            || Vec::new(),
            |es| {
                es.flat_map(|e| e.split(','))
                    .filter(|e| {
                        no_proc_macro = *e == "no-proc-macro";
                        !no_proc_macro
                    })
                    .collect()
            },
        );

        if args.is_present("no-dev-dependencies") {
            config
                .shell()
                .warn("the --no-dev-dependencies flag has changed to -e=no-dev")?;
            kinds.push("no-dev");
        }

        if kinds.is_empty() {
            kinds.extend(&["normal", "build", "dev"]);
        }

        (kinds, no_proc_macro)
    };

    let mut result = HashSet::new();
    let insert_defaults = |result: &mut HashSet<EdgeKind>| {
        result.insert(EdgeKind::Dep(DepKind::Normal));
        result.insert(EdgeKind::Dep(DepKind::Build));
        result.insert(EdgeKind::Dep(DepKind::Development));
    };
    let unknown = |k| {
        bail!(
            "unknown edge kind `{}`, valid values are \
                \"normal\", \"build\", \"dev\", \
                \"no-normal\", \"no-build\", \"no-dev\", \"no-proc-macro\", \
                \"features\", or \"all\"",
            k
        )
    };
    if kinds.iter().any(|k| k.starts_with("no-")) {
        insert_defaults(&mut result);
        for kind in &kinds {
            match *kind {
                "no-normal" => result.remove(&EdgeKind::Dep(DepKind::Normal)),
                "no-build" => result.remove(&EdgeKind::Dep(DepKind::Build)),
                "no-dev" => result.remove(&EdgeKind::Dep(DepKind::Development)),
                "features" => result.insert(EdgeKind::Feature),
                k @ "normal" | k @ "build" | k @ "dev" | k @ "all" => {
                    bail!(
                        "`{}` dependency kind cannot be mixed with \
                            \"no-normal\", \"no-build\", or \"no-dev\" \
                            dependency kinds",
                        k
                    )
                }
                k => return unknown(k),
            };
        }
        return Ok((result, no_proc_macro));
    }
    for kind in &kinds {
        match *kind {
            "all" => {
                insert_defaults(&mut result);
                result.insert(EdgeKind::Feature);
            }
            "features" => {
                result.insert(EdgeKind::Feature);
            }
            "normal" => {
                result.insert(EdgeKind::Dep(DepKind::Normal));
            }
            "build" => {
                result.insert(EdgeKind::Dep(DepKind::Build));
            }
            "dev" => {
                result.insert(EdgeKind::Dep(DepKind::Development));
            }
            k => return unknown(k),
        }
    }
    if kinds.len() == 1 && kinds[0] == "features" {
        insert_defaults(&mut result);
    }
    Ok((result, no_proc_macro))
}
<commit_msg>fix: remove unnecessary match bindings<commit_after>use crate::cli;
use crate::command_prelude::*;
use anyhow::{bail, format_err};
use cargo::core::dependency::DepKind;
use cargo::ops::tree::{self, EdgeKind};
use cargo::ops::Packages;
use cargo::util::print_available_packages;
use cargo::util::CargoResult;
use std::collections::HashSet;
use std::str::FromStr;

pub fn cli() -> App {
    subcommand("tree")
        .about("Display a tree visualization of a dependency graph")
        .arg(opt("quiet", "Suppress status messages").short("q"))
        .arg_manifest_path()
        .arg_package_spec_no_all(
            "Package to be used as the root of the tree",
            "Display the tree for all packages in the workspace",
            "Exclude specific workspace members",
        )
        // Deprecated, use --no-dedupe instead.
        .arg(Arg::with_name("all").long("all").short("a").hidden(true))
        // Deprecated, use --target=all instead.
        .arg(
            Arg::with_name("all-targets")
                .long("all-targets")
                .hidden(true),
        )
        .arg_features()
        .arg_target_triple(
            "Filter dependencies matching the given target-triple (default host platform). \
            Pass `all` to include all targets.",
        )
        // Deprecated, use -e=no-dev instead.
        .arg(
            Arg::with_name("no-dev-dependencies")
                .long("no-dev-dependencies")
                .hidden(true),
        )
        .arg(
            multi_opt(
                "edges",
                "KINDS",
                "The kinds of dependencies to display \
                 (features, normal, build, dev, all, \
                 no-normal, no-build, no-dev, no-proc-macro)",
            )
            .short("e"),
        )
        .arg(
            optional_multi_opt(
                "invert",
                "SPEC",
                "Invert the tree direction and focus on the given package",
            )
            .short("i"),
        )
        // Deprecated, use --prefix=none instead.
        .arg(Arg::with_name("no-indent").long("no-indent").hidden(true))
        // Deprecated, use --prefix=depth instead.
        .arg(
            Arg::with_name("prefix-depth")
                .long("prefix-depth")
                .hidden(true),
        )
        .arg(
            opt(
                "prefix",
                "Change the prefix (indentation) of how each entry is displayed",
            )
            .value_name("PREFIX")
            .possible_values(&["depth", "indent", "none"])
            .default_value("indent"),
        )
        .arg(opt(
            "no-dedupe",
            "Do not de-duplicate (repeats all shared dependencies)",
        ))
        .arg(
            opt(
                "duplicates",
                "Show only dependencies which come in multiple versions (implies -i)",
            )
            .short("d")
            .alias("duplicate"),
        )
        .arg(
            opt("charset", "Character set to use in output: utf8, ascii")
                .value_name("CHARSET")
                .possible_values(&["utf8", "ascii"])
                .default_value("utf8"),
        )
        .arg(
            opt("format", "Format string used for printing dependencies")
                .value_name("FORMAT")
                .short("f")
                .default_value("{p}"),
        )
        .arg(
            // Backwards compatibility with old cargo-tree.
            Arg::with_name("version")
                .long("version")
                .short("V")
                .hidden(true),
        )
        .after_help("Run `cargo help tree` for more detailed information.\n")
}

pub fn exec(config: &mut Config, args: &ArgMatches<'_>) -> CliResult {
    if args.is_present("version") {
        let verbose = args.occurrences_of("verbose") > 0;
        let version = cli::get_version_string(verbose);
        cargo::drop_print!(config, "{}", version);
        return Ok(());
    }
    let prefix = if args.is_present("no-indent") {
        config
            .shell()
            .warn("the --no-indent flag has been changed to --prefix=none")?;
        "none"
    } else if args.is_present("prefix-depth") {
        config
            .shell()
            .warn("the --prefix-depth flag has been changed to --prefix=depth")?;
        "depth"
    } else {
        args.value_of("prefix").unwrap()
    };
    let prefix = tree::Prefix::from_str(prefix).map_err(|e| anyhow::anyhow!("{}", e))?;

    let no_dedupe = args.is_present("no-dedupe") || args.is_present("all");
    if args.is_present("all") {
        config.shell().warn(
            "The `cargo tree` --all flag has been changed to --no-dedupe, \
             and may be removed in a future version.\n\
             If you are looking to display all workspace members, use the --workspace flag.",
        )?;
    }

    let targets = if args.is_present("all-targets") {
        config
            .shell()
            .warn("the --all-targets flag has been changed to --target=all")?;
        vec!["all".to_string()]
    } else {
        args._values_of("target")
    };
    let target = tree::Target::from_cli(targets);

    let (edge_kinds, no_proc_macro) = parse_edge_kinds(config, args)?;
    let graph_features = edge_kinds.contains(&EdgeKind::Feature);

    let packages = args.packages_from_flags()?;
    let mut invert = args
        .values_of("invert")
        .map_or_else(|| Vec::new(), |is| is.map(|s| s.to_string()).collect());
    if args.is_present_with_zero_values("invert") {
        match &packages {
            Packages::Packages(ps) => {
                // Backwards compatibility with old syntax of `cargo tree -i -p foo`.
                invert.extend(ps.clone());
            }
            _ => {
                return Err(format_err!(
                    "The `-i` flag requires a package name.\n\
\n\
The `-i` flag is used to inspect the reverse dependencies of a specific\n\
package. It will invert the tree and display the packages that depend on the\n\
given package.\n\
\n\
Note that in a workspace, by default it will only display the package's\n\
reverse dependencies inside the tree of the workspace member in the current\n\
directory. The --workspace flag can be used to extend it so that it will show\n\
the package's reverse dependencies across the entire workspace. The -p flag\n\
can be used to display the package's reverse dependencies only with the\n\
subtree of the package given to -p.\n\
"
                )
                .into());
            }
        }
    }

    let ws = args.workspace(config)?;

    if args.is_present_with_zero_values("package") {
        print_available_packages(&ws)?;
    }

    let charset = tree::Charset::from_str(args.value_of("charset").unwrap())
        .map_err(|e| anyhow::anyhow!("{}", e))?;
    let opts = tree::TreeOptions {
        cli_features: args.cli_features()?,
        packages,
        target,
        edge_kinds,
        invert,
        prefix,
        no_dedupe,
        duplicates: args.is_present("duplicates"),
        charset,
        format: args.value_of("format").unwrap().to_string(),
        graph_features,
        no_proc_macro,
    };

    if opts.graph_features && opts.duplicates {
        return Err(format_err!("the `-e features` flag does not support `--duplicates`").into());
    }

    tree::build_and_print(&ws, &opts)?;
    Ok(())
}

/// Parses `--edges` option.
///
/// Returns a tuple of `EdgeKind` map and `no_proc_marco` flag.
fn parse_edge_kinds(
    config: &Config,
    args: &ArgMatches<'_>,
) -> CargoResult<(HashSet<EdgeKind>, bool)> {
    let (kinds, no_proc_macro) = {
        let mut no_proc_macro = false;
        let mut kinds = args.values_of("edges").map_or_else(
            || Vec::new(),
            |es| {
                es.flat_map(|e| e.split(','))
                    .filter(|e| {
                        no_proc_macro = *e == "no-proc-macro";
                        !no_proc_macro
                    })
                    .collect()
            },
        );

        if args.is_present("no-dev-dependencies") {
            config
                .shell()
                .warn("the --no-dev-dependencies flag has changed to -e=no-dev")?;
            kinds.push("no-dev");
        }

        if kinds.is_empty() {
            kinds.extend(&["normal", "build", "dev"]);
        }

        (kinds, no_proc_macro)
    };

    let mut result = HashSet::new();
    let insert_defaults = |result: &mut HashSet<EdgeKind>| {
        result.insert(EdgeKind::Dep(DepKind::Normal));
        result.insert(EdgeKind::Dep(DepKind::Build));
        result.insert(EdgeKind::Dep(DepKind::Development));
    };
    let unknown = |k| {
        bail!(
            "unknown edge kind `{}`, valid values are \
                \"normal\", \"build\", \"dev\", \
                \"no-normal\", \"no-build\", \"no-dev\", \"no-proc-macro\", \
                \"features\", or \"all\"",
            k
        )
    };
    if kinds.iter().any(|k| k.starts_with("no-")) {
        insert_defaults(&mut result);
        for kind in &kinds {
            match *kind {
                "no-normal" => result.remove(&EdgeKind::Dep(DepKind::Normal)),
                "no-build" => result.remove(&EdgeKind::Dep(DepKind::Build)),
                "no-dev" => result.remove(&EdgeKind::Dep(DepKind::Development)),
                "features" => result.insert(EdgeKind::Feature),
                "normal" | "build" | "dev" | "all" => {
                    bail!(
                        "`{}` dependency kind cannot be mixed with \
                            \"no-normal\", \"no-build\", or \"no-dev\" \
                            dependency kinds",
                        kind
                    )
                }
                k => return unknown(k),
            };
        }
        return Ok((result, no_proc_macro));
    }
    for kind in &kinds {
        match *kind {
            "all" => {
                insert_defaults(&mut result);
                result.insert(EdgeKind::Feature);
            }
            "features" => {
                result.insert(EdgeKind::Feature);
            }
            "normal" => {
                result.insert(EdgeKind::Dep(DepKind::Normal));
            }
            "build" => {
                result.insert(EdgeKind::Dep(DepKind::Build));
            }
            "dev" => {
                result.insert(EdgeKind::Dep(DepKind::Development));
            }
            k => return unknown(k),
        }
    }
    if kinds.len() == 1 && kinds[0] == "features" {
        insert_defaults(&mut result);
    }
    Ok((result, no_proc_macro))
}
<|endoftext|>


<commit_before>// Copyright 2017 click2stream, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::io;

use std::rc::Rc;
use std::cell::RefCell;
use std::collections::HashMap;

use bytes::{Bytes, BytesMut};

use futures::task;

use futures::{StartSend, Async, AsyncSink, Poll};
use futures::task::Task;
use futures::stream::Stream;
use futures::sink::Sink;

use net::arrow::proto::error::ArrowError;
use net::arrow::proto::msg::ArrowMessage;

const INPUT_BUFFER_LIMIT: usize  = 32768;
const OUTPUT_BUFFER_LIMIT: usize = 4 * 1024 * 1024 * 1024;

/// Session context.
struct SessionContext {
    service_id:   u16,
    session_id:   u32,
    input:        BytesMut,
    output:       BytesMut,
    input_ready:  Option<Task>,
    input_empty:  Option<Task>,
    output_ready: Option<Task>,
    closed:       bool,
}

impl SessionContext {
    /// Create a new session context for a given service ID and session ID.
    fn new(service_id: u16, session_id: u32) -> SessionContext {
        SessionContext {
            service_id:   service_id,
            session_id:   session_id,
            input:        BytesMut::with_capacity(8192),
            output:       BytesMut::with_capacity(8192),
            input_ready:  None,
            input_empty:  None,
            output_ready: None,
            closed:       false,
        }
    }

    /// Extend the output buffer with data from a given Arrow Message. The
    /// method returns:
    /// * `()` on success
    /// * an error if the context has been closed or the output buffer hard
    ///   limit has been exceeded
    fn push_output_message(&mut self, msg: ArrowMessage) -> io::Result<()> {
        if self.closed {
            return Err(io::Error::new(io::ErrorKind::ConnectionReset, "connection has been closed"))
        }

        let data = msg.payload();

        // we cannot backpressure here, so we'll return an error if size
        // of the output buffer gets greater than a given hard limit
        if (self.output.len() + data.len()) > OUTPUT_BUFFER_LIMIT {
            return Err(io::Error::new(io::ErrorKind::Other, "output buffer limit exceeded"))
        }

        self.output.extend(data);

        // we MUST notify any possible task consuming the output buffer that
        // there is some data available again
        if self.output.len() > 0 {
            if let Some(task) = self.output_ready.take() {
                task.unpark();
            }
        }

        Ok(())
    }

    /// Take all the data from the input buffer and return them as an Arrow
    /// Message. The method returns:
    /// * `Async::Ready(Some(_))` if there was some data available
    /// * `Async::Ready(None)` if there was no data available and the context
    ///   has been closed
    /// * `Async::NotReady` if there was no data available
    fn take_input_message(&mut self) -> Poll<Option<ArrowMessage>, io::Error> {
        let data = self.input.take()
            .freeze();

        // we MUST notify any possible task feeding the input buffer that the
        // buffer is empty again
        if let Some(task) = self.input_empty.take() {
            task.unpark();
        }

        if data.len() > 0 {
            let message = ArrowMessage::new(
                self.service_id,
                self.session_id,
                data);

            Ok(Async::Ready(Some(message)))
        } else if self.closed {
            Ok(Async::Ready(None))
        } else {
            // park the current task and wait until there is some data
            // available in the input buffer
            self.input_ready = Some(task::park());

            Ok(Async::NotReady)
        }
    }

    /// Extend the input buffer with given data. The method returns:
    /// * `AsyncSink::NotReady(_)` with remaining data if the input buffer is
    ///   full
    /// * `AsyncSink::Ready` if all the given data has been inserted into the
    ///   input buffer
    /// * an error if the context has been closed
    fn push_input_data(&mut self, mut msg: Bytes) -> StartSend<Bytes, io::Error> {
        if self.closed {
            return Err(io::Error::new(io::ErrorKind::ConnectionReset, "connection has been closed"))
        }

        let mut take = msg.len();

        if (take + self.input.len()) > INPUT_BUFFER_LIMIT {
            take = INPUT_BUFFER_LIMIT - self.input.len();
        }

        self.input.extend(msg.split_to(take));

        // we MUST notify any possible task consuming the input buffer that
        // there is some data available again
        if self.input.len() > 0 {
            if let Some(task) = self.input_ready.take() {
                task.unpark();
            }
        }

        if msg.len() > 0 {
            // park the current task and wait until there is some space in
            // the input buffer again
            self.input_empty = Some(task::park());

            Ok(AsyncSink::NotReady(msg))
        } else {
            Ok(AsyncSink::Ready)
        }
    }

    /// Flush the input buffer. The method returns:
    /// * `Async::Ready(())` if the input buffer is empty
    /// * `Async::NotReady` if the buffer is not empty
    fn flush_input_buffer(&mut self) -> Poll<(), io::Error> {
        if self.input.len() > 0 {
            // park the current task and wait until the input buffer is empty
            self.input_empty = Some(task::park());

            Ok(Async::NotReady)
        } else {
            Ok(Async::Ready(()))
        }
    }

    /// Take data from the output buffer. The method returns:
    /// * `Async::Ready(Some(_))` if there is some data available
    /// * `Async::Ready(None)` if the context has been closed and there is
    ///   not data in the output buffer
    /// * `Async::NotReady` if there is no data available
    fn take_output_data(&mut self) -> Poll<Option<Bytes>, io::Error> {
        let data = self.output.take()
            .freeze();

        if data.len() > 0 {
            Ok(Async::Ready(Some(data)))
        } else if self.closed {
            Ok(Async::Ready(None))
        } else {
            // park the current task and wait until there is some data in
            // the output buffer available again
            self.output_ready = Some(task::park());

            Ok(Async::NotReady)
        }
    }

    /// Mark the context as closed. Note that this method does not flush any
    /// buffer.
    fn close(&mut self) {
        self.closed = true;
    }
}

/// Arrow session (i.e. connection to an external service).
pub struct Session {
    context: Rc<RefCell<SessionContext>>,
}

impl Session {
    /// Create a new session for a given service ID and session ID.
    pub fn new(service_id: u16, session_id: u32) -> Session {
        let context = SessionContext::new(service_id, session_id);

        Session {
            context: Rc::new(RefCell::new(context))
        }
    }

    /// Push a given Arrow Message into the output buffer. The method returns:
    /// * `()` on success
    /// * an error if the session has been closed or the output buffer hard
    ///   limit has been exceeded
    pub fn push(&mut self, msg: ArrowMessage) -> io::Result<()> {
        self.context.borrow_mut()
            .push_output_message(msg)
    }

    /// Take an Arrow Message from the input buffer. The method returns:
    /// * `Async::Ready(Some(_))` if there was some data available
    /// * `Async::Ready(None)` if there was no data available and the context
    ///   has been closed
    /// * `Async::NotReady` if there was no data available
    pub fn take(&mut self) -> Poll<Option<ArrowMessage>, io::Error> {
        self.context.borrow_mut()
            .take_input_message()
    }

    /// Mark the session as closed. The session context won't accept any new
    /// data, however the buffered data can be still processed. It's up to
    /// the corresponding tasks to consume all remaining data.
    pub fn close(&mut self) {
        self.context.borrow_mut()
            .close()
    }

    /// Get session transport.
    pub fn transport(&self) -> SessionTransport {
        SessionTransport {
            context: self.context.clone()
        }
    }
}

/// Session transport.
pub struct SessionTransport {
    context: Rc<RefCell<SessionContext>>,
}

impl Stream for SessionTransport {
    type Item  = Bytes;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Bytes>, io::Error> {
        self.context.borrow_mut()
            .take_output_data()
    }
}

impl Sink for SessionTransport {
    type SinkItem  = Bytes;
    type SinkError = io::Error;

    fn start_send(&mut self, data: Bytes) -> StartSend<Bytes, io::Error> {
        self.context.borrow_mut()
            .push_input_data(data)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        self.context.borrow_mut()
            .flush_input_buffer()
    }

    fn close(&mut self) -> Poll<(), io::Error> {
        let mut context = self.context.borrow_mut();

        // mark the context as closed
        context.close();

        // and wait until the input buffer is fully consumed
        context.flush_input_buffer()
    }
}

/// Arrow session manager.
pub struct SessionManager {
    sessions: HashMap<u32, Session>,
}

impl SessionManager {
    /// Create a new session manager.
    pub fn new() -> SessionManager {
        SessionManager {
            sessions: HashMap::new(),
        }
    }
}

impl Stream for SessionManager {
    type Item  = ArrowMessage;
    type Error = ArrowError;

    fn poll(&mut self) -> Poll<Option<ArrowMessage>, ArrowError> {
        // TODO
        Ok(Async::NotReady)
    }
}

impl Sink for SessionManager {
    type SinkItem  = ArrowMessage;
    type SinkError = ArrowError;

    fn start_send(&mut self, msg: ArrowMessage) -> StartSend<ArrowMessage, ArrowError> {
        // TODO
        Ok(AsyncSink::NotReady(msg))
    }

    fn poll_complete(&mut self) -> Poll<(), ArrowError> {
        // TODO
        Ok(Async::NotReady)
    }

    fn close(&mut self) -> Poll<(), ArrowError> {
        // TODO
        self.poll_complete()
    }
}
<commit_msg>Implementation of the message dispatch in the SessionManager<commit_after>// Copyright 2017 click2stream, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::io;

use std::rc::Rc;
use std::cell::RefCell;
use std::collections::HashMap;
use std::net::ToSocketAddrs;

use bytes::{Bytes, BytesMut};

use futures::task;

use futures::{Async, AsyncSink, Future, Poll, StartSend};
use futures::task::Task;
use futures::stream::Stream;
use futures::sink::Sink;

use tokio_core::net::TcpStream;
use tokio_core::reactor::Handle as TokioCoreHandle;

use tokio_io::AsyncRead;

use futures_ex::StreamEx;

use net::arrow::proto::codec::RawCodec;
use net::arrow::proto::error::ArrowError;
use net::arrow::proto::msg::ArrowMessage;

const INPUT_BUFFER_LIMIT: usize  = 32768;
const OUTPUT_BUFFER_LIMIT: usize = 4 * 1024 * 1024 * 1024;

/// Session context.
struct SessionContext {
    service_id:   u16,
    session_id:   u32,
    input:        BytesMut,
    output:       BytesMut,
    input_ready:  Option<Task>,
    input_empty:  Option<Task>,
    output_ready: Option<Task>,
    closed:       bool,
    error:        Option<io::Error>,
}

impl SessionContext {
    /// Create a new session context for a given service ID and session ID.
    fn new(service_id: u16, session_id: u32) -> SessionContext {
        SessionContext {
            service_id:   service_id,
            session_id:   session_id,
            input:        BytesMut::with_capacity(8192),
            output:       BytesMut::with_capacity(8192),
            input_ready:  None,
            input_empty:  None,
            output_ready: None,
            closed:       false,
            error:        None,
        }
    }

    /// Extend the output buffer with data from a given Arrow Message.
    fn push_output_message(&mut self, msg: ArrowMessage) {
        // ignore all incoming messages after the connection gets closed
        if self.closed {
            return
        }

        let data = msg.payload();

        if (self.output.len() + data.len()) > OUTPUT_BUFFER_LIMIT {
            // we cannot backpressure here, so we'll set an error state
            self.set_error(io::Error::new(io::ErrorKind::Other, "output buffer limit exceeded"));
        } else {
            self.output.extend(data);

            // we MUST notify any possible task consuming the output buffer that
            // there is some data available again
            if self.output.len() > 0 {
                if let Some(task) = self.output_ready.take() {
                    task.unpark();
                }
            }
        }
    }

    /// Take all the data from the input buffer and return them as an Arrow
    /// Message. The method returns:
    /// * `Async::Ready(Some(_))` if there was some data available
    /// * `Async::Ready(None)` if there was no data available and the context
    ///   has been closed
    /// * `Async::NotReady` if there was no data available
    fn take_input_message(&mut self) -> Poll<Option<ArrowMessage>, io::Error> {
        let data = self.input.take()
            .freeze();

        // we MUST notify any possible task feeding the input buffer that the
        // buffer is empty again
        if let Some(task) = self.input_empty.take() {
            task.unpark();
        }

        if data.len() > 0 {
            let message = ArrowMessage::new(
                self.service_id,
                self.session_id,
                data);

            Ok(Async::Ready(Some(message)))
        } else if self.closed {
            match self.error.take() {
                Some(err) => Err(err),
                None      => Ok(Async::Ready(None)),
            }
        } else {
            // park the current task and wait until there is some data
            // available in the input buffer
            self.input_ready = Some(task::park());

            Ok(Async::NotReady)
        }
    }

    /// Extend the input buffer with given data. The method returns:
    /// * `AsyncSink::NotReady(_)` with remaining data if the input buffer is
    ///   full
    /// * `AsyncSink::Ready` if all the given data has been inserted into the
    ///   input buffer
    /// * an error if the context has been closed
    fn push_input_data(&mut self, mut msg: Bytes) -> StartSend<Bytes, io::Error> {
        if self.closed {
            return Err(io::Error::new(io::ErrorKind::ConnectionReset, "connection has been closed"))
        }

        let mut take = msg.len();

        if (take + self.input.len()) > INPUT_BUFFER_LIMIT {
            take = INPUT_BUFFER_LIMIT - self.input.len();
        }

        self.input.extend(msg.split_to(take));

        // we MUST notify any possible task consuming the input buffer that
        // there is some data available again
        if self.input.len() > 0 {
            if let Some(task) = self.input_ready.take() {
                task.unpark();
            }
        }

        if msg.len() > 0 {
            // park the current task and wait until there is some space in
            // the input buffer again
            self.input_empty = Some(task::park());

            Ok(AsyncSink::NotReady(msg))
        } else {
            Ok(AsyncSink::Ready)
        }
    }

    /// Flush the input buffer. The method returns:
    /// * `Async::Ready(())` if the input buffer is empty
    /// * `Async::NotReady` if the buffer is not empty
    fn flush_input_buffer(&mut self) -> Poll<(), io::Error> {
        if self.input.len() > 0 {
            // park the current task and wait until the input buffer is empty
            self.input_empty = Some(task::park());

            Ok(Async::NotReady)
        } else {
            Ok(Async::Ready(()))
        }
    }

    /// Take data from the output buffer. The method returns:
    /// * `Async::Ready(Some(_))` if there is some data available
    /// * `Async::Ready(None)` if the context has been closed and there is
    ///   not data in the output buffer
    /// * `Async::NotReady` if there is no data available
    fn take_output_data(&mut self) -> Poll<Option<Bytes>, io::Error> {
        let data = self.output.take()
            .freeze();

        if data.len() > 0 {
            Ok(Async::Ready(Some(data)))
        } else if self.closed {
            Ok(Async::Ready(None))
        } else {
            // park the current task and wait until there is some data in
            // the output buffer available again
            self.output_ready = Some(task::park());

            Ok(Async::NotReady)
        }
    }

    /// Mark the context as closed. Note that this method does not flush any
    /// buffer.
    fn close(&mut self) {
        self.closed = true;
    }

    /// Mark the context as closed and set a given error. Note that this
    /// method does not flush any buffer.
    fn set_error(&mut self, err: io::Error) {
        // ignore all errors after the connection gets closed
        if !self.closed {
            self.closed = true;
            self.error  = Some(err);
        }
    }
}

/// Arrow session (i.e. connection to an external service).
struct Session {
    context: Rc<RefCell<SessionContext>>,
}

impl Session {
    /// Create a new session for a given service ID and session ID.
    fn new(service_id: u16, session_id: u32) -> Session {
        let context = SessionContext::new(service_id, session_id);

        Session {
            context: Rc::new(RefCell::new(context))
        }
    }

    /// Push a given Arrow Message into the output buffer.
    fn push(&mut self, msg: ArrowMessage) {
        self.context.borrow_mut()
            .push_output_message(msg)
    }

    /// Take an Arrow Message from the input buffer. The method returns:
    /// * `Async::Ready(Some(_))` if there was some data available
    /// * `Async::Ready(None)` if there was no data available and the context
    ///   has been closed
    /// * `Async::NotReady` if there was no data available
    fn take(&mut self) -> Poll<Option<ArrowMessage>, io::Error> {
        self.context.borrow_mut()
            .take_input_message()
    }

    /// Mark the session as closed. The session context won't accept any new
    /// data, however the buffered data can be still processed. It's up to
    /// the corresponding tasks to consume all remaining data.
    fn close(&mut self) {
        self.context.borrow_mut()
            .close()
    }

    /// Get session transport.
    fn transport(&self) -> SessionTransport {
        SessionTransport {
            context: self.context.clone()
        }
    }

    /// Get session error handler.
    fn error_handler(&self) -> SessionErrorHandler {
        SessionErrorHandler {
            context: self.context.clone()
        }
    }
}

/// Session transport.
struct SessionTransport {
    context: Rc<RefCell<SessionContext>>,
}

impl Stream for SessionTransport {
    type Item  = Bytes;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Bytes>, io::Error> {
        self.context.borrow_mut()
            .take_output_data()
    }
}

impl Sink for SessionTransport {
    type SinkItem  = Bytes;
    type SinkError = io::Error;

    fn start_send(&mut self, data: Bytes) -> StartSend<Bytes, io::Error> {
        self.context.borrow_mut()
            .push_input_data(data)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        self.context.borrow_mut()
            .flush_input_buffer()
    }

    fn close(&mut self) -> Poll<(), io::Error> {
        let mut context = self.context.borrow_mut();

        // mark the context as closed
        context.close();

        // and wait until the input buffer is fully consumed
        context.flush_input_buffer()
    }
}

/// Session error handler.
struct SessionErrorHandler {
    context: Rc<RefCell<SessionContext>>,
}

impl SessionErrorHandler {
    /// Save a given transport error into the session context.
    fn set_error(&mut self, err: io::Error) {
        self.context.borrow_mut()
            .set_error(err)
    }
}

/// Arrow session manager.
pub struct SessionManager {
    tc_handle: TokioCoreHandle,
    sessions:  HashMap<u32, Session>,
}

impl SessionManager {
    /// Create a new session manager.
    pub fn new(tc_handle: TokioCoreHandle) -> SessionManager {
        SessionManager {
            tc_handle: tc_handle,
            sessions:  HashMap::new(),
        }
    }

    /// Send a given Arrow Message to the corresponding service using a given
    /// session (as specified by the message). The method returns an error
    /// if the session could not be created for some reason.
    pub fn send(&mut self, msg: ArrowMessage) -> Result<(), ArrowError> {
        let header = *msg.header();

        self.get_session_mut(header.service, header.session)?
            .push(msg);

        Ok(())
    }

    /// Get mutable reference to a given session.
    fn get_session_mut(
        &mut self,
        service_id: u16,
        session_id: u32) -> Result<&mut Session, ArrowError> {
        if !self.sessions.contains_key(&session_id) {
            let session = self.connect(service_id, session_id)?;

            self.sessions.insert(
                session_id,
                session);
        }

        let session = self.sessions.get_mut(&session_id);

        Ok(session.unwrap())
    }

    /// Connect to a given service and create an associated session object
    /// with a given ID.
    fn connect(
        &mut self,
        service_id: u16,
        session_id: u32) -> Result<Session, ArrowError> {
        // TODO: get address of a given service
        let addr = "127.0.0.1:80";
        let addr = addr.to_socket_addrs()?
            .next()
            .ok_or(io::Error::new(io::ErrorKind::Other, "unable to resolve a given address"))?;

        let session = Session::new(service_id, session_id);
        let transport = session.transport();
        let mut err_handler = session.error_handler();

        let client = TcpStream::connect(&addr, &self.tc_handle)
            .and_then(|stream| {
                let framed = stream.framed(RawCodec);
                let (sink, stream) = framed.split();

                let messages = stream.pipe(transport);

                sink.send_all(messages)
            })
            .then(move |res| {
                if let Err(err) = res {
                    err_handler.set_error(err);
                }

                Ok(())
            });

        self.tc_handle.spawn(client);

        Ok(session)
    }
}

impl Stream for SessionManager {
    type Item  = ArrowMessage;
    type Error = ArrowError;

    fn poll(&mut self) -> Poll<Option<ArrowMessage>, ArrowError> {
        // TODO
        Ok(Async::NotReady)
    }
}
<|endoftext|>


<commit_before>use light::Light;
use raytracer::compositor::{ColorRGBA, Surface, SurfaceFactory};
use raytracer::{Intersection, Ray};
use scene::{Camera, Scene};
use std::iter::range;
use std::num::Float;
use std::ops::Deref;
use std::rand::{thread_rng, Rng, SeedableRng, Isaac64Rng};
use std::sync::Arc;
use std::sync::mpsc::channel;
use std::sync::TaskPool;
use vec3::Vec3;

pub static EPSILON: f64 = ::std::f64::EPSILON * 10000.0;

#[derive(Clone)]
pub struct Renderer {
    pub reflect_depth: uint,  // Maximum reflection recursions.
    pub refract_depth: uint,  // Maximum refraction recursions. A sphere takes up 2 recursions.
    pub shadow_samples: uint, // Number of samples for soft shadows and area lights.
    pub pixel_samples: uint,  // The square of this is the number of samples per pixel.
    pub tasks: uint           // Minimum number of tasks to spawn.
}

impl Renderer {
    pub fn render(&self, camera: Camera, shared_scene: Arc<Scene>) -> Surface {

        let mut surface = Surface::new(camera.image_width as uint,
                                       camera.image_height as uint,
                                       ColorRGBA::new_rgb(0, 0, 0));

        let pool = TaskPool::new(self.tasks);

        let (tx, rx) = channel();

        let mut jobs = 0;

        for subsurface_factory in surface.divide(128, 8) {
            jobs += 1;

            let renderer = self.clone();
            let child_tx = tx.clone();
            let scene_local = shared_scene.clone();
            let camera_local = camera.clone();

            pool.execute(move || {
                let _ = child_tx.send(renderer.render_tile(camera_local.clone(),
                    scene_local.deref(), subsurface_factory));
            });
        }

        let start_time = ::time::get_time();

        for (i, subsurface) in rx.iter().take(jobs).enumerate() {
            surface.merge(subsurface);
            ::util::print_progress("Tile", start_time.clone(), (i + 1) as uint, jobs);
        }
        surface
    }

    fn render_tile(&self, camera: Camera, scene: &Scene,
                   tile_factory: SurfaceFactory) -> Box<Surface> {

        let shadow_samples = self.shadow_samples;
        let pixel_samples = self.pixel_samples;
        let reflect_depth = self.reflect_depth;
        let refract_depth = self.refract_depth;

        let mut tile = tile_factory.create();

        let random_data: Vec<u64> = range(0u, 64u).map(|_| {
            thread_rng().next_u64()
        }).collect();
        let mut rng: Isaac64Rng = SeedableRng::from_seed(random_data.as_slice());

        for rel_y in range(0u, tile.height) {
            let abs_y = (camera.image_height as uint) - (tile.y_off + rel_y) - 1;
            for rel_x in range(0u, tile.width) {
                let abs_x = tile.x_off + rel_x;

                // Supersampling, jitter algorithm
                let pixel_width = 1.0 / pixel_samples as f64;
                let mut color = Vec3::zero();

                for y_subpixel in range(0, pixel_samples) {
                    for x_subpixel in range(0, pixel_samples) {
                        // Don't jitter if not antialiasing
                        let (j_x, j_y) = if pixel_samples > 1 {
                            (x_subpixel as f64 * pixel_width + rng.gen::<f64>() * pixel_width,
                             y_subpixel as f64 * pixel_width + rng.gen::<f64>() * pixel_width)
                        } else {
                            (0.0, 0.0)
                        };

                        let ray = camera.get_ray(abs_x as f64 + j_x, abs_y as f64 + j_y);
                        let result = Renderer::trace(scene, &ray, shadow_samples,
                                                     reflect_depth, refract_depth, false);
                        // Clamp subpixels for now to avoid intense aliasing when combined value is clamped later
                        // Should think of a better way to handle this
                        color = color + result.clamp(0.0, 1.0).scale(1.0 / (pixel_samples * pixel_samples) as f64);
                    }
                }
                tile[(rel_x, rel_y)] = ColorRGBA::new_rgb_clamped(color.x, color.y, color.z);
            }
        }

        box tile
    }

    fn trace(scene: &Scene, ray: &Ray, shadow_samples: uint,
             reflect_depth: uint, refract_depth: uint, inside: bool) -> Vec3 {

        if reflect_depth <= 0 || refract_depth <= 0 { return Vec3::zero() }

        match ray.get_nearest_hit(scene) {
            Some(hit) => {
                let n = hit.n.unit();
                let i = (-ray.direction).unit();

                // Local lighting computation: surface shading, shadows
                let mut result = scene.lights.iter().fold(Vec3::zero(), |color_acc, light| {
                    let shadow = Renderer::shadow_intensity(scene, &hit, light, shadow_samples);
                    let l = (light.center() - hit.position).unit();

                    color_acc + light.color() * hit.material.sample(n, i, l, hit.u, hit.v) * shadow
                });

                if hit.material.is_reflective() || hit.material.is_refractive() {
                    let reflect_fresnel = Renderer::fresnel_reflect(hit.material.ior(), &i, &n, inside);
                    let mut refract_fresnel = 1.0 - reflect_fresnel;

                    // Global reflection
                    if hit.material.is_reflective() {
                        let r = Vec3::reflect(&i, &n);
                        let reflect_ray = Ray::new(hit.position, r);
                        let reflection = Renderer::trace(scene, &reflect_ray, shadow_samples,
                                                         reflect_depth - 1, refract_depth, inside);

                        result = result + hit.material.global_specular(&reflection).scale(reflect_fresnel);
                    }

                    // Global refraction
                    if hit.material.is_refractive() {
                        let t = match Vec3::refract(&i, &n, hit.material.ior(), inside) {
                            Some(ref t) => *t,
                            None => {
                                refract_fresnel = 1.0; // Total internal reflection (TODO: verify)
                                Vec3::reflect(&i, &n)
                            }
                        };

                        let refract_ray = Ray::new(hit.position + t.scale(EPSILON), t);
                        let refraction = Renderer::trace(scene, &refract_ray, shadow_samples,
                                                         reflect_depth, refract_depth - 1, !inside);

                        result = result + hit.material.global_transmissive(&refraction).scale(refract_fresnel);
                    }
                }

                result
            },
            None => {
                match scene.skybox {
                    Some(ref skybox) => skybox.color(ray.direction),
                    None => scene.background
                }
            }
        }
    }

    fn shadow_intensity(scene: &Scene, hit: &Intersection,
                        light: &Box<Light+Send+Sync>, shadow_samples: uint) -> Vec3 {

        if shadow_samples <= 0 { return Vec3::one() }

        // Point light speedup (no point in sampling a point light multiple times)
        let shadow_sample_tries = if light.is_point() { 1 } else { shadow_samples };
        let mut shadow = Vec3::zero();

        // Take average shadow color after jittering/sampling light position
        for _ in range(0, shadow_sample_tries) {
            // L has to be a unit vector for t_max 1:1 correspondence to
            // distance to light to work. Shadow feelers only search up
            // until light source.
            let sampled_light_position = light.position();
            let shadow_l = (sampled_light_position - hit.position).unit();
            let shadow_ray = Ray::new(hit.position, shadow_l);
            let distance_to_light = (sampled_light_position - hit.position).len();

            // Check against candidate primitives in scene for occlusion
            // and multiply shadow color by occluders' shadow colors
            let candidate_nodes = scene.octree.get_intersected_objects(&shadow_ray);

            shadow = shadow + candidate_nodes.fold(Vec3::one(), |shadow_acc, prim| {
                let occlusion = prim.intersects(&shadow_ray, EPSILON, distance_to_light);
                match occlusion {
                    Some(occlusion) => shadow_acc * occlusion.material.transmission(),
                    None => shadow_acc
                }
            });
        }

        shadow.scale(1.0 / shadow_sample_tries as f64)
    }

    /// Calculates the fresnel (reflectivity) given the index of refraction and the cos_angle
    /// This uses Schlick's approximation. cos_angle is normal_dot_incoming
    /// http://graphics.stanford.edu/courses/cs148-10-summer/docs/2006--degreve--reflection_refraction.pdf
    fn fresnel_reflect(ior: f64, i: &Vec3, n: &Vec3, inside: bool) -> f64 {
        let (n1, n2) = if inside { (ior, 1.0) } else { (1.0, ior) };
        let actual_n = if inside { -*n } else { *n };

        let r0_sqrt = (n1 - n2) / (n1 + n2);
        let r0 = r0_sqrt * r0_sqrt;

        let cos_angle = if n1 <= n2 {
            i.dot(&actual_n)
        } else {
            let t = match Vec3::refract(i, &-actual_n, ior, inside) {
                Some(x) => x,
                None => return 1.0 // n1 > n2 && TIR
            };

            -actual_n.dot(&t) // n1 > n2 && !TIR
        };

        let cos_term = 1.0 - cos_angle;

        (r0 + ((1.0 - r0) * cos_term * cos_term * cos_term * cos_term * cos_term)).max(0.0).min(1.0)
    }
}
<commit_msg>Add a basic smoke test for Renderer::render<commit_after>use light::Light;
use raytracer::compositor::{ColorRGBA, Surface, SurfaceFactory};
use raytracer::{Intersection, Ray};
use scene::{Camera, Scene};
use std::iter::range;
use std::num::Float;
use std::ops::Deref;
use std::rand::{thread_rng, Rng, SeedableRng, Isaac64Rng};
use std::sync::Arc;
use std::sync::mpsc::channel;
use std::sync::TaskPool;
use vec3::Vec3;

pub static EPSILON: f64 = ::std::f64::EPSILON * 10000.0;

#[derive(Clone)]
pub struct Renderer {
    pub reflect_depth: uint,  // Maximum reflection recursions.
    pub refract_depth: uint,  // Maximum refraction recursions. A sphere takes up 2 recursions.
    pub shadow_samples: uint, // Number of samples for soft shadows and area lights.
    pub pixel_samples: uint,  // The square of this is the number of samples per pixel.
    pub tasks: uint           // Minimum number of tasks to spawn.
}

impl Renderer {
    pub fn render(&self, camera: Camera, shared_scene: Arc<Scene>) -> Surface {

        let mut surface = Surface::new(camera.image_width as uint,
                                       camera.image_height as uint,
                                       ColorRGBA::new_rgb(0, 0, 0));

        let pool = TaskPool::new(self.tasks);

        let (tx, rx) = channel();

        let mut jobs = 0;

        for subsurface_factory in surface.divide(128, 8) {
            jobs += 1;

            let renderer = self.clone();
            let child_tx = tx.clone();
            let scene_local = shared_scene.clone();
            let camera_local = camera.clone();

            pool.execute(move || {
                let _ = child_tx.send(renderer.render_tile(camera_local.clone(),
                    scene_local.deref(), subsurface_factory));
            });
        }

        let start_time = ::time::get_time();

        for (i, subsurface) in rx.iter().take(jobs).enumerate() {
            surface.merge(subsurface);
            ::util::print_progress("Tile", start_time.clone(), (i + 1) as uint, jobs);
        }
        surface
    }

    fn render_tile(&self, camera: Camera, scene: &Scene,
                   tile_factory: SurfaceFactory) -> Box<Surface> {

        let shadow_samples = self.shadow_samples;
        let pixel_samples = self.pixel_samples;
        let reflect_depth = self.reflect_depth;
        let refract_depth = self.refract_depth;

        let mut tile = tile_factory.create();

        let random_data: Vec<u64> = range(0u, 64u).map(|_| {
            thread_rng().next_u64()
        }).collect();
        let mut rng: Isaac64Rng = SeedableRng::from_seed(random_data.as_slice());

        for rel_y in range(0u, tile.height) {
            let abs_y = (camera.image_height as uint) - (tile.y_off + rel_y) - 1;
            for rel_x in range(0u, tile.width) {
                let abs_x = tile.x_off + rel_x;

                // Supersampling, jitter algorithm
                let pixel_width = 1.0 / pixel_samples as f64;
                let mut color = Vec3::zero();

                for y_subpixel in range(0, pixel_samples) {
                    for x_subpixel in range(0, pixel_samples) {
                        // Don't jitter if not antialiasing
                        let (j_x, j_y) = if pixel_samples > 1 {
                            (x_subpixel as f64 * pixel_width + rng.gen::<f64>() * pixel_width,
                             y_subpixel as f64 * pixel_width + rng.gen::<f64>() * pixel_width)
                        } else {
                            (0.0, 0.0)
                        };

                        let ray = camera.get_ray(abs_x as f64 + j_x, abs_y as f64 + j_y);
                        let result = Renderer::trace(scene, &ray, shadow_samples,
                                                     reflect_depth, refract_depth, false);
                        // Clamp subpixels for now to avoid intense aliasing when combined value is clamped later
                        // Should think of a better way to handle this
                        color = color + result.clamp(0.0, 1.0).scale(1.0 / (pixel_samples * pixel_samples) as f64);
                    }
                }
                tile[(rel_x, rel_y)] = ColorRGBA::new_rgb_clamped(color.x, color.y, color.z);
            }
        }

        box tile
    }

    fn trace(scene: &Scene, ray: &Ray, shadow_samples: uint,
             reflect_depth: uint, refract_depth: uint, inside: bool) -> Vec3 {

        if reflect_depth <= 0 || refract_depth <= 0 { return Vec3::zero() }

        match ray.get_nearest_hit(scene) {
            Some(hit) => {
                let n = hit.n.unit();
                let i = (-ray.direction).unit();

                // Local lighting computation: surface shading, shadows
                let mut result = scene.lights.iter().fold(Vec3::zero(), |color_acc, light| {
                    let shadow = Renderer::shadow_intensity(scene, &hit, light, shadow_samples);
                    let l = (light.center() - hit.position).unit();

                    color_acc + light.color() * hit.material.sample(n, i, l, hit.u, hit.v) * shadow
                });

                if hit.material.is_reflective() || hit.material.is_refractive() {
                    let reflect_fresnel = Renderer::fresnel_reflect(hit.material.ior(), &i, &n, inside);
                    let mut refract_fresnel = 1.0 - reflect_fresnel;

                    // Global reflection
                    if hit.material.is_reflective() {
                        let r = Vec3::reflect(&i, &n);
                        let reflect_ray = Ray::new(hit.position, r);
                        let reflection = Renderer::trace(scene, &reflect_ray, shadow_samples,
                                                         reflect_depth - 1, refract_depth, inside);

                        result = result + hit.material.global_specular(&reflection).scale(reflect_fresnel);
                    }

                    // Global refraction
                    if hit.material.is_refractive() {
                        let t = match Vec3::refract(&i, &n, hit.material.ior(), inside) {
                            Some(ref t) => *t,
                            None => {
                                refract_fresnel = 1.0; // Total internal reflection (TODO: verify)
                                Vec3::reflect(&i, &n)
                            }
                        };

                        let refract_ray = Ray::new(hit.position + t.scale(EPSILON), t);
                        let refraction = Renderer::trace(scene, &refract_ray, shadow_samples,
                                                         reflect_depth, refract_depth - 1, !inside);

                        result = result + hit.material.global_transmissive(&refraction).scale(refract_fresnel);
                    }
                }

                result
            },
            None => {
                match scene.skybox {
                    Some(ref skybox) => skybox.color(ray.direction),
                    None => scene.background
                }
            }
        }
    }

    fn shadow_intensity(scene: &Scene, hit: &Intersection,
                        light: &Box<Light+Send+Sync>, shadow_samples: uint) -> Vec3 {

        if shadow_samples <= 0 { return Vec3::one() }

        // Point light speedup (no point in sampling a point light multiple times)
        let shadow_sample_tries = if light.is_point() { 1 } else { shadow_samples };
        let mut shadow = Vec3::zero();

        // Take average shadow color after jittering/sampling light position
        for _ in range(0, shadow_sample_tries) {
            // L has to be a unit vector for t_max 1:1 correspondence to
            // distance to light to work. Shadow feelers only search up
            // until light source.
            let sampled_light_position = light.position();
            let shadow_l = (sampled_light_position - hit.position).unit();
            let shadow_ray = Ray::new(hit.position, shadow_l);
            let distance_to_light = (sampled_light_position - hit.position).len();

            // Check against candidate primitives in scene for occlusion
            // and multiply shadow color by occluders' shadow colors
            let candidate_nodes = scene.octree.get_intersected_objects(&shadow_ray);

            shadow = shadow + candidate_nodes.fold(Vec3::one(), |shadow_acc, prim| {
                let occlusion = prim.intersects(&shadow_ray, EPSILON, distance_to_light);
                match occlusion {
                    Some(occlusion) => shadow_acc * occlusion.material.transmission(),
                    None => shadow_acc
                }
            });
        }

        shadow.scale(1.0 / shadow_sample_tries as f64)
    }

    /// Calculates the fresnel (reflectivity) given the index of refraction and the cos_angle
    /// This uses Schlick's approximation. cos_angle is normal_dot_incoming
    /// http://graphics.stanford.edu/courses/cs148-10-summer/docs/2006--degreve--reflection_refraction.pdf
    fn fresnel_reflect(ior: f64, i: &Vec3, n: &Vec3, inside: bool) -> f64 {
        let (n1, n2) = if inside { (ior, 1.0) } else { (1.0, ior) };
        let actual_n = if inside { -*n } else { *n };

        let r0_sqrt = (n1 - n2) / (n1 + n2);
        let r0 = r0_sqrt * r0_sqrt;

        let cos_angle = if n1 <= n2 {
            i.dot(&actual_n)
        } else {
            let t = match Vec3::refract(i, &-actual_n, ior, inside) {
                Some(x) => x,
                None => return 1.0 // n1 > n2 && TIR
            };

            -actual_n.dot(&t) // n1 > n2 && !TIR
        };

        let cos_term = 1.0 - cos_angle;

        (r0 + ((1.0 - r0) * cos_term * cos_term * cos_term * cos_term * cos_term)).max(0.0).min(1.0)
    }
}

#[test]
fn it_renders_the_background_of_an_empty_scene() {
    let camera = Camera::new(
        Vec3 { x: 0.0, y: 0.0, z: 0.0 },
        Vec3 { x: 0.0, y: 1.0, z: 0.0 },
        Vec3 { x: 0.0, y: 0.0, z: 1.0 },
        45.0,
        32,
        32
    );

    let test_scene = Scene {
        lights: vec!(),
        octree: ::raytracer::Octree::new_from_prims(vec!()),
        background: Vec3 { x: 1.0, y: 0.0, z: 0.0 },
        skybox: None
    };

    let shared_scene = Arc::new(test_scene);

    let renderer = Renderer {
        reflect_depth: 1,
        refract_depth: 1,
        shadow_samples: 1,
        pixel_samples: 1,
        tasks: 2
    };

    let image_data = renderer.render(camera, shared_scene);

    for color in image_data.buffer.iter() {
        assert_eq!(color.r, 255);
        assert_eq!(color.g, 0);
        assert_eq!(color.b, 0);
    }
}
<|endoftext|>


<commit_before>// error-pattern:quux
fn main() {
    let x: int = do { fail "quux" } while (true);
}
<commit_msg>XFAIL do-while-body-fails<commit_after>// xfail-stage1
// xfail-stage2
// xfail-stage3
// error-pattern:quux
fn main() {
    let x: int = do { fail "quux" } while (true);
}
<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::RData,
};

use crate::{
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

async fn acl_lookup<Local, Remote>(
    context: SharedContext,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let qname_in_proxy_list = context.check_query_in_proxy_list(query);
    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match qname_in_proxy_list {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (qname): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (qname): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    // FIXME: spawn(remote_response_fut)
    let local_response = local_response_fut.await;
    for rec in local_response.unwrap_or_else(|_| Message::new()).answers() {
        if rec.record_type() != query.query_type() {
            warn!("local DNS response has inconsistent answer type {} for query {}", rec.record_type(), query);
            break
        }
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::PTR(_) => panic!("PTR records should not reach here"),
            _ => context.is_default_in_proxy_list(),
        };
        if !forward {
            debug!("pick local response (response): {:?}", local_response);
            return (local_response, false);
        }
    }

    let remote_response = remote_response_fut.await;
    debug!("pick remote response (response): {:?}", remote_response);
    (remote_response, true)
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    #[cfg(target_os = "android")]
    let local_upstream = Arc::new(upstream::UnixSocketUpstream {
        path: config.local_dns_path.clone().expect("local query DNS path"),
    });
    #[cfg(not(target_os = "android"))]
    let local_upstream = Arc::new(upstream::UdpUpstream {
        server: config.local_dns_addr.clone().expect("local query DNS address"),
    });
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let local_upstream = Arc::clone(&local_upstream);
        let remote_upstream = Arc::clone(&remote_upstream);
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if request.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.clone(), local_upstream, remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<commit_msg>Refine code style<commit_after>use std::{
    io,
    net::{IpAddr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::RData,
};

use crate::{
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

async fn acl_lookup<Local, Remote>(
    context: SharedContext,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match context.check_query_in_proxy_list(query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (qname): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (qname): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    // FIXME: spawn(remote_response_fut)
    let local_response = local_response_fut.await;
    for rec in local_response.unwrap_or_else(|_| Message::new()).answers() {
        if rec.record_type() != query.query_type() {
            warn!("local DNS response has inconsistent answer type {} for query {}", rec.record_type(), query);
            break
        }
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::PTR(_) => panic!("PTR records should not reach here"),
            _ => context.is_default_in_proxy_list(),
        };
        if !forward {
            debug!("pick local response (response): {:?}", local_response);
            return (local_response, false);
        }
    }

    let remote_response = remote_response_fut.await;
    debug!("pick remote response (response): {:?}", remote_response);
    (remote_response, true)
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    #[cfg(target_os = "android")]
    let local_upstream = Arc::new(upstream::UnixSocketUpstream {
        path: config.local_dns_path.clone().expect("local query DNS path"),
    });
    #[cfg(not(target_os = "android"))]
    let local_upstream = Arc::new(upstream::UdpUpstream {
        server: config.local_dns_addr.clone().expect("local query DNS address"),
    });
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let local_upstream = Arc::clone(&local_upstream);
        let remote_upstream = Arc::clone(&remote_upstream);
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if request.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.clone(), local_upstream, remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>use crate::config::PgpPublicKey;
use crate::dist::dist::TargetTriple;
use crate::dist::manifest::Component;
use crate::dist::temp;
use crate::errors::*;
use crate::utils::notify::NotificationLevel;
use std::fmt::{self, Display};
use std::path::Path;

#[derive(Debug)]
pub enum Notification<'a> {
    Utils(crate::utils::Notification<'a>),
    Temp(temp::Notification<'a>),

    Extracting(&'a Path, &'a Path),
    ComponentAlreadyInstalled(&'a str),
    CantReadUpdateHash(&'a Path),
    NoUpdateHash(&'a Path),
    ChecksumValid(&'a str),
    SignatureValid(&'a str, &'a PgpPublicKey),
    FileAlreadyDownloaded,
    CachedFileChecksumFailed,
    RollingBack,
    ExtensionNotInstalled(&'a str),
    NonFatalError(&'a Error),
    MissingInstalledComponent(&'a str),
    DownloadingComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
    InstallingComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
    RemovingComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
    RemovingOldComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
    DownloadingManifest(&'a str),
    DownloadedManifest(&'a str, Option<&'a str>),
    DownloadingLegacyManifest,
    SkippingNightlyMissingComponent(&'a [Component]),
    ForcingUnavailableComponent(&'a str),
    ManifestChecksumFailedHack,
    ComponentUnavailable(&'a str, Option<&'a TargetTriple>),
    StrayHash(&'a Path),
    SignatureInvalid(&'a str),
}

impl<'a> From<crate::utils::Notification<'a>> for Notification<'a> {
    fn from(n: crate::utils::Notification<'a>) -> Self {
        Notification::Utils(n)
    }
}

impl<'a> From<temp::Notification<'a>> for Notification<'a> {
    fn from(n: temp::Notification<'a>) -> Self {
        Notification::Temp(n)
    }
}

impl<'a> Notification<'a> {
    pub fn level(&self) -> NotificationLevel {
        use self::Notification::*;
        match self {
            Temp(n) => n.level(),
            Utils(n) => n.level(),
            ChecksumValid(_)
            | SignatureValid(_, _)
            | NoUpdateHash(_)
            | FileAlreadyDownloaded
            | DownloadingLegacyManifest => NotificationLevel::Verbose,
            Extracting(_, _)
            | DownloadingComponent(_, _, _)
            | InstallingComponent(_, _, _)
            | RemovingComponent(_, _, _)
            | RemovingOldComponent(_, _, _)
            | ComponentAlreadyInstalled(_)
            | ManifestChecksumFailedHack
            | RollingBack
            | DownloadingManifest(_)
            | SkippingNightlyMissingComponent(_)
            | DownloadedManifest(_, _) => NotificationLevel::Info,
            CantReadUpdateHash(_)
            | ExtensionNotInstalled(_)
            | MissingInstalledComponent(_)
            | CachedFileChecksumFailed
            | ComponentUnavailable(_, _)
            | ForcingUnavailableComponent(_)
            | StrayHash(_) => NotificationLevel::Warn,
            NonFatalError(_) => NotificationLevel::Error,
            SignatureInvalid(_) => NotificationLevel::Warn,
        }
    }
}

impl<'a> Display for Notification<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::result::Result<(), fmt::Error> {
        use self::Notification::*;
        match self {
            Temp(n) => n.fmt(f),
            Utils(n) => n.fmt(f),
            Extracting(_, _) => write!(f, "extracting..."),
            ComponentAlreadyInstalled(c) => write!(f, "component {} is up to date", c),
            CantReadUpdateHash(path) => write!(
                f,
                "can't read update hash file: '{}', can't skip update...",
                path.display()
            ),
            NoUpdateHash(path) => write!(f, "no update hash at: '{}'", path.display()),
            ChecksumValid(_) => write!(f, "checksum passed"),
            SignatureValid(url, key) => (|| {
                writeln!(f, "Good signature from on {} from:", url)?;
                for l in key.show_key().map_err(|_| std::fmt::Error)?.iter() {
                    writeln!(f, "{}", l)?;
                }
                Ok(())
            })(),
            FileAlreadyDownloaded => write!(f, "reusing previously downloaded file"),
            CachedFileChecksumFailed => write!(f, "bad checksum for cached download"),
            RollingBack => write!(f, "rolling back changes"),
            ExtensionNotInstalled(c) => write!(f, "extension '{}' was not installed", c),
            NonFatalError(e) => write!(f, "{}", e),
            MissingInstalledComponent(c) => {
                write!(f, "during uninstall component {} was not found", c)
            }
            DownloadingComponent(c, h, t) => {
                if Some(h) == t.as_ref() || t.is_none() {
                    write!(f, "downloading component '{}'", c)
                } else {
                    write!(f, "downloading component '{}' for '{}'", c, t.unwrap())
                }
            }
            InstallingComponent(c, h, t) => {
                if Some(h) == t.as_ref() || t.is_none() {
                    write!(f, "installing component '{}'", c)
                } else {
                    write!(f, "installing component '{}' for '{}'", c, t.unwrap())
                }
            }
            RemovingComponent(c, h, t) => {
                if Some(h) == t.as_ref() || t.is_none() {
                    write!(f, "removing component '{}'", c)
                } else {
                    write!(f, "removing component '{}' for '{}'", c, t.unwrap())
                }
            }
            RemovingOldComponent(c, h, t) => {
                if Some(h) == t.as_ref() || t.is_none() {
                    write!(f, "removing previous version of component '{}'", c)
                } else {
                    write!(
                        f,
                        "removing previous version of component '{}' for '{}'",
                        c,
                        t.unwrap()
                    )
                }
            }
            DownloadingManifest(t) => write!(f, "syncing channel updates for '{}'", t),
            DownloadedManifest(date, Some(version)) => {
                write!(f, "latest update on {}, rust version {}", date, version)
            }
            DownloadedManifest(date, None) => {
                write!(f, "latest update on {}, no rust version", date)
            }
            DownloadingLegacyManifest => write!(f, "manifest not found. trying legacy manifest"),
            ManifestChecksumFailedHack => {
                write!(f, "update not yet available, sorry! try again later")
            }
            ComponentUnavailable(pkg, toolchain) => {
                if let Some(tc) = toolchain {
                    write!(f, "component '{}' is not available on target '{}'", pkg, tc)
                } else {
                    write!(f, "component '{}' is not available", pkg)
                }
            }
            StrayHash(path) => write!(
                f,
                "removing stray hash found at '{}' in order to continue",
                path.display()
            ),
            SkippingNightlyMissingComponent(components) => write!(
                f,
                "skipping nightly which is missing installed component '{}'",
                components[0].short_name_in_manifest()
            ),
            ForcingUnavailableComponent(component) => {
                write!(f, "Force-skipping unavailable component '{}'", component)
            }
            SignatureInvalid(url) => write!(f, "Signature verification failed for '{}'", url),
        }
    }
}
<commit_msg>notifications: Add RetryingDownload notification<commit_after>use crate::config::PgpPublicKey;
use crate::dist::dist::TargetTriple;
use crate::dist::manifest::Component;
use crate::dist::temp;
use crate::errors::*;
use crate::utils::notify::NotificationLevel;
use std::fmt::{self, Display};
use std::path::Path;

#[derive(Debug)]
pub enum Notification<'a> {
    Utils(crate::utils::Notification<'a>),
    Temp(temp::Notification<'a>),

    Extracting(&'a Path, &'a Path),
    ComponentAlreadyInstalled(&'a str),
    CantReadUpdateHash(&'a Path),
    NoUpdateHash(&'a Path),
    ChecksumValid(&'a str),
    SignatureValid(&'a str, &'a PgpPublicKey),
    FileAlreadyDownloaded,
    CachedFileChecksumFailed,
    RollingBack,
    ExtensionNotInstalled(&'a str),
    NonFatalError(&'a Error),
    MissingInstalledComponent(&'a str),
    DownloadingComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
    InstallingComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
    RemovingComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
    RemovingOldComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
    DownloadingManifest(&'a str),
    DownloadedManifest(&'a str, Option<&'a str>),
    DownloadingLegacyManifest,
    SkippingNightlyMissingComponent(&'a [Component]),
    ForcingUnavailableComponent(&'a str),
    ManifestChecksumFailedHack,
    ComponentUnavailable(&'a str, Option<&'a TargetTriple>),
    StrayHash(&'a Path),
    SignatureInvalid(&'a str),
    RetryingDownload(&'a str),
}

impl<'a> From<crate::utils::Notification<'a>> for Notification<'a> {
    fn from(n: crate::utils::Notification<'a>) -> Self {
        Notification::Utils(n)
    }
}

impl<'a> From<temp::Notification<'a>> for Notification<'a> {
    fn from(n: temp::Notification<'a>) -> Self {
        Notification::Temp(n)
    }
}

impl<'a> Notification<'a> {
    pub fn level(&self) -> NotificationLevel {
        use self::Notification::*;
        match self {
            Temp(n) => n.level(),
            Utils(n) => n.level(),
            ChecksumValid(_)
            | SignatureValid(_, _)
            | NoUpdateHash(_)
            | FileAlreadyDownloaded
            | DownloadingLegacyManifest => NotificationLevel::Verbose,
            Extracting(_, _)
            | DownloadingComponent(_, _, _)
            | InstallingComponent(_, _, _)
            | RemovingComponent(_, _, _)
            | RemovingOldComponent(_, _, _)
            | ComponentAlreadyInstalled(_)
            | ManifestChecksumFailedHack
            | RollingBack
            | DownloadingManifest(_)
            | SkippingNightlyMissingComponent(_)
            | RetryingDownload(_)
            | DownloadedManifest(_, _) => NotificationLevel::Info,
            CantReadUpdateHash(_)
            | ExtensionNotInstalled(_)
            | MissingInstalledComponent(_)
            | CachedFileChecksumFailed
            | ComponentUnavailable(_, _)
            | ForcingUnavailableComponent(_)
            | StrayHash(_) => NotificationLevel::Warn,
            NonFatalError(_) => NotificationLevel::Error,
            SignatureInvalid(_) => NotificationLevel::Warn,
        }
    }
}

impl<'a> Display for Notification<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::result::Result<(), fmt::Error> {
        use self::Notification::*;
        match self {
            Temp(n) => n.fmt(f),
            Utils(n) => n.fmt(f),
            Extracting(_, _) => write!(f, "extracting..."),
            ComponentAlreadyInstalled(c) => write!(f, "component {} is up to date", c),
            CantReadUpdateHash(path) => write!(
                f,
                "can't read update hash file: '{}', can't skip update...",
                path.display()
            ),
            NoUpdateHash(path) => write!(f, "no update hash at: '{}'", path.display()),
            ChecksumValid(_) => write!(f, "checksum passed"),
            SignatureValid(url, key) => (|| {
                writeln!(f, "Good signature from on {} from:", url)?;
                for l in key.show_key().map_err(|_| std::fmt::Error)?.iter() {
                    writeln!(f, "{}", l)?;
                }
                Ok(())
            })(),
            FileAlreadyDownloaded => write!(f, "reusing previously downloaded file"),
            CachedFileChecksumFailed => write!(f, "bad checksum for cached download"),
            RollingBack => write!(f, "rolling back changes"),
            ExtensionNotInstalled(c) => write!(f, "extension '{}' was not installed", c),
            NonFatalError(e) => write!(f, "{}", e),
            MissingInstalledComponent(c) => {
                write!(f, "during uninstall component {} was not found", c)
            }
            DownloadingComponent(c, h, t) => {
                if Some(h) == t.as_ref() || t.is_none() {
                    write!(f, "downloading component '{}'", c)
                } else {
                    write!(f, "downloading component '{}' for '{}'", c, t.unwrap())
                }
            }
            InstallingComponent(c, h, t) => {
                if Some(h) == t.as_ref() || t.is_none() {
                    write!(f, "installing component '{}'", c)
                } else {
                    write!(f, "installing component '{}' for '{}'", c, t.unwrap())
                }
            }
            RemovingComponent(c, h, t) => {
                if Some(h) == t.as_ref() || t.is_none() {
                    write!(f, "removing component '{}'", c)
                } else {
                    write!(f, "removing component '{}' for '{}'", c, t.unwrap())
                }
            }
            RemovingOldComponent(c, h, t) => {
                if Some(h) == t.as_ref() || t.is_none() {
                    write!(f, "removing previous version of component '{}'", c)
                } else {
                    write!(
                        f,
                        "removing previous version of component '{}' for '{}'",
                        c,
                        t.unwrap()
                    )
                }
            }
            DownloadingManifest(t) => write!(f, "syncing channel updates for '{}'", t),
            DownloadedManifest(date, Some(version)) => {
                write!(f, "latest update on {}, rust version {}", date, version)
            }
            DownloadedManifest(date, None) => {
                write!(f, "latest update on {}, no rust version", date)
            }
            DownloadingLegacyManifest => write!(f, "manifest not found. trying legacy manifest"),
            ManifestChecksumFailedHack => {
                write!(f, "update not yet available, sorry! try again later")
            }
            ComponentUnavailable(pkg, toolchain) => {
                if let Some(tc) = toolchain {
                    write!(f, "component '{}' is not available on target '{}'", pkg, tc)
                } else {
                    write!(f, "component '{}' is not available", pkg)
                }
            }
            StrayHash(path) => write!(
                f,
                "removing stray hash found at '{}' in order to continue",
                path.display()
            ),
            SkippingNightlyMissingComponent(components) => write!(
                f,
                "skipping nightly which is missing installed component '{}'",
                components[0].short_name_in_manifest()
            ),
            ForcingUnavailableComponent(component) => {
                write!(f, "Force-skipping unavailable component '{}'", component)
            }
            SignatureInvalid(url) => write!(f, "Signature verification failed for '{}'", url),
            RetryingDownload(url) => write!(f, "Retrying download for '{}'", url),
        }
    }
}
<|endoftext|>