<commit_before>//! Initialization logic for a generic instance of the `Allocate` channel allocation trait.

use std::thread;
use std::io::BufRead;
use getopts;
use std::sync::Arc;

use allocator::{Thread, Process, Generic};
use networking::initialize_networking;

/// Possible configurations for the communication infrastructure.
pub enum Configuration {
    /// Use one thread.
    Thread,
    /// Use one process with an indicated number of threads.
    Process(usize),
    /// Expect multiple processes indicated by `(threads, process, host_list, report)`.
    Cluster(usize, usize, Vec<String>, bool)
}

impl Configuration {

    /// Constructs a new configuration by parsing supplied text arguments.
    ///
    /// Most commonly, this uses `std::env::Args()` as the supplied iterator.
    pub fn from_args<I: Iterator<Item=String>>(args: I) -> Result<Configuration,String> {

        let mut opts = getopts::Options::new();
        opts.optopt("w", "threads", "number of per-process worker threads", "NUM");
        opts.optopt("p", "process", "identity of this process", "IDX");
        opts.optopt("n", "processes", "number of processes", "NUM");
        opts.optopt("h", "hostfile", "text file whose lines are process addresses", "FILE");
        opts.optflag("r", "report", "reports connection progress");

        opts.parse(args)
            .map_err(|e| format!("{:?}", e))
            .map(|matches| {

            // let mut config = Configuration::new(1, 0, Vec::new());
            let threads = matches.opt_str("w").map(|x| x.parse().unwrap_or(1)).unwrap_or(1);
            let process = matches.opt_str("p").map(|x| x.parse().unwrap_or(0)).unwrap_or(0);
            let processes = matches.opt_str("n").map(|x| x.parse().unwrap_or(1)).unwrap_or(1);
            let report = matches.opt_present("report");

            assert!(process < processes);

            if processes > 1 {
                let mut addresses = Vec::new();
                if let Some(hosts) = matches.opt_str("h") {
                    let reader = ::std::io::BufReader::new(::std::fs::File::open(hosts.clone()).unwrap());
                    for x in reader.lines().take(processes) {
                        addresses.push(x.unwrap());
                    }
                    if addresses.len() < processes {
                        panic!("could only read {} addresses from {}, but -n: {}", addresses.len(), hosts, processes);
                    }
                }
                else {
                    for index in 0..processes {
                        addresses.push(format!("localhost:{}", 2101 + index));
                    }
                }

                assert!(processes == addresses.len());
                Configuration::Cluster(threads, process, addresses, report)
            }
            else {
                if threads > 1 { Configuration::Process(threads) }
                else { Configuration::Thread }
            }
        })
    }
}

fn create_allocators(config: Configuration, logger: Arc<Fn(::logging::CommsSetup)->::logging::CommsLogger+Send+Sync>) -> Result<Vec<Generic>,String> {
    match config {
        Configuration::Thread => Ok(vec![Generic::Thread(Thread)]),
        Configuration::Process(threads) => Ok(Process::new_vector(threads).into_iter().map(Generic::Process).collect()),
        Configuration::Cluster(threads, process, addresses, report) => {
            if let Ok(stuff) = initialize_networking(addresses, process, threads, report, logger) {
                Ok(stuff.into_iter().map(Generic::Binary).collect())
            }
            else {
                Err("failed to initialize networking".to_owned())
            }
        },
    }
}

/// Initializes communication and executes a distributed computation.
///
/// This method allocates an `allocator::Generic` for each thread, spawns local worker threads,
/// and invokes the supplied function with the allocator.
/// The method returns a `WorkerGuards<T>` which can be `join`ed to retrieve the return values
/// (or errors) of the workers.
///
///
/// #Examples
/// ```
/// // configure for two threads, just one process.
/// let config = timely_communication::Configuration::Process(2);
///
/// // create a source of inactive loggers.
/// let logger = ::std::sync::Arc::new(|_| timely_communication::logging::BufferingLogger::new_inactive());
///
/// // initializes communication, spawns workers
/// let guards = timely_communication::initialize(config, logger, |mut allocator| {
///     println!("worker {} started", allocator.index());
///
///     // allocates pair of senders list and one receiver.
///     let (mut senders, mut receiver, _) = allocator.allocate();
///
///     // send typed data along each channel
///     senders[0].send(format!("hello, {}", 0));
///     senders[1].send(format!("hello, {}", 1));
///
///     // no support for termination notification,
///     // we have to count down ourselves.
///     let mut expecting = 2;
///     while expecting > 0 {
///         if let Some(message) = receiver.recv() {
///             println!("worker {}: received: <{}>", allocator.index(), message);
///             expecting -= 1;
///         }
///     }
///
///     // optionally, return something
///     allocator.index()
/// });
///
/// // computation runs until guards are joined or dropped.
/// if let Ok(guards) = guards {
///     for guard in guards.join() {
///         println!("result: {:?}", guard);
///     }
/// }
/// else { println!("error in computation"); }
/// ```
///
/// The should produce output like:
///
/// ```ignore
/// worker 0 started
/// worker 1 started
/// worker 0: received: <hello, 0>
/// worker 1: received: <hello, 1>
/// worker 0: received: <hello, 0>
/// worker 1: received: <hello, 1>
/// result: Ok(0)
/// result: Ok(1)
/// ```
pub fn initialize<T:Send+'static, F: Fn(Generic)->T+Send+Sync+'static>(
    config: Configuration,
    log_sender: Arc<Fn(::logging::CommsSetup)->::logging::CommsLogger+Send+Sync>,
    func: F,
) -> Result<WorkerGuards<T>,String> {

    let allocators = try!(create_allocators(config, log_sender));
    let logic = Arc::new(func);

    let mut guards = Vec::new();
    for allocator in allocators.into_iter() {
        let clone = logic.clone();
        guards.push(try!(thread::Builder::new()
                            .name(format!("worker thread {}", allocator.index()))
                            .spawn(move || {
                                (*clone)(allocator)
                            })
                            .map_err(|e| format!("{:?}", e))));
    }

    Ok(WorkerGuards { guards: guards })
}

/// Maintains `JoinHandle`s for worker threads.
pub struct WorkerGuards<T:Send+'static> {
    guards: Vec<::std::thread::JoinHandle<T>>
}

impl<T:Send+'static> WorkerGuards<T> {
    /// Waits on the worker threads and returns the results they produce.
    pub fn join(mut self) -> Vec<Result<T,String>> {
        self.guards.drain(..)
                   .map(|guard| guard.join().map_err(|e| format!("{:?}", e)))
                   .collect()
    }
}

impl<T:Send+'static> Drop for WorkerGuards<T> {
    fn drop(&mut self) {
        for guard in self.guards.drain(..) {
            guard.join().unwrap();
        }
    }
}
<commit_msg>timely-communication: Collapse if statement.<commit_after>//! Initialization logic for a generic instance of the `Allocate` channel allocation trait.

use std::thread;
use std::io::BufRead;
use getopts;
use std::sync::Arc;

use allocator::{Thread, Process, Generic};
use networking::initialize_networking;

/// Possible configurations for the communication infrastructure.
pub enum Configuration {
    /// Use one thread.
    Thread,
    /// Use one process with an indicated number of threads.
    Process(usize),
    /// Expect multiple processes indicated by `(threads, process, host_list, report)`.
    Cluster(usize, usize, Vec<String>, bool)
}

impl Configuration {

    /// Constructs a new configuration by parsing supplied text arguments.
    ///
    /// Most commonly, this uses `std::env::Args()` as the supplied iterator.
    pub fn from_args<I: Iterator<Item=String>>(args: I) -> Result<Configuration,String> {

        let mut opts = getopts::Options::new();
        opts.optopt("w", "threads", "number of per-process worker threads", "NUM");
        opts.optopt("p", "process", "identity of this process", "IDX");
        opts.optopt("n", "processes", "number of processes", "NUM");
        opts.optopt("h", "hostfile", "text file whose lines are process addresses", "FILE");
        opts.optflag("r", "report", "reports connection progress");

        opts.parse(args)
            .map_err(|e| format!("{:?}", e))
            .map(|matches| {

            // let mut config = Configuration::new(1, 0, Vec::new());
            let threads = matches.opt_str("w").map(|x| x.parse().unwrap_or(1)).unwrap_or(1);
            let process = matches.opt_str("p").map(|x| x.parse().unwrap_or(0)).unwrap_or(0);
            let processes = matches.opt_str("n").map(|x| x.parse().unwrap_or(1)).unwrap_or(1);
            let report = matches.opt_present("report");

            assert!(process < processes);

            if processes > 1 {
                let mut addresses = Vec::new();
                if let Some(hosts) = matches.opt_str("h") {
                    let reader = ::std::io::BufReader::new(::std::fs::File::open(hosts.clone()).unwrap());
                    for x in reader.lines().take(processes) {
                        addresses.push(x.unwrap());
                    }
                    if addresses.len() < processes {
                        panic!("could only read {} addresses from {}, but -n: {}", addresses.len(), hosts, processes);
                    }
                }
                else {
                    for index in 0..processes {
                        addresses.push(format!("localhost:{}", 2101 + index));
                    }
                }

                assert!(processes == addresses.len());
                Configuration::Cluster(threads, process, addresses, report)
            }
            else if threads > 1 { Configuration::Process(threads) }
            else { Configuration::Thread }
        })
    }
}

fn create_allocators(config: Configuration, logger: Arc<Fn(::logging::CommsSetup)->::logging::CommsLogger+Send+Sync>) -> Result<Vec<Generic>,String> {
    match config {
        Configuration::Thread => Ok(vec![Generic::Thread(Thread)]),
        Configuration::Process(threads) => Ok(Process::new_vector(threads).into_iter().map(Generic::Process).collect()),
        Configuration::Cluster(threads, process, addresses, report) => {
            if let Ok(stuff) = initialize_networking(addresses, process, threads, report, logger) {
                Ok(stuff.into_iter().map(Generic::Binary).collect())
            }
            else {
                Err("failed to initialize networking".to_owned())
            }
        },
    }
}

/// Initializes communication and executes a distributed computation.
///
/// This method allocates an `allocator::Generic` for each thread, spawns local worker threads,
/// and invokes the supplied function with the allocator.
/// The method returns a `WorkerGuards<T>` which can be `join`ed to retrieve the return values
/// (or errors) of the workers.
///
///
/// #Examples
/// ```
/// // configure for two threads, just one process.
/// let config = timely_communication::Configuration::Process(2);
///
/// // create a source of inactive loggers.
/// let logger = ::std::sync::Arc::new(|_| timely_communication::logging::BufferingLogger::new_inactive());
///
/// // initializes communication, spawns workers
/// let guards = timely_communication::initialize(config, logger, |mut allocator| {
///     println!("worker {} started", allocator.index());
///
///     // allocates pair of senders list and one receiver.
///     let (mut senders, mut receiver, _) = allocator.allocate();
///
///     // send typed data along each channel
///     senders[0].send(format!("hello, {}", 0));
///     senders[1].send(format!("hello, {}", 1));
///
///     // no support for termination notification,
///     // we have to count down ourselves.
///     let mut expecting = 2;
///     while expecting > 0 {
///         if let Some(message) = receiver.recv() {
///             println!("worker {}: received: <{}>", allocator.index(), message);
///             expecting -= 1;
///         }
///     }
///
///     // optionally, return something
///     allocator.index()
/// });
///
/// // computation runs until guards are joined or dropped.
/// if let Ok(guards) = guards {
///     for guard in guards.join() {
///         println!("result: {:?}", guard);
///     }
/// }
/// else { println!("error in computation"); }
/// ```
///
/// The should produce output like:
///
/// ```ignore
/// worker 0 started
/// worker 1 started
/// worker 0: received: <hello, 0>
/// worker 1: received: <hello, 1>
/// worker 0: received: <hello, 0>
/// worker 1: received: <hello, 1>
/// result: Ok(0)
/// result: Ok(1)
/// ```
pub fn initialize<T:Send+'static, F: Fn(Generic)->T+Send+Sync+'static>(
    config: Configuration,
    log_sender: Arc<Fn(::logging::CommsSetup)->::logging::CommsLogger+Send+Sync>,
    func: F,
) -> Result<WorkerGuards<T>,String> {

    let allocators = try!(create_allocators(config, log_sender));
    let logic = Arc::new(func);

    let mut guards = Vec::new();
    for allocator in allocators.into_iter() {
        let clone = logic.clone();
        guards.push(try!(thread::Builder::new()
                            .name(format!("worker thread {}", allocator.index()))
                            .spawn(move || {
                                (*clone)(allocator)
                            })
                            .map_err(|e| format!("{:?}", e))));
    }

    Ok(WorkerGuards { guards: guards })
}

/// Maintains `JoinHandle`s for worker threads.
pub struct WorkerGuards<T:Send+'static> {
    guards: Vec<::std::thread::JoinHandle<T>>
}

impl<T:Send+'static> WorkerGuards<T> {
    /// Waits on the worker threads and returns the results they produce.
    pub fn join(mut self) -> Vec<Result<T,String>> {
        self.guards.drain(..)
                   .map(|guard| guard.join().map_err(|e| format!("{:?}", e)))
                   .collect()
    }
}

impl<T:Send+'static> Drop for WorkerGuards<T> {
    fn drop(&mut self) {
        for guard in self.guards.drain(..) {
            guard.join().unwrap();
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#[macro_export]
macro_rules! make_getter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self) -> DOMString {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            element.get_string_attribute(&Atom::from_slice($htmlname))
        }
    );
    ($attr:ident) => {
        make_getter!($attr, to_lower!(stringify!($attr)));
    }
);

#[macro_export]
macro_rules! make_bool_getter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self) -> bool {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not runtime.
            element.has_attribute(&Atom::from_slice($htmlname))
        }
    );
    ($attr:ident) => {
        make_bool_getter!($attr, to_lower!(stringify!($attr)));
    }
);

#[macro_export]
macro_rules! make_uint_getter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self) -> u32 {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not runtime.
            element.get_uint_attribute(&Atom::from_slice($htmlname))
        }
    );
    ($attr:ident) => {
        make_uint_getter!($attr, to_lower!(stringify!($attr)));
    }
);

#[macro_export]
macro_rules! make_url_getter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self) -> DOMString {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not runtime.
            element.get_url_attribute(&Atom::from_slice($htmlname))
        }
    );
    ($attr:ident) => {
        // FIXME(pcwalton): Do this at compile time, not runtime.
        make_url_getter!($attr, to_lower!(stringify!($attr)));
    }
);

#[macro_export]
macro_rules! make_url_or_base_getter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self) -> DOMString {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            use dom::window::WindowHelpers;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            let url = element.get_url_attribute(&Atom::from_slice($htmlname));
            match &*url {
                "" => {
                    let window = window_from_node(self).root();
                    window.r().get_url().serialize()
                },
                _ => url
            }
        }
    );
    ($attr:ident) => {
        make_url_or_base_getter!($attr, to_lower!(stringify!($attr)));
    }
);

#[macro_export]
macro_rules! make_enumerated_getter(
    ( $attr:ident, $htmlname:expr, $default:expr, $(($choices: pat))|+) => (
        fn $attr(self) -> DOMString {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            use std::borrow::ToOwned;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            let val = element.get_string_attribute(&Atom::from_slice($htmlname))
                             .into_ascii_lowercase();
            // https://html.spec.whatwg.org/multipage/#attr-fs-method
            match &*val {
                $($choices)|+ => val,
                _ => $default.to_owned()
            }
        }
    );
    ($attr:ident, $default:expr, $(($choices: pat))|+) => {
        make_enumerated_getter!($attr, &to_lower!(stringify!($attr)), $default, $(($choices))|+);
    }
);

// concat_idents! doesn't work for function name positions, so
// we have to specify both the content name and the HTML name here
#[macro_export]
macro_rules! make_setter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self, value: DOMString) {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not at runtime.
            element.set_string_attribute(&Atom::from_slice($htmlname), value)
        }
    );
);

#[macro_export]
macro_rules! make_bool_setter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self, value: bool) {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not at runtime.
            element.set_bool_attribute(&Atom::from_slice($htmlname), value)
        }
    );
);

#[macro_export]
macro_rules! make_uint_setter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self, value: u32) {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not at runtime.
            element.set_uint_attribute(&Atom::from_slice($htmlname), value)
        }
    );
);

/// For use on non-jsmanaged types
/// Use #[jstraceable] on JS managed types
macro_rules! no_jsmanaged_fields(
    ($($ty:ident),+) => (
        $(
            impl $crate::dom::bindings::trace::JSTraceable for $ty {
                #[inline]
                fn trace(&self, _: *mut ::js::jsapi::JSTracer) {
                    // Do nothing
                }
            }
        )+
    );
    ($ty:ident<$($gen:ident),+>) => (
        impl<$($gen),+> $crate::dom::bindings::trace::JSTraceable for $ty<$($gen),+> {
            #[inline]
            fn trace(&self, _: *mut ::js::jsapi::JSTracer) {
                // Do nothing
            }
        }
    );
);

/// These are used to generate a event handler which has no special case.
macro_rules! define_event_handler(
    ($handler: ident, $event_type: ident, $getter: ident, $setter: ident) => (
        fn $getter(self) -> Option<$handler> {
            let eventtarget: JSRef<EventTarget> = EventTargetCast::from_ref(self);
            eventtarget.get_event_handler_common(stringify!($event_type))
        }

        fn $setter(self, listener: Option<$handler>) {
            let eventtarget: JSRef<EventTarget> = EventTargetCast::from_ref(self);
            eventtarget.set_event_handler_common(stringify!($event_type), listener)
        }
    )
);

macro_rules! event_handler(
    ($event_type: ident, $getter: ident, $setter: ident) => (
        define_event_handler!(EventHandlerNonNull, $event_type, $getter, $setter);
    )
);

macro_rules! error_event_handler(
    ($event_type: ident, $getter: ident, $setter: ident) => (
        define_event_handler!(OnErrorEventHandlerNonNull, $event_type, $getter, $setter);
    )
);

// https://html.spec.whatwg.org/multipage/#globaleventhandlers
// see webidls/EventHandler.webidl
// As more methods get added, just update them here.
macro_rules! global_event_handlers(
    () => (
        event_handler!(load, GetOnload, SetOnload);
        global_event_handlers!(NoOnload);

    );
    (NoOnload) => (
        event_handler!(click, GetOnclick, SetOnclick);
        event_handler!(input, GetOninput, SetOninput);
        event_handler!(change, GetOnchange, SetOnchange);
        event_handler!(submit, GetOnsubmit, SetOnsubmit);
    )
);
<commit_msg>Simplify the implementation of make_url_or_base_getter.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#[macro_export]
macro_rules! make_getter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self) -> DOMString {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            element.get_string_attribute(&Atom::from_slice($htmlname))
        }
    );
    ($attr:ident) => {
        make_getter!($attr, to_lower!(stringify!($attr)));
    }
);

#[macro_export]
macro_rules! make_bool_getter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self) -> bool {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not runtime.
            element.has_attribute(&Atom::from_slice($htmlname))
        }
    );
    ($attr:ident) => {
        make_bool_getter!($attr, to_lower!(stringify!($attr)));
    }
);

#[macro_export]
macro_rules! make_uint_getter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self) -> u32 {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not runtime.
            element.get_uint_attribute(&Atom::from_slice($htmlname))
        }
    );
    ($attr:ident) => {
        make_uint_getter!($attr, to_lower!(stringify!($attr)));
    }
);

#[macro_export]
macro_rules! make_url_getter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self) -> DOMString {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not runtime.
            element.get_url_attribute(&Atom::from_slice($htmlname))
        }
    );
    ($attr:ident) => {
        // FIXME(pcwalton): Do this at compile time, not runtime.
        make_url_getter!($attr, to_lower!(stringify!($attr)));
    }
);

#[macro_export]
macro_rules! make_url_or_base_getter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self) -> DOMString {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            use dom::window::WindowHelpers;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            let url = element.get_url_attribute(&Atom::from_slice($htmlname));
            if url.is_empty() {
                let window = window_from_node(self).root();
                window.r().get_url().serialize()
            } else {
                url
            }
        }
    );
    ($attr:ident) => {
        make_url_or_base_getter!($attr, to_lower!(stringify!($attr)));
    }
);

#[macro_export]
macro_rules! make_enumerated_getter(
    ( $attr:ident, $htmlname:expr, $default:expr, $(($choices: pat))|+) => (
        fn $attr(self) -> DOMString {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            #[allow(unused_imports)]
            use std::ascii::AsciiExt;
            use std::borrow::ToOwned;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            let val = element.get_string_attribute(&Atom::from_slice($htmlname))
                             .into_ascii_lowercase();
            // https://html.spec.whatwg.org/multipage/#attr-fs-method
            match &*val {
                $($choices)|+ => val,
                _ => $default.to_owned()
            }
        }
    );
    ($attr:ident, $default:expr, $(($choices: pat))|+) => {
        make_enumerated_getter!($attr, &to_lower!(stringify!($attr)), $default, $(($choices))|+);
    }
);

// concat_idents! doesn't work for function name positions, so
// we have to specify both the content name and the HTML name here
#[macro_export]
macro_rules! make_setter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self, value: DOMString) {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not at runtime.
            element.set_string_attribute(&Atom::from_slice($htmlname), value)
        }
    );
);

#[macro_export]
macro_rules! make_bool_setter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self, value: bool) {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not at runtime.
            element.set_bool_attribute(&Atom::from_slice($htmlname), value)
        }
    );
);

#[macro_export]
macro_rules! make_uint_setter(
    ( $attr:ident, $htmlname:expr ) => (
        fn $attr(self, value: u32) {
            use dom::element::{Element, AttributeHandlers};
            use dom::bindings::codegen::InheritTypes::ElementCast;
            let element: JSRef<Element> = ElementCast::from_ref(self);
            // FIXME(pcwalton): Do this at compile time, not at runtime.
            element.set_uint_attribute(&Atom::from_slice($htmlname), value)
        }
    );
);

/// For use on non-jsmanaged types
/// Use #[jstraceable] on JS managed types
macro_rules! no_jsmanaged_fields(
    ($($ty:ident),+) => (
        $(
            impl $crate::dom::bindings::trace::JSTraceable for $ty {
                #[inline]
                fn trace(&self, _: *mut ::js::jsapi::JSTracer) {
                    // Do nothing
                }
            }
        )+
    );
    ($ty:ident<$($gen:ident),+>) => (
        impl<$($gen),+> $crate::dom::bindings::trace::JSTraceable for $ty<$($gen),+> {
            #[inline]
            fn trace(&self, _: *mut ::js::jsapi::JSTracer) {
                // Do nothing
            }
        }
    );
);

/// These are used to generate a event handler which has no special case.
macro_rules! define_event_handler(
    ($handler: ident, $event_type: ident, $getter: ident, $setter: ident) => (
        fn $getter(self) -> Option<$handler> {
            let eventtarget: JSRef<EventTarget> = EventTargetCast::from_ref(self);
            eventtarget.get_event_handler_common(stringify!($event_type))
        }

        fn $setter(self, listener: Option<$handler>) {
            let eventtarget: JSRef<EventTarget> = EventTargetCast::from_ref(self);
            eventtarget.set_event_handler_common(stringify!($event_type), listener)
        }
    )
);

macro_rules! event_handler(
    ($event_type: ident, $getter: ident, $setter: ident) => (
        define_event_handler!(EventHandlerNonNull, $event_type, $getter, $setter);
    )
);

macro_rules! error_event_handler(
    ($event_type: ident, $getter: ident, $setter: ident) => (
        define_event_handler!(OnErrorEventHandlerNonNull, $event_type, $getter, $setter);
    )
);

// https://html.spec.whatwg.org/multipage/#globaleventhandlers
// see webidls/EventHandler.webidl
// As more methods get added, just update them here.
macro_rules! global_event_handlers(
    () => (
        event_handler!(load, GetOnload, SetOnload);
        global_event_handlers!(NoOnload);

    );
    (NoOnload) => (
        event_handler!(click, GetOnclick, SetOnclick);
        event_handler!(input, GetOninput, SetOninput);
        event_handler!(change, GetOnchange, SetOnchange);
        event_handler!(submit, GetOnsubmit, SetOnsubmit);
    )
);
<|endoftext|>


<commit_before>use std::cmp;
use AttackResult::*;
use EffectClass::*;

#[derive(Debug, Copy, Clone)]
struct Character {
    hp: i32,
    mana: i32,
    damage: i32,
    armor: i32
}

#[derive(Debug, Clone)]
enum AttackResult {
    Win,
    Lose,
    Continue(State)
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum EffectClass {
    Shield,
    Poison,
    Recharge
}

#[derive(Debug, Copy, Clone)]
struct Effect {
    turns: i32,
    intensity: i32,
    class: EffectClass
}

#[derive(Debug, Copy, Clone)]
struct Spell {
    cost: i32,
    damage: i32,
    healing: i32,
    effect: Option<Effect>
}

#[derive(Debug, Clone)]
struct State {
    player: Character,
    boss: Character,
    effects: Vec<Effect>
}

fn process_effects(state: &State) -> State {
    let mut result = State {player: state.player, boss: state.boss, effects: Vec::new()};
    let mut armor = 0;

    for mut effect in state.effects.iter().filter(|x| x.turns > 0).cloned() {
        effect.turns -= 1;

        if effect.turns > 0 {
            result.effects.push(effect);
        }

        match effect.class {
            Shield => armor = effect.intensity,
            Poison => result.boss.hp -= effect.intensity,
            Recharge => result.player.mana += effect.intensity
        }
    }

    result.player.armor = armor;
    result
}

fn cast_spell(state: &State, spell: Spell) -> Option<AttackResult> {
    let mut result = process_effects(state);

    // Cast spell

    result.player.mana -= spell.cost;
    result.player.hp += spell.healing;
    result.boss.hp -= spell.damage;

    if result.player.mana <= 0 {
        return Some(Lose);
    }

    // Deal with effects

    if let Some(effect) = spell.effect {
        if result.effects.iter().all(|x| x.class != effect.class) {
            result.effects.push(effect)
        } else {
            return None;
        }
    }

    // Counter attack

    let mut result = process_effects(&result);

    if result.boss.hp > 0 {
        let damage = cmp::max(result.boss.damage - result.player.armor, 1);
        result.player.hp = cmp::max(result.player.hp - damage, 0);

        if result.player.hp <= 0 {
            return Some(Lose);
        }
    } else {
        return Some(Win);
    }

    Some(Continue(result))
}

fn list_strategies(spells: &Vec<Spell>, state: &State) -> Vec<Vec<Spell>> {
    spells.iter().flat_map(|&spell| match cast_spell(state, spell) {
        Some(Continue(next_state)) => {
            list_strategies(spells, &next_state)
            .into_iter()
            .map(|mut strategy| {
                strategy.push(spell);
                strategy
            })
            .collect()
        },
        Some(Win) => vec![vec![spell]],
        _ => vec![]
    }).collect()
}

fn main() {
    let spells = vec![
        Spell {
            cost: 53,
            damage: 4,
            healing: 0,
            effect: None
        },
        Spell {
            cost: 73,
            damage: 2,
            healing: 2,
            effect: None
        },
        Spell {
            cost: 113,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 6, intensity: 7, class: Shield})
        },
        Spell {
            cost: 173,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 6, intensity: 3, class: Poison})
        },
        Spell {
            cost: 229,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 5, intensity: 101, class: Recharge})
        }
    ];

    let boss = Character {
        hp: 51,
        mana: 0,
        damage: 9,
        armor: 0
    };

    let player = Character {
        hp: 50,
        mana: 500,
        damage: 0,
        armor: 0
    };

    let start = State {player, boss, effects: Vec::new()};
    let strategies = list_strategies(&spells, &start);

    println!("{:?}", strategies.len());
}
<commit_msg>process_effects() now mutates State<commit_after>use std::cmp;
use AttackResult::*;
use EffectClass::*;

#[derive(Debug, Copy, Clone)]
struct Character {
    hp: i32,
    mana: i32,
    damage: i32,
    armor: i32
}

#[derive(Debug, Clone)]
enum AttackResult {
    Win,
    Lose,
    Continue(State)
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum EffectClass {
    Shield,
    Poison,
    Recharge
}

#[derive(Debug, Copy, Clone)]
struct Effect {
    turns: i32,
    intensity: i32,
    class: EffectClass
}

#[derive(Debug, Copy, Clone)]
struct Spell {
    cost: i32,
    damage: i32,
    healing: i32,
    effect: Option<Effect>
}

#[derive(Debug, Clone)]
struct State {
    player: Character,
    boss: Character,
    effects: Vec<Effect>
}

fn process_effects(state: &mut State) {
    let mut effects = Vec::new();
    let mut armor = 0;

    for mut effect in state.effects.iter().filter(|x| x.turns > 0).cloned() {
        effect.turns -= 1;

        if effect.turns > 0 {
            effects.push(effect);
        }

        match effect.class {
            Shield => armor = effect.intensity,
            Poison => state.boss.hp -= effect.intensity,
            Recharge => state.player.mana += effect.intensity
        }
    }

    state.player.armor = armor;
    state.effects = effects;
}

fn cast_spell(state: &State, spell: Spell) -> AttackResult {
    let mut result = state.clone();
    process_effects(&mut result);

    // Cast spell

    result.player.mana -= spell.cost;
    result.player.hp += spell.healing;
    result.boss.hp -= spell.damage;

    if result.player.mana <= 0 {
        return Lose;
    }

    // Deal with effects

    if let Some(effect) = spell.effect {
        if result.effects.iter().all(|x| x.class != effect.class) {
            result.effects.push(effect);
        } else {
            return Lose;
        }
    }

    // Counter attack

    process_effects(&mut result);

    if result.boss.hp > 0 {
        let damage = cmp::max(result.boss.damage - result.player.armor, 1);
        result.player.hp = cmp::max(result.player.hp - damage, 0);

        if result.player.hp <= 0 {
            return Lose;
        }
    } else {
        return Win;
    }

    Continue(result)
}

fn list_strategies(spells: &Vec<Spell>, state: &State) -> Vec<Vec<Spell>> {
    spells.iter().flat_map(|&spell| match cast_spell(state, spell) {
        Some(Continue(next_state)) => {
            list_strategies(spells, &next_state)
            .into_iter()
            .map(|mut strategy| {
                strategy.push(spell);
                strategy
            })
            .collect()
        },
        Some(Win) => vec![vec![spell]],
        _ => vec![]
    }).collect()
}

fn main() {
    let spells = vec![
        Spell {
            cost: 53,
            damage: 4,
            healing: 0,
            effect: None
        },
        Spell {
            cost: 73,
            damage: 2,
            healing: 2,
            effect: None
        },
        Spell {
            cost: 113,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 6, intensity: 7, class: Shield})
        },
        Spell {
            cost: 173,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 6, intensity: 3, class: Poison})
        },
        Spell {
            cost: 229,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 5, intensity: 101, class: Recharge})
        }
    ];

    let boss = Character {
        hp: 51,
        mana: 0,
        damage: 9,
        armor: 0
    };

    let player = Character {
        hp: 50,
        mana: 500,
        damage: 0,
        armor: 0
    };

    let start = State {player, boss, effects: Vec::new()};
    let strategies = list_strategies(&spells, &start);

    println!("{:?}", strategies.len());
}
<|endoftext|>


<commit_before>//! An unbounded set of futures.

use std::cell::UnsafeCell;
use std::fmt::{self, Debug};
use std::iter::FromIterator;
use std::marker::{PhantomData, Unpin};
use std::mem::{self, PinMut};
use std::ptr;
use std::sync::atomic::Ordering::SeqCst;
use std::sync::atomic::{AtomicPtr, AtomicBool};
use std::sync::{Arc, Weak};
use std::usize;

use futures_core::{Stream, Future, Poll};
use futures_core::task::{self, AtomicWaker};

mod abort;
mod inner;
mod iter_mut;
mod node;

use self::inner::{Inner, Dequeue};
use self::iter_mut::IterMut;
use self::node::Node;

/// A set of `Future`s which may complete in any order.
///
/// This structure is optimized to manage a large number of futures.
/// Futures managed by `FuturesUnordered` will only be polled when they
/// generate notifications. This reduces the required amount of work needed to
/// poll large numbers of futures.
///
/// `FuturesUnordered` can be filled by `collect`ing an iterator of `Future`s
/// into a `FuturesUnordered`, or by `push`ing `Future`s onto an existing
/// `FuturesUnordered`. When new `Future`s are added, `poll_next` must be
/// called in order to begin receiving wakeups for new `Future`s.
///
/// Note that you can create a ready-made `FuturesUnordered` via the
/// `futures_unordered` function in the `stream` module, or you can start with an
/// empty set with the `FuturesUnordered::new` constructor.
#[must_use = "streams do nothing unless polled"]
pub struct FuturesUnordered<F> {
    inner: Arc<Inner<F>>,
    len: usize,
    head_all: *const Node<F>,
}

unsafe impl<T: Send> Send for FuturesUnordered<T> {}
unsafe impl<T: Sync> Sync for FuturesUnordered<T> {}
impl<T> Unpin for FuturesUnordered<T> {}

// FuturesUnordered is implemented using two linked lists. One which links all
// futures managed by a `FuturesUnordered` and one that tracks futures that have
// been scheduled for polling. The first linked list is not thread safe and is
// only accessed by the thread that owns the `FuturesUnordered` value. The
// second linked list is an implementation of the intrusive MPSC queue algorithm
// described by 1024cores.net.
//
// When a future is submitted to the set a node is allocated and inserted in
// both linked lists. The next call to `poll` will (eventually) see this node
// and call `poll` on the future.
//
// Before a managed future is polled, the current task's `Wake` is replaced
// with one that is aware of the specific future being run. This ensures that
// task notifications generated by that specific future are visible to
// `FuturesUnordered`. When a notification is received, the node is scheduled
// for polling by being inserted into the concurrent linked list.
//
// Each node uses an `AtomicUsize` to track it's state. The node state is the
// reference count (the number of outstanding handles to the node) as well as a
// flag tracking if the node is currently inserted in the atomic queue. When the
// future is notified, it will only insert itself into the linked list if it
// isn't currently inserted.

impl<T> FuturesUnordered<T>
    where T: Future
{
    /// Constructs a new, empty `FuturesUnordered`
    ///
    /// The returned `FuturesUnordered` does not contain any futures and, in this
    /// state, `FuturesUnordered::poll_next` will return `Poll::Ready(None)`.
    pub fn new() -> FuturesUnordered<T> {
        let stub = Arc::new(Node {
            future: UnsafeCell::new(None),
            next_all: UnsafeCell::new(ptr::null()),
            prev_all: UnsafeCell::new(ptr::null()),
            next_readiness: AtomicPtr::new(ptr::null_mut()),
            queued: AtomicBool::new(true),
            queue: Weak::new(),
        });
        let stub_ptr = &*stub as *const Node<T>;
        let inner = Arc::new(Inner {
            parent: AtomicWaker::new(),
            head_readiness: AtomicPtr::new(stub_ptr as *mut _),
            tail_readiness: UnsafeCell::new(stub_ptr),
            stub: stub,
        });

        FuturesUnordered {
            len: 0,
            head_all: ptr::null_mut(),
            inner: inner,
        }
    }
}

impl<T> FuturesUnordered<T> {
    /// Returns the number of futures contained in the set.
    ///
    /// This represents the total number of in-flight futures.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Returns `true` if the set contains no futures
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Push a future into the set.
    ///
    /// This function submits the given future to the set for managing. This
    /// function will not call `poll` on the submitted future. The caller must
    /// ensure that `FuturesUnordered::poll_next` is called in order to receive task
    /// notifications.
    pub fn push(&mut self, future: T) {
        let node = Arc::new(Node {
            future: UnsafeCell::new(Some(future)),
            next_all: UnsafeCell::new(ptr::null_mut()),
            prev_all: UnsafeCell::new(ptr::null_mut()),
            next_readiness: AtomicPtr::new(ptr::null_mut()),
            queued: AtomicBool::new(true),
            queue: Arc::downgrade(&self.inner),
        });

        // Right now our node has a strong reference count of 1. We transfer
        // ownership of this reference count to our internal linked list
        // and we'll reclaim ownership through the `unlink` function below.
        let ptr = self.link(node);

        // We'll need to get the future "into the system" to start tracking it,
        // e.g. getting its unpark notifications going to us tracking which
        // futures are ready. To do that we unconditionally enqueue it for
        // polling here.
        self.inner.enqueue(ptr);
    }

    /// Returns an iterator that allows modifying each future in the set.
    pub fn iter_mut(&mut self) -> IterMut<T> {
        IterMut {
            node: self.head_all,
            len: self.len,
            _marker: PhantomData
        }
    }

    fn release_node(&mut self, node: Arc<Node<T>>) {
        // The future is done, try to reset the queued flag. This will prevent
        // `notify` from doing any work in the future
        let prev = node.queued.swap(true, SeqCst);

        // Drop the future, even if it hasn't finished yet. This is safe
        // because we're dropping the future on the thread that owns
        // `FuturesUnordered`, which correctly tracks T's lifetimes and such.
        unsafe {
            drop((*node.future.get()).take());
        }

        // If the queued flag was previously set then it means that this node
        // is still in our internal mpsc queue. We then transfer ownership
        // of our reference count to the mpsc queue, and it'll come along and
        // free it later, noticing that the future is `None`.
        //
        // If, however, the queued flag was *not* set then we're safe to
        // release our reference count on the internal node. The queued flag
        // was set above so all future `enqueue` operations will not actually
        // enqueue the node, so our node will never see the mpsc queue again.
        // The node itself will be deallocated once all reference counts have
        // been dropped by the various owning tasks elsewhere.
        if prev {
            mem::forget(node);
        }
    }

    /// Insert a new node into the internal linked list.
    fn link(&mut self, node: Arc<Node<T>>) -> *const Node<T> {
        let ptr = Arc::into_raw(node);
        unsafe {
            *(*ptr).next_all.get() = self.head_all;
            if !self.head_all.is_null() {
                *(*self.head_all).prev_all.get() = ptr;
            }
        }

        self.head_all = ptr;
        self.len += 1;
        ptr
    }

    /// Remove the node from the linked list tracking all nodes currently
    /// managed by `FuturesUnordered`.
    /// This function is unsafe because it has be guaranteed that `node` is a
    /// valid pointer.
    unsafe fn unlink(&mut self, node: *const Node<T>) -> Arc<Node<T>> {
        let node = Arc::from_raw(node);
        let next = *node.next_all.get();
        let prev = *node.prev_all.get();
        *node.next_all.get() = ptr::null_mut();
        *node.prev_all.get() = ptr::null_mut();

        if !next.is_null() {
            *(*next).prev_all.get() = prev;
        }

        if !prev.is_null() {
            *(*prev).next_all.get() = next;
        } else {
            self.head_all = next;
        }
        self.len -= 1;
        node
    }
}

impl<T> Stream for FuturesUnordered<T>
    where T: Future
{
    type Item = T::Output;

    fn poll_next(mut self: PinMut<Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
        // Ensure `parent` is correctly set.
        self.inner.parent.register(cx.waker());

        loop {
            // Safety: &mut self guarantees mutual exclusion
            let node = match unsafe { self.inner.dequeue() } {
                Dequeue::Empty => {
                    if self.is_empty() {
                        return Poll::Ready(None);
                    } else {
                        return Poll::Pending;
                    }
                }
                Dequeue::Inconsistent => {
                    // At this point, it may be worth yielding the thread &
                    // spinning a few times... but for now, just yield using the
                    // task system.
                    cx.waker().wake();
                    return Poll::Pending;
                }
                Dequeue::Data(node) => node,
            };

            debug_assert!(node != self.inner.stub());

            // Safety: Node is a valid pointer. We are the only thread that
            // accesses the `UnsafeCell` that contains the future
            let future = match unsafe { &mut *(*node).future.get() } {
                Some(future) => future,

                // If the future has already gone away then we're just
                // cleaning out this node. See the comment in
                // `release_node` for more information, but we're basically
                // just taking ownership of our reference count here.
                None => {
                    // Saftey: `release_node` uses `mem::forget`
                    let node = unsafe { Arc::from_raw(node) };
                    assert!(node.next_all.get().is_null());
                    assert!(node.prev_all.get().is_null());
                    continue
                }
            };

            // Safety: `node` is a valid pointer
            let node = unsafe { self.unlink(node) };

            // Unset queued flag... this must be done before
            // polling. This ensures that the future gets
            // rescheduled if it is notified **during** a call
            // to `poll`.
            let prev = node.queued.swap(false, SeqCst);
            assert!(prev);

            // We're going to need to be very careful if the `poll`
            // function below panics. We need to (a) not leak memory and
            // (b) ensure that we still don't have any use-after-frees. To
            // manage this we do a few things:
            //
            // * When the "bomb" is dropped abnormally, it will call
            //   `release_node`. That way we'll be sure the memory management
            //   of the `node` is managed correctly. In particular
            //   `release_node` will drop the future. This ensures that it is
            //   dropped on this thread and not accidentally on a different
            //   thread (bad).
            // * We unlink the node from our internal queue to preemptively
            //   assume it'll panic, in which case we'll want to discard it
            //   regardless.
            struct Bomb<'a, T: 'a> {
                queue: &'a mut FuturesUnordered<T>,
                node: Option<Arc<Node<T>>>,
            }

            impl<'a, T> Drop for Bomb<'a, T> {
                fn drop(&mut self) {
                    if let Some(node) = self.node.take() {
                        self.queue.release_node(node);
                    }
                }
            }

            let mut bomb = Bomb {
                node: Some(node),
                queue: &mut *self,
            };

            // Poll the underlying future with the appropriate `notify`
            // implementation. This is where a large bit of the unsafety
            // starts to stem from internally. The `notify` instance itself
            // is basically just our `Arc<Node<T>>` and tracks the mpsc
            // queue of ready futures.
            //
            // Critically though `Node<T>` won't actually access `T`, the
            // future, while it's floating around inside of `Task`
            // instances. These structs will basically just use `T` to size
            // the internal allocation, appropriately accessing fields and
            // deallocating the node if need be.
            let res = {
                let local_waker = bomb.node.as_ref().unwrap().local_waker();
                let mut cx = cx.with_waker(&local_waker);

                // Safety: We won't move the future ever again
                let future = unsafe { PinMut::new_unchecked(future) };
                future.poll(&mut cx)
            };

            let ret = match res {
                Poll::Pending => {
                    let node = bomb.node.take().unwrap();
                    bomb.queue.link(node);
                    continue
                }
                Poll::Ready(result) => Poll::Ready(Some(result)),
            };
            return ret
        }
    }
}

impl<T: Debug> Debug for FuturesUnordered<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "FuturesUnordered {{ ... }}")
    }
}

impl<T> Drop for FuturesUnordered<T> {
    fn drop(&mut self) {
        // When a `FuturesUnordered` is dropped we want to drop all futures associated
        // with it. At the same time though there may be tons of `Task` handles
        // flying around which contain `Node<T>` references inside them. We'll
        // let those naturally get deallocated when the `Task` itself goes out
        // of scope or gets notified.
        unsafe {
            while !self.head_all.is_null() {
                let head = self.head_all;
                let node = self.unlink(head);
                self.release_node(node);
            }
        }

        // Note that at this point we could still have a bunch of nodes in the
        // mpsc queue. None of those nodes, however, have futures associated
        // with them so they're safe to destroy on any thread. At this point
        // the `FuturesUnordered` struct, the owner of the one strong reference
        // to `Inner<T>` will drop the strong reference. At that point
        // whichever thread releases the strong refcount last (be it this
        // thread or some other thread as part of an `upgrade`) will clear out
        // the mpsc queue and free all remaining nodes.
        //
        // While that freeing operation isn't guaranteed to happen here, it's
        // guaranteed to happen "promptly" as no more "blocking work" will
        // happen while there's a strong refcount held.
    }
}

impl<F: Future> FromIterator<F> for FuturesUnordered<F> {
    fn from_iter<T>(iter: T) -> Self
    where
        T: IntoIterator<Item = F>,
    {
        let acc = FuturesUnordered::new();
        iter.into_iter().fold(acc, |mut acc, item| { acc.push(item); acc })
    }
}

/// Converts a list of futures into a `Stream` of results from the futures.
///
/// This function will take an list of futures (e.g. a vector, an iterator,
/// etc), and return a stream. The stream will yield items as they become
/// available on the futures internally, in the order that they become
/// available. This function is similar to `buffer_unordered` in that it may
/// return items in a different order than in the list specified.
///
/// Note that the returned set can also be used to dynamically push more
/// futures into the set as they become available.
pub fn futures_unordered<I>(futures: I) -> FuturesUnordered<I::Item>
where
    I: IntoIterator,
    I::Item: Future,
{
    futures.into_iter().collect()
}
<commit_msg>Mostly updates to comments<commit_after>//! An unbounded set of futures.

use std::cell::UnsafeCell;
use std::fmt::{self, Debug};
use std::iter::FromIterator;
use std::marker::{PhantomData, Unpin};
use std::mem::{self, PinMut};
use std::ptr;
use std::sync::atomic::Ordering::SeqCst;
use std::sync::atomic::{AtomicPtr, AtomicBool};
use std::sync::{Arc, Weak};
use std::usize;

use futures_core::{Stream, Future, Poll};
use futures_core::task::{self, AtomicWaker};

mod abort;
mod inner;
mod iter_mut;
mod node;

use self::inner::{Inner, Dequeue};
use self::iter_mut::IterMut;
use self::node::Node;

/// A set of `Future`s which may complete in any order.
///
/// This structure is optimized to manage a large number of futures.
/// Futures managed by `FuturesUnordered` will only be polled when they
/// generate notifications. This reduces the required amount of work needed to
/// poll large numbers of futures.
///
/// `FuturesUnordered` can be filled by `collect`ing an iterator of `Future`s
/// into a `FuturesUnordered`, or by `push`ing `Future`s onto an existing
/// `FuturesUnordered`. When new `Future`s are added, `poll_next` must be
/// called in order to begin receiving wakeups for new `Future`s.
///
/// Note that you can create a ready-made `FuturesUnordered` via the
/// `futures_unordered` function in the `stream` module, or you can start with
/// an empty set with the `FuturesUnordered::new` constructor.
#[must_use = "streams do nothing unless polled"]
pub struct FuturesUnordered<F> {
    inner: Arc<Inner<F>>,
    len: usize,
    head_all: *const Node<F>,
}

unsafe impl<T: Send> Send for FuturesUnordered<T> {}
unsafe impl<T: Sync> Sync for FuturesUnordered<T> {}
impl<T> Unpin for FuturesUnordered<T> {}

// FuturesUnordered is implemented using two linked lists. One which links all
// futures managed by a `FuturesUnordered` and one that tracks futures that have
// been scheduled for polling. The first linked list is not thread safe and is
// only accessed by the thread that owns the `FuturesUnordered` value. The
// second linked list is an implementation of the intrusive MPSC queue algorithm
// described by 1024cores.net.
//
// When a future is submitted to the set a node is allocated and inserted in
// both linked lists. The next call to `poll_next` will (eventually) see this
// node and call `poll` on the future.
//
// Before a managed future is polled, the current task's `Waker` is replaced
// with one that is aware of the specific future being run. This ensures that
// task notifications generated by that specific future are visible to
// `FuturesUnordered`. When a notification is received, the node is scheduled
// for polling by being inserted into the concurrent linked list.
//
// Each node is wrapped in an `Arc` and thereby atomically reference counted.
// Also, each node contains an `AtomicBool` which acts as a flag that indicates
// whether the node is currently inserted in the atomic queue. When the future
// is notified, it will only insert itself into the linked list if it isn't
// currently inserted.

impl<T> FuturesUnordered<T>
    where T: Future
{
    /// Constructs a new, empty `FuturesUnordered`
    ///
    /// The returned `FuturesUnordered` does not contain any futures.
    /// In this state, `FuturesUnordered::poll_next` will return
    /// `Poll::Ready(None)`.
    pub fn new() -> FuturesUnordered<T> {
        let stub = Arc::new(Node {
            future: UnsafeCell::new(None),
            next_all: UnsafeCell::new(ptr::null()),
            prev_all: UnsafeCell::new(ptr::null()),
            next_readiness: AtomicPtr::new(ptr::null_mut()),
            queued: AtomicBool::new(true),
            queue: Weak::new(),
        });
        let stub_ptr = &*stub as *const Node<T>;
        let inner = Arc::new(Inner {
            parent: AtomicWaker::new(),
            head_readiness: AtomicPtr::new(stub_ptr as *mut _),
            tail_readiness: UnsafeCell::new(stub_ptr),
            stub: stub,
        });

        FuturesUnordered {
            len: 0,
            head_all: ptr::null_mut(),
            inner: inner,
        }
    }
}

impl<T> FuturesUnordered<T> {
    /// Returns the number of futures contained in the set.
    ///
    /// This represents the total number of in-flight futures.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Returns `true` if the set contains no futures
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Push a future into the set.
    ///
    /// This function submits the given future to the set for managing. This
    /// function will not call `poll` on the submitted future. The caller must
    /// ensure that `FuturesUnordered::poll_next` is called in order to receive
    /// task notifications.
    pub fn push(&mut self, future: T) {
        let node = Arc::new(Node {
            future: UnsafeCell::new(Some(future)),
            next_all: UnsafeCell::new(ptr::null_mut()),
            prev_all: UnsafeCell::new(ptr::null_mut()),
            next_readiness: AtomicPtr::new(ptr::null_mut()),
            queued: AtomicBool::new(true),
            queue: Arc::downgrade(&self.inner),
        });

        // Right now our node has a strong reference count of 1. We transfer
        // ownership of this reference count to our internal linked list
        // and we'll reclaim ownership through the `unlink` function below.
        let ptr = self.link(node);

        // We'll need to get the future "into the system" to start tracking it,
        // e.g. getting its unpark notifications going to us tracking which
        // futures are ready. To do that we unconditionally enqueue it for
        // polling here.
        self.inner.enqueue(ptr);
    }

    /// Returns an iterator that allows modifying each future in the set.
    pub fn iter_mut(&mut self) -> IterMut<T> {
        IterMut {
            node: self.head_all,
            len: self.len,
            _marker: PhantomData
        }
    }

    fn release_node(&mut self, node: Arc<Node<T>>) {
        // The future is done, try to reset the queued flag. This will prevent
        // `notify` from doing any work in the future
        let prev = node.queued.swap(true, SeqCst);

        // Drop the future, even if it hasn't finished yet. This is safe
        // because we're dropping the future on the thread that owns
        // `FuturesUnordered`, which correctly tracks T's lifetimes and such.
        unsafe {
            drop((*node.future.get()).take());
        }

        // If the queued flag was previously set then it means that this node
        // is still in our internal mpsc queue. We then transfer ownership
        // of our reference count to the mpsc queue, and it'll come along and
        // free it later, noticing that the future is `None`.
        //
        // If, however, the queued flag was *not* set then we're safe to
        // release our reference count on the internal node. The queued flag
        // was set above so all future `enqueue` operations will not actually
        // enqueue the node, so our node will never see the mpsc queue again.
        // The node itself will be deallocated once all reference counts have
        // been dropped by the various owning tasks elsewhere.
        if prev {
            mem::forget(node);
        }
    }

    /// Insert a new node into the internal linked list.
    fn link(&mut self, node: Arc<Node<T>>) -> *const Node<T> {
        let ptr = Arc::into_raw(node);
        unsafe {
            *(*ptr).next_all.get() = self.head_all;
            if !self.head_all.is_null() {
                *(*self.head_all).prev_all.get() = ptr;
            }
        }

        self.head_all = ptr;
        self.len += 1;
        ptr
    }

    /// Remove the node from the linked list tracking all nodes currently
    /// managed by `FuturesUnordered`.
    /// This function is unsafe because it has be guaranteed that `node` is a
    /// valid pointer.
    unsafe fn unlink(&mut self, node: *const Node<T>) -> Arc<Node<T>> {
        let node = Arc::from_raw(node);
        let next = *node.next_all.get();
        let prev = *node.prev_all.get();
        *node.next_all.get() = ptr::null_mut();
        *node.prev_all.get() = ptr::null_mut();

        if !next.is_null() {
            *(*next).prev_all.get() = prev;
        }

        if !prev.is_null() {
            *(*prev).next_all.get() = next;
        } else {
            self.head_all = next;
        }
        self.len -= 1;
        node
    }
}

impl<T> Stream for FuturesUnordered<T>
    where T: Future
{
    type Item = T::Output;

    fn poll_next(mut self: PinMut<Self>, cx: &mut task::Context)
        -> Poll<Option<Self::Item>>
    {
        // Ensure `parent` is correctly set.
        self.inner.parent.register(cx.waker());

        loop {
            // Safety: &mut self guarantees the mutual exclusion `dequeue`
            // expects
            let node = match unsafe { self.inner.dequeue() } {
                Dequeue::Empty => {
                    if self.is_empty() {
                        return Poll::Ready(None);
                    } else {
                        return Poll::Pending;
                    }
                }
                Dequeue::Inconsistent => {
                    // At this point, it may be worth yielding the thread &
                    // spinning a few times... but for now, just yield using the
                    // task system.
                    cx.waker().wake();
                    return Poll::Pending;
                }
                Dequeue::Data(node) => node,
            };

            debug_assert!(node != self.inner.stub());

            // Safety:
            // - Node is a valid pointer.
            // - We are the only thread that accesses the `UnsafeCell` that
            //   contains the future
            let future = match unsafe { &mut *(*node).future.get() } {
                Some(future) => future,

                // If the future has already gone away then we're just
                // cleaning out this node. See the comment in
                // `release_node` for more information, but we're basically
                // just taking ownership of our reference count here.
                None => {
                    // Saftey: `release_node` uses `mem::forget`
                    let node = unsafe { Arc::from_raw(node) };
                    assert!(node.next_all.get().is_null());
                    assert!(node.prev_all.get().is_null());
                    continue
                }
            };

            // Safety: `node` is a valid pointer
            let node = unsafe { self.unlink(node) };

            // Unset queued flag... this must be done before
            // polling. This ensures that the future gets
            // rescheduled if it is notified **during** a call
            // to `poll`.
            let prev = node.queued.swap(false, SeqCst);
            assert!(prev);

            let local_waker = node.local_waker();

            // We're going to need to be very careful if the `poll`
            // function below panics. We need to (a) not leak memory and
            // (b) ensure that we still don't have any use-after-frees. To
            // manage this we do a few things:
            //
            // * A "bomb" is created which if dropped abnormally will call
            //   `release_node`. That way we'll be sure the memory management
            //   of the `node` is managed correctly. In particular
            //   `release_node` will drop the future. This ensures that it is
            //   dropped on this thread and not accidentally on a different
            //   thread (bad).
            // * We unlink the node from our internal queue to preemptively
            //   assume it'll panic, in which case we'll want to discard it
            //   regardless.
            struct Bomb<'a, T: 'a> {
                queue: &'a mut FuturesUnordered<T>,
                node: Option<Arc<Node<T>>>,
            }

            impl<'a, T> Drop for Bomb<'a, T> {
                fn drop(&mut self) {
                    if let Some(node) = self.node.take() {
                        self.queue.release_node(node);
                    }
                }
            }

            let mut bomb = Bomb {
                node: Some(node),
                queue: &mut *self,
            };

            // Poll the underlying future with the appropriate waker
            // implementation. This is where a large bit of the unsafety
            // starts to stem from internally. The waker is basically just
            // our `Arc<Node<T>>` and tracks the mpsc queue of ready futures.
            //
            // Critically though `Node<T>` won't actually access `T`, the
            // future, while it's floating around inside of `Task`
            // instances. These structs will basically just use `T` to size
            // the internal allocation, appropriately accessing fields and
            // deallocating the node if need be.

            // Safety: We won't move the future ever again
            let future = unsafe { PinMut::new_unchecked(future) };

            let mut cx = cx.with_waker(&local_waker);
            let res = future.poll(&mut cx);

            let ret = match res {
                Poll::Pending => {
                    let node = bomb.node.take().unwrap();
                    bomb.queue.link(node);
                    continue
                }
                Poll::Ready(result) => Poll::Ready(Some(result)),
            };
            return ret
        }
    }
}

impl<T: Debug> Debug for FuturesUnordered<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "FuturesUnordered {{ ... }}")
    }
}

impl<T> Drop for FuturesUnordered<T> {
    fn drop(&mut self) {
        // When a `FuturesUnordered` is dropped we want to drop all futures
        // associated with it. At the same time though there may be tons of
        // `Task` handles flying around which contain `Node<T>` references
        // inside them. We'll let those naturally get deallocated when the
        // `Task` itself goes out of scope or gets notified.
        unsafe {
            while !self.head_all.is_null() {
                let head = self.head_all;
                let node = self.unlink(head);
                self.release_node(node);
            }
        }

        // Note that at this point we could still have a bunch of nodes in the
        // mpsc queue. None of those nodes, however, have futures associated
        // with them so they're safe to destroy on any thread. At this point
        // the `FuturesUnordered` struct, the owner of the one strong reference
        // to `Inner<T>` will drop the strong reference. At that point
        // whichever thread releases the strong refcount last (be it this
        // thread or some other thread as part of an `upgrade`) will clear out
        // the mpsc queue and free all remaining nodes.
        //
        // While that freeing operation isn't guaranteed to happen here, it's
        // guaranteed to happen "promptly" as no more "blocking work" will
        // happen while there's a strong refcount held.
    }
}

impl<F: Future> FromIterator<F> for FuturesUnordered<F> {
    fn from_iter<T>(iter: T) -> Self
    where
        T: IntoIterator<Item = F>,
    {
        let acc = FuturesUnordered::new();
        iter.into_iter().fold(acc, |mut acc, item| { acc.push(item); acc })
    }
}

/// Converts a list of futures into a `Stream` of results from the futures.
///
/// This function will take an list of futures (e.g. a vector, an iterator,
/// etc), and return a stream. The stream will yield items as they become
/// available on the futures internally, in the order that they become
/// available. This function is similar to `buffer_unordered` in that it may
/// return items in a different order than in the list specified.
///
/// Note that the returned set can also be used to dynamically push more
/// futures into the set as they become available.
pub fn futures_unordered<I>(futures: I) -> FuturesUnordered<I::Item>
where
    I: IntoIterator,
    I::Item: Future,
{
    futures.into_iter().collect()
}
<|endoftext|>


<commit_before>//! Low-level bindings for the c-ares library.
//!
//! In most cases this crate should not be used directly.  The c-ares crate
//! provides a safe wrapper and should be preferred wherever possible.

#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
extern crate libc;
extern crate c_types;

#[cfg(windows)]
extern crate winapi;

mod ffi;
mod constants;

pub use constants::*;
pub use ffi::*;
<commit_msg>Remove unnecessary annotations<commit_after>//! Low-level bindings for the c-ares library.
//!
//! In most cases this crate should not be used directly.  The c-ares crate
//! provides a safe wrapper and should be preferred wherever possible.

extern crate libc;
extern crate c_types;

#[cfg(windows)]
extern crate winapi;

mod ffi;
mod constants;

pub use constants::*;
pub use ffi::*;
<|endoftext|>


<commit_before>use decoder::{Decoder, MAX_COMPONENTS};
use error::{Error, Result, UnsupportedFeature};
use huffman::{fill_default_mjpeg_tables, HuffmanDecoder, HuffmanTable};
use marker::Marker;
use parser::Predictor;
use parser::{
    parse_app, parse_com, parse_dht, parse_dqt, parse_dri, parse_sof, parse_sos,
    AdobeColorTransform, AppData, CodingProcess, Component, Dimensions, EntropyCoding, FrameInfo,
    IccChunk, ScanInfo,
};
use std::io::Read;

impl<R: Read> Decoder<R> {
    /// decode_scan_lossless
    pub fn decode_scan_lossless(
        &mut self,
        frame: &FrameInfo,
        scan: &ScanInfo,
    ) -> Result<(Option<Marker>, Vec<Vec<u16>>)> {
        let ncomp = scan.component_indices.len();
        let npixel = frame.image_size.height as usize * frame.image_size.width as usize;
        assert!(ncomp <= MAX_COMPONENTS);
        let mut results = vec![vec!(0;npixel); ncomp];

        let components: Vec<Component> = scan
            .component_indices
            .iter()
            .map(|&i| frame.components[i].clone())
            .collect();

        // Verify that all required huffman tables has been set.
        if scan
            .dc_table_indices
            .iter()
            .any(|&i| self.dc_huffman_tables[i].is_none())
        {
            return Err(Error::Format(
                "scan makes use of unset dc huffman table".to_owned(),
            ));
        }

        let mut huffman = HuffmanDecoder::new();
        let reader = &mut self.reader;
        let mut mcus_left_until_restart = self.restart_interval;
        let mut expected_rst_num = 0;
        let mut ra = [0u16; MAX_COMPONENTS];
        let mut rb = [0u16; MAX_COMPONENTS];
        let mut rc = [0u16; MAX_COMPONENTS];

        let width = frame.image_size.width as usize;
        let height = frame.image_size.height as usize;

        let mut differences = vec![Vec::with_capacity(npixel); ncomp];
        for mcu_y in 0..height {
            for mcu_x in 0..width {
                if self.restart_interval > 0 {
                    if mcus_left_until_restart == 0 {
                        match huffman.take_marker(reader)? {
                            Some(Marker::RST(n)) => {
                                if n != expected_rst_num {
                                    return Err(Error::Format(format!(
                                        "found RST{} where RST{} was expected",
                                        n, expected_rst_num
                                    )));
                                }

                                huffman.reset();

                                expected_rst_num = (expected_rst_num + 1) % 8;
                                mcus_left_until_restart = self.restart_interval;
                            }
                            Some(marker) => {
                                return Err(Error::Format(format!(
                                    "found marker {:?} inside scan where RST{} was expected",
                                    marker, expected_rst_num
                                )))
                            }
                            None => {
                                return Err(Error::Format(format!(
                                    "no marker found where RST{} was expected",
                                    expected_rst_num
                                )))
                            }
                        }
                    }

                    mcus_left_until_restart -= 1;
                }

                for (i, _component) in components.iter().enumerate() {
                    let dc_table = self.dc_huffman_tables[scan.dc_table_indices[i]]
                        .as_ref()
                        .unwrap();
                    let value = huffman.decode(reader, dc_table)?;
                    let diff = match value {
                        0 => 0,
                        1..=15 => huffman.receive_extend(reader, value)? as i32,
                        16 => 32768,
                        _ => {
                            // Section F.1.2.1.1
                            // Table F.1
                            return Err(Error::Format(
                                "invalid DC difference magnitude category".to_owned(),
                            ));
                        }
                    };
                    differences[i].push(diff);
                }
            }
        }

        if scan.predictor_selection == Predictor::Ra {
            for (i, _component) in components.iter().enumerate() {
                // calculate the top left pixel
                let diff = differences[i][0];
                let prediction = 1 << (frame.precision - scan.point_transform - 1) as i32;
                let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                let result = result << scan.point_transform;
                results[i][0] = result;

                // calculate leftmost column, using top pixel as predictor
                let mut previous = result;
                for mcu_y in 1..height {
                    let diff = differences[i][mcu_y * width];
                    let prediction = previous as i32;
                    let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                    let result = result << scan.point_transform;
                    results[i][mcu_y * width] = result;
                    previous = result;
                }

                // calculate rows, using left pixel as predictor
                for mcu_x in 1..width {
                    for mcu_y in 0..height {
                        let diff = differences[i][mcu_y * width + mcu_x];
                        let prediction = results[i][mcu_y * width + mcu_x - 1] as i32;
                        let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                        let result = result << scan.point_transform;
                        results[i][mcu_y * width + mcu_x] = result;
                    }
                }

            }
        } else {
            for mcu_y in 0..height {
                for mcu_x in 0..width {
                    for (i, _component) in components.iter().enumerate() {
                        let diff = differences[i][mcu_y * width + mcu_x];

                        // The following lines could be further optimized, e.g. moving the checks
                        // and updates of the previous values into the prediction function or
                        // iterating such that diagonals with mcu_x + mcu_y = const are computed at
                        // the same time to exploit independent predictions in this case
                        if mcu_x > 0 {
                            ra[i] = results[i][mcu_y * frame.image_size.width as usize + mcu_x - 1];
                        }
                        if mcu_y > 0 {
                            rb[i] = results[i][(mcu_y - 1) * frame.image_size.width as usize + mcu_x];
                            if mcu_x > 0 {
                                rc[i] = results[i]
                                    [(mcu_y - 1) * frame.image_size.width as usize + (mcu_x - 1)];
                            }
                        }
                        let prediction = predict(
                            ra[i] as i32,
                            rb[i] as i32,
                            rc[i] as i32,
                            scan.predictor_selection,
                            scan.point_transform,
                            frame.precision,
                            mcu_x,
                            mcu_y,
                            self.restart_interval > 0
                                && mcus_left_until_restart == self.restart_interval - 1,
                        );
                        let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                        results[i][mcu_y * width + mcu_x] = result << scan.point_transform;
                    }
                }
            }
        }

        let mut marker = huffman.take_marker(&mut self.reader)?;
        while let Some(Marker::RST(_)) = marker {
            marker = self.read_marker().ok();
        }
        Ok((marker, results))
    }
}

/// H.1.2.1
fn predict(
    ra: i32,
    rb: i32,
    rc: i32,
    predictor: Predictor,
    point_transform: u8,
    input_precision: u8,
    ix: usize,
    iy: usize,
    restart: bool,
) -> i32 {
    let result = if (ix == 0 && iy == 0) || restart {
        // start of first line or restart
        1 << (input_precision - point_transform - 1)
    } else if iy == 0 {
        // rest of first line
        ra
    } else if ix == 0 {
        // start of other line
        rb
    } else {
        // use predictor Table H.1
        match predictor {
            Predictor::NoPrediction => 0,
            Predictor::Ra => ra,
            Predictor::Rb => rb,
            Predictor::Rc => rc,
            Predictor::RaRbRc1 => ra + rb - rc,
            Predictor::RaRbRc2 => ra + ((rb - rc) >> 1),
            Predictor::RaRbRc3 => rb + ((ra - rc) >> 1),
            Predictor::RaRb => (ra + rb) / 2,
        }
    };
    result
}

pub fn compute_image_lossless(
    frame: &FrameInfo,
    mut data: Vec<Vec<u16>>
) -> Result<Vec<u8>> {
    if data.is_empty() || data.iter().any(Vec::is_empty) {
        return Err(Error::Format("not all components have data".to_owned()));
    }
    let output_size = frame.output_size;
    let components = &frame.components;
    let ncomp = components.len();

    if ncomp == 1 {
        let decoded = convert_to_u8(frame, data.remove(0));
        Ok(decoded)
    } else {
        let mut decoded: Vec<u16> = vec![
            0u16;
            ncomp * output_size.width as usize * output_size.height as usize
        ];
        for (x, chunk) in decoded.chunks_mut(ncomp).enumerate() {
            for (i, (component_data, _)) in data.iter().zip(components.iter()).enumerate() {
                chunk[i] = component_data[x];
            }
        }
        let decoded = convert_to_u8(frame, decoded);
        Ok(decoded)
    }
}

fn convert_to_u8(frame: &FrameInfo, data: Vec<u16>) -> Vec<u8> {
    if frame.precision == 8 {
        data.iter().map(|x| *x as u8).collect()
    } else {
        // we use big endian to conform with PNG
        let out: Vec<[u8; 2]> = data.iter().map(|x| x.to_ne_bytes()).collect();
        out.iter().flatten().map(|x| *x).collect()
    }
}
<commit_msg>better to iterate columns in inner loop<commit_after>use decoder::{Decoder, MAX_COMPONENTS};
use error::{Error, Result, UnsupportedFeature};
use huffman::{fill_default_mjpeg_tables, HuffmanDecoder, HuffmanTable};
use marker::Marker;
use parser::Predictor;
use parser::{
    parse_app, parse_com, parse_dht, parse_dqt, parse_dri, parse_sof, parse_sos,
    AdobeColorTransform, AppData, CodingProcess, Component, Dimensions, EntropyCoding, FrameInfo,
    IccChunk, ScanInfo,
};
use std::io::Read;

impl<R: Read> Decoder<R> {
    /// decode_scan_lossless
    pub fn decode_scan_lossless(
        &mut self,
        frame: &FrameInfo,
        scan: &ScanInfo,
    ) -> Result<(Option<Marker>, Vec<Vec<u16>>)> {
        let ncomp = scan.component_indices.len();
        let npixel = frame.image_size.height as usize * frame.image_size.width as usize;
        assert!(ncomp <= MAX_COMPONENTS);
        let mut results = vec![vec!(0;npixel); ncomp];

        let components: Vec<Component> = scan
            .component_indices
            .iter()
            .map(|&i| frame.components[i].clone())
            .collect();

        // Verify that all required huffman tables has been set.
        if scan
            .dc_table_indices
            .iter()
            .any(|&i| self.dc_huffman_tables[i].is_none())
        {
            return Err(Error::Format(
                "scan makes use of unset dc huffman table".to_owned(),
            ));
        }

        let mut huffman = HuffmanDecoder::new();
        let reader = &mut self.reader;
        let mut mcus_left_until_restart = self.restart_interval;
        let mut expected_rst_num = 0;
        let mut ra = [0u16; MAX_COMPONENTS];
        let mut rb = [0u16; MAX_COMPONENTS];
        let mut rc = [0u16; MAX_COMPONENTS];

        let width = frame.image_size.width as usize;
        let height = frame.image_size.height as usize;

        let mut differences = vec![Vec::with_capacity(npixel); ncomp];
        for mcu_y in 0..height {
            for mcu_x in 0..width {
                if self.restart_interval > 0 {
                    if mcus_left_until_restart == 0 {
                        match huffman.take_marker(reader)? {
                            Some(Marker::RST(n)) => {
                                if n != expected_rst_num {
                                    return Err(Error::Format(format!(
                                        "found RST{} where RST{} was expected",
                                        n, expected_rst_num
                                    )));
                                }

                                huffman.reset();

                                expected_rst_num = (expected_rst_num + 1) % 8;
                                mcus_left_until_restart = self.restart_interval;
                            }
                            Some(marker) => {
                                return Err(Error::Format(format!(
                                    "found marker {:?} inside scan where RST{} was expected",
                                    marker, expected_rst_num
                                )))
                            }
                            None => {
                                return Err(Error::Format(format!(
                                    "no marker found where RST{} was expected",
                                    expected_rst_num
                                )))
                            }
                        }
                    }

                    mcus_left_until_restart -= 1;
                }

                for (i, _component) in components.iter().enumerate() {
                    let dc_table = self.dc_huffman_tables[scan.dc_table_indices[i]]
                        .as_ref()
                        .unwrap();
                    let value = huffman.decode(reader, dc_table)?;
                    let diff = match value {
                        0 => 0,
                        1..=15 => huffman.receive_extend(reader, value)? as i32,
                        16 => 32768,
                        _ => {
                            // Section F.1.2.1.1
                            // Table F.1
                            return Err(Error::Format(
                                "invalid DC difference magnitude category".to_owned(),
                            ));
                        }
                    };
                    differences[i].push(diff);
                }
            }
        }

        if scan.predictor_selection == Predictor::Ra {
            for (i, _component) in components.iter().enumerate() {
                // calculate the top left pixel
                let diff = differences[i][0];
                let prediction = 1 << (frame.precision - scan.point_transform - 1) as i32;
                let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                let result = result << scan.point_transform;
                results[i][0] = result;

                // calculate leftmost column, using top pixel as predictor
                let mut previous = result;
                for mcu_y in 1..height {
                    let diff = differences[i][mcu_y * width];
                    let prediction = previous as i32;
                    let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                    let result = result << scan.point_transform;
                    results[i][mcu_y * width] = result;
                    previous = result;
                }

                // calculate rows, using left pixel as predictor
                for mcu_y in 0..height {
                    for mcu_x in 1..width {
                        let diff = differences[i][mcu_y * width + mcu_x];
                        let prediction = results[i][mcu_y * width + mcu_x - 1] as i32;
                        let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                        let result = result << scan.point_transform;
                        results[i][mcu_y * width + mcu_x] = result;
                    }
                }

            }
        } else {
            for mcu_y in 0..height {
                for mcu_x in 0..width {
                    for (i, _component) in components.iter().enumerate() {
                        let diff = differences[i][mcu_y * width + mcu_x];

                        // The following lines could be further optimized, e.g. moving the checks
                        // and updates of the previous values into the prediction function or
                        // iterating such that diagonals with mcu_x + mcu_y = const are computed at
                        // the same time to exploit independent predictions in this case
                        if mcu_x > 0 {
                            ra[i] = results[i][mcu_y * frame.image_size.width as usize + mcu_x - 1];
                        }
                        if mcu_y > 0 {
                            rb[i] = results[i][(mcu_y - 1) * frame.image_size.width as usize + mcu_x];
                            if mcu_x > 0 {
                                rc[i] = results[i]
                                    [(mcu_y - 1) * frame.image_size.width as usize + (mcu_x - 1)];
                            }
                        }
                        let prediction = predict(
                            ra[i] as i32,
                            rb[i] as i32,
                            rc[i] as i32,
                            scan.predictor_selection,
                            scan.point_transform,
                            frame.precision,
                            mcu_x,
                            mcu_y,
                            self.restart_interval > 0
                                && mcus_left_until_restart == self.restart_interval - 1,
                        );
                        let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                        results[i][mcu_y * width + mcu_x] = result << scan.point_transform;
                    }
                }
            }
        }

        let mut marker = huffman.take_marker(&mut self.reader)?;
        while let Some(Marker::RST(_)) = marker {
            marker = self.read_marker().ok();
        }
        Ok((marker, results))
    }
}

/// H.1.2.1
fn predict(
    ra: i32,
    rb: i32,
    rc: i32,
    predictor: Predictor,
    point_transform: u8,
    input_precision: u8,
    ix: usize,
    iy: usize,
    restart: bool,
) -> i32 {
    let result = if (ix == 0 && iy == 0) || restart {
        // start of first line or restart
        1 << (input_precision - point_transform - 1)
    } else if iy == 0 {
        // rest of first line
        ra
    } else if ix == 0 {
        // start of other line
        rb
    } else {
        // use predictor Table H.1
        match predictor {
            Predictor::NoPrediction => 0,
            Predictor::Ra => ra,
            Predictor::Rb => rb,
            Predictor::Rc => rc,
            Predictor::RaRbRc1 => ra + rb - rc,
            Predictor::RaRbRc2 => ra + ((rb - rc) >> 1),
            Predictor::RaRbRc3 => rb + ((ra - rc) >> 1),
            Predictor::RaRb => (ra + rb) / 2,
        }
    };
    result
}

pub fn compute_image_lossless(
    frame: &FrameInfo,
    mut data: Vec<Vec<u16>>
) -> Result<Vec<u8>> {
    if data.is_empty() || data.iter().any(Vec::is_empty) {
        return Err(Error::Format("not all components have data".to_owned()));
    }
    let output_size = frame.output_size;
    let components = &frame.components;
    let ncomp = components.len();

    if ncomp == 1 {
        let decoded = convert_to_u8(frame, data.remove(0));
        Ok(decoded)
    } else {
        let mut decoded: Vec<u16> = vec![
            0u16;
            ncomp * output_size.width as usize * output_size.height as usize
        ];
        for (x, chunk) in decoded.chunks_mut(ncomp).enumerate() {
            for (i, (component_data, _)) in data.iter().zip(components.iter()).enumerate() {
                chunk[i] = component_data[x];
            }
        }
        let decoded = convert_to_u8(frame, decoded);
        Ok(decoded)
    }
}

fn convert_to_u8(frame: &FrameInfo, data: Vec<u16>) -> Vec<u8> {
    if frame.precision == 8 {
        data.iter().map(|x| *x as u8).collect()
    } else {
        // we use big endian to conform with PNG
        let out: Vec<[u8; 2]> = data.iter().map(|x| x.to_ne_bytes()).collect();
        out.iter().flatten().map(|x| *x).collect()
    }
}
<|endoftext|>