<commit_before>
use std::collections::HashMap;
use parse::*;
use clap::ArgMatches;
use intercom_common::utils;
use intercom_common::guid::*;
use error::*;
use std::io;

/// Runs the 'cpp' subcommand.
pub fn run( cpp_params : &ArgMatches ) -> AppResult {

    // Parse the sources and convert the result into an IDL.
    let ( renames, result ) = parse_crate(
            cpp_params.value_of( "path" ).unwrap() )?;
    result_to_cpp( &result, &renames, &mut io::stdout() );

    Ok(())
}

/// Converts the parse result into an header  that gets written to stdout.
fn result_to_cpp(
    r : &ParseResult,
    rn : &HashMap<String, String>,
    out : &mut io::Write,
) {
    // Introduce all interfaces so we don't get errors on undeclared items.
    let itf_forward_declarations = r.interfaces.iter().map(|itf| {
        format!( "{}struct {};", get_indentation( 1 ), try_rename( rn, &itf.name ) )
    } ).collect::<Vec<_>>().join( "\n" );

    // Raw interfaces that give direct access to the components implemented in Rust.
    let itfs = r.interfaces.iter().map(|itf| {
        generate_raw_interface( rn, itf, 1 )
    } ).collect::<Vec<_>>().join( "\n" );

    // Generate class descriptors.
    let class_descriptors = r.class_names.iter().map(|class_name| {

        // Get the class details by matching the name.
        let class  = r.classes.iter().find(|cls| &cls.name == class_name )
            .unwrap();

        generate_class_descriptor( rn, &r.libname, class_name, class, 1 )
    } ).collect::<Vec<_>>().join( "\n" );

    // Generate using statements that flatten the raw interfaces
    // to mimic the results of a Microsoft MIDL compiler.
    let flat_interface_declarations = r.interfaces.iter().map(|itf| {
        flatten_interface( rn, &r.libname, itf, 1 )
    } ).collect::<Vec<_>>().join( "\n" );

    // Generate using statements that flatten the class descriptors
    // to mimic the results of a Microsoft MIDL compiler.
    let flat_class_descriptor_declarations = r.class_names.iter().map(|class_name| {

                // Get the class details by matching the name.
        let class  = r.classes.iter().find(|cls| &cls.name == class_name )
            .unwrap();

        flatten_class_descriptor( class_name, class, 1 )
    } ).collect::<Vec<_>>().join( "\n" );

    // We got the interfaces and classes. We can format and output the raw interfaces.
    let raw_namespace: String = format!( r###"
>       #include <array>
>       #include <intercom.h>
>       namespace {}
>       {{
>       namespace raw
>       {{
>       {}
>       {}
>       {}
>       }}
>       }}
>       #ifdef INTERCOM_FLATTEN_DECLARATIONS
>       {}
>       {}
>       #endif
>       "###,
        &r.libname, itf_forward_declarations, itfs, class_descriptors,
        flat_interface_declarations, flat_class_descriptor_declarations );
    let raw_namespace = raw_namespace.replace( ">       ", "" );
    writeln!( out, "{}", raw_namespace ).unwrap();
}

fn generate_raw_interface(
    rn : &HashMap<String, String>,
    interface: &Interface,
    base_indentation: usize
) -> String {

    // Now that we have methods sorted out, we can construct the final
    // interface definition.
    let indentation = get_indentation( base_indentation );
    let interface_text = format!( "{}struct {} : IUnknown\n{}{{\n{}static constexpr intercom::IID ID = {};\n\n{}\n{}}};\n",
            indentation, try_rename( rn, &interface.name ), indentation,
            get_indentation( base_indentation + 1 ), guid_to_binary( &interface.iid ),
            generate_raw_methods( interface, base_indentation + 1 ), indentation );
    interface_text
}

fn generate_raw_methods(
    itf: &Interface,
    indentation: usize
) -> String {

    // Get the method definitions for the current interface.
    itf.methods.iter().enumerate().map(|(_i,m)| {

        // Construct the argument list.
        let args = m.arguments.iter().map(|a| {

            // Argument direction affects both the argument attribute and
            // whether the argument is passed by pointer or value.
            let out_ptr = match a.dir {
                ArgDirection::In => "",
                ArgDirection::Out => "*",
                ArgDirection::Return => "*",
            };
            format!( "{}{} {}", to_cpp_type( &a.ty ), out_ptr, a.name )

        } ).collect::<Vec<_>>().join( ", " );

        // Format the method.
        // To maintain backwards compatibility all new methods should
        // be added to the end of the traits.
        format!( r###"{}virtual {} INTERCOM_CC {}( {} ) = 0;"###, get_indentation( indentation ), to_cpp_type( &m.rvalue ), utils::pascal_case( &m.name ), args )

    } ).collect::<Vec<_>>().join( "\n" )
}

fn flatten_interface(
    rn : &HashMap<String, String>,
    libname: &str,
    interface: &Interface,
    indentation: usize
) -> String {
    let interface_name = try_rename( rn, &interface.name );
    let using_iid = format!( "static constexpr intercom::IID IID_{} = {};",
            utils::pascal_case( &interface_name ), guid_to_binary( &interface.iid ) );
    let using_interface = format!( "using {} = {}::raw::{};", utils::pascal_case( &interface_name ), libname,  &interface_name );
    let indentation = get_indentation( indentation );
    format!("{}{}\n{}{}", indentation, using_iid, indentation, using_interface )
}

fn flatten_class_descriptor(
    class_name: &str,
    coclass: &CoClass,
    indentation: usize
) -> String {
    let using_clsid = format!( "static constexpr CLSID CLSID_{} = {};",
            utils::pascal_case( class_name ), guid_to_binary( &coclass.clsid ) );
    let indentation = get_indentation( indentation );
    format!("{}{}", indentation, using_clsid )
}

fn generate_class_descriptor(
    rn : &HashMap<String, String>,
    libname: &str,
    class_name: &str,
    coclass: &CoClass,
    indentation: usize
) -> String {

    // Create a list of interfaces to be declared in the class descriptor.
    let interfaces =  coclass.interfaces.iter().map(|itf| {
            format!( r###"{}{}::raw::{}::ID"###,
                &get_indentation( indentation + 2 ), libname, try_rename( rn, itf ) )
        } ).collect::<Vec<_>>().join( ",\n" );

    // Class descriptors hold information about COM classes implemented by the library.
    let clsid = guid_to_binary( &coclass.clsid );
    let class_descriptor: String = format!( r###"
>       class {}Descriptor
>       {{
>           static constexpr CLSID ID = {};

>           static constexpr std::array<intercom::IID, {}> INTERFACES = {{
{}
>           }};

>           {}Descriptor() = delete;
>           ~{}Descriptor() = delete;

>       }};
>       "###,
    class_name, clsid, coclass.interfaces.len(), interfaces, class_name, class_name );
    class_descriptor.replace( ">       ", &get_indentation( indentation ) )
}

/// Gets an empty string for ind One level of indentation is four spaces.
fn get_indentation(
    level: usize
) -> String {
    let spaces = level * 4;
    let indentation = format!( r###"{: <1$}"###, "", spaces );
    indentation
}

/// Converts a Rust type into applicable C++ type.
fn to_cpp_type(
    ty: &str
) -> &str {

    match ty {
        "int8" => "int8_t",
        "uint8" => "uint8_t",
        "int16" => "int16_t",
        "uint16" => "uint16_t",
        "int32" => "int32_t",
        "uint32" => "uint32_t",
        "int64" => "int64_t",
        "uint64" => "uint64_t",
        "BStr" => "intercom::BSTR",
        "HRESULT" => "intercom::HRESULT",
        _ => ty,
    }
}

/// Converts a guid to binarys representation.
fn guid_to_binary(
    g: &GUID
) -> String {

    format!( "{{0x{:08x},0x{:04x},0x{:04x},{{0x{:02x},0x{:02x},0x{:02x},0x{:02x},0x{:02x},0x{:02x},0x{:02x},0x{:02x}}}}}",
            g.data1, g.data2, g.data3,
            g.data4[0], g.data4[1], g.data4[2], g.data4[3],
            g.data4[4], g.data4[5], g.data4[6], g.data4[7] )
}
<commit_msg>Fixed a warning due to duplicate match arms by clippy<commit_after>
use std::collections::HashMap;
use parse::*;
use clap::ArgMatches;
use intercom_common::utils;
use intercom_common::guid::*;
use error::*;
use std::io;

/// Runs the 'cpp' subcommand.
pub fn run( cpp_params : &ArgMatches ) -> AppResult {

    // Parse the sources and convert the result into an IDL.
    let ( renames, result ) = parse_crate(
            cpp_params.value_of( "path" ).unwrap() )?;
    result_to_cpp( &result, &renames, &mut io::stdout() );

    Ok(())
}

/// Converts the parse result into an header  that gets written to stdout.
fn result_to_cpp(
    r : &ParseResult,
    rn : &HashMap<String, String>,
    out : &mut io::Write,
) {
    // Introduce all interfaces so we don't get errors on undeclared items.
    let itf_forward_declarations = r.interfaces.iter().map(|itf| {
        format!( "{}struct {};", get_indentation( 1 ), try_rename( rn, &itf.name ) )
    } ).collect::<Vec<_>>().join( "\n" );

    // Raw interfaces that give direct access to the components implemented in Rust.
    let itfs = r.interfaces.iter().map(|itf| {
        generate_raw_interface( rn, itf, 1 )
    } ).collect::<Vec<_>>().join( "\n" );

    // Generate class descriptors.
    let class_descriptors = r.class_names.iter().map(|class_name| {

        // Get the class details by matching the name.
        let class  = r.classes.iter().find(|cls| &cls.name == class_name )
            .unwrap();

        generate_class_descriptor( rn, &r.libname, class_name, class, 1 )
    } ).collect::<Vec<_>>().join( "\n" );

    // Generate using statements that flatten the raw interfaces
    // to mimic the results of a Microsoft MIDL compiler.
    let flat_interface_declarations = r.interfaces.iter().map(|itf| {
        flatten_interface( rn, &r.libname, itf, 1 )
    } ).collect::<Vec<_>>().join( "\n" );

    // Generate using statements that flatten the class descriptors
    // to mimic the results of a Microsoft MIDL compiler.
    let flat_class_descriptor_declarations = r.class_names.iter().map(|class_name| {

                // Get the class details by matching the name.
        let class  = r.classes.iter().find(|cls| &cls.name == class_name )
            .unwrap();

        flatten_class_descriptor( class_name, class, 1 )
    } ).collect::<Vec<_>>().join( "\n" );

    // We got the interfaces and classes. We can format and output the raw interfaces.
    let raw_namespace: String = format!( r###"
>       #include <array>
>       #include <intercom.h>
>       namespace {}
>       {{
>       namespace raw
>       {{
>       {}
>       {}
>       {}
>       }}
>       }}
>       #ifdef INTERCOM_FLATTEN_DECLARATIONS
>       {}
>       {}
>       #endif
>       "###,
        &r.libname, itf_forward_declarations, itfs, class_descriptors,
        flat_interface_declarations, flat_class_descriptor_declarations );
    let raw_namespace = raw_namespace.replace( ">       ", "" );
    writeln!( out, "{}", raw_namespace ).unwrap();
}

fn generate_raw_interface(
    rn : &HashMap<String, String>,
    interface: &Interface,
    base_indentation: usize
) -> String {

    // Now that we have methods sorted out, we can construct the final
    // interface definition.
    let indentation = get_indentation( base_indentation );
    let interface_text = format!( "{}struct {} : IUnknown\n{}{{\n{}static constexpr intercom::IID ID = {};\n\n{}\n{}}};\n",
            indentation, try_rename( rn, &interface.name ), indentation,
            get_indentation( base_indentation + 1 ), guid_to_binary( &interface.iid ),
            generate_raw_methods( interface, base_indentation + 1 ), indentation );
    interface_text
}

fn generate_raw_methods(
    itf: &Interface,
    indentation: usize
) -> String {

    // Get the method definitions for the current interface.
    itf.methods.iter().enumerate().map(|(_i,m)| {

        // Construct the argument list.
        let args = m.arguments.iter().map(|a| {

            // Argument direction affects both the argument attribute and
            // whether the argument is passed by pointer or value.
            let out_ptr = match a.dir {
                ArgDirection::In => "",
                ArgDirection::Out | ArgDirection::Return => "*",
            };
            format!( "{}{} {}", to_cpp_type( &a.ty ), out_ptr, a.name )

        } ).collect::<Vec<_>>().join( ", " );

        // Format the method.
        // To maintain backwards compatibility all new methods should
        // be added to the end of the traits.
        format!( r###"{}virtual {} INTERCOM_CC {}( {} ) = 0;"###, get_indentation( indentation ), to_cpp_type( &m.rvalue ), utils::pascal_case( &m.name ), args )

    } ).collect::<Vec<_>>().join( "\n" )
}

fn flatten_interface(
    rn : &HashMap<String, String>,
    libname: &str,
    interface: &Interface,
    indentation: usize
) -> String {
    let interface_name = try_rename( rn, &interface.name );
    let using_iid = format!( "static constexpr intercom::IID IID_{} = {};",
            utils::pascal_case( &interface_name ), guid_to_binary( &interface.iid ) );
    let using_interface = format!( "using {} = {}::raw::{};", utils::pascal_case( &interface_name ), libname,  &interface_name );
    let indentation = get_indentation( indentation );
    format!("{}{}\n{}{}", indentation, using_iid, indentation, using_interface )
}

fn flatten_class_descriptor(
    class_name: &str,
    coclass: &CoClass,
    indentation: usize
) -> String {
    let using_clsid = format!( "static constexpr CLSID CLSID_{} = {};",
            utils::pascal_case( class_name ), guid_to_binary( &coclass.clsid ) );
    let indentation = get_indentation( indentation );
    format!("{}{}", indentation, using_clsid )
}

fn generate_class_descriptor(
    rn : &HashMap<String, String>,
    libname: &str,
    class_name: &str,
    coclass: &CoClass,
    indentation: usize
) -> String {

    // Create a list of interfaces to be declared in the class descriptor.
    let interfaces =  coclass.interfaces.iter().map(|itf| {
            format!( r###"{}{}::raw::{}::ID"###,
                &get_indentation( indentation + 2 ), libname, try_rename( rn, itf ) )
        } ).collect::<Vec<_>>().join( ",\n" );

    // Class descriptors hold information about COM classes implemented by the library.
    let clsid = guid_to_binary( &coclass.clsid );
    let class_descriptor: String = format!( r###"
>       class {}Descriptor
>       {{
>           static constexpr CLSID ID = {};

>           static constexpr std::array<intercom::IID, {}> INTERFACES = {{
{}
>           }};

>           {}Descriptor() = delete;
>           ~{}Descriptor() = delete;

>       }};
>       "###,
    class_name, clsid, coclass.interfaces.len(), interfaces, class_name, class_name );
    class_descriptor.replace( ">       ", &get_indentation( indentation ) )
}

/// Gets an empty string for ind One level of indentation is four spaces.
fn get_indentation(
    level: usize
) -> String {
    let spaces = level * 4;
    let indentation = format!( r###"{: <1$}"###, "", spaces );
    indentation
}

/// Converts a Rust type into applicable C++ type.
fn to_cpp_type(
    ty: &str
) -> &str {

    match ty {
        "int8" => "int8_t",
        "uint8" => "uint8_t",
        "int16" => "int16_t",
        "uint16" => "uint16_t",
        "int32" => "int32_t",
        "uint32" => "uint32_t",
        "int64" => "int64_t",
        "uint64" => "uint64_t",
        "BStr" => "intercom::BSTR",
        "HRESULT" => "intercom::HRESULT",
        _ => ty,
    }
}

/// Converts a guid to binarys representation.
fn guid_to_binary(
    g: &GUID
) -> String {

    format!( "{{0x{:08x},0x{:04x},0x{:04x},{{0x{:02x},0x{:02x},0x{:02x},0x{:02x},0x{:02x},0x{:02x},0x{:02x},0x{:02x}}}}}",
            g.data1, g.data2, g.data3,
            g.data4[0], g.data4[1], g.data4[2], g.data4[3],
            g.data4[4], g.data4[5], g.data4[6], g.data4[7] )
}
<|endoftext|>


<commit_before>//! 非同期タスクを管理する。

use crate::sysmod::SystemModules;
use anyhow::Result;
use chrono::prelude::*;
use log::{error, info, trace};
use std::future::Future;
use std::sync::Arc;
use tokio::signal::unix::{signal, SignalKind};
use tokio::sync::watch;

/// システムシャットダウン開始通知送信側 (単数)
type CancelTx = watch::Sender<bool>;
/// システムシャットダウン開始通知受信側 (複数)
type CancelRx = watch::Receiver<bool>;

/// [TaskServer] と各非同期タスク間で共有されるデータ。
///
/// インスタンスは1つだけ生成され、[Arc] により共有される。
struct InternalControl {
    /// Tokio ランタイム。
    rt: tokio::runtime::Runtime,
    /// 全システムモジュールのリスト。
    sysmods: SystemModules,
}

/// 各非同期タスクに clone され渡されるコントロールハンドル。
///
/// 各フィールドは private で、メソッドで機能を提供する。
/// [Clone::clone] 可能で、各タスクごとに複製したものが渡される。
#[derive(Clone)]
pub struct Control {
    /// private で [InternalControl] への [Arc] を持つ。
    internal: Arc<InternalControl>,
    /// システムシャットダウン時、true が設定送信される。
    ///
    /// また、シャットダウンシーケンスにおいて、全タスクの完了待ちのためにも使う。
    /// サーバ側は clone されたこれがすべて drop されるまで待機する。
    cancel_rx: Option<CancelRx>,
}

/// [TaskServer::run] の返す実行終了種別。
pub enum RunResult {
    Shutdown,
    Reboot,
}

/// タスクサーバ本体。
pub struct TaskServer {
    /// 各タスクに clone して渡すオリジナルの [Control]。
    ctrl: Control,
    /// システムシャットダウン時の中断リクエストの送信側。
    /// <https://tokio.rs/tokio/topics/shutdown>
    /// "Telling things to shut down" + "Waiting for things to finish shutting down"
    cancel_tx: CancelTx,
}

impl Control {
    /// [crate::sysmod::SystemModule] リストへの参照を取得する。
    pub fn sysmods(&self) -> &SystemModules {
        &self.internal.sysmods
    }

    /// キャンセル通知の受信オブジェクトを取得する。
    pub fn cancel_rx(&mut self) -> &mut CancelRx {
        self.cancel_rx.as_mut().unwrap()
    }

    /// 1回限りのタスクを生成して実行開始する。
    ///
    /// F: [Control] を引数に、T を返す関数。
    /// T: Future<Output = anyhow::Result<()> かつスレッド間移動可能。
    ///
    /// つまり、F は [Control] を引数に、anyhow::Result<()> を返す async function。
    pub fn spawn_oneshot_task<F, T>(&self, name: &str, f: F)
    where
        F: Fn(Control) -> T + Send + Sync + 'static,
        T: Future<Output = Result<()>> + Send,
    {
        // move するデータを準備する
        let name = name.to_string();
        let ctrl = self.clone();

        self.internal.rt.spawn(async move {
            info!("[{}] start (one-shot)", name);

            // ctrl を clone して move する
            let future = f(ctrl.clone());
            let result = future.await;
            // drop clone of ctrl

            if let Err(e) = result {
                error!("[{}] finish (error): {:?}", name, e);
            } else {
                info!("[{}] finish (success)", name);
            }
            // drop ctrl
        });
    }

    /// 周期タスクを生成する。
    ///
    /// wakeup_list: 起動時刻。以下を満たさないと panic する。
    /// * second 以下が 0 である。(分単位)
    /// * 昇順ソート済みである。
    ///
    /// F: [Control] を引数に、T を返す関数。
    /// T: Future<Output = anyhow::Result<()> かつスレッド間移動可能。
    ///
    /// つまり、F は [Control] を引数に、anyhow::Result<()> を返す async function。
    pub fn spawn_periodic_task<F, T>(&self, name: &str, wakeup_list: &[NaiveTime], f: F)
    where
        F: Fn(Control) -> T + Send + Sync + 'static,
        T: Future<Output = Result<()>> + Send + Sync + 'static,
    {
        // move するデータを準備する
        let name = name.to_string();
        let ctrl = self.clone();

        // 空でなくソート済み、秒以下がゼロなのを確認後
        // 今日のその時刻からなる Local DateTime に変換する
        // TODO: is_sorted() がまだ unstable
        assert!(!wakeup_list.is_empty(), "time list is empty");
        let sorted = wakeup_list.windows(2).all(|t| t[0] <= t[1]);
        assert!(sorted, "time list is not sorted");
        let today = Local::today();
        let mut dt_list: Vec<_> = wakeup_list
            .iter()
            .map(|time| {
                assert_eq!(time.second(), 0);
                assert_eq!(time.nanosecond(), 0);
                today.and_time(*time).unwrap()
            })
            .collect();

        // wakeup time list を最初の LOG_LIMIT 個までログに出力する
        const LOG_LIMIT: usize = 5;
        let log_iter = wakeup_list.iter().take(LOG_LIMIT);
        let mut str = log_iter.enumerate().fold(String::new(), |sum, (i, v)| {
            let str = if i == 0 {
                format!("{}", v)
            } else {
                format!(", {}", v)
            };
            sum + &str
        });
        if wakeup_list.len() > LOG_LIMIT {
            str += &format!(", ... ({} items)", wakeup_list.len());
        }
        info!("[{}] registered as a periodic task", name);
        info!("[{}] wakeup time: {}", name, str);

        // spawn async task
        self.internal.rt.spawn(async move {
            type CDuration = chrono::Duration;
            type TDuration = tokio::time::Duration;

            loop {
                // 現在時刻を取得して分までに切り捨てる
                let now = Local::now();
                let now_hmd = now.date().and_hms(now.hour(), now.minute(), 0);
                let next_min = now_hmd + CDuration::minutes(1);
                trace!("[{}] periodic task check: {}", name, now_hmd);

                // 起動時刻リスト内で二分探索
                match dt_list.binary_search(&now_hmd) {
                    Ok(ind) => {
                        // 一致するものを発見したので続行
                        trace!("[{}] hit in time list: {}", name, now_hmd);
                    }
                    Err(ind) => {
                        // ind = insertion point
                        trace!("[{}] not found in time list: {}", name, now_hmd);
                        // 起きるべき時刻は dt_list[ind]
                        if ind < dt_list.len() {
                            let target_dt = dt_list[ind] + CDuration::seconds(1);
                            let sleep_duration = target_dt - Local::now();
                            let sleep_sec = sleep_duration.num_seconds().clamp(0, i64::MAX) as u64;
                            trace!("[{}] target: {}, sleep_sec: {}", name, target_dt, sleep_sec);
                            tokio::time::sleep(TDuration::from_secs(sleep_sec)).await;
                            trace!("[{}] wake up", name);
                        } else {
                            // 一番後ろよりも現在時刻が後
                            // 起動時刻リストをすべて1日ずつ後ろにずらす
                            for dt in dt_list.iter_mut() {
                                let tomorrow = dt.date() + CDuration::days(1);
                                let time = dt.time();
                                *dt = tomorrow.and_time(time).unwrap();
                                trace!("[{}] advance time list by 1 day", name);
                            }
                        }
                        continue;
                    }
                }

                // ctrl を clone して future 内に move する
                let future = f(ctrl.clone());
                info!("[{}] start (periodic)", name);
                let result = future.await;
                // drop clone of ctrl
                if let Err(e) = result {
                    error!("[{}] finish (error): {:?}", name, e);
                } else {
                    info!("[{}] finish (success)", name);
                }

                // 次の "分" を狙って sleep する
                // 目標は安全のため hh:mm:05 を狙う
                let target_dt = next_min + CDuration::seconds(5);
                // タスクの実行に1分以上かかると負になるが、
                // chrono::Duration は負数を許している
                // その場合は 0 に補正する
                let sleep_duration = target_dt - Local::now();
                let sleep_sec = sleep_duration.num_seconds().clamp(0, i64::MAX) as u64;
                trace!("[{}] target: {}, sleep_sec: {}", name, target_dt, sleep_sec);
                tokio::time::sleep(tokio::time::Duration::from_secs(sleep_sec)).await;
                trace!("[{}] wake up", name);
            }
            // drop ctrl
        });
    }
}

impl TaskServer {
    /// タスクサーバを生成して初期化する。
    pub fn new(sysmods: SystemModules) -> Self {
        let rt = tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .build()
            .unwrap();

        let (cancel_tx, cancel_rx) = watch::channel(false);

        let internal = InternalControl { rt, sysmods };
        let ctrl = Control {
            internal: Arc::new(internal),
            cancel_rx: Some(cancel_rx),
        };

        TaskServer { ctrl, cancel_tx }
    }

    /// [Control::spawn_oneshot_task] と同じ。
    pub fn spawn_oneshot_task<F, T>(&self, name: &str, f: F)
    where
        F: Fn(Control) -> T + Send + Sync + 'static,
        T: Future<Output = Result<()>> + Send,
    {
        self.ctrl.spawn_oneshot_task(name, f);
    }

    /// 実行を開始し、完了するまでブロックする。
    ///
    /// self の所有権はを consume するため、一度しか実行できない。
    pub fn run(mut self) -> RunResult {
        // async block へ move するためのコピーを作る
        let ctrl = self.ctrl.clone();
        // オリジナルの cancel_rx を self から奪って drop しておく
        drop(self.ctrl.cancel_rx.take().unwrap());

        self.ctrl.internal.rt.block_on(async move {
            // SystemModule 全体に on_start イベントを配送
            ctrl.internal.sysmods.on_start(&ctrl).await;

            // この async block をシグナル処理に使う
            let mut sigint = signal(SignalKind::interrupt()).unwrap();
            let mut sigterm = signal(SignalKind::terminate()).unwrap();
            let mut sighup = signal(SignalKind::hangup()).unwrap();

            let run_result;
            tokio::select! {
                _ = sigint.recv() => {
                    info!("[signal] SIGINT");
                    run_result = RunResult::Shutdown;
                },
                _ = sigterm.recv() => {
                    info!("[signal] SIGTERM");
                    run_result = RunResult::Shutdown;
                },
                _ = sighup.recv() => {
                    info!("[signal] SIGHUP");
                    run_result = RunResult::Reboot;
                },
            }

            // 値を true に設定して全タスクにキャンセルリクエストを通知する
            self.cancel_tx.send_replace(true);
            // 全タスク完了待ち
            // この async block で持っている分の cancel_rx を ctrl ごと drop する
            // オリジナルの self.ctrl.cancel_rx も drop 済みなので
            // 残りは各 async task に clone された Control
            drop(ctrl);
            // 全 cancel_rx が drop されるまで待つ
            info!("waiting for all tasks to be completed....");
            self.cancel_tx.closed().await;
            info!("OK: all tasks are completed");

            run_result
        })
    }
}
<commit_msg>周期タスクの sleep をキャンセル可能にする<commit_after>//! 非同期タスクを管理する。

use crate::sysmod::SystemModules;
use anyhow::Result;
use chrono::prelude::*;
use log::{error, info, trace};
use std::future::Future;
use std::sync::Arc;
use tokio::select;
use tokio::signal::unix::{signal, SignalKind};
use tokio::sync::watch;

/// システムシャットダウン開始通知送信側 (単数)
type CancelTx = watch::Sender<bool>;
/// システムシャットダウン開始通知受信側 (複数)
type CancelRx = watch::Receiver<bool>;

/// [TaskServer] と各非同期タスク間で共有されるデータ。
///
/// インスタンスは1つだけ生成され、[Arc] により共有される。
struct InternalControl {
    /// Tokio ランタイム。
    rt: tokio::runtime::Runtime,
    /// 全システムモジュールのリスト。
    sysmods: SystemModules,
}

/// 各非同期タスクに clone され渡されるコントロールハンドル。
///
/// 各フィールドは private で、メソッドで機能を提供する。
/// [Clone::clone] 可能で、各タスクごとに複製したものが渡される。
#[derive(Clone)]
pub struct Control {
    /// private で [InternalControl] への [Arc] を持つ。
    internal: Arc<InternalControl>,
    /// システムシャットダウン時、true が設定送信される。
    ///
    /// また、シャットダウンシーケンスにおいて、全タスクの完了待ちのためにも使う。
    /// サーバ側は clone されたこれがすべて drop されるまで待機する。
    cancel_rx: Option<CancelRx>,
}

/// [TaskServer::run] の返す実行終了種別。
pub enum RunResult {
    Shutdown,
    Reboot,
}

/// タスクサーバ本体。
pub struct TaskServer {
    /// 各タスクに clone して渡すオリジナルの [Control]。
    ctrl: Control,
    /// システムシャットダウン時の中断リクエストの送信側。
    /// <https://tokio.rs/tokio/topics/shutdown>
    /// "Telling things to shut down" + "Waiting for things to finish shutting down"
    cancel_tx: CancelTx,
}

impl Control {
    /// [crate::sysmod::SystemModule] リストへの参照を取得する。
    pub fn sysmods(&self) -> &SystemModules {
        &self.internal.sysmods
    }

    /// キャンセル通知の受信オブジェクトを取得する。
    pub fn cancel_rx(&mut self) -> &mut CancelRx {
        self.cancel_rx.as_mut().unwrap()
    }

    /// 1回限りのタスクを生成して実行開始する。
    ///
    /// F: [Control] を引数に、T を返す関数。
    /// T: Future<Output = anyhow::Result<()> かつスレッド間移動可能。
    ///
    /// つまり、F は [Control] を引数に、anyhow::Result<()> を返す async function。
    pub fn spawn_oneshot_task<F, T>(&self, name: &str, f: F)
    where
        F: Fn(Control) -> T + Send + Sync + 'static,
        T: Future<Output = Result<()>> + Send,
    {
        // move するデータを準備する
        let name = name.to_string();
        let ctrl = self.clone();

        self.internal.rt.spawn(async move {
            info!("[{}] start (one-shot)", name);

            // ctrl を clone して move する
            let future = f(ctrl.clone());
            let result = future.await;
            // drop clone of ctrl

            if let Err(e) = result {
                error!("[{}] finish (error): {:?}", name, e);
            } else {
                info!("[{}] finish (success)", name);
            }
            // drop ctrl
        });
    }

    /// 周期タスクを生成する。
    ///
    /// wakeup_list: 起動時刻。以下を満たさないと panic する。
    /// * second 以下が 0 である。(分単位)
    /// * 昇順ソート済みである。
    ///
    /// F: [Control] を引数に、T を返す関数。
    /// T: Future<Output = anyhow::Result<()> かつスレッド間移動可能。
    ///
    /// つまり、F は [Control] を引数に、anyhow::Result<()> を返す async function。
    pub fn spawn_periodic_task<F, T>(&self, name: &str, wakeup_list: &[NaiveTime], f: F)
    where
        F: Fn(Control) -> T + Send + Sync + 'static,
        T: Future<Output = Result<()>> + Send + Sync + 'static,
    {
        // move するデータを準備する
        let name = name.to_string();
        let mut ctrl = self.clone();

        // 空でなくソート済み、秒以下がゼロなのを確認後
        // 今日のその時刻からなる Local DateTime に変換する
        // TODO: is_sorted() がまだ unstable
        assert!(!wakeup_list.is_empty(), "time list is empty");
        let sorted = wakeup_list.windows(2).all(|t| t[0] <= t[1]);
        assert!(sorted, "time list is not sorted");
        let today = Local::today();
        let mut dt_list: Vec<_> = wakeup_list
            .iter()
            .map(|time| {
                assert_eq!(time.second(), 0);
                assert_eq!(time.nanosecond(), 0);
                today.and_time(*time).unwrap()
            })
            .collect();

        // wakeup time list を最初の LOG_LIMIT 個までログに出力する
        const LOG_LIMIT: usize = 5;
        let log_iter = wakeup_list.iter().take(LOG_LIMIT);
        let mut str = log_iter.enumerate().fold(String::new(), |sum, (i, v)| {
            let str = if i == 0 {
                format!("{}", v)
            } else {
                format!(", {}", v)
            };
            sum + &str
        });
        if wakeup_list.len() > LOG_LIMIT {
            str += &format!(", ... ({} items)", wakeup_list.len());
        }
        info!("[{}] registered as a periodic task", name);
        info!("[{}] wakeup time: {}", name, str);

        // spawn async task
        self.internal.rt.spawn(async move {
            type CDuration = chrono::Duration;
            type TDuration = tokio::time::Duration;

            loop {
                // 現在時刻を取得して分までに切り捨てる
                let now = Local::now();
                let now_hmd = now.date().and_hms(now.hour(), now.minute(), 0);
                let next_min = now_hmd + CDuration::minutes(1);
                trace!("[{}] periodic task check: {}", name, now_hmd);

                // 起動時刻リスト内で二分探索
                match dt_list.binary_search(&now_hmd) {
                    Ok(ind) => {
                        // 一致するものを発見したので続行
                        trace!("[{}] hit in time list: {}", name, now_hmd);
                    }
                    Err(ind) => {
                        // ind = insertion point
                        trace!("[{}] not found in time list: {}", name, now_hmd);
                        // 起きるべき時刻は dt_list[ind]
                        if ind < dt_list.len() {
                            let target_dt = dt_list[ind] + CDuration::seconds(1);
                            let sleep_duration = target_dt - Local::now();
                            let sleep_sec = sleep_duration.num_seconds().clamp(0, i64::MAX) as u64;
                            trace!("[{}] target: {}, sleep_sec: {}", name, target_dt, sleep_sec);
                            select! {
                                _ = tokio::time::sleep(TDuration::from_secs(sleep_sec)) => {}
                                _ = ctrl.cancel_rx().changed() => {
                                    info!("[{}] cancel periodic task", name);
                                    return;
                                }
                            }

                            trace!("[{}] wake up", name);
                        } else {
                            // 一番後ろよりも現在時刻が後
                            // 起動時刻リストをすべて1日ずつ後ろにずらす
                            for dt in dt_list.iter_mut() {
                                let tomorrow = dt.date() + CDuration::days(1);
                                let time = dt.time();
                                *dt = tomorrow.and_time(time).unwrap();
                                trace!("[{}] advance time list by 1 day", name);
                            }
                        }
                        continue;
                    }
                }

                // ctrl を clone して future 内に move する
                let future = f(ctrl.clone());
                info!("[{}] start (periodic)", name);
                let result = future.await;
                // drop clone of ctrl
                if let Err(e) = result {
                    error!("[{}] finish (error): {:?}", name, e);
                } else {
                    info!("[{}] finish (success)", name);
                }

                // 次の "分" を狙って sleep する
                // 目標は安全のため hh:mm:05 を狙う
                let target_dt = next_min + CDuration::seconds(5);
                // タスクの実行に1分以上かかると負になるが、
                // chrono::Duration は負数を許している
                // その場合は 0 に補正する
                let sleep_duration = target_dt - Local::now();
                let sleep_sec = sleep_duration.num_seconds().clamp(0, i64::MAX) as u64;
                trace!("[{}] target: {}, sleep_sec: {}", name, target_dt, sleep_sec);
                select! {
                    _ = tokio::time::sleep(tokio::time::Duration::from_secs(sleep_sec)) => {}
                    _ = ctrl.cancel_rx().changed() => {
                        info!("[{}] cancel periodic task", name);
                        return;
                    }
                }
                trace!("[{}] wake up", name);
            }
            // drop ctrl
        });
    }
}

impl TaskServer {
    /// タスクサーバを生成して初期化する。
    pub fn new(sysmods: SystemModules) -> Self {
        let rt = tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .build()
            .unwrap();

        let (cancel_tx, cancel_rx) = watch::channel(false);

        let internal = InternalControl { rt, sysmods };
        let ctrl = Control {
            internal: Arc::new(internal),
            cancel_rx: Some(cancel_rx),
        };

        TaskServer { ctrl, cancel_tx }
    }

    /// [Control::spawn_oneshot_task] と同じ。
    pub fn spawn_oneshot_task<F, T>(&self, name: &str, f: F)
    where
        F: Fn(Control) -> T + Send + Sync + 'static,
        T: Future<Output = Result<()>> + Send,
    {
        self.ctrl.spawn_oneshot_task(name, f);
    }

    /// 実行を開始し、完了するまでブロックする。
    ///
    /// self の所有権はを consume するため、一度しか実行できない。
    pub fn run(mut self) -> RunResult {
        // async block へ move するためのコピーを作る
        let ctrl = self.ctrl.clone();
        // オリジナルの cancel_rx を self から奪って drop しておく
        drop(self.ctrl.cancel_rx.take().unwrap());

        self.ctrl.internal.rt.block_on(async move {
            // SystemModule 全体に on_start イベントを配送
            ctrl.internal.sysmods.on_start(&ctrl).await;

            // この async block をシグナル処理に使う
            let mut sigint = signal(SignalKind::interrupt()).unwrap();
            let mut sigterm = signal(SignalKind::terminate()).unwrap();
            let mut sighup = signal(SignalKind::hangup()).unwrap();

            let run_result;
            tokio::select! {
                _ = sigint.recv() => {
                    info!("[signal] SIGINT");
                    run_result = RunResult::Shutdown;
                },
                _ = sigterm.recv() => {
                    info!("[signal] SIGTERM");
                    run_result = RunResult::Shutdown;
                },
                _ = sighup.recv() => {
                    info!("[signal] SIGHUP");
                    run_result = RunResult::Reboot;
                },
            }

            // 値を true に設定して全タスクにキャンセルリクエストを通知する
            self.cancel_tx.send_replace(true);
            // 全タスク完了待ち
            // この async block で持っている分の cancel_rx を ctrl ごと drop する
            // オリジナルの self.ctrl.cancel_rx も drop 済みなので
            // 残りは各 async task に clone された Control
            drop(ctrl);
            // 全 cancel_rx が drop されるまで待つ
            info!("waiting for all tasks to be completed....");
            self.cancel_tx.closed().await;
            info!("OK: all tasks are completed");

            run_result
        })
    }
}
<|endoftext|>


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

#![allow(missing_docs)]
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]

use prelude::v1::*;

use ffi::OsStr;
use io::{self, ErrorKind};
use libc;
use mem;
use old_io::{self, IoResult, IoError};
use num::Int;
use os::windows::OsStrExt;
use sync::{Once, ONCE_INIT};

macro_rules! helper_init { (static $name:ident: Helper<$m:ty>) => (
    static $name: Helper<$m> = Helper {
        lock: ::sync::MUTEX_INIT,
        cond: ::sync::CONDVAR_INIT,
        chan: ::cell::UnsafeCell { value: 0 as *mut ::sync::mpsc::Sender<$m> },
        signal: ::cell::UnsafeCell { value: 0 },
        initialized: ::cell::UnsafeCell { value: false },
        shutdown: ::cell::UnsafeCell { value: false },
    };
) }

pub mod backtrace;
pub mod c;
pub mod condvar;
pub mod ext;
pub mod fs;
pub mod fs2;
pub mod handle;
pub mod helper_signal;
pub mod mutex;
pub mod net;
pub mod os;
pub mod os_str;
pub mod pipe;
pub mod pipe2;
pub mod process;
pub mod process2;
pub mod rwlock;
pub mod stack_overflow;
pub mod sync;
pub mod tcp;
pub mod thread;
pub mod thread_local;
pub mod time;
pub mod timer;
pub mod tty;
pub mod udp;

pub mod addrinfo {
    pub use sys_common::net::get_host_addresses;
    pub use sys_common::net::get_address_name;
}

// FIXME: move these to c module
pub type sock_t = libc::SOCKET;
pub type wrlen = libc::c_int;
pub type msglen_t = libc::c_int;
pub unsafe fn close_sock(sock: sock_t) { let _ = libc::closesocket(sock); }

// windows has zero values as errors
fn mkerr_winbool(ret: libc::c_int) -> IoResult<()> {
    if ret == 0 {
        Err(last_error())
    } else {
        Ok(())
    }
}

pub fn last_error() -> IoError {
    let errno = os::errno() as i32;
    let mut err = decode_error(errno);
    err.detail = Some(os::error_string(errno));
    err
}

pub fn last_net_error() -> IoError {
    let errno = unsafe { c::WSAGetLastError() as i32 };
    let mut err = decode_error(errno);
    err.detail = Some(os::error_string(errno));
    err
}

pub fn last_gai_error(_errno: i32) -> IoError {
    last_net_error()
}

/// Convert an `errno` value into a high-level error variant and description.
pub fn decode_error(errno: i32) -> IoError {
    let (kind, desc) = match errno {
        libc::EOF => (old_io::EndOfFile, "end of file"),
        libc::ERROR_NO_DATA => (old_io::BrokenPipe, "the pipe is being closed"),
        libc::ERROR_FILE_NOT_FOUND => (old_io::FileNotFound, "file not found"),
        libc::ERROR_INVALID_NAME => (old_io::InvalidInput, "invalid file name"),
        libc::WSAECONNREFUSED => (old_io::ConnectionRefused, "connection refused"),
        libc::WSAECONNRESET => (old_io::ConnectionReset, "connection reset"),
        libc::ERROR_ACCESS_DENIED | libc::WSAEACCES =>
            (old_io::PermissionDenied, "permission denied"),
        libc::WSAEWOULDBLOCK => {
            (old_io::ResourceUnavailable, "resource temporarily unavailable")
        }
        libc::WSAENOTCONN => (old_io::NotConnected, "not connected"),
        libc::WSAECONNABORTED => (old_io::ConnectionAborted, "connection aborted"),
        libc::WSAEADDRNOTAVAIL => (old_io::ConnectionRefused, "address not available"),
        libc::WSAEADDRINUSE => (old_io::ConnectionRefused, "address in use"),
        libc::ERROR_BROKEN_PIPE => (old_io::EndOfFile, "the pipe has ended"),
        libc::ERROR_OPERATION_ABORTED =>
            (old_io::TimedOut, "operation timed out"),
        libc::WSAEINVAL => (old_io::InvalidInput, "invalid argument"),
        libc::ERROR_CALL_NOT_IMPLEMENTED =>
            (old_io::IoUnavailable, "function not implemented"),
        libc::ERROR_INVALID_HANDLE =>
            (old_io::MismatchedFileTypeForOperation,
             "invalid handle provided to function"),
        libc::ERROR_NOTHING_TO_TERMINATE =>
            (old_io::InvalidInput, "no process to kill"),
        libc::ERROR_ALREADY_EXISTS =>
            (old_io::PathAlreadyExists, "path already exists"),

        // libuv maps this error code to EISDIR. we do too. if it is found
        // to be incorrect, we can add in some more machinery to only
        // return this message when ERROR_INVALID_FUNCTION after certain
        // Windows calls.
        libc::ERROR_INVALID_FUNCTION => (old_io::InvalidInput,
                                         "illegal operation on a directory"),

        _ => (old_io::OtherIoError, "unknown error")
    };
    IoError { kind: kind, desc: desc, detail: None }
}

pub fn decode_error_detailed(errno: i32) -> IoError {
    let mut err = decode_error(errno);
    err.detail = Some(os::error_string(errno));
    err
}

pub fn decode_error_kind(errno: i32) -> ErrorKind {
    match errno as libc::c_int {
        libc::ERROR_ACCESS_DENIED => ErrorKind::PermissionDenied,
        libc::ERROR_ALREADY_EXISTS => ErrorKind::PathAlreadyExists,
        libc::ERROR_BROKEN_PIPE => ErrorKind::BrokenPipe,
        libc::ERROR_FILE_NOT_FOUND => ErrorKind::FileNotFound,
        libc::ERROR_INVALID_FUNCTION => ErrorKind::InvalidInput,
        libc::ERROR_INVALID_HANDLE => ErrorKind::MismatchedFileTypeForOperation,
        libc::ERROR_INVALID_NAME => ErrorKind::InvalidInput,
        libc::ERROR_NOTHING_TO_TERMINATE => ErrorKind::InvalidInput,
        libc::ERROR_NO_DATA => ErrorKind::BrokenPipe,
        libc::ERROR_OPERATION_ABORTED => ErrorKind::TimedOut,

        libc::WSAEACCES => ErrorKind::PermissionDenied,
        libc::WSAEADDRINUSE => ErrorKind::ConnectionRefused,
        libc::WSAEADDRNOTAVAIL => ErrorKind::ConnectionRefused,
        libc::WSAECONNABORTED => ErrorKind::ConnectionAborted,
        libc::WSAECONNREFUSED => ErrorKind::ConnectionRefused,
        libc::WSAECONNRESET => ErrorKind::ConnectionReset,
        libc::WSAEINVAL => ErrorKind::InvalidInput,
        libc::WSAENOTCONN => ErrorKind::NotConnected,
        libc::WSAEWOULDBLOCK => ErrorKind::ResourceUnavailable,

        _ => ErrorKind::Other,
    }
}

#[inline]
pub fn retry<I, F>(f: F) -> I where F: FnOnce() -> I { f() } // PR rust-lang/rust/#17020

pub fn ms_to_timeval(ms: u64) -> libc::timeval {
    libc::timeval {
        tv_sec: (ms / 1000) as libc::c_long,
        tv_usec: ((ms % 1000) * 1000) as libc::c_long,
    }
}

pub fn wouldblock() -> bool {
    let err = os::errno();
    err == libc::WSAEWOULDBLOCK as i32
}

pub fn set_nonblocking(fd: sock_t, nb: bool) {
    let mut set = nb as libc::c_ulong;
    (if unsafe { c::ioctlsocket(fd, c::FIONBIO, &mut set) } != 0 {
        Err(last_error())
    } else {
        Ok(())
    }).unwrap();
}

pub fn init_net() {
    unsafe {
        static START: Once = ONCE_INIT;

        START.call_once(|| {
            let mut data: c::WSADATA = mem::zeroed();
            let ret = c::WSAStartup(0x202, // version 2.2
                                    &mut data);
            assert_eq!(ret, 0);
        });
    }
}

pub fn unimpl() -> IoError {
    IoError {
        kind: old_io::IoUnavailable,
        desc: "operation is not implemented",
        detail: None,
    }
}

fn to_utf16(s: Option<&str>) -> IoResult<Vec<u16>> {
    match s {
        Some(s) => Ok(to_utf16_os(OsStr::from_str(s))),
        None => Err(IoError {
            kind: old_io::InvalidInput,
            desc: "valid unicode input required",
            detail: None,
        }),
    }
}

fn to_utf16_os(s: &OsStr) -> Vec<u16> {
    let mut v: Vec<_> = s.encode_wide().collect();
    v.push(0);
    v
}

// Many Windows APIs follow a pattern of where we hand the a buffer and then
// they will report back to us how large the buffer should be or how many bytes
// currently reside in the buffer. This function is an abstraction over these
// functions by making them easier to call.
//
// The first callback, `f1`, is yielded a (pointer, len) pair which can be
// passed to a syscall. The `ptr` is valid for `len` items (u16 in this case).
// The closure is expected to return what the syscall returns which will be
// interpreted by this function to determine if the syscall needs to be invoked
// again (with more buffer space).
//
// Once the syscall has completed (errors bail out early) the second closure is
// yielded the data which has been read from the syscall. The return value
// from this closure is then the return value of the function.
fn fill_utf16_buf_base<F1, F2, T>(mut f1: F1, f2: F2) -> Result<T, ()>
    where F1: FnMut(*mut u16, libc::DWORD) -> libc::DWORD,
          F2: FnOnce(&[u16]) -> T
{
    // Start off with a stack buf but then spill over to the heap if we end up
    // needing more space.
    let mut stack_buf = [0u16; 512];
    let mut heap_buf = Vec::new();
    unsafe {
        let mut n = stack_buf.len();
        loop {
            let buf = if n <= stack_buf.len() {
                &mut stack_buf[..]
            } else {
                let extra = n - heap_buf.len();
                heap_buf.reserve(extra);
                heap_buf.set_len(n);
                &mut heap_buf[..]
            };

            // This function is typically called on windows API functions which
            // will return the correct length of the string, but these functions
            // also return the `0` on error. In some cases, however, the
            // returned "correct length" may actually be 0!
            //
            // To handle this case we call `SetLastError` to reset it to 0 and
            // then check it again if we get the "0 error value". If the "last
            // error" is still 0 then we interpret it as a 0 length buffer and
            // not an actual error.
            c::SetLastError(0);
            let k = match f1(buf.as_mut_ptr(), n as libc::DWORD) {
                0 if libc::GetLastError() == 0 => 0,
                0 => return Err(()),
                n => n,
            } as usize;
            if k == n && libc::GetLastError() ==
                            libc::ERROR_INSUFFICIENT_BUFFER as libc::DWORD {
                n *= 2;
            } else if k >= n {
                n = k;
            } else {
                return Ok(f2(&buf[..k]))
            }
        }
    }
}

fn fill_utf16_buf<F1, F2, T>(f1: F1, f2: F2) -> IoResult<T>
    where F1: FnMut(*mut u16, libc::DWORD) -> libc::DWORD,
          F2: FnOnce(&[u16]) -> T
{
    fill_utf16_buf_base(f1, f2).map_err(|()| IoError::last_error())
}

fn fill_utf16_buf_new<F1, F2, T>(f1: F1, f2: F2) -> io::Result<T>
    where F1: FnMut(*mut u16, libc::DWORD) -> libc::DWORD,
          F2: FnOnce(&[u16]) -> T
{
    fill_utf16_buf_base(f1, f2).map_err(|()| io::Error::last_os_error())
}

fn os2path(s: &[u16]) -> Path {
    // FIXME: this should not be a panicking conversion (aka path reform)
    Path::new(String::from_utf16(s).unwrap())
}

pub fn truncate_utf16_at_nul<'a>(v: &'a [u16]) -> &'a [u16] {
    match v.iter().position(|c| *c == 0) {
        // don't include the 0
        Some(i) => &v[..i],
        None => v
    }
}

fn cvt<I: Int>(i: I) -> io::Result<I> {
    if i == Int::zero() {
        Err(io::Error::last_os_error())
    } else {
        Ok(i)
    }
}

fn ms_to_filetime(ms: u64) -> libc::FILETIME {
    // A FILETIME is a count of 100 nanosecond intervals, so we multiply by
    // 10000 b/c there are 10000 intervals in 1 ms
    let ms = ms * 10000;
    libc::FILETIME {
        dwLowDateTime: ms as u32,
        dwHighDateTime: (ms >> 32) as u32,
    }
}
<commit_msg>Improve readability of an error check in `set_non_blocking`<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![allow(missing_docs)]
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]

use prelude::v1::*;

use ffi::OsStr;
use io::{self, ErrorKind};
use libc;
use mem;
use old_io::{self, IoResult, IoError};
use num::Int;
use os::windows::OsStrExt;
use sync::{Once, ONCE_INIT};

macro_rules! helper_init { (static $name:ident: Helper<$m:ty>) => (
    static $name: Helper<$m> = Helper {
        lock: ::sync::MUTEX_INIT,
        cond: ::sync::CONDVAR_INIT,
        chan: ::cell::UnsafeCell { value: 0 as *mut ::sync::mpsc::Sender<$m> },
        signal: ::cell::UnsafeCell { value: 0 },
        initialized: ::cell::UnsafeCell { value: false },
        shutdown: ::cell::UnsafeCell { value: false },
    };
) }

pub mod backtrace;
pub mod c;
pub mod condvar;
pub mod ext;
pub mod fs;
pub mod fs2;
pub mod handle;
pub mod helper_signal;
pub mod mutex;
pub mod net;
pub mod os;
pub mod os_str;
pub mod pipe;
pub mod pipe2;
pub mod process;
pub mod process2;
pub mod rwlock;
pub mod stack_overflow;
pub mod sync;
pub mod tcp;
pub mod thread;
pub mod thread_local;
pub mod time;
pub mod timer;
pub mod tty;
pub mod udp;

pub mod addrinfo {
    pub use sys_common::net::get_host_addresses;
    pub use sys_common::net::get_address_name;
}

// FIXME: move these to c module
pub type sock_t = libc::SOCKET;
pub type wrlen = libc::c_int;
pub type msglen_t = libc::c_int;
pub unsafe fn close_sock(sock: sock_t) { let _ = libc::closesocket(sock); }

// windows has zero values as errors
fn mkerr_winbool(ret: libc::c_int) -> IoResult<()> {
    if ret == 0 {
        Err(last_error())
    } else {
        Ok(())
    }
}

pub fn last_error() -> IoError {
    let errno = os::errno() as i32;
    let mut err = decode_error(errno);
    err.detail = Some(os::error_string(errno));
    err
}

pub fn last_net_error() -> IoError {
    let errno = unsafe { c::WSAGetLastError() as i32 };
    let mut err = decode_error(errno);
    err.detail = Some(os::error_string(errno));
    err
}

pub fn last_gai_error(_errno: i32) -> IoError {
    last_net_error()
}

/// Convert an `errno` value into a high-level error variant and description.
pub fn decode_error(errno: i32) -> IoError {
    let (kind, desc) = match errno {
        libc::EOF => (old_io::EndOfFile, "end of file"),
        libc::ERROR_NO_DATA => (old_io::BrokenPipe, "the pipe is being closed"),
        libc::ERROR_FILE_NOT_FOUND => (old_io::FileNotFound, "file not found"),
        libc::ERROR_INVALID_NAME => (old_io::InvalidInput, "invalid file name"),
        libc::WSAECONNREFUSED => (old_io::ConnectionRefused, "connection refused"),
        libc::WSAECONNRESET => (old_io::ConnectionReset, "connection reset"),
        libc::ERROR_ACCESS_DENIED | libc::WSAEACCES =>
            (old_io::PermissionDenied, "permission denied"),
        libc::WSAEWOULDBLOCK => {
            (old_io::ResourceUnavailable, "resource temporarily unavailable")
        }
        libc::WSAENOTCONN => (old_io::NotConnected, "not connected"),
        libc::WSAECONNABORTED => (old_io::ConnectionAborted, "connection aborted"),
        libc::WSAEADDRNOTAVAIL => (old_io::ConnectionRefused, "address not available"),
        libc::WSAEADDRINUSE => (old_io::ConnectionRefused, "address in use"),
        libc::ERROR_BROKEN_PIPE => (old_io::EndOfFile, "the pipe has ended"),
        libc::ERROR_OPERATION_ABORTED =>
            (old_io::TimedOut, "operation timed out"),
        libc::WSAEINVAL => (old_io::InvalidInput, "invalid argument"),
        libc::ERROR_CALL_NOT_IMPLEMENTED =>
            (old_io::IoUnavailable, "function not implemented"),
        libc::ERROR_INVALID_HANDLE =>
            (old_io::MismatchedFileTypeForOperation,
             "invalid handle provided to function"),
        libc::ERROR_NOTHING_TO_TERMINATE =>
            (old_io::InvalidInput, "no process to kill"),
        libc::ERROR_ALREADY_EXISTS =>
            (old_io::PathAlreadyExists, "path already exists"),

        // libuv maps this error code to EISDIR. we do too. if it is found
        // to be incorrect, we can add in some more machinery to only
        // return this message when ERROR_INVALID_FUNCTION after certain
        // Windows calls.
        libc::ERROR_INVALID_FUNCTION => (old_io::InvalidInput,
                                         "illegal operation on a directory"),

        _ => (old_io::OtherIoError, "unknown error")
    };
    IoError { kind: kind, desc: desc, detail: None }
}

pub fn decode_error_detailed(errno: i32) -> IoError {
    let mut err = decode_error(errno);
    err.detail = Some(os::error_string(errno));
    err
}

pub fn decode_error_kind(errno: i32) -> ErrorKind {
    match errno as libc::c_int {
        libc::ERROR_ACCESS_DENIED => ErrorKind::PermissionDenied,
        libc::ERROR_ALREADY_EXISTS => ErrorKind::PathAlreadyExists,
        libc::ERROR_BROKEN_PIPE => ErrorKind::BrokenPipe,
        libc::ERROR_FILE_NOT_FOUND => ErrorKind::FileNotFound,
        libc::ERROR_INVALID_FUNCTION => ErrorKind::InvalidInput,
        libc::ERROR_INVALID_HANDLE => ErrorKind::MismatchedFileTypeForOperation,
        libc::ERROR_INVALID_NAME => ErrorKind::InvalidInput,
        libc::ERROR_NOTHING_TO_TERMINATE => ErrorKind::InvalidInput,
        libc::ERROR_NO_DATA => ErrorKind::BrokenPipe,
        libc::ERROR_OPERATION_ABORTED => ErrorKind::TimedOut,

        libc::WSAEACCES => ErrorKind::PermissionDenied,
        libc::WSAEADDRINUSE => ErrorKind::ConnectionRefused,
        libc::WSAEADDRNOTAVAIL => ErrorKind::ConnectionRefused,
        libc::WSAECONNABORTED => ErrorKind::ConnectionAborted,
        libc::WSAECONNREFUSED => ErrorKind::ConnectionRefused,
        libc::WSAECONNRESET => ErrorKind::ConnectionReset,
        libc::WSAEINVAL => ErrorKind::InvalidInput,
        libc::WSAENOTCONN => ErrorKind::NotConnected,
        libc::WSAEWOULDBLOCK => ErrorKind::ResourceUnavailable,

        _ => ErrorKind::Other,
    }
}

#[inline]
pub fn retry<I, F>(f: F) -> I where F: FnOnce() -> I { f() } // PR rust-lang/rust/#17020

pub fn ms_to_timeval(ms: u64) -> libc::timeval {
    libc::timeval {
        tv_sec: (ms / 1000) as libc::c_long,
        tv_usec: ((ms % 1000) * 1000) as libc::c_long,
    }
}

pub fn wouldblock() -> bool {
    let err = os::errno();
    err == libc::WSAEWOULDBLOCK as i32
}

pub fn set_nonblocking(fd: sock_t, nb: bool) {
    let mut set = nb as libc::c_ulong;
    if unsafe { c::ioctlsocket(fd, c::FIONBIO, &mut set) } != 0 {
        // The above function should not return an error unless we passed it
        // invalid parameters. Panic on errors.
        Err(last_error()).unwrap();
    }
}

pub fn init_net() {
    unsafe {
        static START: Once = ONCE_INIT;

        START.call_once(|| {
            let mut data: c::WSADATA = mem::zeroed();
            let ret = c::WSAStartup(0x202, // version 2.2
                                    &mut data);
            assert_eq!(ret, 0);
        });
    }
}

pub fn unimpl() -> IoError {
    IoError {
        kind: old_io::IoUnavailable,
        desc: "operation is not implemented",
        detail: None,
    }
}

fn to_utf16(s: Option<&str>) -> IoResult<Vec<u16>> {
    match s {
        Some(s) => Ok(to_utf16_os(OsStr::from_str(s))),
        None => Err(IoError {
            kind: old_io::InvalidInput,
            desc: "valid unicode input required",
            detail: None,
        }),
    }
}

fn to_utf16_os(s: &OsStr) -> Vec<u16> {
    let mut v: Vec<_> = s.encode_wide().collect();
    v.push(0);
    v
}

// Many Windows APIs follow a pattern of where we hand the a buffer and then
// they will report back to us how large the buffer should be or how many bytes
// currently reside in the buffer. This function is an abstraction over these
// functions by making them easier to call.
//
// The first callback, `f1`, is yielded a (pointer, len) pair which can be
// passed to a syscall. The `ptr` is valid for `len` items (u16 in this case).
// The closure is expected to return what the syscall returns which will be
// interpreted by this function to determine if the syscall needs to be invoked
// again (with more buffer space).
//
// Once the syscall has completed (errors bail out early) the second closure is
// yielded the data which has been read from the syscall. The return value
// from this closure is then the return value of the function.
fn fill_utf16_buf_base<F1, F2, T>(mut f1: F1, f2: F2) -> Result<T, ()>
    where F1: FnMut(*mut u16, libc::DWORD) -> libc::DWORD,
          F2: FnOnce(&[u16]) -> T
{
    // Start off with a stack buf but then spill over to the heap if we end up
    // needing more space.
    let mut stack_buf = [0u16; 512];
    let mut heap_buf = Vec::new();
    unsafe {
        let mut n = stack_buf.len();
        loop {
            let buf = if n <= stack_buf.len() {
                &mut stack_buf[..]
            } else {
                let extra = n - heap_buf.len();
                heap_buf.reserve(extra);
                heap_buf.set_len(n);
                &mut heap_buf[..]
            };

            // This function is typically called on windows API functions which
            // will return the correct length of the string, but these functions
            // also return the `0` on error. In some cases, however, the
            // returned "correct length" may actually be 0!
            //
            // To handle this case we call `SetLastError` to reset it to 0 and
            // then check it again if we get the "0 error value". If the "last
            // error" is still 0 then we interpret it as a 0 length buffer and
            // not an actual error.
            c::SetLastError(0);
            let k = match f1(buf.as_mut_ptr(), n as libc::DWORD) {
                0 if libc::GetLastError() == 0 => 0,
                0 => return Err(()),
                n => n,
            } as usize;
            if k == n && libc::GetLastError() ==
                            libc::ERROR_INSUFFICIENT_BUFFER as libc::DWORD {
                n *= 2;
            } else if k >= n {
                n = k;
            } else {
                return Ok(f2(&buf[..k]))
            }
        }
    }
}

fn fill_utf16_buf<F1, F2, T>(f1: F1, f2: F2) -> IoResult<T>
    where F1: FnMut(*mut u16, libc::DWORD) -> libc::DWORD,
          F2: FnOnce(&[u16]) -> T
{
    fill_utf16_buf_base(f1, f2).map_err(|()| IoError::last_error())
}

fn fill_utf16_buf_new<F1, F2, T>(f1: F1, f2: F2) -> io::Result<T>
    where F1: FnMut(*mut u16, libc::DWORD) -> libc::DWORD,
          F2: FnOnce(&[u16]) -> T
{
    fill_utf16_buf_base(f1, f2).map_err(|()| io::Error::last_os_error())
}

fn os2path(s: &[u16]) -> Path {
    // FIXME: this should not be a panicking conversion (aka path reform)
    Path::new(String::from_utf16(s).unwrap())
}

pub fn truncate_utf16_at_nul<'a>(v: &'a [u16]) -> &'a [u16] {
    match v.iter().position(|c| *c == 0) {
        // don't include the 0
        Some(i) => &v[..i],
        None => v
    }
}

fn cvt<I: Int>(i: I) -> io::Result<I> {
    if i == Int::zero() {
        Err(io::Error::last_os_error())
    } else {
        Ok(i)
    }
}

fn ms_to_filetime(ms: u64) -> libc::FILETIME {
    // A FILETIME is a count of 100 nanosecond intervals, so we multiply by
    // 10000 b/c there are 10000 intervals in 1 ms
    let ms = ms * 10000;
    libc::FILETIME {
        dwLowDateTime: ms as u32,
        dwHighDateTime: (ms >> 32) as u32,
    }
}
<|endoftext|>


<commit_before>//! Infers how each function uses every register

use analysis::inst_combine;
use frontend::radeco_containers::{RadecoFunction, RadecoModule};
use middle::dce;
use middle::ir;
use middle::regfile::*;
use middle::ssa::cfg_traits::*;
use middle::ssa::ssa_traits::*;
use middle::ssa::ssastorage::SSAStorage;
use middle::ssa::utils;

use petgraph::visit::{DfsPostOrder, Walker};

use std::collections::{BTreeMap, HashSet};

/// For every function, patch all of its call sites to ignore registers that the
/// callee doesn't read and to preserve register values that the callee
/// preserves. Then, record which registers it reads and which registers it
/// preserves.
///
/// After this, all functions should have mutable [`regusage`][RadecoFunction::regusage]s.
///
/// This analysis is super conservative; for example, if a function preserves a
/// register by pushing it onto the stack and popping it back right before
/// returning, it is considered to be read and not preserved because we can't
/// guarantee that that stack location is never subsequently read or modified.
/// See #147 for further discussion
pub fn run(rmod: &mut RadecoModule, reginfo: &SubRegisterFile) -> () {
    Inferer::new().run(rmod, reginfo);
}

struct Inferer {
    /// Addresses of the functions we've already analyzed
    analyzed: HashSet<u64>,
}

impl Inferer {
    fn new() -> Inferer {
        Inferer {
            analyzed: HashSet::new(),
        }
    }

    /// Calls `patch_fn`, `dce::collect`, and `analyze_fn` on every function,
    /// callees first
    fn run(&mut self, rmod: &mut RadecoModule, reginfo: &SubRegisterFile) -> () {
        // for imports, *ASSUME* that the callconv that r2 says is correct
        {
            let imp_ru_iter = rmod.imports.iter().filter_map(|(&imp_addr, imp_info)| {
                let imp_rfn = imp_info.rfn.borrow();
                let regusage = reginfo.r2callconv_to_register_usage(
                    imp_rfn.callconv.as_ref()?, // ignore imports without callconvs
                    &*imp_rfn.callconv_name,
                )?;
                Some((imp_addr, regusage))
            });
            for (imp_addr, imp_ru) in imp_ru_iter {
                rmod.functions.get_mut(&imp_addr).unwrap().regusage = imp_ru;
                self.analyzed.insert(imp_addr);
            }
        }

        let mut dfs_wi = DfsPostOrder::empty(&rmod.callgraph).iter(&rmod.callgraph);
        // pick a function ...
        for fn_ni in rmod.callgraph.node_indices() {
            // ... and start a dfs on it
            dfs_wi.inner_mut().move_to(fn_ni);
            while let Some(fn_to_anal) = dfs_wi.next() {
                let fn_addr = rmod.callgraph[fn_to_anal];

                // ignore functions already in `call_convs` (probably because its an import)
                if !self.analyzed.contains(&fn_addr) {
                    self.patch_fn(fn_addr, &mut rmod.functions);

                    let rfn = &mut rmod.functions.get_mut(&fn_addr).unwrap();
                    dce::collect(rfn.ssa_mut());
                    // TODO preserve NodeIndex for recovering local variables
                    // inst_combine::run(rfn.ssa_mut());

                    let ru = self.analyze_fn(rfn, reginfo).unwrap_or_else(|| {
                        radeco_err!("Failed to analyze fn: {:?} (@ {:#X})", rfn.name, fn_addr);
                        // if analysis failed, default to "reads and clobbers everything"
                        reginfo.new_register_usage()
                    });

                    rfn.regusage = ru;
                    self.analyzed.insert(fn_addr);
                }
            }
        }
    }

    /// Using the callconv info we've gathered so far, patch-up call sites to
    /// to remove arguments that the callee doesn't read and make values in
    /// callee-saved registers be preserved across the call.
    fn patch_fn(&self, fn_addr: u64, fn_map: &mut BTreeMap<u64, RadecoFunction>) -> () {
        radeco_trace!("patching calls in fn: {}", fn_map[&fn_addr].name);
        for node in fn_map[&fn_addr].ssa().inorder_walk() {
            if let Ok(NodeType::Op(ir::MOpcode::OpCall)) =
                fn_map[&fn_addr].ssa().node_data(node).map(|nd| nd.nt)
            {
                self.patch_call_node(node, fn_addr, fn_map)
                    .unwrap_or_else(|| {
                        radeco_warn!(
                            "failed to remove unused args for call at {:#X}",
                            fn_map[&fn_addr].ssa().address(node).unwrap()
                        );
                    });
            }
        }
    }

    fn patch_call_node(
        &self,
        call_node: <SSAStorage as SSA>::ValueRef,
        fn_addr: u64,
        fn_map: &mut BTreeMap<u64, RadecoFunction>,
    ) -> Option<()> {
        // bail on indirect or weird call
        let (call_tgt_addr, call_reg_map) = direct_call_info(fn_map[&fn_addr].ssa(), call_node)?;

        // remove unread args
        for (regid, &op_node) in &call_reg_map {
            if fn_map[&call_tgt_addr].regusage.is_ignored(regid) {
                fn_map
                    .get_mut(&fn_addr)
                    .unwrap()
                    .ssa_mut()
                    .op_unuse(call_node, op_node);
            }
        }

        // bridge preserved registers
        for (regid, (use_node, _)) in utils::call_rets(call_node, fn_map[&fn_addr].ssa()) {
            if fn_map[&call_tgt_addr].regusage.is_preserved(regid) {
                fn_map
                    .get_mut(&fn_addr)
                    .unwrap()
                    .ssa_mut()
                    .replace_value(use_node, call_reg_map[regid]);
            }
        }

        Some(())
    }

    fn analyze_fn(&self, rfn: &RadecoFunction, reginfo: &SubRegisterFile) -> Option<RegisterUsage> {
        radeco_trace!("analyzing fn: {}", rfn.name);
        let ssa = rfn.ssa();
        let entry_regstate_node = ssa.registers_in(ssa.entry_node()?)?;
        let exit_regstate_node = ssa.registers_in(ssa.exit_node()?)?;
        // some registers may not be present in the entry node;
        // this means that the function neither reads nor preserves that register
        let entry_regstate = utils::register_state_info(entry_regstate_node, ssa);
        let exit_regstate = utils::register_state_info(exit_regstate_node, ssa);

        let mut ret = reginfo.new_register_usage();
        ret.set_all_ignored();

        for regid in ssa.regfile.iter_register_ids() {
            // ignore registers not in entry regstate
            if let Some(&(reg_val_entry, _)) = entry_regstate.get(regid) {
                // bail if a register isn't present in exit regstate
                let &(reg_val_exit, _) = exit_regstate.get(regid)?;

                if reg_val_exit == reg_val_entry {
                    ret.set_preserved(regid);
                }

                // find all uses, ignoring entry/exit register state
                let mut uses_iter = ssa
                    .uses_of(reg_val_entry)
                    .into_iter()
                    .filter(|&n| n != entry_regstate_node && n != exit_regstate_node);
                if uses_iter.next().is_some() {
                    ret.set_read(regid);
                }
            }
        }

        Some(ret)
    }
}

fn direct_call_info(
    ssa: &SSAStorage,
    call_node: <SSAStorage as SSA>::ValueRef,
) -> Option<(u64, RegisterMap<<SSAStorage as SSA>::ValueRef>)> {
    let callinfo = utils::call_info(call_node, ssa)?;
    Some((ssa.constant(callinfo.target)?, callinfo.register_args))
}
<commit_msg>Revert "Comment `inst_combine::run` out"<commit_after>//! Infers how each function uses every register

use analysis::inst_combine;
use frontend::radeco_containers::{RadecoFunction, RadecoModule};
use middle::dce;
use middle::ir;
use middle::regfile::*;
use middle::ssa::cfg_traits::*;
use middle::ssa::ssa_traits::*;
use middle::ssa::ssastorage::SSAStorage;
use middle::ssa::utils;

use petgraph::visit::{DfsPostOrder, Walker};

use std::collections::{BTreeMap, HashSet};

/// For every function, patch all of its call sites to ignore registers that the
/// callee doesn't read and to preserve register values that the callee
/// preserves. Then, record which registers it reads and which registers it
/// preserves.
///
/// After this, all functions should have mutable [`regusage`][RadecoFunction::regusage]s.
///
/// This analysis is super conservative; for example, if a function preserves a
/// register by pushing it onto the stack and popping it back right before
/// returning, it is considered to be read and not preserved because we can't
/// guarantee that that stack location is never subsequently read or modified.
/// See #147 for further discussion
pub fn run(rmod: &mut RadecoModule, reginfo: &SubRegisterFile) -> () {
    Inferer::new().run(rmod, reginfo);
}

struct Inferer {
    /// Addresses of the functions we've already analyzed
    analyzed: HashSet<u64>,
}

impl Inferer {
    fn new() -> Inferer {
        Inferer {
            analyzed: HashSet::new(),
        }
    }

    /// Calls `patch_fn`, `dce::collect`, and `analyze_fn` on every function,
    /// callees first
    fn run(&mut self, rmod: &mut RadecoModule, reginfo: &SubRegisterFile) -> () {
        // for imports, *ASSUME* that the callconv that r2 says is correct
        {
            let imp_ru_iter = rmod.imports.iter().filter_map(|(&imp_addr, imp_info)| {
                let imp_rfn = imp_info.rfn.borrow();
                let regusage = reginfo.r2callconv_to_register_usage(
                    imp_rfn.callconv.as_ref()?, // ignore imports without callconvs
                    &*imp_rfn.callconv_name,
                )?;
                Some((imp_addr, regusage))
            });
            for (imp_addr, imp_ru) in imp_ru_iter {
                rmod.functions.get_mut(&imp_addr).unwrap().regusage = imp_ru;
                self.analyzed.insert(imp_addr);
            }
        }

        let mut dfs_wi = DfsPostOrder::empty(&rmod.callgraph).iter(&rmod.callgraph);
        // pick a function ...
        for fn_ni in rmod.callgraph.node_indices() {
            // ... and start a dfs on it
            dfs_wi.inner_mut().move_to(fn_ni);
            while let Some(fn_to_anal) = dfs_wi.next() {
                let fn_addr = rmod.callgraph[fn_to_anal];

                // ignore functions already in `call_convs` (probably because its an import)
                if !self.analyzed.contains(&fn_addr) {
                    self.patch_fn(fn_addr, &mut rmod.functions);

                    let rfn = &mut rmod.functions.get_mut(&fn_addr).unwrap();
                    dce::collect(rfn.ssa_mut());
                    inst_combine::run(rfn.ssa_mut());

                    let ru = self.analyze_fn(rfn, reginfo).unwrap_or_else(|| {
                        radeco_err!("Failed to analyze fn: {:?} (@ {:#X})", rfn.name, fn_addr);
                        // if analysis failed, default to "reads and clobbers everything"
                        reginfo.new_register_usage()
                    });

                    rfn.regusage = ru;
                    self.analyzed.insert(fn_addr);
                }
            }
        }
    }

    /// Using the callconv info we've gathered so far, patch-up call sites to
    /// to remove arguments that the callee doesn't read and make values in
    /// callee-saved registers be preserved across the call.
    fn patch_fn(&self, fn_addr: u64, fn_map: &mut BTreeMap<u64, RadecoFunction>) -> () {
        radeco_trace!("patching calls in fn: {}", fn_map[&fn_addr].name);
        for node in fn_map[&fn_addr].ssa().inorder_walk() {
            if let Ok(NodeType::Op(ir::MOpcode::OpCall)) =
                fn_map[&fn_addr].ssa().node_data(node).map(|nd| nd.nt)
            {
                self.patch_call_node(node, fn_addr, fn_map)
                    .unwrap_or_else(|| {
                        radeco_warn!(
                            "failed to remove unused args for call at {:#X}",
                            fn_map[&fn_addr].ssa().address(node).unwrap()
                        );
                    });
            }
        }
    }

    fn patch_call_node(
        &self,
        call_node: <SSAStorage as SSA>::ValueRef,
        fn_addr: u64,
        fn_map: &mut BTreeMap<u64, RadecoFunction>,
    ) -> Option<()> {
        // bail on indirect or weird call
        let (call_tgt_addr, call_reg_map) = direct_call_info(fn_map[&fn_addr].ssa(), call_node)?;

        // remove unread args
        for (regid, &op_node) in &call_reg_map {
            if fn_map[&call_tgt_addr].regusage.is_ignored(regid) {
                fn_map
                    .get_mut(&fn_addr)
                    .unwrap()
                    .ssa_mut()
                    .op_unuse(call_node, op_node);
            }
        }

        // bridge preserved registers
        for (regid, (use_node, _)) in utils::call_rets(call_node, fn_map[&fn_addr].ssa()) {
            if fn_map[&call_tgt_addr].regusage.is_preserved(regid) {
                fn_map
                    .get_mut(&fn_addr)
                    .unwrap()
                    .ssa_mut()
                    .replace_value(use_node, call_reg_map[regid]);
            }
        }

        Some(())
    }

    fn analyze_fn(&self, rfn: &RadecoFunction, reginfo: &SubRegisterFile) -> Option<RegisterUsage> {
        radeco_trace!("analyzing fn: {}", rfn.name);
        let ssa = rfn.ssa();
        let entry_regstate_node = ssa.registers_in(ssa.entry_node()?)?;
        let exit_regstate_node = ssa.registers_in(ssa.exit_node()?)?;
        // some registers may not be present in the entry node;
        // this means that the function neither reads nor preserves that register
        let entry_regstate = utils::register_state_info(entry_regstate_node, ssa);
        let exit_regstate = utils::register_state_info(exit_regstate_node, ssa);

        let mut ret = reginfo.new_register_usage();
        ret.set_all_ignored();

        for regid in ssa.regfile.iter_register_ids() {
            // ignore registers not in entry regstate
            if let Some(&(reg_val_entry, _)) = entry_regstate.get(regid) {
                // bail if a register isn't present in exit regstate
                let &(reg_val_exit, _) = exit_regstate.get(regid)?;

                if reg_val_exit == reg_val_entry {
                    ret.set_preserved(regid);
                }

                // find all uses, ignoring entry/exit register state
                let mut uses_iter = ssa
                    .uses_of(reg_val_entry)
                    .into_iter()
                    .filter(|&n| n != entry_regstate_node && n != exit_regstate_node);
                if uses_iter.next().is_some() {
                    ret.set_read(regid);
                }
            }
        }

        Some(ret)
    }
}

fn direct_call_info(
    ssa: &SSAStorage,
    call_node: <SSAStorage as SSA>::ValueRef,
) -> Option<(u64, RegisterMap<<SSAStorage as SSA>::ValueRef>)> {
    let callinfo = utils::call_info(call_node, ssa)?;
    Some((ssa.constant(callinfo.target)?, callinfo.register_args))
}
<|endoftext|>


<commit_before>use std::f64;
use rand::Rng;
use rand::distributions::{Sample, IndependentSample};
use function::gamma;
use statistics::*;
use distribution::{Univariate, Continuous, Distribution};
use result::Result;
use error::StatsError;

/// Implements the [Inverse Gamma](https://en.wikipedia.org/wiki/Inverse-gamma_distribution) distribution
///
/// # Examples
///
/// ```
/// ```
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct InverseGamma {
    shape: f64,
    rate: f64,
}

impl InverseGamma {
    /// Constructs a new inverse gamma distribution with a shape (α)
    /// of `shape` and a rate (β) of `rate`
    ///
    /// # Errors
    ///
    /// Returns an error if `shape` or `rate` are `NaN`.
    /// Also returns an error if `shape <= 0.0` or `rate <= 0.0`
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::InverseGamma;
    ///
    /// let mut result = InverseGamma::new(3.0, 1.0);
    /// assert!(result.is_ok());
    ///
    /// result = InverseGamma::new(0.0, 0.0);
    /// assert!(result.is_err());
    /// ```
    pub fn new(shape: f64, rate: f64) -> Result<InverseGamma> {
        let is_nan = shape.is_nan() || rate.is_nan();
        match (shape, rate, is_nan) {
            (_, _, true) => Err(StatsError::BadParams),
            (_, _, false) if shape <= 0.0 || rate <= 0.0 => Err(StatsError::BadParams),
            (_, _, false) => {
                Ok(InverseGamma {
                    shape: shape,
                    rate: rate,
                })
            }
        }
    }

    /// Returns the shape (α) of the inverse gamma distribution
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::InverseGamma;
    ///
    /// let n = InverseGamma::new(3.0, 1.0).unwrap();
    /// assert_eq!(n.shape(), 3.0);
    /// ```
    pub fn shape(&self) -> f64 {
        self.shape
    }

    /// Returns the rate (β) of the inverse gamma distribution
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::InverseGamma;
    ///
    /// let n = InverseGamma::new(3.0, 1.0).unwrap();
    /// assert_eq!(n.rate(), 1.0);
    /// ```
    pub fn rate(&self) -> f64 {
        self.rate
    }
}

impl Sample<f64> for InverseGamma {
    /// Generate a random sample from an inverse gamma
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn sample<R: Rng>(&mut self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl IndependentSample<f64> for InverseGamma {
    /// Generate a random independent sample from an inverse gamma
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn ind_sample<R: Rng>(&self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl Distribution<f64> for InverseGamma {
    /// Generate a random sample from an inverse gamma distribution
    /// using `r` as the source of randomness.
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate rand;
    /// # extern crate statrs;
    /// use rand::StdRng;
    /// use statrs::distribution::{InverseGamma, Distribution};
    ///
    /// # fn main() {
    /// let mut r = rand::StdRng::new().unwrap();
    /// let n = InverseGamma::new(3.0, 1.0).unwrap();
    /// print!("{}", n.sample::<StdRng>(&mut r));
    /// # }
    /// ```
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        1.0 / super::gamma::sample_unchecked(r, self.shape, self.rate)
    }
}

impl Univariate<f64, f64> for InverseGamma {
    /// Calculates the cumulative distribution function for the inverse gamma
    /// distribution at `x`
    ///
    /// # Formula
    ///
    /// ```ignore
    /// Γ(α, β / x) / Γ(α)
    /// ```
    ///
    /// where the numerator is the upper incomplete gamma function,
    /// the denominator is the gamma function, `α` is the shape,
    /// and `β` is the rate
    fn cdf(&self, x: f64) -> f64 {
        gamma::gamma_ur(self.shape, self.rate / x)
    }
}

impl Min<f64> for InverseGamma {
    /// Returns the minimum value in the domain of the
    /// inverse gamma distribution representable by a double precision
    /// float
    ///
    /// # Formula
    ///
    /// ```ignore
    /// 0
    /// ```
    fn min(&self) -> f64 {
        0.0
    }
}

impl Max<f64> for InverseGamma {
    /// Returns the maximum value in the domain of the
    /// inverse gamma distribution representable by a double precision
    /// float
    ///
    /// # Formula
    ///
    /// ```ignore
    /// INF
    /// ```
    fn max(&self) -> f64 {
        f64::INFINITY
    }
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[cfg(test)]
mod test {
    use std::f64;
    use statistics::*;
    use distribution::{Univariate, Continuous, InverseGamma};

    fn try_create(shape: f64, rate: f64) -> InverseGamma {
        let n = InverseGamma::new(shape, rate);
        assert!(n.is_ok());
        n.unwrap()
    }

    fn create_case(shape: f64, rate: f64) {
        let n = try_create(shape, rate);
        assert_eq!(shape, n.shape());
        assert_eq!(rate, n.rate());
    }

    fn bad_create_case(shape: f64, rate: f64) {
        let n = InverseGamma::new(shape, rate);
        assert!(n.is_err());
    }

    fn get_value<F>(shape: f64, rate: f64, eval: F) -> f64
        where F: Fn(InverseGamma) -> f64
    {
        let n = try_create(shape, rate);
        eval(n)
    }

    fn test_case<F>(shape: f64, rate: f64, expected: f64, eval: F)
        where F: Fn(InverseGamma) -> f64
    {
        let x = get_value(shape, rate, eval);
        assert_eq!(expected, x);
    }

    fn test_almost<F>(shape: f64, rate: f64, expected: f64, acc: f64, eval: F)
        where F: Fn(InverseGamma) -> f64
    {
        let x = get_value(shape, rate, eval);
        assert_almost_eq!(expected, x, acc);
    }

    fn test_is_nan<F>(shape: f64, rate: f64, eval: F)
        where F: Fn(InverseGamma) -> f64
    {
        let x = get_value(shape, rate, eval);
        assert!(x.is_nan());
    }

    #[test]
    fn test_create() {
        create_case(0.1, 0.1);
        create_case(1.0, 1.0);
    }

    #[test]
    fn test_bad_create() {
        bad_create_case(0.0, 1.0);
        bad_create_case(-1.0, 1.0);
        bad_create_case(-100.0, 1.0);
        bad_create_case(f64::NEG_INFINITY, 1.0);
        bad_create_case(f64::NAN, 1.0);
        bad_create_case(1.0, 0.0);
        bad_create_case(1.0, -1.0);
        bad_create_case(1.0, -100.0);
        bad_create_case(1.0, f64::NEG_INFINITY);
        bad_create_case(1.0, f64::NAN);
    }

    #[test]
    fn test_cdf() {
        test_almost(0.1, 0.1, 0.1862151961946054271994, 1e-14, |x| x.cdf(1.2));
        test_almost(0.1, 1.0, 0.05859755410986647796141, 1e-14, |x| x.cdf(2.0));
        test_case(0.1, f64::INFINITY, 0.0, |x| x.cdf(1.1));
        test_case(1.0, 0.1, 0.9355069850316177377304, |x| x.cdf(1.5));
        test_almost(1.0, 1.0, 0.4345982085070782231613, 1e-14, |x| x.cdf(1.2));
        test_case(1.0, f64::INFINITY, 0.0, |x| x.cdf(1.5));
        test_is_nan(f64::INFINITY, 0.1, |x| x.cdf(5.0));
        test_is_nan(f64::INFINITY, 1.0, |x| x.cdf(2.5));
        test_is_nan(f64::INFINITY, f64::INFINITY, |x| x.cdf(1.0));
    }

    #[test]
    fn test_min_max() {
        test_case(1.0, 1.0, 0.0, |x| x.min());
        test_case(1.0, 1.0, f64::INFINITY, |x| x.max());
    }
}<commit_msg>wip: impl mean and variance for inverse gamma<commit_after>use std::f64;
use rand::Rng;
use rand::distributions::{Sample, IndependentSample};
use function::gamma;
use statistics::*;
use distribution::{Univariate, Continuous, Distribution};
use result::Result;
use error::StatsError;

/// Implements the [Inverse Gamma](https://en.wikipedia.org/wiki/Inverse-gamma_distribution) distribution
///
/// # Examples
///
/// ```
/// ```
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct InverseGamma {
    shape: f64,
    rate: f64,
}

impl InverseGamma {
    /// Constructs a new inverse gamma distribution with a shape (α)
    /// of `shape` and a rate (β) of `rate`
    ///
    /// # Errors
    ///
    /// Returns an error if `shape` or `rate` are `NaN`.
    /// Also returns an error if `shape <= 0.0` or `rate <= 0.0`
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::InverseGamma;
    ///
    /// let mut result = InverseGamma::new(3.0, 1.0);
    /// assert!(result.is_ok());
    ///
    /// result = InverseGamma::new(0.0, 0.0);
    /// assert!(result.is_err());
    /// ```
    pub fn new(shape: f64, rate: f64) -> Result<InverseGamma> {
        let is_nan = shape.is_nan() || rate.is_nan();
        match (shape, rate, is_nan) {
            (_, _, true) => Err(StatsError::BadParams),
            (_, _, false) if shape <= 0.0 || rate <= 0.0 => Err(StatsError::BadParams),
            (_, _, false) => {
                Ok(InverseGamma {
                    shape: shape,
                    rate: rate,
                })
            }
        }
    }

    /// Returns the shape (α) of the inverse gamma distribution
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::InverseGamma;
    ///
    /// let n = InverseGamma::new(3.0, 1.0).unwrap();
    /// assert_eq!(n.shape(), 3.0);
    /// ```
    pub fn shape(&self) -> f64 {
        self.shape
    }

    /// Returns the rate (β) of the inverse gamma distribution
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::InverseGamma;
    ///
    /// let n = InverseGamma::new(3.0, 1.0).unwrap();
    /// assert_eq!(n.rate(), 1.0);
    /// ```
    pub fn rate(&self) -> f64 {
        self.rate
    }
}

impl Sample<f64> for InverseGamma {
    /// Generate a random sample from an inverse gamma
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn sample<R: Rng>(&mut self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl IndependentSample<f64> for InverseGamma {
    /// Generate a random independent sample from an inverse gamma
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn ind_sample<R: Rng>(&self, r: &mut R) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl Distribution<f64> for InverseGamma {
    /// Generate a random sample from an inverse gamma distribution
    /// using `r` as the source of randomness.
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate rand;
    /// # extern crate statrs;
    /// use rand::StdRng;
    /// use statrs::distribution::{InverseGamma, Distribution};
    ///
    /// # fn main() {
    /// let mut r = rand::StdRng::new().unwrap();
    /// let n = InverseGamma::new(3.0, 1.0).unwrap();
    /// print!("{}", n.sample::<StdRng>(&mut r));
    /// # }
    /// ```
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        1.0 / super::gamma::sample_unchecked(r, self.shape, self.rate)
    }
}

impl Univariate<f64, f64> for InverseGamma {
    /// Calculates the cumulative distribution function for the inverse gamma
    /// distribution at `x`
    ///
    /// # Formula
    ///
    /// ```ignore
    /// Γ(α, β / x) / Γ(α)
    /// ```
    ///
    /// where the numerator is the upper incomplete gamma function,
    /// the denominator is the gamma function, `α` is the shape,
    /// and `β` is the rate
    fn cdf(&self, x: f64) -> f64 {
        gamma::gamma_ur(self.shape, self.rate / x)
    }
}

impl Min<f64> for InverseGamma {
    /// Returns the minimum value in the domain of the
    /// inverse gamma distribution representable by a double precision
    /// float
    ///
    /// # Formula
    ///
    /// ```ignore
    /// 0
    /// ```
    fn min(&self) -> f64 {
        0.0
    }
}

impl Max<f64> for InverseGamma {
    /// Returns the maximum value in the domain of the
    /// inverse gamma distribution representable by a double precision
    /// float
    ///
    /// # Formula
    ///
    /// ```ignore
    /// INF
    /// ```
    fn max(&self) -> f64 {
        f64::INFINITY
    }
}

impl Mean<f64> for InverseGamma {
    /// Returns the mean of the inverse distribution
    ///
    /// # Panics
    ///
    /// If `shape <= 1.0`
    ///
    /// # Formula
    ///
    /// ```ignore
    /// β / (α - 1)
    /// ```
    ///
    /// where `α` is the shape and `β` is the rate
    fn mean(&self) -> f64 {
        assert!(self.shape > 1.0,
                format!("{}", StatsError::ArgGt("shape", 1.0)));
        self.rate / (self.shape - 1.0)
    }
}

impl Variance<f64> for InverseGamma {
    /// Returns the variance of the inverse gamma distribution
    ///
    /// # Panics
    ///
    /// If `shape <= 2.0`
    ///
    /// # Formula
    ///
    /// ```ignore
    /// β^2 / ((α - 1)^2 * (α - 2))
    /// ```
    ///
    /// where `α` is the shape and `β` is the rate
    fn variance(&self) -> f64 {
        assert!(self.shape > 2.0,
                format!("{}", StatsError::ArgGt("shape", 2.0)));
        self.rate * self.rate / ((self.shape - 1.0) * (self.shape - 1.0) * (self.shape - 2.0))
    }

    /// Returns the standard deviation of the inverse gamma distribution
    ///
    /// # Panics
    ///
    /// If `shape <= 2.0`
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sqrt(β^2 / ((α - 1)^2 * (α - 2)))
    /// ```
    ///
    /// where `α` is the shape and `β` is the rate
    fn std_dev(&self) -> f64 {
        self.variance().sqrt()
    }
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[cfg(test)]
mod test {
    use std::f64;
    use statistics::*;
    use distribution::{Univariate, Continuous, InverseGamma};

    fn try_create(shape: f64, rate: f64) -> InverseGamma {
        let n = InverseGamma::new(shape, rate);
        assert!(n.is_ok());
        n.unwrap()
    }

    fn create_case(shape: f64, rate: f64) {
        let n = try_create(shape, rate);
        assert_eq!(shape, n.shape());
        assert_eq!(rate, n.rate());
    }

    fn bad_create_case(shape: f64, rate: f64) {
        let n = InverseGamma::new(shape, rate);
        assert!(n.is_err());
    }

    fn get_value<F>(shape: f64, rate: f64, eval: F) -> f64
        where F: Fn(InverseGamma) -> f64
    {
        let n = try_create(shape, rate);
        eval(n)
    }

    fn test_case<F>(shape: f64, rate: f64, expected: f64, eval: F)
        where F: Fn(InverseGamma) -> f64
    {
        let x = get_value(shape, rate, eval);
        assert_eq!(expected, x);
    }

    fn test_almost<F>(shape: f64, rate: f64, expected: f64, acc: f64, eval: F)
        where F: Fn(InverseGamma) -> f64
    {
        let x = get_value(shape, rate, eval);
        assert_almost_eq!(expected, x, acc);
    }

    fn test_is_nan<F>(shape: f64, rate: f64, eval: F)
        where F: Fn(InverseGamma) -> f64
    {
        let x = get_value(shape, rate, eval);
        assert!(x.is_nan());
    }

    #[test]
    fn test_create() {
        create_case(0.1, 0.1);
        create_case(1.0, 1.0);
    }

    #[test]
    fn test_bad_create() {
        bad_create_case(0.0, 1.0);
        bad_create_case(-1.0, 1.0);
        bad_create_case(-100.0, 1.0);
        bad_create_case(f64::NEG_INFINITY, 1.0);
        bad_create_case(f64::NAN, 1.0);
        bad_create_case(1.0, 0.0);
        bad_create_case(1.0, -1.0);
        bad_create_case(1.0, -100.0);
        bad_create_case(1.0, f64::NEG_INFINITY);
        bad_create_case(1.0, f64::NAN);
    }

    #[test]
    fn test_cdf() {
        test_almost(0.1, 0.1, 0.1862151961946054271994, 1e-14, |x| x.cdf(1.2));
        test_almost(0.1, 1.0, 0.05859755410986647796141, 1e-14, |x| x.cdf(2.0));
        test_case(0.1, f64::INFINITY, 0.0, |x| x.cdf(1.1));
        test_case(1.0, 0.1, 0.9355069850316177377304, |x| x.cdf(1.5));
        test_almost(1.0, 1.0, 0.4345982085070782231613, 1e-14, |x| x.cdf(1.2));
        test_case(1.0, f64::INFINITY, 0.0, |x| x.cdf(1.5));
        test_is_nan(f64::INFINITY, 0.1, |x| x.cdf(5.0));
        test_is_nan(f64::INFINITY, 1.0, |x| x.cdf(2.5));
        test_is_nan(f64::INFINITY, f64::INFINITY, |x| x.cdf(1.0));
    }

    #[test]
    #[should_panic]
    fn test_cdf_panics() {
        get_value(0.1, 0.1, |x| x.cdf(-1.0));
    }

    #[test]
    fn test_min_max() {
        test_case(1.0, 1.0, 0.0, |x| x.min());
        test_case(1.0, 1.0, f64::INFINITY, |x| x.max());
    }

    #[test]
    fn test_mean() {
        test_almost(1.1, 0.1, 1.0, 1e-14, |x| x.mean());
        test_almost(1.1, 1.0, 10.0, 1e-14, |x| x.mean());
        test_is_nan(f64::INFINITY, f64::INFINITY, |x| x.mean());
    }

    #[test]
    #[should_panic]
    fn test_mean_panics() {
        get_value(0.1, 0.1, |x| x.mean());
    }

    #[test]
    fn test_variance() {
        test_almost(2.1, 0.1, 0.08264462809917355371901, 1e-15, |x| x.variance());
        test_almost(2.1, 1.0, 8.264462809917355371901, 1e-13, |x| x.variance());
        test_is_nan(f64::INFINITY, f64::INFINITY, |x| x.variance());
    }

    #[test]
    #[should_panic]
    fn test_variance_panics() {
        get_value(0.1, 0.1, |x| x.variance());
    }

    #[test]
    fn test_std_dev() {
        test_almost(2.1, 0.1, 0.2874797872880344847272, 1e-15, |x| x.std_dev());
        test_almost(2.1, 1.0, 2.874797872880344847272, 1e-14, |x| x.std_dev());
        test_is_nan(f64::INFINITY, f64::INFINITY, |x| x.std_dev());
    }

    #[test]
    #[should_panic]
    fn test_std_dev_panics() {
        get_value(0.1, 0.1, |x| x.std_dev());
    }
}<|endoftext|>


<commit_before>use crate::core::prelude::*;

use itertools::Either;

fn is_event_owned(event: &Event, owned_tags: &[String]) -> bool {
    // Exclusive ownership of events is determined by the associated
    // tags.
    owned_tags
        .iter()
        .any(|ref tag| event.tags.iter().any(|ref e_tag| e_tag == tag))
}

pub fn strip_event_details_on_search(
    event_iter: impl Iterator<Item = Event>,
    owned_tags: Vec<String>,
) -> impl Iterator<Item = Event> {
    event_iter.map(move |e| {
        // Hide the created_by e-mail address if the event is not owned.
        if is_event_owned(&e, &owned_tags) {
            // Include created_by
            e
        } else {
            // Exclude created_by
            Event {
                created_by: None,
                ..e
            }
        }
    })
}

pub fn strip_event_details_on_export(
    event_iter: impl Iterator<Item = Event>,
    role: Role,
    owned_tags: Vec<String>,
) -> impl Iterator<Item = Event> {
    if role == Role::Admin {
        // Admin sees everything even if no owned tags are provided
        Either::Left(event_iter)
    } else {
        Either::Right(event_iter.map(move |e| {
            // Contact details are only visible for scouts and admins
            if role >= Role::Scout {
                // Include contact details
                let owned = is_event_owned(&e, &owned_tags);
                if owned {
                    // Include created_by
                    e
                } else {
                    // Exclude created_by
                    Event {
                        created_by: None,
                        ..e
                    }
                }
            } else {
                // Exclude both contact details and created_by
                Event {
                    contact: None,
                    created_by: None,
                    ..e
                }
            }
        }))
    }
}
<commit_msg>Make admin filter future proof for new roles<commit_after>use crate::core::prelude::*;

use itertools::Either;

fn is_event_owned(event: &Event, owned_tags: &[String]) -> bool {
    // Exclusive ownership of events is determined by the associated
    // tags.
    owned_tags
        .iter()
        .any(|ref tag| event.tags.iter().any(|ref e_tag| e_tag == tag))
}

pub fn strip_event_details_on_search(
    event_iter: impl Iterator<Item = Event>,
    owned_tags: Vec<String>,
) -> impl Iterator<Item = Event> {
    event_iter.map(move |e| {
        // Hide the created_by e-mail address if the event is not owned.
        if is_event_owned(&e, &owned_tags) {
            // Include created_by
            e
        } else {
            // Exclude created_by
            Event {
                created_by: None,
                ..e
            }
        }
    })
}

pub fn strip_event_details_on_export(
    event_iter: impl Iterator<Item = Event>,
    role: Role,
    owned_tags: Vec<String>,
) -> impl Iterator<Item = Event> {
    if role >= Role::Admin {
        // Admin sees everything even if no owned tags are provided
        Either::Left(event_iter)
    } else {
        Either::Right(event_iter.map(move |e| {
            // Contact details are only visible for scouts and admins
            if role >= Role::Scout {
                // Include contact details
                let owned = is_event_owned(&e, &owned_tags);
                if owned {
                    // Include created_by
                    e
                } else {
                    // Exclude created_by
                    Event {
                        created_by: None,
                        ..e
                    }
                }
            } else {
                // Exclude both contact details and created_by
                Event {
                    contact: None,
                    created_by: None,
                    ..e
                }
            }
        }))
    }
}
<|endoftext|>