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

use std::collections::HashMap;
use std::io;
use std::path::{Path, PathBuf};
use std::sync::{
    mpsc::{
        channel, sync_channel, Receiver, SendError, Sender, SyncSender, TryRecvError, TrySendError,
    },
    Arc, Mutex,
};
use std::thread::{self, Builder as ThreadBuilder};
use std::time::Duration;

use super::file_watcher::{default_file_watcher_with_no_initial_event, Callbacks};

use hcore::fs::USER_CONFIG_FILE;
use hcore::service::ServiceGroup;
use manager::service::Service;
use manager::service::UserConfigPath;

static LOGKEY: &'static str = "UCW";

// This trait exists to ease the testing of functions that receive a Service. Creating Services
// requires a lot of ceremony, so we work around this with this trait.
pub trait Serviceable {
    fn name(&self) -> &str;
    fn user_config_path(&self) -> &UserConfigPath;
    fn service_group(&self) -> &ServiceGroup;
}

impl Serviceable for Service {
    fn name(&self) -> &str {
        &self.pkg.name
    }

    fn user_config_path(&self) -> &UserConfigPath {
        &self.cfg.user_config_path
    }

    fn service_group(&self) -> &ServiceGroup {
        &self.service_group
    }
}

// WorkerState contains the channels the worker uses to communicate
// with the Watcher.
struct WorkerState {
    // This receiver is used by the watcher to be notified when a
    // worker has events.  The channel is a SyncChannel with buffer
    // size 1, as we are only interested in the fact that there were
    // events, not how many there were.
    have_events: Receiver<()>,
    // This sender is used by the watcher to notify a worker to stop
    // running.  It is an async channel because we never want the
    // UserConfigWatcher to block, even if the receiver end of the
    // channel somehow dies and/or fails to consume the message.
    stop_running: Sender<()>,
    // This receiver is used by the watcher tests to be notified when
    // a worker finished setting up the watcher and is about to
    // starting looping it.
    //
    // Silence the dead code warnings from rustc, because it is only
    // used in tests for synchronization purposes.
    #[allow(dead_code)]
    started_watching: Receiver<()>,
}

type ServiceName = String;
pub struct UserConfigWatcher {
    // We use Arc/Mutex here, because this needs to be shareable
    // across threads so we can remove watchers from futures; the
    // Sender and Receiver members of WorkerState aren't themselves
    // shareable.
    states: Arc<Mutex<HashMap<ServiceName, WorkerState>>>,
}

impl UserConfigWatcher {
    pub fn new() -> Self {
        Self {
            states: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Adds a service to the User Config Watcher, thereby starting a watcher thread.
    pub fn add<T: Serviceable>(&mut self, service: &T) -> io::Result<()> {
        // It isn't possible to use the `or_insert_with` function here because it can't have a
        // return value, which we need to return the error from `Worker::run`.
        let mut states = self.states.lock().expect("states lock was poisoned");
        if states.get(service.name()).is_none() {
            let user_toml_path = match service.user_config_path() {
                &UserConfigPath::Recommended(ref p) => p.join(USER_CONFIG_FILE),
                &UserConfigPath::Deprecated(ref p) => {
                    outputln!(
                        preamble service.service_group(),
                        "Not watching {}, because it is located in deprecated path ({}).",
                        USER_CONFIG_FILE,
                        p.display(),
                    );
                    return Ok(());
                }
            };
            // Establish bi-directional communication with the worker by creating two channels.
            // The sync_channel's buffer size is 1 because we want to use it as a boolean, i.e. we
            // are not interested in the events themselves, but only whether at least one has
            // happened.
            let (events_tx, events_rx) = sync_channel(1);
            let (running_tx, running_rx) = channel();
            let (watching_tx, watching_rx) = sync_channel(1);

            Worker::run(user_toml_path, events_tx, running_rx, watching_tx)?;

            outputln!(preamble service.service_group(), "Watching {}", USER_CONFIG_FILE);

            let state = WorkerState {
                have_events: events_rx,
                stop_running: running_tx,
                started_watching: watching_rx,
            };

            states.insert(service.name().to_owned(), state);
        }

        Ok(())
    }

    /// Removes a service from the User Config Watcher, and sends a message to the watcher thread
    /// to stop running.
    pub fn remove<T: Serviceable>(&mut self, service: &T) -> Result<(), SendError<()>> {
        if let Some(state) = self
            .states
            .lock()
            .expect("states lock was poisoned")
            .remove(service.name())
        {
            state.stop_running.send(())?;
        }

        Ok(())
    }

    /// Checks whether the watcher for the specified service has observed any events.
    ///
    /// This also consumes the events.
    pub fn have_events_for<T: Serviceable>(&self, service: &T) -> bool {
        if let Some(state) = self
            .states
            .lock()
            .expect("states lock was poisoned")
            .get(service.name())
        {
            let rx = &state.have_events;

            match rx.try_recv() {
                Ok(_) => {
                    return true;
                }
                Err(TryRecvError::Empty) => return false,
                Err(TryRecvError::Disconnected) => {
                    debug!("UserConfigWatcher worker has died!");
                    return false;
                }
            }
        }

        false
    }
}

struct UserConfigCallbacks {
    have_events: SyncSender<()>,
}

impl Callbacks for UserConfigCallbacks {
    fn file_appeared(&mut self, _: &Path) {
        self.perform();
    }

    fn file_modified(&mut self, _: &Path) {
        self.perform();
    }

    fn file_disappeared(&mut self, _: &Path) {
        self.perform();
    }
}

impl UserConfigCallbacks {
    fn perform(&self) {
        if let Err(TrySendError::Disconnected(_)) = self.have_events.try_send(()) {
            debug!("Worker could not notify Manager of event");
        }
    }
}

struct Worker;

impl Worker {
    // starts a new thread with the file watcher tracking the service's user-config file
    pub fn run(
        path: PathBuf,
        have_events: SyncSender<()>,
        stop_running: Receiver<()>,
        started_watching: SyncSender<()>,
    ) -> io::Result<()> {
        ThreadBuilder::new()
            .name(format!("user-config-watcher-{}", path.display()))
            .spawn(move || {
                debug!(
                    "UserConfigWatcher({}) worker thread starting",
                    path.display(),
                );
                let callbacks = UserConfigCallbacks {
                    have_events: have_events,
                };
                let mut file_watcher =
                    match default_file_watcher_with_no_initial_event(&path, callbacks) {
                        Ok(w) => w,
                        Err(e) => {
                            outputln!(
                                "UserConfigWatcher({}) could not start notifier, \
                                 ending thread ({})",
                                path.display(),
                                e,
                            );
                            return;
                        }
                    };

                let _ = started_watching.try_send(());

                loop {
                    match stop_running.try_recv() {
                        // As long as the `stop_running` channel is
                        // empty, this branch will execute on every
                        // iteration.
                        Err(TryRecvError::Empty) => {
                            if let Err(e) = file_watcher.single_iteration() {
                                outputln!(
                                    "UserConfigWatcher({}) could not run notifier, \
                                     ending thread ({})",
                                    path.display(),
                                    e,
                                );
                                return;
                            };
                        }

                        // If we receive a message on the channel, we stop.
                        Ok(_) => break,

                        // If the channel is disconnected, we stop as well.
                        Err(TryRecvError::Disconnected) => {
                            debug!(
                                "UserConfigWatcher({}) worker thread failed to receive on channel",
                                path.display(),
                            );
                            break;
                        }
                    }

                    thread::sleep(Duration::from_secs(1));
                }
            })?;

        Ok(())
    }
}

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

    use std::fs::{remove_file, File};
    use std::io::Write;
    use std::str::FromStr;
    use std::thread;
    use std::time::{Duration, Instant};

    use manager::file_watcher::WATCHER_DELAY_MS;

    use tempfile::TempDir;

    #[test]
    fn no_events_at_first() {
        let service = TestService::default();
        let mut ucm = UserConfigWatcher::new();
        ucm.add(&service).expect("adding service");

        assert!(!ucm.have_events_for(&service));
    }

    #[test]
    fn events_present_after_adding_config() {
        let service = TestService::default();
        let mut ucm = UserConfigWatcher::new();
        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));

        File::create(service.user_config_path().get_path().join(USER_CONFIG_FILE))
            .expect("creating file");

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_changing_config() {
        let service = TestService::default();
        let file_path = service.user_config_path().get_path().join(USER_CONFIG_FILE);
        let mut ucm = UserConfigWatcher::new();

        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));
        let mut file = File::create(&file_path).expect("creating file");

        file.write_all(b"42").expect(USER_CONFIG_FILE);

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_removing_config() {
        let service = TestService::default();
        let file_path = service.user_config_path().get_path().join(USER_CONFIG_FILE);
        let mut ucm = UserConfigWatcher::new();

        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));
        File::create(&file_path).expect("creating file");

        // Allow the watcher to notice that a file was created.
        wait_for_events(&ucm, &service);

        remove_file(&file_path).expect("removing file");

        assert!(wait_for_events(&ucm, &service));
    }

    fn wait_for_watcher<T: Serviceable>(ucm: &UserConfigWatcher, service: &T) -> bool {
        let start = Instant::now();
        let timeout = Duration::from_secs(10);

        while start.elapsed() < timeout {
            let states = ucm.states.lock().expect("states lock was poisoned");
            let state = states.get(service.name()).expect("service added");
            match state.started_watching.try_recv() {
                Ok(_) => {
                    println!("Received data on the start_watching channel. Returning true.");
                    return true;
                }
                Err(TryRecvError::Empty) => {
                    println!("Received nothing on the start_watching channel. Returning ().");
                    ()
                }
                Err(TryRecvError::Disconnected) => {
                    println!("The start_watching channel was disconnected. Returning false.");
                    return false;
                }
            }

            thread::sleep(Duration::from_millis(100));
        }

        false
    }

    fn wait_for_events<T: Serviceable>(ucm: &UserConfigWatcher, service: &T) -> bool {
        let start = Instant::now();
        let timeout = Duration::from_millis(WATCHER_DELAY_MS * 2);

        while start.elapsed() < timeout {
            if ucm.have_events_for(service) {
                return true;
            }

            thread::sleep(Duration::from_millis(100));
        }

        false
    }

    struct TestService {
        // It is only used as a way to remove the temporary directory
        // when the test finishes. Not used otherwise.
        #[allow(dead_code)]
        tmp: TempDir,
        name: String,
        user_config_path: UserConfigPath,
        service_group: ServiceGroup,
    }

    impl Serviceable for TestService {
        fn name(&self) -> &str {
            &self.name
        }

        fn user_config_path(&self) -> &UserConfigPath {
            &self.user_config_path
        }

        fn service_group(&self) -> &ServiceGroup {
            &self.service_group
        }
    }

    impl Default for TestService {
        fn default() -> Self {
            let tmp = TempDir::new().expect("creating temp dir");
            let path = UserConfigPath::Recommended(tmp.path().to_path_buf());
            Self {
                tmp: tmp,
                name: String::from("foo"),
                user_config_path: path,
                service_group: ServiceGroup::from_str("foo.bar@yoyodine").unwrap(),
            }
        }
    }
}
<commit_msg>Increase wait timeout for UserConfigWatcher tests<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::io;
use std::path::{Path, PathBuf};
use std::sync::{
    mpsc::{
        channel, sync_channel, Receiver, SendError, Sender, SyncSender, TryRecvError, TrySendError,
    },
    Arc, Mutex,
};
use std::thread::{self, Builder as ThreadBuilder};
use std::time::Duration;

use super::file_watcher::{default_file_watcher_with_no_initial_event, Callbacks};

use hcore::fs::USER_CONFIG_FILE;
use hcore::service::ServiceGroup;
use manager::service::Service;
use manager::service::UserConfigPath;

static LOGKEY: &'static str = "UCW";

// This trait exists to ease the testing of functions that receive a Service. Creating Services
// requires a lot of ceremony, so we work around this with this trait.
pub trait Serviceable {
    fn name(&self) -> &str;
    fn user_config_path(&self) -> &UserConfigPath;
    fn service_group(&self) -> &ServiceGroup;
}

impl Serviceable for Service {
    fn name(&self) -> &str {
        &self.pkg.name
    }

    fn user_config_path(&self) -> &UserConfigPath {
        &self.cfg.user_config_path
    }

    fn service_group(&self) -> &ServiceGroup {
        &self.service_group
    }
}

// WorkerState contains the channels the worker uses to communicate
// with the Watcher.
struct WorkerState {
    // This receiver is used by the watcher to be notified when a
    // worker has events.  The channel is a SyncChannel with buffer
    // size 1, as we are only interested in the fact that there were
    // events, not how many there were.
    have_events: Receiver<()>,
    // This sender is used by the watcher to notify a worker to stop
    // running.  It is an async channel because we never want the
    // UserConfigWatcher to block, even if the receiver end of the
    // channel somehow dies and/or fails to consume the message.
    stop_running: Sender<()>,
    // This receiver is used by the watcher tests to be notified when
    // a worker finished setting up the watcher and is about to
    // starting looping it.
    //
    // Silence the dead code warnings from rustc, because it is only
    // used in tests for synchronization purposes.
    #[allow(dead_code)]
    started_watching: Receiver<()>,
}

type ServiceName = String;
pub struct UserConfigWatcher {
    // We use Arc/Mutex here, because this needs to be shareable
    // across threads so we can remove watchers from futures; the
    // Sender and Receiver members of WorkerState aren't themselves
    // shareable.
    states: Arc<Mutex<HashMap<ServiceName, WorkerState>>>,
}

impl UserConfigWatcher {
    pub fn new() -> Self {
        Self {
            states: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Adds a service to the User Config Watcher, thereby starting a watcher thread.
    pub fn add<T: Serviceable>(&mut self, service: &T) -> io::Result<()> {
        // It isn't possible to use the `or_insert_with` function here because it can't have a
        // return value, which we need to return the error from `Worker::run`.
        let mut states = self.states.lock().expect("states lock was poisoned");
        if states.get(service.name()).is_none() {
            let user_toml_path = match service.user_config_path() {
                &UserConfigPath::Recommended(ref p) => p.join(USER_CONFIG_FILE),
                &UserConfigPath::Deprecated(ref p) => {
                    outputln!(
                        preamble service.service_group(),
                        "Not watching {}, because it is located in deprecated path ({}).",
                        USER_CONFIG_FILE,
                        p.display(),
                    );
                    return Ok(());
                }
            };
            // Establish bi-directional communication with the worker by creating two channels.
            // The sync_channel's buffer size is 1 because we want to use it as a boolean, i.e. we
            // are not interested in the events themselves, but only whether at least one has
            // happened.
            let (events_tx, events_rx) = sync_channel(1);
            let (running_tx, running_rx) = channel();
            let (watching_tx, watching_rx) = sync_channel(1);

            Worker::run(user_toml_path, events_tx, running_rx, watching_tx)?;

            outputln!(preamble service.service_group(), "Watching {}", USER_CONFIG_FILE);

            let state = WorkerState {
                have_events: events_rx,
                stop_running: running_tx,
                started_watching: watching_rx,
            };

            states.insert(service.name().to_owned(), state);
        }

        Ok(())
    }

    /// Removes a service from the User Config Watcher, and sends a message to the watcher thread
    /// to stop running.
    pub fn remove<T: Serviceable>(&mut self, service: &T) -> Result<(), SendError<()>> {
        if let Some(state) = self
            .states
            .lock()
            .expect("states lock was poisoned")
            .remove(service.name())
        {
            state.stop_running.send(())?;
        }

        Ok(())
    }

    /// Checks whether the watcher for the specified service has observed any events.
    ///
    /// This also consumes the events.
    pub fn have_events_for<T: Serviceable>(&self, service: &T) -> bool {
        if let Some(state) = self
            .states
            .lock()
            .expect("states lock was poisoned")
            .get(service.name())
        {
            let rx = &state.have_events;

            match rx.try_recv() {
                Ok(_) => {
                    return true;
                }
                Err(TryRecvError::Empty) => return false,
                Err(TryRecvError::Disconnected) => {
                    debug!("UserConfigWatcher worker has died!");
                    return false;
                }
            }
        }

        false
    }
}

struct UserConfigCallbacks {
    have_events: SyncSender<()>,
}

impl Callbacks for UserConfigCallbacks {
    fn file_appeared(&mut self, _: &Path) {
        self.perform();
    }

    fn file_modified(&mut self, _: &Path) {
        self.perform();
    }

    fn file_disappeared(&mut self, _: &Path) {
        self.perform();
    }
}

impl UserConfigCallbacks {
    fn perform(&self) {
        if let Err(TrySendError::Disconnected(_)) = self.have_events.try_send(()) {
            debug!("Worker could not notify Manager of event");
        }
    }
}

struct Worker;

impl Worker {
    // starts a new thread with the file watcher tracking the service's user-config file
    pub fn run(
        path: PathBuf,
        have_events: SyncSender<()>,
        stop_running: Receiver<()>,
        started_watching: SyncSender<()>,
    ) -> io::Result<()> {
        ThreadBuilder::new()
            .name(format!("user-config-watcher-{}", path.display()))
            .spawn(move || {
                debug!(
                    "UserConfigWatcher({}) worker thread starting",
                    path.display(),
                );
                let callbacks = UserConfigCallbacks {
                    have_events: have_events,
                };
                let mut file_watcher =
                    match default_file_watcher_with_no_initial_event(&path, callbacks) {
                        Ok(w) => w,
                        Err(e) => {
                            outputln!(
                                "UserConfigWatcher({}) could not start notifier, \
                                 ending thread ({})",
                                path.display(),
                                e,
                            );
                            return;
                        }
                    };

                let _ = started_watching.try_send(());

                loop {
                    match stop_running.try_recv() {
                        // As long as the `stop_running` channel is
                        // empty, this branch will execute on every
                        // iteration.
                        Err(TryRecvError::Empty) => {
                            if let Err(e) = file_watcher.single_iteration() {
                                outputln!(
                                    "UserConfigWatcher({}) could not run notifier, \
                                     ending thread ({})",
                                    path.display(),
                                    e,
                                );
                                return;
                            };
                        }

                        // If we receive a message on the channel, we stop.
                        Ok(_) => break,

                        // If the channel is disconnected, we stop as well.
                        Err(TryRecvError::Disconnected) => {
                            debug!(
                                "UserConfigWatcher({}) worker thread failed to receive on channel",
                                path.display(),
                            );
                            break;
                        }
                    }

                    thread::sleep(Duration::from_secs(1));
                }
            })?;

        Ok(())
    }
}

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

    use std::fs::{remove_file, File};
    use std::io::Write;
    use std::str::FromStr;
    use std::thread;
    use std::time::{Duration, Instant};

    use manager::file_watcher::WATCHER_DELAY_MS;

    use tempfile::TempDir;

    #[test]
    fn no_events_at_first() {
        let service = TestService::default();
        let mut ucm = UserConfigWatcher::new();
        ucm.add(&service).expect("adding service");

        assert!(!ucm.have_events_for(&service));
    }

    #[test]
    fn events_present_after_adding_config() {
        let service = TestService::default();
        let mut ucm = UserConfigWatcher::new();
        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));

        File::create(service.user_config_path().get_path().join(USER_CONFIG_FILE))
            .expect("creating file");

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_changing_config() {
        let service = TestService::default();
        let file_path = service.user_config_path().get_path().join(USER_CONFIG_FILE);
        let mut ucm = UserConfigWatcher::new();

        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));
        let mut file = File::create(&file_path).expect("creating file");

        file.write_all(b"42").expect(USER_CONFIG_FILE);

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_removing_config() {
        let service = TestService::default();
        let file_path = service.user_config_path().get_path().join(USER_CONFIG_FILE);
        let mut ucm = UserConfigWatcher::new();

        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));
        File::create(&file_path).expect("creating file");

        // Allow the watcher to notice that a file was created.
        wait_for_events(&ucm, &service);

        remove_file(&file_path).expect("removing file");

        assert!(wait_for_events(&ucm, &service));
    }

    fn wait_for_watcher<T: Serviceable>(ucm: &UserConfigWatcher, service: &T) -> bool {
        let start = Instant::now();
        let timeout = Duration::from_secs(10);

        while start.elapsed() < timeout {
            let states = ucm.states.lock().expect("states lock was poisoned");
            let state = states.get(service.name()).expect("service added");
            match state.started_watching.try_recv() {
                Ok(_) => {
                    println!("Received data on the start_watching channel. Returning true.");
                    return true;
                }
                Err(TryRecvError::Empty) => {
                    println!("Received nothing on the start_watching channel. Returning ().");
                    ()
                }
                Err(TryRecvError::Disconnected) => {
                    println!("The start_watching channel was disconnected. Returning false.");
                    return false;
                }
            }

            thread::sleep(Duration::from_millis(100));
        }

        false
    }

    fn wait_for_events<T: Serviceable>(ucm: &UserConfigWatcher, service: &T) -> bool {
        let start = Instant::now();
        let timeout = Duration::from_millis(WATCHER_DELAY_MS * 5);

        while start.elapsed() < timeout {
            if ucm.have_events_for(service) {
                return true;
            }

            thread::sleep(Duration::from_millis(100));
        }

        false
    }

    struct TestService {
        // It is only used as a way to remove the temporary directory
        // when the test finishes. Not used otherwise.
        #[allow(dead_code)]
        tmp: TempDir,
        name: String,
        user_config_path: UserConfigPath,
        service_group: ServiceGroup,
    }

    impl Serviceable for TestService {
        fn name(&self) -> &str {
            &self.name
        }

        fn user_config_path(&self) -> &UserConfigPath {
            &self.user_config_path
        }

        fn service_group(&self) -> &ServiceGroup {
            &self.service_group
        }
    }

    impl Default for TestService {
        fn default() -> Self {
            let tmp = TempDir::new().expect("creating temp dir");
            let path = UserConfigPath::Recommended(tmp.path().to_path_buf());
            Self {
                tmp: tmp,
                name: String::from("foo"),
                user_config_path: path,
                service_group: ServiceGroup::from_str("foo.bar@yoyodine").unwrap(),
            }
        }
    }
}
<|endoftext|>


<commit_before>use crate::event::{Event, EventArg, EventReceiver};
use crate::item::ItemGroup;
use std::sync::mpsc::SyncSender;
use std::time::Duration;

// sender is a cache of reader
pub struct CachedSender {
    items: Vec<ItemGroup>, // cache
    rx_sender: EventReceiver,
    tx_item: SyncSender<(Event, EventArg)>,
}

const WAIT_TIMEOUT: Duration = Duration::from_millis(10);

impl CachedSender {
    pub fn new(rx_sender: EventReceiver, tx_item: SyncSender<(Event, EventArg)>) -> Self {
        CachedSender {
            items: Vec::new(),
            rx_sender,
            tx_item,
        }
    }

    pub fn run(&mut self) {
        // main loop for sending objects

        // if sender is not running, no need to send the items.
        let mut am_i_runing = false;
        // if the reader stopped, no need to wait for more items.
        let mut reader_stopped = false;
        let mut index = 0;

        loop {
            // try to read a bunch of items first
            if let Ok((ev, arg)) = self.rx_sender.recv_timeout(WAIT_TIMEOUT) {
                match ev {
                    Event::EvReaderStarted => {
                        reader_stopped = false;
                        self.items.clear();
                        index = 0;
                        am_i_runing = true;
                    }

                    Event::EvReaderStopped => {
                        // send the total number that reader read.
                        let total_num: usize = self.items.iter().map(|group| group.len()).sum();
                        let _ = self.tx_item.send((ev, Box::new(total_num)));

                        reader_stopped = true;
                    }

                    Event::EvSenderRestart => {
                        // pass the event to matcher, it includes the query
                        let _ = self.tx_item.send((Event::EvMatcherRestart, arg));

                        if !reader_stopped {
                            // pass the event to matcher
                            let _ = self.tx_item.send((Event::EvReaderStarted, Box::new(true)));
                        }

                        am_i_runing = true;
                        index = 0;
                    }

                    Event::EvReaderNewItem => {
                        self.items.push(
                            *arg.downcast::<ItemGroup>()
                                .expect("sender:EvReaderNewItem: failed to get argument"),
                        );
                    }

                    ev @ Event::EvActAccept | ev @ Event::EvActAbort => {
                        // pass the event to sender
                        let _ = self.tx_item.send((ev, Box::new(true)));

                        // quit the loop
                        break;
                    }

                    _ => {}
                }
            }

            if am_i_runing && index < self.items.len() {
                // try to send a bunch of items:
                if self
                    .tx_item
                    .try_send((Event::EvMatcherNewItem, Box::new(self.items[index].clone())))
                    .is_ok()
                {
                    index += 1;
                }
            } else if am_i_runing && reader_stopped {
                // stop running if reader had stopped
                let _ = self.tx_item.send((Event::EvSenderStopped, Box::new(true)));
                am_i_runing = false;
            }
        }
    }
}
<commit_msg>[sender] remove sender, no longer needed<commit_after><|endoftext|>


<commit_before>//! Serde integration support.

use std::fmt;
use std::marker::PhantomData;

use serde::*;
use serde::de::{Visitor, Error, Unexpected};

fn i64_to_u64<'d, V: Visitor<'d>, E: Error>(v: V, n: i64) -> Result<V::Value, E> {
	if n >= 0 {
		v.visit_u64(n as u64)
	} else {
		Err(E::invalid_value(Unexpected::Signed(n), &v))
	}
}

/// Ignore deserialization errors and revert to default.
pub fn ignore_errors<'d, T: Deserialize<'d> + Default, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
	Ok(T::deserialize(d).ok().unwrap_or_default())
}

/// Deserialize a maybe-string ID into a u64.
pub fn deserialize_id<'d, D: Deserializer<'d>>(d: D) -> Result<u64, D::Error> {
	struct IdVisitor;
	impl<'d> Visitor<'d> for IdVisitor {
		type Value = u64;

		fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
			write!(fmt, "a u64 or parseable string")
		}

		fn visit_i64<E: Error>(self, v: i64) -> Result<u64, E> {
			i64_to_u64(self, v)
		}

		fn visit_u64<E: Error>(self, v: u64) -> Result<u64, E> {
			Ok(v)
		}

		fn visit_str<E: Error>(self, v: &str) -> Result<u64, E> {
			v.parse::<u64>().map_err(|_| E::invalid_value(Unexpected::Str(v), &self))
		}
	}

	d.deserialize_u64(IdVisitor)
}

/// Deserialize a maybe-string discriminator into a u16.
/// Also enforces 0 <= N <= 9999.
#[allow(unused_comparisons)]
pub fn deserialize_discrim<'d, D: Deserializer<'d>>(d: D) -> Result<u16, D::Error> {
	macro_rules! check {
		($self:ident, $v:ident, $wrong:expr) => {
			if $v >= 0 && $v <= 9999 {
				Ok($v as u16)
			} else {
				Err(E::invalid_value($wrong, &$self))
			}
		}
	}

	struct DiscrimVisitor;
	impl<'d> Visitor<'d> for DiscrimVisitor {
		type Value = u16;

		fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
			write!(fmt, "a u16 in [0, 9999] or parseable string")
		}

		fn visit_i64<E: Error>(self, v: i64) -> Result<u16, E> {
			check!(self, v, Unexpected::Signed(v))
		}

		fn visit_u64<E: Error>(self, v: u64) -> Result<u16, E> {
			check!(self, v, Unexpected::Unsigned(v))
		}

		fn visit_str<E: Error>(self, v: &str) -> Result<u16, E> {
			v.parse::<u16>()
				.map_err(|_| E::invalid_value(Unexpected::Str(v), &self))
				.and_then(|v| self.visit_u16(v))
		}
	}

	d.deserialize_u16(DiscrimVisitor)
}

/// Deserialize a single-field struct like a newtype struct.
macro_rules! serial_single_field {
	($typ:ident as $field:ident: $inner:path) => {
		impl ::serde::Serialize for $typ {
			fn serialize<S: ::serde::ser::Serializer>(&self, s: S) -> ::std::result::Result<S::Ok, S::Error> {
				self.$field.serialize(s)
			}
		}

		impl<'d> ::serde::Deserialize<'d> for $typ {
			fn deserialize<D: ::serde::de::Deserializer<'d>>(d: D) -> ::std::result::Result<$typ, D::Error> {
				<$inner as ::serde::de::Deserialize>::deserialize(d).map(|v| $typ { $field: v })
			}
		}
	}
}

/// Special support for the oddly complex ReactionEmoji.
pub mod reaction_emoji {
	use super::*;
	use model::{ReactionEmoji, EmojiId};

	#[derive(Serialize, Deserialize)]
	struct RawEmoji<'s> {
		name: &'s str,
		id: Option<EmojiId>,
	}

	pub fn serialize<S: Serializer>(v: &ReactionEmoji, s: S) -> Result<S::Ok, S::Error> {
		(match *v {
			ReactionEmoji::Unicode(ref name) => RawEmoji { name: name, id: None },
			ReactionEmoji::Custom { ref name, id } => RawEmoji { name: name, id: Some(id) },
		}).serialize(s)
	}

	pub fn deserialize<'d, D: Deserializer<'d>>(d: D) -> Result<ReactionEmoji, D::Error> {
		Ok(match try!(RawEmoji::deserialize(d)) {
			RawEmoji { name, id: None } => ReactionEmoji::Unicode(name.to_owned()),
			RawEmoji { name, id: Some(id) } => ReactionEmoji::Custom { name: name.to_owned(), id: id },
		})
	}
}

/// Support for named enums.
pub mod named {
	use super::*;

	pub trait NamedEnum: Sized {
		fn name(&self) -> &'static str;
		fn from_name(name: &str) -> Option<Self>;
		fn typename() -> &'static str;
	}

	pub fn serialize<T: NamedEnum, S: Serializer>(v: &T, s: S) -> Result<S::Ok, S::Error> {
		v.name().serialize(s)
	}

	pub fn deserialize<'d, T: NamedEnum, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
		struct NameVisitor<T>(PhantomData<T>);
		impl<'d, T: NamedEnum> Visitor<'d> for NameVisitor<T> {
			type Value = T;

			fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
				write!(fmt, "a valid {} name", T::typename())
			}

			fn visit_str<E: Error>(self, v: &str) -> Result<T, E> {
				T::from_name(v).ok_or_else(|| E::invalid_value(Unexpected::Str(v), &self))
			}
		}

		d.deserialize_string(NameVisitor(PhantomData))
	}
}
macro_rules! serial_names {
	($typ:ident; $($entry:ident, $value:expr;)*) => {
		impl $typ {
			pub fn name(&self) -> &'static str {
				match *self {
					$($typ::$entry => $value,)*
				}
			}

			pub fn from_name(name: &str) -> Option<Self> {
				match name {
					$($value => Some($typ::$entry),)*
					_ => None,
				}
			}
		}

		impl ::serial::named::NamedEnum for $typ {
			fn name(&self) -> &'static str {
				self.name()
			}

			fn from_name(name: &str) -> Option<Self> {
				Self::from_name(name)
			}

			fn typename() -> &'static str {
				stringify!($typ)
			}
		}
	}
}

/// Support for numeric enums.
pub mod numeric {
	use super::*;

	pub trait NumericEnum: Sized {
		fn num(&self) -> u64;
		fn from_num(num: u64) -> Option<Self>;
		fn typename() -> &'static str;
	}

	pub fn serialize<T: NumericEnum, S: Serializer>(v: &T, s: S) -> Result<S::Ok, S::Error> {
		v.num().serialize(s)
	}

	pub fn deserialize<'d, T: NumericEnum, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
		struct NumVisitor<T>(PhantomData<T>);
		impl<'d, T: NumericEnum> Visitor<'d> for NumVisitor<T> {
			type Value = T;

			fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
				write!(fmt, "a valid {} number", T::typename())
			}

			fn visit_i64<E: Error>(self, v: i64) -> Result<T, E> {
				i64_to_u64(self, v)
			}

			fn visit_u64<E: Error>(self, v: u64) -> Result<T, E> {
				T::from_num(v).ok_or_else(|| E::invalid_value(Unexpected::Unsigned(v), &self))
			}
		}

		d.deserialize_string(NumVisitor(PhantomData))
	}
}
macro_rules! serial_numbers {
	($typ:ident; $($entry:ident, $value:expr;)*) => {
		impl $typ {
			pub fn num(&self) -> u64 {
				match *self {
					$($typ::$entry => $value,)*
				}
			}

			pub fn from_num(num: u64) -> Option<Self> {
				match num {
					$($value => Some($typ::$entry),)*
					_ => None,
				}
			}
		}
		impl ::serial::numeric::NumericEnum for $typ {
			fn num(&self) -> u64 {
				self.num()
			}

			fn from_num(num: u64) -> Option<Self> {
				Self::from_num(num)
			}

			fn typename() -> &'static str {
				stringify!($typ)
			}
		}
	}
}

/// Support for using "named" or "numeric" as the default ser/de impl.
macro_rules! serial_use_mapping {
	($typ:ident, $which:ident) => {
		impl ::serde::Serialize for $typ {
			#[inline]
			fn serialize<S: ::serde::ser::Serializer>(&self, s: S) -> ::std::result::Result<S::Ok, S::Error> {
				::serial::$which::serialize(self, s)
			}
		}

		impl<'d> ::serde::Deserialize<'d> for $typ {
			#[inline]
			fn deserialize<D: ::serde::de::Deserializer<'d>>(d: D) -> ::std::result::Result<$typ, D::Error> {
				::serial::$which::deserialize(d)
			}
		}
	}
}
<commit_msg>Fix broken ReactionEmoji deserialization<commit_after>//! Serde integration support.

use std::fmt;
use std::marker::PhantomData;

use serde::*;
use serde::de::{Visitor, Error, Unexpected};

fn i64_to_u64<'d, V: Visitor<'d>, E: Error>(v: V, n: i64) -> Result<V::Value, E> {
	if n >= 0 {
		v.visit_u64(n as u64)
	} else {
		Err(E::invalid_value(Unexpected::Signed(n), &v))
	}
}

/// Ignore deserialization errors and revert to default.
pub fn ignore_errors<'d, T: Deserialize<'d> + Default, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
	Ok(T::deserialize(d).ok().unwrap_or_default())
}

/// Deserialize a maybe-string ID into a u64.
pub fn deserialize_id<'d, D: Deserializer<'d>>(d: D) -> Result<u64, D::Error> {
	struct IdVisitor;
	impl<'d> Visitor<'d> for IdVisitor {
		type Value = u64;

		fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
			write!(fmt, "a u64 or parseable string")
		}

		fn visit_i64<E: Error>(self, v: i64) -> Result<u64, E> {
			i64_to_u64(self, v)
		}

		fn visit_u64<E: Error>(self, v: u64) -> Result<u64, E> {
			Ok(v)
		}

		fn visit_str<E: Error>(self, v: &str) -> Result<u64, E> {
			v.parse::<u64>().map_err(|_| E::invalid_value(Unexpected::Str(v), &self))
		}
	}

	d.deserialize_u64(IdVisitor)
}

/// Deserialize a maybe-string discriminator into a u16.
/// Also enforces 0 <= N <= 9999.
#[allow(unused_comparisons)]
pub fn deserialize_discrim<'d, D: Deserializer<'d>>(d: D) -> Result<u16, D::Error> {
	macro_rules! check {
		($self:ident, $v:ident, $wrong:expr) => {
			if $v >= 0 && $v <= 9999 {
				Ok($v as u16)
			} else {
				Err(E::invalid_value($wrong, &$self))
			}
		}
	}

	struct DiscrimVisitor;
	impl<'d> Visitor<'d> for DiscrimVisitor {
		type Value = u16;

		fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
			write!(fmt, "a u16 in [0, 9999] or parseable string")
		}

		fn visit_i64<E: Error>(self, v: i64) -> Result<u16, E> {
			check!(self, v, Unexpected::Signed(v))
		}

		fn visit_u64<E: Error>(self, v: u64) -> Result<u16, E> {
			check!(self, v, Unexpected::Unsigned(v))
		}

		fn visit_str<E: Error>(self, v: &str) -> Result<u16, E> {
			v.parse::<u16>()
				.map_err(|_| E::invalid_value(Unexpected::Str(v), &self))
				.and_then(|v| self.visit_u16(v))
		}
	}

	d.deserialize_u16(DiscrimVisitor)
}

/// Deserialize a single-field struct like a newtype struct.
macro_rules! serial_single_field {
	($typ:ident as $field:ident: $inner:path) => {
		impl ::serde::Serialize for $typ {
			fn serialize<S: ::serde::ser::Serializer>(&self, s: S) -> ::std::result::Result<S::Ok, S::Error> {
				self.$field.serialize(s)
			}
		}

		impl<'d> ::serde::Deserialize<'d> for $typ {
			fn deserialize<D: ::serde::de::Deserializer<'d>>(d: D) -> ::std::result::Result<$typ, D::Error> {
				<$inner as ::serde::de::Deserialize>::deserialize(d).map(|v| $typ { $field: v })
			}
		}
	}
}

/// Special support for the oddly complex ReactionEmoji.
pub mod reaction_emoji {
	use super::*;
	use model::{ReactionEmoji, EmojiId};

	#[derive(Serialize)]
	struct EmojiSer<'s> {
		name: &'s str,
		id: Option<EmojiId>
	}

	#[derive(Deserialize)]
	struct EmojiDe {
		name: String,
		id: Option<EmojiId>,
	}

	pub fn serialize<S: Serializer>(v: &ReactionEmoji, s: S) -> Result<S::Ok, S::Error> {
		(match *v {
			ReactionEmoji::Unicode(ref name) => EmojiSer { name: name, id: None },
			ReactionEmoji::Custom { ref name, id } => EmojiSer { name: name, id: Some(id) },
		}).serialize(s)
	}

	pub fn deserialize<'d, D: Deserializer<'d>>(d: D) -> Result<ReactionEmoji, D::Error> {
		Ok(match try!(EmojiDe::deserialize(d)) {
			EmojiDe { name, id: None } => ReactionEmoji::Unicode(name),
			EmojiDe { name, id: Some(id) } => ReactionEmoji::Custom { name: name, id: id },
		})
	}
}

/// Support for named enums.
pub mod named {
	use super::*;

	pub trait NamedEnum: Sized {
		fn name(&self) -> &'static str;
		fn from_name(name: &str) -> Option<Self>;
		fn typename() -> &'static str;
	}

	pub fn serialize<T: NamedEnum, S: Serializer>(v: &T, s: S) -> Result<S::Ok, S::Error> {
		v.name().serialize(s)
	}

	pub fn deserialize<'d, T: NamedEnum, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
		struct NameVisitor<T>(PhantomData<T>);
		impl<'d, T: NamedEnum> Visitor<'d> for NameVisitor<T> {
			type Value = T;

			fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
				write!(fmt, "a valid {} name", T::typename())
			}

			fn visit_str<E: Error>(self, v: &str) -> Result<T, E> {
				T::from_name(v).ok_or_else(|| E::invalid_value(Unexpected::Str(v), &self))
			}
		}

		d.deserialize_string(NameVisitor(PhantomData))
	}
}
macro_rules! serial_names {
	($typ:ident; $($entry:ident, $value:expr;)*) => {
		impl $typ {
			pub fn name(&self) -> &'static str {
				match *self {
					$($typ::$entry => $value,)*
				}
			}

			pub fn from_name(name: &str) -> Option<Self> {
				match name {
					$($value => Some($typ::$entry),)*
					_ => None,
				}
			}
		}

		impl ::serial::named::NamedEnum for $typ {
			fn name(&self) -> &'static str {
				self.name()
			}

			fn from_name(name: &str) -> Option<Self> {
				Self::from_name(name)
			}

			fn typename() -> &'static str {
				stringify!($typ)
			}
		}
	}
}

/// Support for numeric enums.
pub mod numeric {
	use super::*;

	pub trait NumericEnum: Sized {
		fn num(&self) -> u64;
		fn from_num(num: u64) -> Option<Self>;
		fn typename() -> &'static str;
	}

	pub fn serialize<T: NumericEnum, S: Serializer>(v: &T, s: S) -> Result<S::Ok, S::Error> {
		v.num().serialize(s)
	}

	pub fn deserialize<'d, T: NumericEnum, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
		struct NumVisitor<T>(PhantomData<T>);
		impl<'d, T: NumericEnum> Visitor<'d> for NumVisitor<T> {
			type Value = T;

			fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
				write!(fmt, "a valid {} number", T::typename())
			}

			fn visit_i64<E: Error>(self, v: i64) -> Result<T, E> {
				i64_to_u64(self, v)
			}

			fn visit_u64<E: Error>(self, v: u64) -> Result<T, E> {
				T::from_num(v).ok_or_else(|| E::invalid_value(Unexpected::Unsigned(v), &self))
			}
		}

		d.deserialize_string(NumVisitor(PhantomData))
	}
}
macro_rules! serial_numbers {
	($typ:ident; $($entry:ident, $value:expr;)*) => {
		impl $typ {
			pub fn num(&self) -> u64 {
				match *self {
					$($typ::$entry => $value,)*
				}
			}

			pub fn from_num(num: u64) -> Option<Self> {
				match num {
					$($value => Some($typ::$entry),)*
					_ => None,
				}
			}
		}
		impl ::serial::numeric::NumericEnum for $typ {
			fn num(&self) -> u64 {
				self.num()
			}

			fn from_num(num: u64) -> Option<Self> {
				Self::from_num(num)
			}

			fn typename() -> &'static str {
				stringify!($typ)
			}
		}
	}
}

/// Support for using "named" or "numeric" as the default ser/de impl.
macro_rules! serial_use_mapping {
	($typ:ident, $which:ident) => {
		impl ::serde::Serialize for $typ {
			#[inline]
			fn serialize<S: ::serde::ser::Serializer>(&self, s: S) -> ::std::result::Result<S::Ok, S::Error> {
				::serial::$which::serialize(self, s)
			}
		}

		impl<'d> ::serde::Deserialize<'d> for $typ {
			#[inline]
			fn deserialize<D: ::serde::de::Deserializer<'d>>(d: D) -> ::std::result::Result<$typ, D::Error> {
				::serial::$which::deserialize(d)
			}
		}
	}
}
<|endoftext|>


<commit_before>extern crate portaudio;

static SECONDS: uint = 1;

fn main()
{
    portaudio::pa::initialize().unwrap();
    print_devs();
    println!("{}", demo());
    portaudio::pa::terminate().unwrap();
}

fn print_devs()
{
    for i in range(0, portaudio::device::get_count().unwrap())
    {
        match portaudio::device::get_info(i)
        {
            None => {},
            Some(info) => println!("{}: {}", i, info.name),
        }
    }
}

fn demo() -> portaudio::pa::PaResult
{
    let stream = try!(portaudio::stream::Stream::open_default(1, 1, 44100.0, portaudio::stream::FRAMES_PER_BUFFER_UNSPECIFIED, None));

    try!(stream.start());

    let input = try!(stream.read(44100));

    let mut phase = 0.0f32;
    let mut buffer = Vec::with_capacity(44100 * SECONDS);
    for _i in range(0u, 44100 * SECONDS)
    {
        buffer.push(phase);

        phase += 0.007;
        if phase > 1.0 { phase -= 2.0; }
    }

    let mut timer = match std::io::timer::Timer::new()
    {
        Err(e) => { fail!("{}", e); },
        Ok(t) => t,
    };
    let waiter = timer.oneshot(std::time::duration::Duration::seconds(SECONDS as i64));

    match stream.write(buffer.as_slice())
    {
        Err(e) => { println!("write 1: Err({})", e); },
        Ok(()) => {},
    }

    match stream.write(input.as_slice())
    {
        Err(e) => { println!("write 2: Err({})", e); },
        Ok(()) => {},
    }

    waiter.recv();

    Ok(())
}
<commit_msg>fail -> panic<commit_after>extern crate portaudio;

static SECONDS: uint = 1;

fn main()
{
    portaudio::pa::initialize().unwrap();
    print_devs();
    println!("{}", demo());
    portaudio::pa::terminate().unwrap();
}

fn print_devs()
{
    for i in range(0, portaudio::device::get_count().unwrap())
    {
        match portaudio::device::get_info(i)
        {
            None => {},
            Some(info) => println!("{}: {}", i, info.name),
        }
    }
}

fn demo() -> portaudio::pa::PaResult
{
    let stream = try!(portaudio::stream::Stream::open_default(1, 1, 44100.0, portaudio::stream::FRAMES_PER_BUFFER_UNSPECIFIED, None));

    try!(stream.start());

    let input = try!(stream.read(44100));

    let mut phase = 0.0f32;
    let mut buffer = Vec::with_capacity(44100 * SECONDS);
    for _i in range(0u, 44100 * SECONDS)
    {
        buffer.push(phase);

        phase += 0.007;
        if phase > 1.0 { phase -= 2.0; }
    }

    let mut timer = match std::io::timer::Timer::new()
    {
        Err(e) => { panic!("{}", e); },
        Ok(t) => t,
    };
    let waiter = timer.oneshot(std::time::duration::Duration::seconds(SECONDS as i64));

    match stream.write(buffer.as_slice())
    {
        Err(e) => { println!("write 1: Err({})", e); },
        Ok(()) => {},
    }

    match stream.write(input.as_slice())
    {
        Err(e) => { println!("write 2: Err({})", e); },
        Ok(()) => {},
    }

    waiter.recv();

    Ok(())
}
<|endoftext|>


<commit_before>#![feature(rustc_private)]

extern crate rustc_middle;
extern crate rustc_driver;
extern crate rustc_hir;
extern crate rustc_interface;
extern crate rustc_session;

use std::convert::TryFrom;
use std::env;
use std::str::FromStr;

use hex::FromHexError;
use log::debug;

use rustc_session::CtfeBacktrace;
use rustc_driver::Compilation;
use rustc_hir::def_id::LOCAL_CRATE;
use rustc_middle::ty::TyCtxt;

struct MiriCompilerCalls {
    miri_config: miri::MiriConfig,
}

impl rustc_driver::Callbacks for MiriCompilerCalls {
    fn after_analysis<'tcx>(
        &mut self,
        compiler: &rustc_interface::interface::Compiler,
        queries: &'tcx rustc_interface::Queries<'tcx>,
    ) -> Compilation {
        compiler.session().abort_if_errors();

        queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
            init_late_loggers(tcx);
            let (entry_def_id, _) = tcx.entry_fn(LOCAL_CRATE).expect("no main function found!");
            let mut config = self.miri_config.clone();

            // Add filename to `miri` arguments.
            config.args.insert(0, compiler.input().filestem().to_string());

            if let Some(return_code) = miri::eval_main(tcx, entry_def_id.to_def_id(), config) {
                std::process::exit(
                    i32::try_from(return_code).expect("Return value was too large!"),
                );
            }
        });

        compiler.session().abort_if_errors();

        Compilation::Stop
    }
}

fn init_early_loggers() {
    // Note that our `extern crate log` is *not* the same as rustc's; as a result, we have to
    // initialize them both, and we always initialize `miri`'s first.
    let env = env_logger::Env::new().filter("MIRI_LOG").write_style("MIRI_LOG_STYLE");
    env_logger::init_from_env(env);
    // We only initialize `rustc` if the env var is set (so the user asked for it).
    // If it is not set, we avoid initializing now so that we can initialize
    // later with our custom settings, and *not* log anything for what happens before
    // `miri` gets started.
    if env::var_os("RUSTC_LOG").is_some() {
        rustc_driver::init_rustc_env_logger();
    }
}

fn init_late_loggers(tcx: TyCtxt<'_>) {
    // We initialize loggers right before we start evaluation. We overwrite the `RUSTC_LOG`
    // env var if it is not set, control it based on `MIRI_LOG`.
    // (FIXME: use `var_os`, but then we need to manually concatenate instead of `format!`.)
    if let Ok(var) = env::var("MIRI_LOG") {
        if env::var_os("RUSTC_LOG").is_none() {
            // We try to be a bit clever here: if `MIRI_LOG` is just a single level
            // used for everything, we only apply it to the parts of rustc that are
            // CTFE-related. Otherwise, we use it verbatim for `RUSTC_LOG`.
            // This way, if you set `MIRI_LOG=trace`, you get only the right parts of
            // rustc traced, but you can also do `MIRI_LOG=miri=trace,rustc_mir::interpret=debug`.
            if log::Level::from_str(&var).is_ok() {
                env::set_var(
                    "RUSTC_LOG",
                    &format!("rustc_middle::mir::interpret={0},rustc_mir::interpret={0}", var),
                );
            } else {
                env::set_var("RUSTC_LOG", &var);
            }
            rustc_driver::init_rustc_env_logger();
        }
    }

    // If `MIRI_BACKTRACE` is set and `RUSTC_CTFE_BACKTRACE` is not, set `RUSTC_CTFE_BACKTRACE`.
    // Do this late, so we ideally only apply this to Miri's errors.
    if let Some(val) = env::var_os("MIRI_BACKTRACE") {
        let ctfe_backtrace = match &*val.to_string_lossy() {
            "immediate" => CtfeBacktrace::Immediate,
            "0" => CtfeBacktrace::Disabled,
            _ => CtfeBacktrace::Capture,
        };
        *tcx.sess.ctfe_backtrace.borrow_mut() = ctfe_backtrace;
    }
}

/// Returns the "default sysroot" that Miri will use if no `--sysroot` flag is set.
/// Should be a compile-time constant.
fn compile_time_sysroot() -> Option<String> {
    if option_env!("RUSTC_STAGE").is_some() {
        // This is being built as part of rustc, and gets shipped with rustup.
        // We can rely on the sysroot computation in librustc_session.
        return None;
    }
    // For builds outside rustc, we need to ensure that we got a sysroot
    // that gets used as a default.  The sysroot computation in librustc_session would
    // end up somewhere in the build dir (see `get_or_default_sysroot`).
    // Taken from PR <https://github.com/Manishearth/rust-clippy/pull/911>.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    Some(match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => option_env!("RUST_SYSROOT")
            .expect("To build Miri without rustup, set the `RUST_SYSROOT` env var at build time")
            .to_owned(),
    })
}

/// Execute a compiler with the given CLI arguments and callbacks.
fn run_compiler(mut args: Vec<String>, callbacks: &mut (dyn rustc_driver::Callbacks + Send)) -> ! {
    // Make sure we use the right default sysroot. The default sysroot is wrong,
    // because `get_or_default_sysroot` in `librustc_session` bases that on `current_exe`.
    //
    // Make sure we always call `compile_time_sysroot` as that also does some sanity-checks
    // of the environment we were built in.
    // FIXME: Ideally we'd turn a bad build env into a compile-time error via CTFE or so.
    if let Some(sysroot) = compile_time_sysroot() {
        let sysroot_flag = "--sysroot";
        if !args.iter().any(|e| e == sysroot_flag) {
            // We need to overwrite the default that librustc_session would compute.
            args.push(sysroot_flag.to_owned());
            args.push(sysroot);
        }
    }

    // Some options have different defaults in Miri than in plain rustc; apply those by making
    // them the first arguments after the binary name (but later arguments can overwrite them).
    args.splice(1..1, miri::miri_default_args().iter().map(ToString::to_string));

    // Invoke compiler, and handle return code.
    let exit_code = rustc_driver::catch_with_exit_code(move || {
        rustc_driver::run_compiler(&args, callbacks, None, None)
    });
    std::process::exit(exit_code)
}

fn main() {
    rustc_driver::install_ice_hook();

    // If the environment asks us to actually be rustc, then do that.
    if env::var_os("MIRI_BE_RUSTC").is_some() {
        rustc_driver::init_rustc_env_logger();
        // We cannot use `rustc_driver::main` as we need to adjust the CLI arguments.
        let mut callbacks = rustc_driver::TimePassesCallbacks::default();
        run_compiler(env::args().collect(), &mut callbacks)
    }

    // Init loggers the Miri way.
    init_early_loggers();

    // Parse our arguments and split them across `rustc` and `miri`.
    let mut validate = true;
    let mut stacked_borrows = true;
    let mut check_alignment = true;
    let mut communicate = false;
    let mut ignore_leaks = false;
    let mut seed: Option<u64> = None;
    let mut tracked_pointer_tag: Option<miri::PtrId> = None;
    let mut tracked_alloc_id: Option<miri::AllocId> = None;
    let mut rustc_args = vec![];
    let mut crate_args = vec![];
    let mut after_dashdash = false;
    let mut excluded_env_vars = vec![];
    for arg in env::args() {
        if rustc_args.is_empty() {
            // Very first arg: binary name.
            rustc_args.push(arg);
        } else if after_dashdash {
            // Everything that comes after `--` is forwarded to the interpreted crate.
            crate_args.push(arg);
        } else {
            match arg.as_str() {
                "-Zmiri-disable-validation" => {
                    validate = false;
                }
                "-Zmiri-disable-stacked-borrows" => {
                    stacked_borrows = false;
                }
                "-Zmiri-disable-alignment-check" => {
                    check_alignment = false;
                }
                "-Zmiri-disable-isolation" => {
                    communicate = true;
                }
                "-Zmiri-ignore-leaks" => {
                    ignore_leaks = true;
                }
                "--" => {
                    after_dashdash = true;
                }
                arg if arg.starts_with("-Zmiri-seed=") => {
                    if seed.is_some() {
                        panic!("Cannot specify -Zmiri-seed multiple times!");
                    }
                    let seed_raw = hex::decode(arg.trim_start_matches("-Zmiri-seed="))
                        .unwrap_or_else(|err| match err {
                            FromHexError::InvalidHexCharacter { .. } => panic!(
                                "-Zmiri-seed should only contain valid hex digits [0-9a-fA-F]"
                            ),
                            FromHexError::OddLength =>
                                panic!("-Zmiri-seed should have an even number of digits"),
                            err => panic!("Unknown error decoding -Zmiri-seed as hex: {:?}", err),
                        });
                    if seed_raw.len() > 8 {
                        panic!(format!(
                            "-Zmiri-seed must be at most 8 bytes, was {}",
                            seed_raw.len()
                        ));
                    }

                    let mut bytes = [0; 8];
                    bytes[..seed_raw.len()].copy_from_slice(&seed_raw);
                    seed = Some(u64::from_be_bytes(bytes));
                }
                arg if arg.starts_with("-Zmiri-env-exclude=") => {
                    excluded_env_vars
                        .push(arg.trim_start_matches("-Zmiri-env-exclude=").to_owned());
                }
                arg if arg.starts_with("-Zmiri-track-pointer-tag=") => {
                    let id: u64 = match arg.trim_start_matches("-Zmiri-track-pointer-tag=").parse()
                    {
                        Ok(id) => id,
                        Err(err) => panic!(
                            "-Zmiri-track-pointer-tag requires a valid `u64` as the argument: {}",
                            err
                        ),
                    };
                    if let Some(id) = miri::PtrId::new(id) {
                        tracked_pointer_tag = Some(id);
                    } else {
                        panic!("-Zmiri-track-pointer-tag must be a nonzero id");
                    }
                }
                arg if arg.starts_with("-Zmiri-track-alloc-id=") => {
                    let id: u64 = match arg.trim_start_matches("-Zmiri-track-alloc-id=").parse()
                    {
                        Ok(id) => id,
                        Err(err) => panic!(
                            "-Zmiri-track-alloc-id requires a valid `u64` as the argument: {}",
                            err
                        ),
                    };
                    tracked_alloc_id = Some(miri::AllocId(id));
                }
                _ => {
                    // Forward to rustc.
                    rustc_args.push(arg);
                }
            }
        }
    }

    debug!("rustc arguments: {:?}", rustc_args);
    debug!("crate arguments: {:?}", crate_args);
    let miri_config = miri::MiriConfig {
        validate,
        stacked_borrows,
        check_alignment,
        communicate,
        ignore_leaks,
        excluded_env_vars,
        seed,
        args: crate_args,
        tracked_pointer_tag,
        tracked_alloc_id,
    };
    run_compiler(rustc_args, &mut MiriCompilerCalls { miri_config })
}
<commit_msg>Auto merge of #1437 - RalfJung:strip, r=RalfJung<commit_after>#![feature(rustc_private)]
#![feature(str_strip)]

extern crate rustc_middle;
extern crate rustc_driver;
extern crate rustc_hir;
extern crate rustc_interface;
extern crate rustc_session;

use std::convert::TryFrom;
use std::env;
use std::str::FromStr;

use hex::FromHexError;
use log::debug;

use rustc_session::CtfeBacktrace;
use rustc_driver::Compilation;
use rustc_hir::def_id::LOCAL_CRATE;
use rustc_middle::ty::TyCtxt;

struct MiriCompilerCalls {
    miri_config: miri::MiriConfig,
}

impl rustc_driver::Callbacks for MiriCompilerCalls {
    fn after_analysis<'tcx>(
        &mut self,
        compiler: &rustc_interface::interface::Compiler,
        queries: &'tcx rustc_interface::Queries<'tcx>,
    ) -> Compilation {
        compiler.session().abort_if_errors();

        queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
            init_late_loggers(tcx);
            let (entry_def_id, _) = tcx.entry_fn(LOCAL_CRATE).expect("no main function found!");
            let mut config = self.miri_config.clone();

            // Add filename to `miri` arguments.
            config.args.insert(0, compiler.input().filestem().to_string());

            if let Some(return_code) = miri::eval_main(tcx, entry_def_id.to_def_id(), config) {
                std::process::exit(
                    i32::try_from(return_code).expect("Return value was too large!"),
                );
            }
        });

        compiler.session().abort_if_errors();

        Compilation::Stop
    }
}

fn init_early_loggers() {
    // Note that our `extern crate log` is *not* the same as rustc's; as a result, we have to
    // initialize them both, and we always initialize `miri`'s first.
    let env = env_logger::Env::new().filter("MIRI_LOG").write_style("MIRI_LOG_STYLE");
    env_logger::init_from_env(env);
    // We only initialize `rustc` if the env var is set (so the user asked for it).
    // If it is not set, we avoid initializing now so that we can initialize
    // later with our custom settings, and *not* log anything for what happens before
    // `miri` gets started.
    if env::var_os("RUSTC_LOG").is_some() {
        rustc_driver::init_rustc_env_logger();
    }
}

fn init_late_loggers(tcx: TyCtxt<'_>) {
    // We initialize loggers right before we start evaluation. We overwrite the `RUSTC_LOG`
    // env var if it is not set, control it based on `MIRI_LOG`.
    // (FIXME: use `var_os`, but then we need to manually concatenate instead of `format!`.)
    if let Ok(var) = env::var("MIRI_LOG") {
        if env::var_os("RUSTC_LOG").is_none() {
            // We try to be a bit clever here: if `MIRI_LOG` is just a single level
            // used for everything, we only apply it to the parts of rustc that are
            // CTFE-related. Otherwise, we use it verbatim for `RUSTC_LOG`.
            // This way, if you set `MIRI_LOG=trace`, you get only the right parts of
            // rustc traced, but you can also do `MIRI_LOG=miri=trace,rustc_mir::interpret=debug`.
            if log::Level::from_str(&var).is_ok() {
                env::set_var(
                    "RUSTC_LOG",
                    &format!("rustc_middle::mir::interpret={0},rustc_mir::interpret={0}", var),
                );
            } else {
                env::set_var("RUSTC_LOG", &var);
            }
            rustc_driver::init_rustc_env_logger();
        }
    }

    // If `MIRI_BACKTRACE` is set and `RUSTC_CTFE_BACKTRACE` is not, set `RUSTC_CTFE_BACKTRACE`.
    // Do this late, so we ideally only apply this to Miri's errors.
    if let Some(val) = env::var_os("MIRI_BACKTRACE") {
        let ctfe_backtrace = match &*val.to_string_lossy() {
            "immediate" => CtfeBacktrace::Immediate,
            "0" => CtfeBacktrace::Disabled,
            _ => CtfeBacktrace::Capture,
        };
        *tcx.sess.ctfe_backtrace.borrow_mut() = ctfe_backtrace;
    }
}

/// Returns the "default sysroot" that Miri will use if no `--sysroot` flag is set.
/// Should be a compile-time constant.
fn compile_time_sysroot() -> Option<String> {
    if option_env!("RUSTC_STAGE").is_some() {
        // This is being built as part of rustc, and gets shipped with rustup.
        // We can rely on the sysroot computation in librustc_session.
        return None;
    }
    // For builds outside rustc, we need to ensure that we got a sysroot
    // that gets used as a default.  The sysroot computation in librustc_session would
    // end up somewhere in the build dir (see `get_or_default_sysroot`).
    // Taken from PR <https://github.com/Manishearth/rust-clippy/pull/911>.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    Some(match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => option_env!("RUST_SYSROOT")
            .expect("To build Miri without rustup, set the `RUST_SYSROOT` env var at build time")
            .to_owned(),
    })
}

/// Execute a compiler with the given CLI arguments and callbacks.
fn run_compiler(mut args: Vec<String>, callbacks: &mut (dyn rustc_driver::Callbacks + Send)) -> ! {
    // Make sure we use the right default sysroot. The default sysroot is wrong,
    // because `get_or_default_sysroot` in `librustc_session` bases that on `current_exe`.
    //
    // Make sure we always call `compile_time_sysroot` as that also does some sanity-checks
    // of the environment we were built in.
    // FIXME: Ideally we'd turn a bad build env into a compile-time error via CTFE or so.
    if let Some(sysroot) = compile_time_sysroot() {
        let sysroot_flag = "--sysroot";
        if !args.iter().any(|e| e == sysroot_flag) {
            // We need to overwrite the default that librustc_session would compute.
            args.push(sysroot_flag.to_owned());
            args.push(sysroot);
        }
    }

    // Some options have different defaults in Miri than in plain rustc; apply those by making
    // them the first arguments after the binary name (but later arguments can overwrite them).
    args.splice(1..1, miri::miri_default_args().iter().map(ToString::to_string));

    // Invoke compiler, and handle return code.
    let exit_code = rustc_driver::catch_with_exit_code(move || {
        rustc_driver::run_compiler(&args, callbacks, None, None)
    });
    std::process::exit(exit_code)
}

fn main() {
    rustc_driver::install_ice_hook();

    // If the environment asks us to actually be rustc, then do that.
    if env::var_os("MIRI_BE_RUSTC").is_some() {
        rustc_driver::init_rustc_env_logger();
        // We cannot use `rustc_driver::main` as we need to adjust the CLI arguments.
        let mut callbacks = rustc_driver::TimePassesCallbacks::default();
        run_compiler(env::args().collect(), &mut callbacks)
    }

    // Init loggers the Miri way.
    init_early_loggers();

    // Parse our arguments and split them across `rustc` and `miri`.
    let mut validate = true;
    let mut stacked_borrows = true;
    let mut check_alignment = true;
    let mut communicate = false;
    let mut ignore_leaks = false;
    let mut seed: Option<u64> = None;
    let mut tracked_pointer_tag: Option<miri::PtrId> = None;
    let mut tracked_alloc_id: Option<miri::AllocId> = None;
    let mut rustc_args = vec![];
    let mut crate_args = vec![];
    let mut after_dashdash = false;
    let mut excluded_env_vars = vec![];
    for arg in env::args() {
        if rustc_args.is_empty() {
            // Very first arg: binary name.
            rustc_args.push(arg);
        } else if after_dashdash {
            // Everything that comes after `--` is forwarded to the interpreted crate.
            crate_args.push(arg);
        } else {
            match arg.as_str() {
                "-Zmiri-disable-validation" => {
                    validate = false;
                }
                "-Zmiri-disable-stacked-borrows" => {
                    stacked_borrows = false;
                }
                "-Zmiri-disable-alignment-check" => {
                    check_alignment = false;
                }
                "-Zmiri-disable-isolation" => {
                    communicate = true;
                }
                "-Zmiri-ignore-leaks" => {
                    ignore_leaks = true;
                }
                "--" => {
                    after_dashdash = true;
                }
                arg if arg.starts_with("-Zmiri-seed=") => {
                    if seed.is_some() {
                        panic!("Cannot specify -Zmiri-seed multiple times!");
                    }
                    let seed_raw = hex::decode(arg.strip_prefix("-Zmiri-seed=").unwrap())
                        .unwrap_or_else(|err| match err {
                            FromHexError::InvalidHexCharacter { .. } => panic!(
                                "-Zmiri-seed should only contain valid hex digits [0-9a-fA-F]"
                            ),
                            FromHexError::OddLength =>
                                panic!("-Zmiri-seed should have an even number of digits"),
                            err => panic!("Unknown error decoding -Zmiri-seed as hex: {:?}", err),
                        });
                    if seed_raw.len() > 8 {
                        panic!(format!(
                            "-Zmiri-seed must be at most 8 bytes, was {}",
                            seed_raw.len()
                        ));
                    }

                    let mut bytes = [0; 8];
                    bytes[..seed_raw.len()].copy_from_slice(&seed_raw);
                    seed = Some(u64::from_be_bytes(bytes));
                }
                arg if arg.starts_with("-Zmiri-env-exclude=") => {
                    excluded_env_vars
                        .push(arg.strip_prefix("-Zmiri-env-exclude=").unwrap().to_owned());
                }
                arg if arg.starts_with("-Zmiri-track-pointer-tag=") => {
                    let id: u64 = match arg.strip_prefix("-Zmiri-track-pointer-tag=").unwrap().parse()
                    {
                        Ok(id) => id,
                        Err(err) => panic!(
                            "-Zmiri-track-pointer-tag requires a valid `u64` as the argument: {}",
                            err
                        ),
                    };
                    if let Some(id) = miri::PtrId::new(id) {
                        tracked_pointer_tag = Some(id);
                    } else {
                        panic!("-Zmiri-track-pointer-tag must be a nonzero id");
                    }
                }
                arg if arg.starts_with("-Zmiri-track-alloc-id=") => {
                    let id: u64 = match arg.strip_prefix("-Zmiri-track-alloc-id=").unwrap().parse()
                    {
                        Ok(id) => id,
                        Err(err) => panic!(
                            "-Zmiri-track-alloc-id requires a valid `u64` as the argument: {}",
                            err
                        ),
                    };
                    tracked_alloc_id = Some(miri::AllocId(id));
                }
                _ => {
                    // Forward to rustc.
                    rustc_args.push(arg);
                }
            }
        }
    }

    debug!("rustc arguments: {:?}", rustc_args);
    debug!("crate arguments: {:?}", crate_args);
    let miri_config = miri::MiriConfig {
        validate,
        stacked_borrows,
        check_alignment,
        communicate,
        ignore_leaks,
        excluded_env_vars,
        seed,
        args: crate_args,
        tracked_pointer_tag,
        tracked_alloc_id,
    };
    run_compiler(rustc_args, &mut MiriCompilerCalls { miri_config })
}
<|endoftext|>


<commit_before>use std::path::Path;
use std::io::prelude::*;
use std::fs::File;
#[cfg(not(test))] use std::env::home_dir;
use yaml_rust::{YamlLoader, YamlEmitter, Yaml};

static MUXED_FOLDER: &'static str = "muxed";

pub fn open(project_name: String) -> Vec<Yaml> {
    load_yaml(&read(&format!("./{}{}", MUXED_FOLDER, project_name)))
}

#[cfg(not(test))] fn homedir_string() -> String {
    match home_dir() {
        Some(dir) => format!("{}", dir.display()),
        None      => panic!("Impossible to get your home dir!")
    }
}

#[cfg(test)] fn homedir_string() -> String {
  String::from_str("/tmp")
}

fn parse_config() {
}

fn read(config_str: &String) -> String {
    let path = Path::new(config_str);
    let mut s = String::new();
    File::open(path).expect("Config Read error").read_to_string(&mut s);

    return s
}

fn load_yaml(yaml_string: &String) -> Vec<Yaml> {
    return YamlLoader::load_from_str(yaml_string).unwrap();
}

#[test]
pub fn path_returns_muxed_inside_homedir() {
    let path = format!("{}", path().display());
    let new  = format!("{}", Path::new("/tmp/.muxed").display());
    assert_eq!(path, new)
}
<commit_msg>Pass strings instead of Path<commit_after>use std::path::Path;
use std::io::prelude::*;
use std::fs::File;
#[cfg(not(test))] use std::env::home_dir;
use yaml_rust::{YamlLoader, YamlEmitter, Yaml};

static MUXED_FOLDER: &'static str = "muxed";

pub fn open(project_name: String) -> Vec<Yaml> {
    YamlLoader::load_from_str(&read(path_string(project_name))).unwrap()
}

pub fn path_string(project_name: String) -> String {
    format!("{}/.{}/{}", homedir_string(), &MUXED_FOLDER.to_string(), project_name)
}

fn read(config_str: String) -> String {
    let path = Path::new(&config_str);
    let mut s = String::new();
    File::open(path).expect("Config Read error").read_to_string(&mut s);

    return s
}

fn parse_config() {
}

#[cfg(not(test))] fn homedir_string() -> String {
    match home_dir() {
        Some(dir) => format!("{}", dir.display()),
        None      => panic!("Impossible to get your home dir!")
    }
}

#[cfg(test)] fn homedir_string() -> String {
  String::from("/tmp")
}

#[test]
pub fn path_string_returns_muxed_inside_homedir() {
    let path = format!("{}", path_string("".to_string()));
    let new  = format!("{}", Path::new("/tmp/.muxed/").display());
    assert_eq!(path, new)
}
<|endoftext|>


<commit_before>use std::env;
use std::fs::File;
use std::io::{stderr, Read, Write};
use std::process;
use std::str;

fn main() {
    if let Some(url) = env::args().nth(1) {
        let (scheme, reference) = url.split_at(url.find(':').unwrap_or(0));
        if scheme == "http" {
            let mut parts = reference.split('/').skip(2); //skip first two slashes
            let remote = parts.next().unwrap_or("");
            let path = parts.next().unwrap_or("/");

            let mut remote_parts = remote.split(':');
            let host = remote_parts.next().unwrap_or("127.0.0.1");
            let port = remote_parts.next().unwrap_or("80");

            let tcp = format!("tcp:{}:{}", host, port);
            let mut stream = File::open(tcp).unwrap();

            let request = format!("GET {} HTTP/1.0\r\n\r\n", path);
            stream.write(request.as_bytes()).unwrap();
            stream.flush().unwrap();

            let mut bytes = [0; 8192];
            let count = stream.read(&mut bytes).unwrap();
            println!("{}", unsafe { str::from_utf8_unchecked(&bytes[.. count]) });
        } else {
            write!(stderr(), "wget: unknown scheme '{}'\n", scheme).unwrap();
            process::exit(1);
        }
    } else {
        write!(stderr(), "wget: http://host:port/path\n").unwrap();
        process::exit(1);
    }
}
<commit_msg>Parse out headers in wget<commit_after>use std::env;
use std::fs::File;
use std::io::{stderr, Read, Write};
use std::process;
use std::str;

fn main() {
    if let Some(url) = env::args().nth(1) {
        let (scheme, reference) = url.split_at(url.find(':').unwrap_or(0));
        if scheme == "http" {
            let mut parts = reference.split('/').skip(2); //skip first two slashes
            let remote = parts.next().unwrap_or("");
            let path = parts.next().unwrap_or("/");

            let mut remote_parts = remote.split(':');
            let host = remote_parts.next().unwrap_or("127.0.0.1");
            let port = remote_parts.next().unwrap_or("80");

            write!(stderr(), "* Connecting to {}:{}\n", host, port).unwrap();

            let tcp = format!("tcp:{}:{}", host, port);
            let mut stream = File::open(tcp).unwrap();

            write!(stderr(), "* Requesting {}\n", path).unwrap();

            let request = format!("GET {} HTTP/1.0\r\n\r\n", path);
            stream.write(request.as_bytes()).unwrap();
            stream.flush().unwrap();

            write!(stderr(), "* Waiting for response\n").unwrap();

            let mut response = [0; 8192];
            let count = stream.read(&mut response).unwrap();

            let mut headers = true;
            for line in unsafe { str::from_utf8_unchecked(&response[.. count]) }.lines() {
                if headers {
                    if line.is_empty() {
                        headers = false;
                    } else {
                        write!(stderr(), "> {}\n", line).unwrap();
                    }
                } else {
                    println!("{}", line);
                }
            }
        } else {
            write!(stderr(), "wget: unknown scheme '{}'\n", scheme).unwrap();
            process::exit(1);
        }
    } else {
        write!(stderr(), "wget: http://host:port/path\n").unwrap();
        process::exit(1);
    }
}
<|endoftext|>


<commit_before>use std::mem;
use na::{self, Cross, Dot, zero, Norm, clamp};

use ::{Vector, Point2f};
use colour::Colourf;
use intersection::Intersection;
use interaction::SurfaceInteraction;

bitflags! {
    pub flags BxDFType: u32 {
        const REFLECTION   = 0b_0000_0001,
        const TRANSMISSION = 0b_0000_0010,
        const DIFFUSE      = 0b_0000_0100,
        const GLOSSY       = 0b_0000_1000,
        const SPECULAR     = 0b_0001_0000,
    }
}

/// Represents the Bidirectional Scattering Distribution Function.
/// It represents the properties of a material at a given point.
#[derive(Copy, Clone, Debug)]
pub struct BSDF {
    /// Index of refraction of the surface
    eta: f32,
    /// Shading normal (i.e. potentially affected by bump-mapping)
    ns: Vector,
    /// Geometry normal
    ng: Vector,
    ss: Vector,
    ts: Vector, // bxdfs: BxDFType,
}

impl BSDF {
    pub fn new(isect: &Intersection, eta: f32) -> Self {
        let n = isect.dg.nhit;
        let ss = isect.dg.dpdu.normalize();
        BSDF {
            eta: eta,
            ns: n,
            ng: n,
            ss: ss,
            ts: n.cross(&ss),
        }
    }

    pub fn new2(isect: &SurfaceInteraction, eta: f32) -> Self {
        let ss = isect.dpdu.normalize();
        BSDF {
            eta: eta,
            ns: isect.shading.n,
            ng: isect.n,
            ss: ss,
            ts: isect.shading.n.cross(&ss),
        }
    }

    /// Evaluate the BSDF for the given incoming light direction and outgoing light direction.
    pub fn f(&self, _wi_w: &Vector, _wo_w: &Vector) -> Colourf {
        Colourf::black()
    }

    pub fn sample_f(&self,
                    wo_w: &Vector,
                    sample: (f32, f32),
                    flags: BxDFType)
                    -> (Colourf, Vector, f32) {
        if !flags.contains(SPECULAR) {
            unimplemented!();
        }

        if flags.contains(REFLECTION) {
            let wo = self.world_to_local(&wo_w);
            let wi = Vector::new(-wo.x, -wo.y, wo.z);
            let cos_theta_i = wi.z;
            let kr = fr_dielectric(cos_theta_i, 1.0, self.eta);
            let colour = Colourf::rgb(1.0, 1.0, 1.0) * kr / cos_theta_i.abs();

            assert!(!colour.has_nan());
            return (colour, self.local_to_world(&wi), 1.0);
        } else if flags.contains(TRANSMISSION) {
            let wo = self.world_to_local(&wo_w);
            let entering = wo.z > 0.0;
            let (eta_i, eta_t) = if entering {
                (1.0, self.eta)
            } else {
                (self.eta, 1.0)
            };
            let n = if wo.z < 0.0 {
                -Vector::z()
            } else {
                Vector::z()
            };
            return refract(&wo, &n, eta_i / eta_t)
                .map(|wi| {
                    let cos_theta_i = wi.z;
                    let kr = fr_dielectric(cos_theta_i, 1.0, self.eta);
                    let colour = Colourf::rgb(1.0, 1.0, 1.0) * (1.0 - kr) / cos_theta_i.abs();

                    assert!(!colour.has_nan());
                    (colour, self.local_to_world(&wi), 1.0)
                })
                .unwrap_or((Colourf::black(), zero(), 0.0));
        }

        (Colourf::black(), zero(), 0.0)
    }

    fn world_to_local(&self, v: &Vector) -> Vector {
        Vector::new(v.dot(&self.ss), v.dot(&self.ts), v.dot(&self.ns))
    }

    fn local_to_world(&self, v: &Vector) -> Vector {
        Vector::new(self.ss.x * v.x + self.ts.x * v.y + self.ns.x * v.z,
                    self.ss.y * v.x + self.ts.y * v.y + self.ns.y * v.z,
                    self.ss.z * v.z + self.ts.z * v.y + self.ns.z * v.z)
    }
}

// Common geometric functions
#[inline]
fn cos_theta(w: &Vector) -> f32 {
    w.z
}

#[inline]
fn cos2_theta(w: &Vector) -> f32 {
    w.z * w.z
}

#[inline]
fn abs_cos_theta(w: &Vector) -> f32 {
    w.z.abs()
}

#[inline]
fn sin2_theta(w: &Vector) -> f32 {
    (1.0 - cos2_theta(w)).max(0.0)
}

#[inline]
fn sin_theta(w: &Vector) -> f32 {
    sin2_theta(w).sqrt()
}

#[inline]
fn tan_theta(w: &Vector) -> f32 {
    sin_theta(w) / cos_theta(w)
}

#[inline]
fn tan2_theta(w: &Vector) -> f32 {
    sin2_theta(w) / cos2_theta(w)
}

#[inline]
fn cos_phi(w: &Vector) -> f32 {
    let sin_theta = sin_theta(w);
    if sin_theta == 0.0 {
        0.0
    } else {
        na::clamp(w.x / sin_theta, -1.0, 1.0)
    }
}

#[inline]
fn sin_phi(w: &Vector) -> f32 {
    let sin_theta = sin_theta(w);
    if sin_theta == 0.0 {
        0.0
    } else {
        na::clamp(w.y / sin_theta, -1.0, 1.0)
    }
}

#[inline]
fn cos2_phi(w: &Vector) -> f32 {
    cos_phi(w) / cos_phi(w)
}

#[inline]
fn sin2_phi(w: &Vector) -> f32 {
    sin_phi(w) / sin_phi(w)
}

#[inline]
fn cos_d_phi(wa: &Vector, wb: &Vector) -> f32 {
    na::clamp((wa.x * wb.x + wa.y * wa.y) /
              ((wa.x * wa.x + wa.y * wa.y) * (wb.x * wb.x + wb.y * wb.y)).sqrt(),
              -1.0,
              1.0)
}

trait BxDF {
    fn f(&self, wo: &Vector, wi: &Vector) -> Colourf;
    fn sample_f(&self, wo: &Vector, sample: &Point2f) -> (Vector, f32, BxDFType, Colourf);
    fn rho(&self, wo: &Vector, n_samples: u32) -> (Point2f, Colourf);
    fn rho_hh(&self, n_samples: u32) -> (Point2f, Point2f, Colourf);
    fn matches(&self, flags: BxDFType) -> bool;
}

struct ScaledBxDF {
    bxdf: Box<BxDF>,
    scale: Colourf,
}

impl ScaledBxDF {
    fn new(bxdf: Box<BxDF>, scale: Colourf) -> ScaledBxDF {
        ScaledBxDF {
            bxdf: bxdf,
            scale: scale,
        }
    }
}

impl BxDF for ScaledBxDF {
    fn f(&self, wo: &Vector, wi: &Vector) -> Colourf {
        self.bxdf.f(wo, wi) * self.scale
    }
    fn sample_f(&self, wo: &Vector, sample: &Point2f) -> (Vector, f32, BxDFType, Colourf) {
        let (wi, pdf, bxdftype, spectrum) = self.bxdf.sample_f(wo, sample);
        (wi, pdf, bxdftype, spectrum * self.scale)
    }
    fn rho(&self, wo: &Vector, n_samples: u32) -> (Point2f, Colourf) {
        let (sample, spectrum) = self.bxdf.rho(wo, n_samples);
        (sample, spectrum * self.scale)
    }
    fn rho_hh(&self, n_samples: u32) -> (Point2f, Point2f, Colourf) {
        let (sample1, sample2, spectrum) = self.bxdf.rho_hh(n_samples);
        (sample1, sample2, spectrum * self.scale)
    }
    fn matches(&self, flags: BxDFType) -> bool {
        self.bxdf.matches(flags)
    }
}

/// Compute the reflection direction
fn reflect(wo: &Vector, n: &Vector) -> Vector {
    (-(*wo) + *n * 2.0 * wo.dot(n)).normalize()
}

/// Compute the refraction direction
fn refract(i: &Vector, n: &Vector, eta: f32) -> Option<Vector> {
    let cos_theta_i = n.dot(i);
    let sin2theta_i = (1.0 - cos_theta_i * cos_theta_i).max(0.0);
    let sin2theta_t = eta * eta * sin2theta_i;

    if sin2theta_t >= 1.0 {
        None
    } else {
        let cos_theta_t = (1.0 - sin2theta_t).sqrt();
        Some(eta * -*i + (eta * cos_theta_i - cos_theta_t) * *n)
    }
}

/// Compute the Fresnel coefficient for dielectric materials
fn fr_dielectric(cos_theta_i: f32, eta_i: f32, eta_t: f32) -> f32 {
    let mut cos_theta_i = clamp(cos_theta_i, -1.0, 1.0);
    let (mut eta_i, mut eta_t) = (eta_i, eta_t);
    if cos_theta_i <= 0.0 {
        // If leaving the surface, swap the indices of refraction
        mem::swap(&mut eta_i, &mut eta_t);
        cos_theta_i = cos_theta_i.abs();
    }

    let sin_theta_i = (1.0 - cos_theta_i * cos_theta_i).max(0.0).sqrt();
    let sin_theta_t = eta_i / eta_t * sin_theta_i;
    if sin_theta_t >= 1.0 {
        // Total internal reflection
        1.0
    } else {
        let cos_theta_t = (1.0 - sin_theta_t * sin_theta_t).max(0.0).sqrt();
        // Reflectance for parallel polarized light
        let r_parl = ((eta_t * cos_theta_i) - (eta_i * cos_theta_t)) /
                     ((eta_t * cos_theta_i) + (eta_i * cos_theta_t));
        // Reflectance for perpendicular polarized light
        let r_perp = ((eta_i * cos_theta_i) - (eta_t * cos_theta_t)) /
                     ((eta_i * cos_theta_i) + (eta_t * cos_theta_t));
        // Total reflectance for unpolarized light
        0.5 * (r_parl * r_parl + r_perp * r_perp)
    }
}

fn fr_conductor(cos_theta_i: f32, eta_i: &Colourf, eta_t: &Colourf, k: &Colourf) -> Colourf {
    let mut cos_theta_i = clamp(cos_theta_i, -1.0, 1.0);
    let eta = *eta_t / *eta_i;
    let eta_k = *k / *eta_i;

    let cos2_theta_i = cos_theta_i * cos_theta_i;
    let sin2_theta_i = 1.0 - cos2_theta_i;
    let eta2 = eta * eta;
    let eta_k2 = eta_k * eta_k;

    let t0 = eta2 - eta_k2 - sin2_theta_i;
    let a2plusb2 = (t0 * t0 + 4.0 * eta2 * eta_k2).sqrt();
    let t1 = a2plusb2 + cos2_theta_i;
    let a = (0.5 * (a2plusb2 + t0)).sqrt();
    let t2 = 2.0 * cos_theta_i * a;
    let r_s = (t1 - t2) / (t1 + t2);

    let t3 = cos2_theta_i * a2plusb2 + sin2_theta_i * sin2_theta_i;
    let t4 = t2 * sin2_theta_i;
    let r_p = r_s * (t3 - t4) / (t3 + t4);

    0.5 * (r_p + r_s)
}

/// Trait for Fresnel materials
trait Fresnel {
    fn evaluate(&self, cos_theta_i: f32) -> Colourf;
}

/// Fresnel for conductor materials
struct FresnelConductor {
    eta_i: Colourf,
    eta_t: Colourf,
    k: Colourf,
}

impl Fresnel for FresnelConductor {
    fn evaluate(&self, cos_theta_i: f32) -> Colourf {
        fr_conductor(cos_theta_i.abs(), &self.eta_i, &self.eta_t, &self.k)
    }
}

/// Fresnel for dielectric materials
struct FresnelDielectric {
    eta_i: f32,
    eta_t: f32,
}

impl Fresnel for FresnelDielectric {
    fn evaluate(&self, cos_theta_i: f32) -> Colourf {
        Colourf::grey(fr_dielectric(cos_theta_i.abs(), self.eta_i, self.eta_t))
    }
}
<commit_msg>Add BxDF impl for SpecularReflection<commit_after>use std::mem;
use na::{self, Cross, Dot, zero, Norm, clamp};

use ::{Vector, Point2f};
use colour::Colourf;
use intersection::Intersection;
use interaction::SurfaceInteraction;

bitflags! {
    pub flags BxDFType: u32 {
        const REFLECTION   = 0b_0000_0001,
        const TRANSMISSION = 0b_0000_0010,
        const DIFFUSE      = 0b_0000_0100,
        const GLOSSY       = 0b_0000_1000,
        const SPECULAR     = 0b_0001_0000,
    }
}

/// Represents the Bidirectional Scattering Distribution Function.
/// It represents the properties of a material at a given point.
#[derive(Copy, Clone, Debug)]
pub struct BSDF {
    /// Index of refraction of the surface
    eta: f32,
    /// Shading normal (i.e. potentially affected by bump-mapping)
    ns: Vector,
    /// Geometry normal
    ng: Vector,
    ss: Vector,
    ts: Vector, // bxdfs: BxDFType,
}

impl BSDF {
    pub fn new(isect: &Intersection, eta: f32) -> Self {
        let n = isect.dg.nhit;
        let ss = isect.dg.dpdu.normalize();
        BSDF {
            eta: eta,
            ns: n,
            ng: n,
            ss: ss,
            ts: n.cross(&ss),
        }
    }

    pub fn new2(isect: &SurfaceInteraction, eta: f32) -> Self {
        let ss = isect.dpdu.normalize();
        BSDF {
            eta: eta,
            ns: isect.shading.n,
            ng: isect.n,
            ss: ss,
            ts: isect.shading.n.cross(&ss),
        }
    }

    /// Evaluate the BSDF for the given incoming light direction and outgoing light direction.
    pub fn f(&self, _wi_w: &Vector, _wo_w: &Vector) -> Colourf {
        Colourf::black()
    }

    pub fn sample_f(&self,
                    wo_w: &Vector,
                    sample: (f32, f32),
                    flags: BxDFType)
                    -> (Colourf, Vector, f32) {
        if !flags.contains(SPECULAR) {
            unimplemented!();
        }

        if flags.contains(REFLECTION) {
            let wo = self.world_to_local(&wo_w);
            let wi = Vector::new(-wo.x, -wo.y, wo.z);
            let cos_theta_i = wi.z;
            let kr = fr_dielectric(cos_theta_i, 1.0, self.eta);
            let colour = Colourf::rgb(1.0, 1.0, 1.0) * kr / cos_theta_i.abs();

            assert!(!colour.has_nan());
            return (colour, self.local_to_world(&wi), 1.0);
        } else if flags.contains(TRANSMISSION) {
            let wo = self.world_to_local(&wo_w);
            let entering = wo.z > 0.0;
            let (eta_i, eta_t) = if entering {
                (1.0, self.eta)
            } else {
                (self.eta, 1.0)
            };
            let n = if wo.z < 0.0 {
                -Vector::z()
            } else {
                Vector::z()
            };
            return refract(&wo, &n, eta_i / eta_t)
                .map(|wi| {
                    let cos_theta_i = wi.z;
                    let kr = fr_dielectric(cos_theta_i, 1.0, self.eta);
                    let colour = Colourf::rgb(1.0, 1.0, 1.0) * (1.0 - kr) / cos_theta_i.abs();

                    assert!(!colour.has_nan());
                    (colour, self.local_to_world(&wi), 1.0)
                })
                .unwrap_or((Colourf::black(), zero(), 0.0));
        }

        (Colourf::black(), zero(), 0.0)
    }

    fn world_to_local(&self, v: &Vector) -> Vector {
        Vector::new(v.dot(&self.ss), v.dot(&self.ts), v.dot(&self.ns))
    }

    fn local_to_world(&self, v: &Vector) -> Vector {
        Vector::new(self.ss.x * v.x + self.ts.x * v.y + self.ns.x * v.z,
                    self.ss.y * v.x + self.ts.y * v.y + self.ns.y * v.z,
                    self.ss.z * v.z + self.ts.z * v.y + self.ns.z * v.z)
    }
}

// Common geometric functions
#[inline]
fn cos_theta(w: &Vector) -> f32 {
    w.z
}

#[inline]
fn cos2_theta(w: &Vector) -> f32 {
    w.z * w.z
}

#[inline]
fn abs_cos_theta(w: &Vector) -> f32 {
    w.z.abs()
}

#[inline]
fn sin2_theta(w: &Vector) -> f32 {
    (1.0 - cos2_theta(w)).max(0.0)
}

#[inline]
fn sin_theta(w: &Vector) -> f32 {
    sin2_theta(w).sqrt()
}

#[inline]
fn tan_theta(w: &Vector) -> f32 {
    sin_theta(w) / cos_theta(w)
}

#[inline]
fn tan2_theta(w: &Vector) -> f32 {
    sin2_theta(w) / cos2_theta(w)
}

#[inline]
fn cos_phi(w: &Vector) -> f32 {
    let sin_theta = sin_theta(w);
    if sin_theta == 0.0 {
        0.0
    } else {
        na::clamp(w.x / sin_theta, -1.0, 1.0)
    }
}

#[inline]
fn sin_phi(w: &Vector) -> f32 {
    let sin_theta = sin_theta(w);
    if sin_theta == 0.0 {
        0.0
    } else {
        na::clamp(w.y / sin_theta, -1.0, 1.0)
    }
}

#[inline]
fn cos2_phi(w: &Vector) -> f32 {
    cos_phi(w) / cos_phi(w)
}

#[inline]
fn sin2_phi(w: &Vector) -> f32 {
    sin_phi(w) / sin_phi(w)
}

#[inline]
fn cos_d_phi(wa: &Vector, wb: &Vector) -> f32 {
    na::clamp((wa.x * wb.x + wa.y * wa.y) /
              ((wa.x * wa.x + wa.y * wa.y) * (wb.x * wb.x + wb.y * wb.y)).sqrt(),
              -1.0,
              1.0)
}

trait BxDF {
    fn f(&self, wo: &Vector, wi: &Vector) -> Colourf;
    fn sample_f(&self, wo: &Vector, sample: &Point2f) -> (Vector, f32, Option<BxDFType>, Colourf);
    // fn rho(&self, wo: &Vector, n_samples: u32) -> (Point2f, Colourf);
    // fn rho_hh(&self, n_samples: u32) -> (Point2f, Point2f, Colourf);
    fn matches(&self, flags: BxDFType) -> bool;
}

struct ScaledBxDF {
    bxdf: Box<BxDF>,
    scale: Colourf,
}

impl ScaledBxDF {
    fn new(bxdf: Box<BxDF>, scale: Colourf) -> ScaledBxDF {
        ScaledBxDF {
            bxdf: bxdf,
            scale: scale,
        }
    }
}

impl BxDF for ScaledBxDF {
    fn f(&self, wo: &Vector, wi: &Vector) -> Colourf {
        self.bxdf.f(wo, wi) * self.scale
    }
    fn sample_f(&self, wo: &Vector, sample: &Point2f) -> (Vector, f32, Option<BxDFType>, Colourf) {
        let (wi, pdf, bxdftype, spectrum) = self.bxdf.sample_f(wo, sample);
        (wi, pdf, bxdftype, spectrum * self.scale)
    }
    // fn rho(&self, wo: &Vector, n_samples: u32) -> (Point2f, Colourf) {
    //     let (sample, spectrum) = self.bxdf.rho(wo, n_samples);
    //     (sample, spectrum * self.scale)
    // }
    // fn rho_hh(&self, n_samples: u32) -> (Point2f, Point2f, Colourf) {
    //     let (sample1, sample2, spectrum) = self.bxdf.rho_hh(n_samples);
    //     (sample1, sample2, spectrum * self.scale)
    // }
    fn matches(&self, flags: BxDFType) -> bool {
        self.bxdf.matches(flags)
    }
}

/// Compute the reflection direction
fn reflect(wo: &Vector, n: &Vector) -> Vector {
    (-(*wo) + *n * 2.0 * wo.dot(n)).normalize()
}

/// Compute the refraction direction
fn refract(i: &Vector, n: &Vector, eta: f32) -> Option<Vector> {
    let cos_theta_i = n.dot(i);
    let sin2theta_i = (1.0 - cos_theta_i * cos_theta_i).max(0.0);
    let sin2theta_t = eta * eta * sin2theta_i;

    if sin2theta_t >= 1.0 {
        None
    } else {
        let cos_theta_t = (1.0 - sin2theta_t).sqrt();
        Some(eta * -*i + (eta * cos_theta_i - cos_theta_t) * *n)
    }
}

/// Compute the Fresnel coefficient for dielectric materials
fn fr_dielectric(cos_theta_i: f32, eta_i: f32, eta_t: f32) -> f32 {
    let mut cos_theta_i = clamp(cos_theta_i, -1.0, 1.0);
    let (mut eta_i, mut eta_t) = (eta_i, eta_t);
    if cos_theta_i <= 0.0 {
        // If leaving the surface, swap the indices of refraction
        mem::swap(&mut eta_i, &mut eta_t);
        cos_theta_i = cos_theta_i.abs();
    }

    let sin_theta_i = (1.0 - cos_theta_i * cos_theta_i).max(0.0).sqrt();
    let sin_theta_t = eta_i / eta_t * sin_theta_i;
    if sin_theta_t >= 1.0 {
        // Total internal reflection
        1.0
    } else {
        let cos_theta_t = (1.0 - sin_theta_t * sin_theta_t).max(0.0).sqrt();
        // Reflectance for parallel polarized light
        let r_parl = ((eta_t * cos_theta_i) - (eta_i * cos_theta_t)) /
                     ((eta_t * cos_theta_i) + (eta_i * cos_theta_t));
        // Reflectance for perpendicular polarized light
        let r_perp = ((eta_i * cos_theta_i) - (eta_t * cos_theta_t)) /
                     ((eta_i * cos_theta_i) + (eta_t * cos_theta_t));
        // Total reflectance for unpolarized light
        0.5 * (r_parl * r_parl + r_perp * r_perp)
    }
}

fn fr_conductor(cos_theta_i: f32, eta_i: &Colourf, eta_t: &Colourf, k: &Colourf) -> Colourf {
    let mut cos_theta_i = clamp(cos_theta_i, -1.0, 1.0);
    let eta = *eta_t / *eta_i;
    let eta_k = *k / *eta_i;

    let cos2_theta_i = cos_theta_i * cos_theta_i;
    let sin2_theta_i = 1.0 - cos2_theta_i;
    let eta2 = eta * eta;
    let eta_k2 = eta_k * eta_k;

    let t0 = eta2 - eta_k2 - sin2_theta_i;
    let a2plusb2 = (t0 * t0 + 4.0 * eta2 * eta_k2).sqrt();
    let t1 = a2plusb2 + cos2_theta_i;
    let a = (0.5 * (a2plusb2 + t0)).sqrt();
    let t2 = 2.0 * cos_theta_i * a;
    let r_s = (t1 - t2) / (t1 + t2);

    let t3 = cos2_theta_i * a2plusb2 + sin2_theta_i * sin2_theta_i;
    let t4 = t2 * sin2_theta_i;
    let r_p = r_s * (t3 - t4) / (t3 + t4);

    0.5 * (r_p + r_s)
}

/// Trait for Fresnel materials
trait Fresnel {
    fn evaluate(&self, cos_theta_i: f32) -> Colourf;
}

/// Fresnel for conductor materials
struct FresnelConductor {
    eta_i: Colourf,
    eta_t: Colourf,
    k: Colourf,
}

impl Fresnel for FresnelConductor {
    fn evaluate(&self, cos_theta_i: f32) -> Colourf {
        fr_conductor(cos_theta_i.abs(), &self.eta_i, &self.eta_t, &self.k)
    }
}

/// Fresnel for dielectric materials
struct FresnelDielectric {
    eta_i: f32,
    eta_t: f32,
}

impl Fresnel for FresnelDielectric {
    fn evaluate(&self, cos_theta_i: f32) -> Colourf {
        Colourf::grey(fr_dielectric(cos_theta_i.abs(), self.eta_i, self.eta_t))
    }
}

/// BRDF for perfect specular reflection
struct SpecularReflection {
    typ: BxDFType,
    r: Colourf,
    fresnel: Box<Fresnel>,
}

impl SpecularReflection {
    fn new(r: Colourf, fresnel: Box<Fresnel>) -> SpecularReflection {
        SpecularReflection {
            typ: SPECULAR | REFLECTION,
            r: r,
            fresnel: fresnel,
        }
    }
}

impl BxDF for SpecularReflection {
    fn f(&self, wo: &Vector, wi: &Vector) -> Colourf {
        // The probability to call f() with the exact (wo, wi) for specular reflection is 0, so we
        // return black here. Use sample_f() instead.
        Colourf::black()
    }

    fn sample_f(&self, wo: &Vector, _sample: &Point2f) -> (Vector, f32, Option<BxDFType>, Colourf) {
        // There's only one possible wi for a given wo, so we always return it with a pdf of 1.
        let wi = Vector::new(-wo.x, -wo.y, wo.z);
        let spectrum = self.fresnel.evaluate(cos_theta(&wi)) * self.r / abs_cos_theta(&wi);
        (wi, 1.0, None, spectrum)
    }

    fn matches(&self, flags: BxDFType) -> bool {
        self.typ & flags == self.typ
    }
}
<|endoftext|>