<commit_before>extern crate glib;
extern crate gstreamer as gst;
extern crate gtk;

#[cfg(any(feature = "profiling-listener", feature = "profiling-tracker"))]
use chrono::Utc;

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

use std::path::PathBuf;

use std::sync::mpsc::{channel, Receiver};

use gtk::prelude::*;

use media::{PlaybackContext, ContextMessage};
use media::ContextMessage::*;

use super::{AudioController, ExportController, InfoController, VideoController};

#[derive(Clone, PartialEq)]
pub enum ControllerState {
    EOS,
    Ready,
    Paused,
    PendingExportToc,
    PendingSelectMedia,
    Playing,
    Stopped,
}

const LISTENER_PERIOD: u32 = 250; // 250 ms (4 Hz)
const TRACKER_PERIOD: u32 = 17; //  17 ms (60 Hz)

pub struct MainController {
    window: gtk::ApplicationWindow,
    header_bar: gtk::HeaderBar,
    play_pause_btn: gtk::ToolButton,
    export_toc_btn: gtk::Button,

    video_ctrl: VideoController,
    info_ctrl: Rc<RefCell<InfoController>>,
    audio_ctrl: Rc<RefCell<AudioController>>,
    export_ctrl: Rc<RefCell<ExportController>>,

    context: Option<PlaybackContext>,
    state: ControllerState,
    duration: Option<u64>, // duration is accurately known at eos
    seeking: bool,

    requires_async_dialog: bool,    // when pipeline contains video, dialogs must wait for
                                    // asyncdone before opening a dialog otherwise the listener
                                    // may borrow the MainController while the dialog is already
                                    // using it leading to a borrowing conflict

    this_opt: Option<Rc<RefCell<MainController>>>,
    keep_going: bool,
    listener_src: Option<glib::SourceId>,
    tracker_src: Option<glib::SourceId>,
}

impl MainController {
    pub fn new(builder: gtk::Builder) -> Rc<RefCell<Self>> {
        let export_toc_btn: gtk::Button = builder.get_object("export_toc-btn").unwrap();
        export_toc_btn.set_sensitive(false);

        let this = Rc::new(RefCell::new(MainController {
            window: builder.get_object("application-window").unwrap(),
            header_bar: builder.get_object("header-bar").unwrap(),
            play_pause_btn: builder.get_object("play_pause-toolbutton").unwrap(),
            export_toc_btn: export_toc_btn,

            video_ctrl: VideoController::new(&builder),
            info_ctrl: InfoController::new(&builder),
            audio_ctrl: AudioController::new(&builder),
            export_ctrl: ExportController::new(&builder),

            context: None,
            state: ControllerState::Stopped,
            duration: None,
            seeking: false,

            requires_async_dialog: false,

            this_opt: None,
            keep_going: true,
            listener_src: None,
            tracker_src: None,
        }));

        {
            let mut this_mut = this.borrow_mut();

            let this_rc = Rc::clone(&this);
            this_mut.this_opt = Some(this_rc);

            this_mut.window.connect_delete_event(|_, _| {
                gtk::main_quit();
                Inhibit(false)
            });
            this_mut.window.set_titlebar(&this_mut.header_bar);

            let this_rc = Rc::clone(&this);
            this_mut.play_pause_btn.connect_clicked(move |_| {
                this_rc.borrow_mut().play_pause();
            });

            // TODO: add key bindings to seek by steps
            // play/pause, etc.

            let this_rc = Rc::clone(&this);
            this_mut.export_toc_btn.connect_clicked(move |_| {
                let mut this = this_rc.borrow_mut();

                if this.requires_async_dialog && this.state == ControllerState::Playing {
                    this.hold();
                    this.state = ControllerState::PendingExportToc;
                } else {
                    this.hold();
                    this.export_toc();
                }
            });

            this_mut.video_ctrl.register_callbacks(&this);
            InfoController::register_callbacks(&this_mut.info_ctrl, &this);
            AudioController::register_callbacks(&this_mut.audio_ctrl, &this);
            ExportController::register_callbacks(&this_mut.export_ctrl, &this);
        }

        let open_btn: gtk::Button = builder.get_object("open-btn").unwrap();
        let this_rc = Rc::clone(&this);
        open_btn.connect_clicked(move |_| {
            let mut this = this_rc.borrow_mut();

            if this.requires_async_dialog && this.state == ControllerState::Playing {
                this.hold();
                this.state = ControllerState::PendingSelectMedia;
            } else {
                this.hold();
                this.select_media();
            }
        });

        this
    }

    pub fn show_all(&self) {
        self.window.show_all();
    }

    pub fn get_state(&self) -> &ControllerState {
        &self.state
    }

    fn play_pause(&mut self) {
        let context = match self.context.take() {
            Some(context) => context,
            None => {
                self.select_media();
                return;
            }
        };

        if self.state != ControllerState::EOS {
            match context.get_state() {
                gst::State::Paused => {
                    self.register_tracker();
                    self.play_pause_btn.set_icon_name("media-playback-pause");
                    self.state = ControllerState::Playing;
                    context.play().unwrap();
                    self.context = Some(context);
                }
                gst::State::Playing => {
                    context.pause().unwrap();
                    self.play_pause_btn.set_icon_name("media-playback-start");
                    self.remove_tracker();
                    self.state = ControllerState::Paused;
                    self.context = Some(context);
                }
                _ => {
                    self.context = Some(context);
                    self.select_media();
                }
            };
        } else {
            // Restart the stream from the begining
            self.context = Some(context);
            self.seek(0, true); // accurate (slow)
        }
    }

    pub fn seek(&mut self, position: u64, accurate: bool) {
        if self.state != ControllerState::Stopped {
            self.seeking = true;

            self.context
                .as_ref()
                .expect("MainController::seek no context")
                .seek(position, accurate);

            if self.state == ControllerState::EOS || self.state == ControllerState::Ready {
                if self.state == ControllerState::Ready {
                    self.context
                        .as_ref()
                        .expect("MainController::seek no context")
                        .play()
                        .unwrap();
                }
                self.register_tracker();
                self.play_pause_btn.set_icon_name("media-playback-pause");
                self.state = ControllerState::Playing;
            }

            self.info_ctrl.borrow_mut().seek(position, &self.state);
            self.audio_ctrl.borrow_mut().seek(position, &self.state);
        }
    }

    pub fn get_position(&mut self) -> u64 {
        self.context
            .as_mut()
            .expect("MainController::get_position no context")
            .get_position()
    }

    fn hold(&mut self) {
        self.remove_tracker();
        if let Some(context) = self.context.as_mut() {
            context.pause().unwrap();
        };

        self.play_pause_btn.set_icon_name("media-playback-start");
    }

    fn export_toc(&mut self) {
        match self.context.take() {
            Some(mut context) => {
                self.info_ctrl.borrow().export_chapters(&mut context);
                self.export_ctrl.borrow_mut().open(context);
                self.state = ControllerState::Paused;
            }
            None => (),
        }
    }

    pub fn restore_context(&mut self, context: PlaybackContext) {
        self.context = Some(context);
        self.state = ControllerState::Paused;
    }

    fn handle_eos(&mut self) {
        #[cfg(feature = "trace-main-controller")]
        println!("MainController::handle_eos");

        self.play_pause_btn.set_icon_name("media-playback-start");
        self.state = ControllerState::EOS;
    }

    fn remove_listener(&mut self) {
        if let Some(source_id) = self.listener_src.take() {
            glib::source_remove(source_id);
        }
    }

    fn register_listener(&mut self, timeout: u32, ui_rx: Receiver<ContextMessage>) {
        if self.listener_src.is_some() {
            return;
        }

        let this_rc = Rc::clone(self.this_opt.as_ref().unwrap());

        self.listener_src = Some(gtk::timeout_add(timeout, move || {
            let mut keep_going = true;

            for message in ui_rx.try_iter() {
                match message {
                    AsyncDone => {
                        let mut this = this_rc.borrow_mut();
                        this.seeking = false;
                        match this.state {
                            ControllerState::PendingSelectMedia => this.select_media(),
                            ControllerState::PendingExportToc => this.export_toc(),
                            _ => (),
                        }
                    }
                    InitDone => {
                        let mut this = this_rc.borrow_mut();

                        let context = this.context.take().expect(
                            "MainController: InitDone but no context available",
                        );

                        this.requires_async_dialog = context.info.lock()
                            .expect(
                                "MainController:: failed to lock media info in InitDone",
                            )
                            .video_best.is_some();

                        this.header_bar.set_subtitle(
                            Some(context.file_name.as_str()),
                        );

                        this.video_ctrl.new_media(&context);
                        this.info_ctrl.borrow_mut().new_media(&context);
                        this.audio_ctrl.borrow_mut().new_media(&context);

                        this.context = Some(context);

                        this.state = ControllerState::Ready;
                    }
                    Eos => {
                        let mut this = this_rc.borrow_mut();

                        let position = this.get_position();
                        this.duration = Some(position);

                        {
                            let mut info_ctrl = this.info_ctrl.borrow_mut();
                            info_ctrl.update_duration(position);
                            info_ctrl.tick(position, true);
                        }

                        this.audio_ctrl.borrow_mut().tick();

                        this.handle_eos();

                        // Remove listener and tracker.
                        // Note: tracker will be register again in case of
                        // a seek. Listener is of no use anymore because
                        // the context won't send any more Eos nor AsyncDone
                        // after an EOS.
                        keep_going = false;
                    }
                    FailedToOpenMedia => {
                        eprintln!("ERROR: failed to open media");

                        let mut this = this_rc.borrow_mut();

                        this.context = None;
                        this.keep_going = false;
                        keep_going = false;
                    }
                };

                if !keep_going {
                    break;
                }
            }

            if !keep_going {
                let mut this = this_rc.borrow_mut();
                this.listener_src = None;
                this.tracker_src = None;
            }

            glib::Continue(keep_going)
        }));
    }

    fn remove_tracker(&mut self) {
        if let Some(source_id) = self.tracker_src.take() {
            glib::source_remove(source_id);
        }
    }

    fn register_tracker(&mut self) {
        if self.tracker_src.is_some() {
            return;
        }

        let this_rc = Rc::clone(self.this_opt.as_ref().unwrap());

        self.tracker_src = Some(gtk::timeout_add(TRACKER_PERIOD, move || {
            #[cfg(feature = "profiling-tracker")]
            let start = Utc::now();

            let mut keep_going = true;

            let mut this = this_rc.borrow_mut();

            #[cfg(feature = "profiling-tracker")]
            let before_tick = Utc::now();

            let position = this
                .context
                .as_mut()
                .expect("MainController::tracker no context while getting position")
                .get_position();

            let is_eos = if let Some(duration) = this.duration {
                if position >= duration {
                    if !this.seeking {
                        // this check is necessary as EOS is not sent
                        // in case of a seek after EOS
                        this.handle_eos();
                        this.tracker_src = None;
                        keep_going = false;
                    }
                    true
                } else if this.seeking {
                    // this check is necessary as AsyncDone is not sent
                    // in case of a seek after EOS
                    this.seeking = false;
                    false
                } else {
                    false
                }
            } else {
                false
            };

            if !this.seeking {
                this.info_ctrl.borrow_mut().tick(position, is_eos);
                this.audio_ctrl.borrow_mut().tick();
            }

            #[cfg(feature = "profiling-tracker")]
            let end = Utc::now();

            #[cfg(feature = "profiling-tracker")]
            println!(
                "tracker,{},{},{}",
                start.time().format("%H:%M:%S%.6f"),
                before_tick.time().format("%H:%M:%S%.6f"),
                end.time().format("%H:%M:%S%.6f"),
            );

            glib::Continue(this.keep_going && keep_going)
        }));
    }

    fn select_media(&mut self) {
        let file_dlg = gtk::FileChooserDialog::new(
            Some("Open a media file"),
            Some(&self.window),
            gtk::FileChooserAction::Open,
        );
        // Note: couldn't find equivalents for STOCK_OK
        file_dlg.add_button("Open", gtk::ResponseType::Ok.into());

        if file_dlg.run() == gtk::ResponseType::Ok.into() {
            if let Some(ref context) = self.context {
                context.stop();
            }
            self.open_media(file_dlg.get_filename().unwrap());
        } else {
            self.state = ControllerState::Paused;;
        }

        file_dlg.close();
    }

    fn open_media(&mut self, filepath: PathBuf) {
        self.remove_listener();

        self.info_ctrl.borrow_mut().cleanup();
        self.audio_ctrl.borrow_mut().cleanup();
        self.header_bar.set_subtitle("");
        self.export_toc_btn.set_sensitive(false);
        self.duration = None;

        let (ctx_tx, ui_rx) = channel();

        self.seeking = false;
        self.keep_going = true;
        self.register_listener(LISTENER_PERIOD, ui_rx);

        let dbl_buffer_mtx = self.audio_ctrl.borrow().get_dbl_buffer_mtx().clone();
        match PlaybackContext::new(filepath, dbl_buffer_mtx, ctx_tx) {
            Ok(context) => {
                self.context = Some(context);
                self.export_toc_btn.set_sensitive(true);
            }
            Err(error) => {
                self.remove_listener();
                eprintln!("Error opening media: {}", error);
            }
        };
    }
}
<commit_msg>MainWindow: set unsensitive while running in the background<commit_after>extern crate glib;
extern crate gstreamer as gst;
extern crate gtk;

#[cfg(any(feature = "profiling-listener", feature = "profiling-tracker"))]
use chrono::Utc;

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

use std::path::PathBuf;

use std::sync::mpsc::{channel, Receiver};

use gtk::prelude::*;

use media::{PlaybackContext, ContextMessage};
use media::ContextMessage::*;

use super::{AudioController, ExportController, InfoController, VideoController};

#[derive(Clone, PartialEq)]
pub enum ControllerState {
    EOS,
    Ready,
    Paused,
    PendingExportToc,
    PendingSelectMedia,
    Playing,
    Stopped,
}

const LISTENER_PERIOD: u32 = 250; // 250 ms (4 Hz)
const TRACKER_PERIOD: u32 = 17; //  17 ms (60 Hz)

pub struct MainController {
    window: gtk::ApplicationWindow,
    header_bar: gtk::HeaderBar,
    play_pause_btn: gtk::ToolButton,
    export_toc_btn: gtk::Button,

    video_ctrl: VideoController,
    info_ctrl: Rc<RefCell<InfoController>>,
    audio_ctrl: Rc<RefCell<AudioController>>,
    export_ctrl: Rc<RefCell<ExportController>>,

    context: Option<PlaybackContext>,
    state: ControllerState,
    duration: Option<u64>, // duration is accurately known at eos
    seeking: bool,

    requires_async_dialog: bool,    // when pipeline contains video, dialogs must wait for
                                    // asyncdone before opening a dialog otherwise the listener
                                    // may borrow the MainController while the dialog is already
                                    // using it leading to a borrowing conflict

    this_opt: Option<Rc<RefCell<MainController>>>,
    keep_going: bool,
    listener_src: Option<glib::SourceId>,
    tracker_src: Option<glib::SourceId>,
}

impl MainController {
    pub fn new(builder: gtk::Builder) -> Rc<RefCell<Self>> {
        let export_toc_btn: gtk::Button = builder.get_object("export_toc-btn").unwrap();
        export_toc_btn.set_sensitive(false);

        let this = Rc::new(RefCell::new(MainController {
            window: builder.get_object("application-window").unwrap(),
            header_bar: builder.get_object("header-bar").unwrap(),
            play_pause_btn: builder.get_object("play_pause-toolbutton").unwrap(),
            export_toc_btn: export_toc_btn,

            video_ctrl: VideoController::new(&builder),
            info_ctrl: InfoController::new(&builder),
            audio_ctrl: AudioController::new(&builder),
            export_ctrl: ExportController::new(&builder),

            context: None,
            state: ControllerState::Stopped,
            duration: None,
            seeking: false,

            requires_async_dialog: false,

            this_opt: None,
            keep_going: true,
            listener_src: None,
            tracker_src: None,
        }));

        {
            let mut this_mut = this.borrow_mut();

            let this_rc = Rc::clone(&this);
            this_mut.this_opt = Some(this_rc);

            this_mut.window.connect_delete_event(|_, _| {
                gtk::main_quit();
                Inhibit(false)
            });
            this_mut.window.set_titlebar(&this_mut.header_bar);

            let this_rc = Rc::clone(&this);
            this_mut.play_pause_btn.connect_clicked(move |_| {
                this_rc.borrow_mut().play_pause();
            });

            // TODO: add key bindings to seek by steps
            // play/pause, etc.

            let this_rc = Rc::clone(&this);
            this_mut.export_toc_btn.connect_clicked(move |_| {
                let mut this = this_rc.borrow_mut();

                if this.requires_async_dialog && this.state == ControllerState::Playing {
                    this.hold();
                    this.state = ControllerState::PendingExportToc;
                } else {
                    this.hold();
                    this.export_toc();
                }
            });

            this_mut.video_ctrl.register_callbacks(&this);
            InfoController::register_callbacks(&this_mut.info_ctrl, &this);
            AudioController::register_callbacks(&this_mut.audio_ctrl, &this);
            ExportController::register_callbacks(&this_mut.export_ctrl, &this);
        }

        let open_btn: gtk::Button = builder.get_object("open-btn").unwrap();
        let this_rc = Rc::clone(&this);
        open_btn.connect_clicked(move |_| {
            let mut this = this_rc.borrow_mut();

            if this.requires_async_dialog && this.state == ControllerState::Playing {
                this.hold();
                this.state = ControllerState::PendingSelectMedia;
            } else {
                this.hold();
                this.select_media();
            }
        });

        this
    }

    pub fn show_all(&self) {
        self.window.show_all();
    }

    pub fn get_state(&self) -> &ControllerState {
        &self.state
    }

    fn play_pause(&mut self) {
        let context = match self.context.take() {
            Some(context) => context,
            None => {
                self.select_media();
                return;
            }
        };

        if self.state != ControllerState::EOS {
            match context.get_state() {
                gst::State::Paused => {
                    self.register_tracker();
                    self.play_pause_btn.set_icon_name("media-playback-pause");
                    self.state = ControllerState::Playing;
                    context.play().unwrap();
                    self.context = Some(context);
                }
                gst::State::Playing => {
                    context.pause().unwrap();
                    self.play_pause_btn.set_icon_name("media-playback-start");
                    self.remove_tracker();
                    self.state = ControllerState::Paused;
                    self.context = Some(context);
                }
                _ => {
                    self.context = Some(context);
                    self.select_media();
                }
            };
        } else {
            // Restart the stream from the begining
            self.context = Some(context);
            self.seek(0, true); // accurate (slow)
        }
    }

    pub fn seek(&mut self, position: u64, accurate: bool) {
        if self.state != ControllerState::Stopped {
            self.seeking = true;

            self.context
                .as_ref()
                .expect("MainController::seek no context")
                .seek(position, accurate);

            if self.state == ControllerState::EOS || self.state == ControllerState::Ready {
                if self.state == ControllerState::Ready {
                    self.context
                        .as_ref()
                        .expect("MainController::seek no context")
                        .play()
                        .unwrap();
                }
                self.register_tracker();
                self.play_pause_btn.set_icon_name("media-playback-pause");
                self.state = ControllerState::Playing;
            }

            self.info_ctrl.borrow_mut().seek(position, &self.state);
            self.audio_ctrl.borrow_mut().seek(position, &self.state);
        }
    }

    pub fn get_position(&mut self) -> u64 {
        self.context
            .as_mut()
            .expect("MainController::get_position no context")
            .get_position()
    }

    fn hold(&mut self) {
        self.remove_tracker();
        if let Some(context) = self.context.as_mut() {
            context.pause().unwrap();
        };

        self.play_pause_btn.set_icon_name("media-playback-start");
        self.window.set_sensitive(false);
    }

    fn export_toc(&mut self) {
        match self.context.take() {
            Some(mut context) => {
                self.info_ctrl.borrow().export_chapters(&mut context);
                self.export_ctrl.borrow_mut().open(context);
                self.state = ControllerState::Paused;
            }
            None => (),
        }
    }

    pub fn restore_context(&mut self, context: PlaybackContext) {
        self.context = Some(context);
        self.state = ControllerState::Paused;
        self.window.set_sensitive(true);
    }

    fn handle_eos(&mut self) {
        #[cfg(feature = "trace-main-controller")]
        println!("MainController::handle_eos");

        self.play_pause_btn.set_icon_name("media-playback-start");
        self.state = ControllerState::EOS;
    }

    fn remove_listener(&mut self) {
        if let Some(source_id) = self.listener_src.take() {
            glib::source_remove(source_id);
        }
    }

    fn register_listener(&mut self, timeout: u32, ui_rx: Receiver<ContextMessage>) {
        if self.listener_src.is_some() {
            return;
        }

        let this_rc = Rc::clone(self.this_opt.as_ref().unwrap());

        self.listener_src = Some(gtk::timeout_add(timeout, move || {
            let mut keep_going = true;

            for message in ui_rx.try_iter() {
                match message {
                    AsyncDone => {
                        let mut this = this_rc.borrow_mut();
                        this.seeking = false;
                        match this.state {
                            ControllerState::PendingSelectMedia => this.select_media(),
                            ControllerState::PendingExportToc => this.export_toc(),
                            _ => (),
                        }
                    }
                    InitDone => {
                        let mut this = this_rc.borrow_mut();

                        let context = this.context.take().expect(
                            "MainController: InitDone but no context available",
                        );

                        this.requires_async_dialog = context.info.lock()
                            .expect(
                                "MainController:: failed to lock media info in InitDone",
                            )
                            .video_best.is_some();

                        this.header_bar.set_subtitle(
                            Some(context.file_name.as_str()),
                        );

                        this.video_ctrl.new_media(&context);
                        this.info_ctrl.borrow_mut().new_media(&context);
                        this.audio_ctrl.borrow_mut().new_media(&context);

                        this.context = Some(context);

                        this.state = ControllerState::Ready;
                    }
                    Eos => {
                        let mut this = this_rc.borrow_mut();

                        let position = this.get_position();
                        this.duration = Some(position);

                        {
                            let mut info_ctrl = this.info_ctrl.borrow_mut();
                            info_ctrl.update_duration(position);
                            info_ctrl.tick(position, true);
                        }

                        this.audio_ctrl.borrow_mut().tick();

                        this.handle_eos();

                        // Remove listener and tracker.
                        // Note: tracker will be register again in case of
                        // a seek. Listener is of no use anymore because
                        // the context won't send any more Eos nor AsyncDone
                        // after an EOS.
                        keep_going = false;
                    }
                    FailedToOpenMedia => {
                        eprintln!("ERROR: failed to open media");

                        let mut this = this_rc.borrow_mut();

                        this.context = None;
                        this.keep_going = false;
                        keep_going = false;
                    }
                };

                if !keep_going {
                    break;
                }
            }

            if !keep_going {
                let mut this = this_rc.borrow_mut();
                this.listener_src = None;
                this.tracker_src = None;
            }

            glib::Continue(keep_going)
        }));
    }

    fn remove_tracker(&mut self) {
        if let Some(source_id) = self.tracker_src.take() {
            glib::source_remove(source_id);
        }
    }

    fn register_tracker(&mut self) {
        if self.tracker_src.is_some() {
            return;
        }

        let this_rc = Rc::clone(self.this_opt.as_ref().unwrap());

        self.tracker_src = Some(gtk::timeout_add(TRACKER_PERIOD, move || {
            #[cfg(feature = "profiling-tracker")]
            let start = Utc::now();

            let mut keep_going = true;

            let mut this = this_rc.borrow_mut();

            #[cfg(feature = "profiling-tracker")]
            let before_tick = Utc::now();

            let position = this
                .context
                .as_mut()
                .expect("MainController::tracker no context while getting position")
                .get_position();

            let is_eos = if let Some(duration) = this.duration {
                if position >= duration {
                    if !this.seeking {
                        // this check is necessary as EOS is not sent
                        // in case of a seek after EOS
                        this.handle_eos();
                        this.tracker_src = None;
                        keep_going = false;
                    }
                    true
                } else if this.seeking {
                    // this check is necessary as AsyncDone is not sent
                    // in case of a seek after EOS
                    this.seeking = false;
                    false
                } else {
                    false
                }
            } else {
                false
            };

            if !this.seeking {
                this.info_ctrl.borrow_mut().tick(position, is_eos);
                this.audio_ctrl.borrow_mut().tick();
            }

            #[cfg(feature = "profiling-tracker")]
            let end = Utc::now();

            #[cfg(feature = "profiling-tracker")]
            println!(
                "tracker,{},{},{}",
                start.time().format("%H:%M:%S%.6f"),
                before_tick.time().format("%H:%M:%S%.6f"),
                end.time().format("%H:%M:%S%.6f"),
            );

            glib::Continue(this.keep_going && keep_going)
        }));
    }

    fn select_media(&mut self) {
        let file_dlg = gtk::FileChooserDialog::new(
            Some("Open a media file"),
            Some(&self.window),
            gtk::FileChooserAction::Open,
        );
        // Note: couldn't find equivalents for STOCK_OK
        file_dlg.add_button("Open", gtk::ResponseType::Ok.into());

        if file_dlg.run() == gtk::ResponseType::Ok.into() {
            if let Some(ref context) = self.context {
                context.stop();
            }
            self.open_media(file_dlg.get_filename().unwrap());
        } else {
            self.state = ControllerState::Paused;
            self.window.set_sensitive(true);
        }

        file_dlg.close();
    }

    fn open_media(&mut self, filepath: PathBuf) {
        self.remove_listener();

        self.info_ctrl.borrow_mut().cleanup();
        self.audio_ctrl.borrow_mut().cleanup();
        self.header_bar.set_subtitle("");
        self.export_toc_btn.set_sensitive(false);
        self.duration = None;

        let (ctx_tx, ui_rx) = channel();

        self.seeking = false;
        self.keep_going = true;
        self.register_listener(LISTENER_PERIOD, ui_rx);

        let dbl_buffer_mtx = self.audio_ctrl.borrow().get_dbl_buffer_mtx().clone();
        match PlaybackContext::new(filepath, dbl_buffer_mtx, ctx_tx) {
            Ok(context) => {
                self.context = Some(context);
                self.export_toc_btn.set_sensitive(true);
            }
            Err(error) => {
                eprintln!("Error opening media: {}", error);
            }
        };

        self.window.set_sensitive(true);
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use super::FileGroup;
use super::WatchmanFile;
use crate::compiler_state::{ProjectSet, SourceSet};
use crate::config::{Config, SchemaLocation};
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::ffi::OsStr;
use std::path::{Component, PathBuf};

/// The watchman query returns a list of files, but for the compiler we
/// need to categorize these files into multiple groups of files like
/// schema files, extensions and sources by their source set name.
///
/// See `FileGroup` for all groups of files.
pub fn categorize_files(
    config: &Config,
    files: &[WatchmanFile],
) -> HashMap<FileGroup, Vec<WatchmanFile>> {
    let categorizer = FileCategorizer::from_config(config);
    let mut categorized = HashMap::new();
    for file in files {
        categorized
            .entry(categorizer.categorize(&file.name))
            .or_insert_with(Vec::new)
            .push(file.clone());
    }
    categorized
}

/// The FileCategorizer is created from a Config and categorizes files found by
/// Watchman into what kind of files they are, such as source files of a
/// specific source file group or generated files from some project.
pub struct FileCategorizer {
    extensions_mapping: PathMapping<ProjectSet>,
    default_generated_dir: &'static OsStr,
    generated_dir_paths: HashSet<PathBuf>,
    source_mapping: PathMapping<SourceSet>,
    schema_file_mapping: HashMap<PathBuf, ProjectSet>,
    schema_dir_mapping: PathMapping<ProjectSet>,
}

impl FileCategorizer {
    pub fn from_config(config: &Config) -> Self {
        let mut source_mapping = vec![];
        for (path, source_set) in &config.sources {
            source_mapping.push((path.clone(), source_set.clone()));
        }
        // Sort so that more specific paths come first, i.e.
        // - foo/bar -> A
        // - foo -> B
        // which ensures we categorize foo/bar/x.js as A.
        source_mapping.sort_by(|(path_a, _), (path_b, _)| path_b.cmp(path_a));

        let mut extensions_map: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            for extension_dir in &project_config.extensions {
                match extensions_map.entry(extension_dir.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }
        let mut extensions_mapping = vec![];
        for (path, project_set) in extensions_map {
            extensions_mapping.push((path, project_set));
        }

        let mut schema_file_mapping: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            if let SchemaLocation::File(schema_file) = &project_config.schema_location {
                match schema_file_mapping.entry(schema_file.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }

        let mut schema_dir_mapping_map: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            if let SchemaLocation::Directory(directory) = &project_config.schema_location {
                match schema_dir_mapping_map.entry(directory.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }

        let mut schema_dir_mapping = vec![];
        for (path, project_set) in schema_dir_mapping_map {
            schema_dir_mapping.push((path, project_set));
        }

        let default_generated_dir = OsStr::new("__generated__");
        let generated_dir_paths: HashSet<PathBuf> = config
            .projects
            .iter()
            .filter_map(|(_, project_config)| project_config.output.clone())
            .collect();

        Self {
            extensions_mapping: PathMapping(extensions_mapping),
            default_generated_dir,
            generated_dir_paths,
            schema_file_mapping,
            schema_dir_mapping: PathMapping(schema_dir_mapping),
            source_mapping: PathMapping(source_mapping),
        }
    }

    pub fn categorize(&self, path: &PathBuf) -> FileGroup {
        let extension = path
            .extension()
            .unwrap_or_else(|| panic!("Got unexpected path without extension: `{:?}`.", path));
        if extension == "js" {
            if self.in_generated_dir(path) {
                FileGroup::Generated
            } else {
                FileGroup::Source {
                    source_set: self.source_mapping.get(path),
                }
            }
        } else if extension == "graphql" {
            if let Some(project_set) = self.schema_file_mapping.get(path) {
                FileGroup::Schema {
                    project_set: project_set.clone(),
                }
            } else if let Some(project_set) = self.extensions_mapping.find(path) {
                FileGroup::Extension { project_set }
            } else if let Some(project_set) = self.schema_dir_mapping.find(path) {
                FileGroup::Schema { project_set }
            } else {
                panic!(
                    "Expected *.graphql file `{:?}` to be either a schema or extension.",
                    path
                )
            }
        } else {
            panic!(
                "Received file {:?} from watchman with unexpected extension.",
                path
            )
        }
    }

    fn in_generated_dir(&self, path: &PathBuf) -> bool {
        self.in_absolute_generated_dir(path) || self.in_relative_generated_dir(path)
    }

    fn in_absolute_generated_dir(&self, path: &PathBuf) -> bool {
        self.generated_dir_paths
            .iter()
            .any(|generated_dir_path| path.starts_with(generated_dir_path))
    }

    fn in_relative_generated_dir(&self, path: &PathBuf) -> bool {
        path.components().any(|comp| match comp {
            Component::Normal(comp) => comp == self.default_generated_dir,
            _ => false,
        })
    }
}

struct PathMapping<T: Clone>(Vec<(PathBuf, T)>);
impl<T: Clone> PathMapping<T> {
    fn get(&self, path: &PathBuf) -> T {
        self.find(path).unwrap_or_else(|| {
            panic!("Path `{:?}` not in of the the expected directories.", path);
        })
    }
    fn find(&self, path: &PathBuf) -> Option<T> {
        self.0.iter().find_map(|(prefix, item)| {
            if path.starts_with(prefix) {
                Some(item.clone())
            } else {
                None
            }
        })
    }
}

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

    #[test]
    fn test_categorize() {
        let config = Config::from_string_for_test(
            r#"
                {
                    "sources": {
                        "src/js": "public",
                        "src/js/internal": "internal",
                        "src/vendor": "public",
                        "src/custom": "with_custom_generated_dir"
                    },
                    "projects": {
                        "public": {
                            "schema": "graphql/public.graphql"
                        },
                        "internal": {
                            "schema": "graphql/__generated__/internal.graphql"
                        },
                        "with_custom_generated_dir": {
                            "schema": "graphql/__generated__/custom.graphql",
                            "output": "graphql/custom-generated"
                        }
                    }
                }
            "#,
        )
        .unwrap();
        let categorizer = FileCategorizer::from_config(&config);

        assert_eq!(
            categorizer.categorize(&"src/js/a.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("public".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/nested/b.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("public".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/internal/nested/c.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("internal".intern()),
            },
        );
        assert_eq!(
            // When custom output dir is provided, path is correctly categorized
            // even if it has same dirname in path as custom output folder.
            // Path is only categorized as generated if it matches the absolute path
            // of the provided custom output.
            categorizer.categorize(&"src/custom/custom-generated/c.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("with_custom_generated_dir".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/internal/nested/__generated__/c.js".into()),
            FileGroup::Generated,
        );
        assert_eq!(
            categorizer.categorize(&"graphql/custom-generated/c.js".into()),
            FileGroup::Generated,
        );
        assert_eq!(
            categorizer.categorize(&"graphql/public.graphql".into()),
            FileGroup::Schema {
                project_set: ProjectSet::ProjectName("public".intern())
            },
        );
        assert_eq!(
            categorizer.categorize(&"graphql/__generated__/internal.graphql".into()),
            FileGroup::Schema {
                project_set: ProjectSet::ProjectName("internal".intern())
            },
        );
    }
}
<commit_msg>add comment to FileCategorizer<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use super::FileGroup;
use super::WatchmanFile;
use crate::compiler_state::{ProjectSet, SourceSet};
use crate::config::{Config, SchemaLocation};
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::ffi::OsStr;
use std::path::{Component, PathBuf};

/// The watchman query returns a list of files, but for the compiler we
/// need to categorize these files into multiple groups of files like
/// schema files, extensions and sources by their source set name.
///
/// See `FileGroup` for all groups of files.
pub fn categorize_files(
    config: &Config,
    files: &[WatchmanFile],
) -> HashMap<FileGroup, Vec<WatchmanFile>> {
    let categorizer = FileCategorizer::from_config(config);
    let mut categorized = HashMap::new();
    for file in files {
        categorized
            .entry(categorizer.categorize(&file.name))
            .or_insert_with(Vec::new)
            .push(file.clone());
    }
    categorized
}

/// The FileCategorizer is created from a Config and categorizes files found by
/// Watchman into what kind of files they are, such as source files of a
/// specific source file group or generated files from some project.
pub struct FileCategorizer {
    extensions_mapping: PathMapping<ProjectSet>,
    default_generated_dir: &'static OsStr,
    generated_dir_paths: HashSet<PathBuf>,
    source_mapping: PathMapping<SourceSet>,
    schema_file_mapping: HashMap<PathBuf, ProjectSet>,
    schema_dir_mapping: PathMapping<ProjectSet>,
}

impl FileCategorizer {
    pub fn from_config(config: &Config) -> Self {
        let mut source_mapping = vec![];
        for (path, source_set) in &config.sources {
            source_mapping.push((path.clone(), source_set.clone()));
        }
        // Sort so that more specific paths come first, i.e.
        // - foo/bar -> A
        // - foo -> B
        // which ensures we categorize foo/bar/x.js as A.
        source_mapping.sort_by(|(path_a, _), (path_b, _)| path_b.cmp(path_a));

        let mut extensions_map: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            for extension_dir in &project_config.extensions {
                match extensions_map.entry(extension_dir.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }
        let mut extensions_mapping = vec![];
        for (path, project_set) in extensions_map {
            extensions_mapping.push((path, project_set));
        }

        let mut schema_file_mapping: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            if let SchemaLocation::File(schema_file) = &project_config.schema_location {
                match schema_file_mapping.entry(schema_file.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }

        let mut schema_dir_mapping_map: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            if let SchemaLocation::Directory(directory) = &project_config.schema_location {
                match schema_dir_mapping_map.entry(directory.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }

        let mut schema_dir_mapping = vec![];
        for (path, project_set) in schema_dir_mapping_map {
            schema_dir_mapping.push((path, project_set));
        }

        let default_generated_dir = OsStr::new("__generated__");
        let generated_dir_paths: HashSet<PathBuf> = config
            .projects
            .iter()
            .filter_map(|(_, project_config)| project_config.output.clone())
            .collect();

        Self {
            extensions_mapping: PathMapping(extensions_mapping),
            default_generated_dir,
            generated_dir_paths,
            schema_file_mapping,
            schema_dir_mapping: PathMapping(schema_dir_mapping),
            source_mapping: PathMapping(source_mapping),
        }
    }

    /// Categorizes a file. This method should be kept as cheap as possible by
    /// preprocessing the config in `from_config` and then re-using the
    /// `FileCategorizer`.
    pub fn categorize(&self, path: &PathBuf) -> FileGroup {
        let extension = path
            .extension()
            .unwrap_or_else(|| panic!("Got unexpected path without extension: `{:?}`.", path));
        if extension == "js" {
            if self.in_generated_dir(path) {
                FileGroup::Generated
            } else {
                FileGroup::Source {
                    source_set: self.source_mapping.get(path),
                }
            }
        } else if extension == "graphql" {
            if let Some(project_set) = self.schema_file_mapping.get(path) {
                FileGroup::Schema {
                    project_set: project_set.clone(),
                }
            } else if let Some(project_set) = self.extensions_mapping.find(path) {
                FileGroup::Extension { project_set }
            } else if let Some(project_set) = self.schema_dir_mapping.find(path) {
                FileGroup::Schema { project_set }
            } else {
                panic!(
                    "Expected *.graphql file `{:?}` to be either a schema or extension.",
                    path
                )
            }
        } else {
            panic!(
                "Received file {:?} from watchman with unexpected extension.",
                path
            )
        }
    }

    fn in_generated_dir(&self, path: &PathBuf) -> bool {
        self.in_absolute_generated_dir(path) || self.in_relative_generated_dir(path)
    }

    fn in_absolute_generated_dir(&self, path: &PathBuf) -> bool {
        self.generated_dir_paths
            .iter()
            .any(|generated_dir_path| path.starts_with(generated_dir_path))
    }

    fn in_relative_generated_dir(&self, path: &PathBuf) -> bool {
        path.components().any(|comp| match comp {
            Component::Normal(comp) => comp == self.default_generated_dir,
            _ => false,
        })
    }
}

struct PathMapping<T: Clone>(Vec<(PathBuf, T)>);
impl<T: Clone> PathMapping<T> {
    fn get(&self, path: &PathBuf) -> T {
        self.find(path).unwrap_or_else(|| {
            panic!("Path `{:?}` not in of the the expected directories.", path);
        })
    }
    fn find(&self, path: &PathBuf) -> Option<T> {
        self.0.iter().find_map(|(prefix, item)| {
            if path.starts_with(prefix) {
                Some(item.clone())
            } else {
                None
            }
        })
    }
}

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

    #[test]
    fn test_categorize() {
        let config = Config::from_string_for_test(
            r#"
                {
                    "sources": {
                        "src/js": "public",
                        "src/js/internal": "internal",
                        "src/vendor": "public",
                        "src/custom": "with_custom_generated_dir"
                    },
                    "projects": {
                        "public": {
                            "schema": "graphql/public.graphql"
                        },
                        "internal": {
                            "schema": "graphql/__generated__/internal.graphql"
                        },
                        "with_custom_generated_dir": {
                            "schema": "graphql/__generated__/custom.graphql",
                            "output": "graphql/custom-generated"
                        }
                    }
                }
            "#,
        )
        .unwrap();
        let categorizer = FileCategorizer::from_config(&config);

        assert_eq!(
            categorizer.categorize(&"src/js/a.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("public".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/nested/b.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("public".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/internal/nested/c.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("internal".intern()),
            },
        );
        assert_eq!(
            // When custom output dir is provided, path is correctly categorized
            // even if it has same dirname in path as custom output folder.
            // Path is only categorized as generated if it matches the absolute path
            // of the provided custom output.
            categorizer.categorize(&"src/custom/custom-generated/c.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("with_custom_generated_dir".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/internal/nested/__generated__/c.js".into()),
            FileGroup::Generated,
        );
        assert_eq!(
            categorizer.categorize(&"graphql/custom-generated/c.js".into()),
            FileGroup::Generated,
        );
        assert_eq!(
            categorizer.categorize(&"graphql/public.graphql".into()),
            FileGroup::Schema {
                project_set: ProjectSet::ProjectName("public".intern())
            },
        );
        assert_eq!(
            categorizer.categorize(&"graphql/__generated__/internal.graphql".into()),
            FileGroup::Schema {
                project_set: ProjectSet::ProjectName("internal".intern())
            },
        );
    }
}
<|endoftext|>


<commit_before>extern crate cc;

mod build_linux;
mod build_windows;
mod build_macos;

use std::env;

fn main() {
    if env::var("TARGET").unwrap().contains("windows") {
        build_windows::build_windows();
    }
    else if env::var("TARGET").unwrap().contains("linux") {
        build_linux::build_linux();
    } if env::var("TARGET").unwrap().contains("linux") {
        build_macos::build_macos();
    } else {
        panic!("Unsupported platform!");
    }

}
<commit_msg>WIP macos build<commit_after>extern crate cc;

mod build_linux;
mod build_windows;
mod build_macos;

use std::env;

fn main() {
    if env::var("TARGET").unwrap().contains("windows") {
        build_windows::build_windows();
    } else if env::var("TARGET").unwrap().contains("linux") {
        build_linux::build_linux();
    } else if env::var("TARGET").unwrap().contains("mac") {
        build_macos::build_macos();
    } else {
        panic!("Unsupported platform!");
    }

}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use super::FileGroup;
use super::WatchmanFile;
use crate::compiler_state::{ProjectName, ProjectSet, SourceSet};
use crate::config::{Config, SchemaLocation};
use fnv::FnvHashSet;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::ffi::OsStr;
use std::path::{Component, PathBuf};

/// The watchman query returns a list of files, but for the compiler we
/// need to categorize these files into multiple groups of files like
/// schema files, extensions and sources by their source set name.
///
/// See `FileGroup` for all groups of files.
pub fn categorize_files(
    config: &Config,
    files: &[WatchmanFile],
) -> HashMap<FileGroup, Vec<WatchmanFile>> {
    let categorizer = FileCategorizer::from_config(config);
    let mut categorized = HashMap::new();

    let mut has_disabled = false;
    let mut relevant_projects = FnvHashSet::default();
    for (project_name, project_config) in &config.projects {
        if project_config.enabled {
            relevant_projects.insert(project_name);
            if let Some(base_project) = &project_config.base {
                relevant_projects.insert(base_project);
            }
        } else {
            has_disabled = true;
        }
    }

    for file in files {
        let file_group = categorizer.categorize(&file.name);
        let should_skip = has_disabled
            && match &file_group {
                FileGroup::Source {
                    source_set: SourceSet::SourceSetName(name),
                }
                | FileGroup::Schema {
                    project_set: ProjectSet::ProjectName(name),
                }
                | FileGroup::Extension {
                    project_set: ProjectSet::ProjectName(name),
                }
                | FileGroup::Generated { project_name: name } => !relevant_projects.contains(name),
                FileGroup::Source {
                    source_set: SourceSet::SourceSetNames(names),
                }
                | FileGroup::Schema {
                    project_set: ProjectSet::ProjectNames(names),
                }
                | FileGroup::Extension {
                    project_set: ProjectSet::ProjectNames(names),
                } => !names.iter().any(|name| relevant_projects.contains(name)),
            };
        if !should_skip {
            categorized
                .entry(file_group)
                .or_insert_with(Vec::new)
                .push(file.clone());
        }
    }
    categorized
}

/// The FileCategorizer is created from a Config and categorizes files found by
/// Watchman into what kind of files they are, such as source files of a
/// specific source file group or generated files from some project.
pub struct FileCategorizer {
    extensions_mapping: PathMapping<ProjectSet>,
    default_generated_dir: &'static OsStr,
    generated_dir_mapping: PathMapping<ProjectName>,
    source_mapping: PathMapping<SourceSet>,
    schema_file_mapping: HashMap<PathBuf, ProjectSet>,
    schema_dir_mapping: PathMapping<ProjectSet>,
}

impl FileCategorizer {
    pub fn from_config(config: &Config) -> Self {
        let mut source_mapping = vec![];
        for (path, source_set) in &config.sources {
            source_mapping.push((path.clone(), source_set.clone()));
        }

        let mut extensions_map: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            for extension_dir in &project_config.extensions {
                match extensions_map.entry(extension_dir.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }

        let mut schema_file_mapping: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            if let SchemaLocation::File(schema_file) = &project_config.schema_location {
                match schema_file_mapping.entry(schema_file.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }

        let mut schema_dir_mapping_map: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            if let SchemaLocation::Directory(directory) = &project_config.schema_location {
                match schema_dir_mapping_map.entry(directory.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }

        let mut schema_dir_mapping = vec![];
        for (path, project_set) in schema_dir_mapping_map {
            schema_dir_mapping.push((path, project_set));
        }

        let mut generated_dir_mapping = Vec::new();
        for (&project_name, project_config) in &config.projects {
            if let Some(output) = &project_config.output {
                generated_dir_mapping.push((output.clone(), project_name));
            }
            if let Some(extra_artifacts_output) = &project_config.extra_artifacts_output {
                generated_dir_mapping.push((extra_artifacts_output.clone(), project_name));
            }
        }

        Self {
            extensions_mapping: PathMapping::new(extensions_map.into_iter().collect()),
            default_generated_dir: OsStr::new("__generated__"),
            generated_dir_mapping: PathMapping::new(generated_dir_mapping),
            schema_file_mapping,
            schema_dir_mapping: PathMapping::new(schema_dir_mapping),
            source_mapping: PathMapping::new(source_mapping),
        }
    }

    /// Categorizes a file. This method should be kept as cheap as possible by
    /// preprocessing the config in `from_config` and then re-using the
    /// `FileCategorizer`.
    pub fn categorize(&self, path: &PathBuf) -> FileGroup {
        if let Some(project_name) = self.generated_dir_mapping.find(path) {
            return FileGroup::Generated { project_name };
        }
        let extension = path
            .extension()
            .unwrap_or_else(|| panic!("Got unexpected path without extension: `{:?}`.", path));
        if extension == "js" {
            let source_set = self.source_mapping.get(path);
            if self.in_relative_generated_dir(path) {
                if let SourceSet::SourceSetName(source_set_name) = source_set {
                    FileGroup::Generated {
                        project_name: source_set_name,
                    }
                } else {
                    panic!(
                        "Overlapping input sources are incompatible with relative generated \
                        directories. Got `{:?}` in a relative generated directory with source set {:?}",
                        path, source_set
                    );
                }
            } else {
                FileGroup::Source { source_set }
            }
        } else if extension == "graphql" {
            if let Some(project_set) = self.schema_file_mapping.get(path) {
                FileGroup::Schema {
                    project_set: project_set.clone(),
                }
            } else if let Some(project_set) = self.extensions_mapping.find(path) {
                FileGroup::Extension { project_set }
            } else if let Some(project_set) = self.schema_dir_mapping.find(path) {
                FileGroup::Schema { project_set }
            } else {
                panic!(
                    "Expected *.graphql file `{:?}` to be either a schema or extension.",
                    path
                )
            }
        } else {
            panic!(
                "Received file {:?} from watchman with unexpected extension.",
                path
            )
        }
    }

    fn in_relative_generated_dir(&self, path: &PathBuf) -> bool {
        path.components().any(|comp| match comp {
            Component::Normal(comp) => comp == self.default_generated_dir,
            _ => false,
        })
    }
}

struct PathMapping<T: Clone>(Vec<(PathBuf, T)>);
impl<T: Clone> PathMapping<T> {
    fn new(mut entries: Vec<(PathBuf, T)>) -> Self {
        // Sort so that more specific paths come first, i.e.
        // - foo/bar -> A
        // - foo -> B
        // which ensures we categorize foo/bar/x.js as A.
        entries.sort_by(|(path_a, _), (path_b, _)| path_b.cmp(path_a));
        Self(entries)
    }

    fn get(&self, path: &PathBuf) -> T {
        self.find(path).unwrap_or_else(|| {
            panic!("Path `{:?}` not in of the the expected directories.", path);
        })
    }
    fn find(&self, path: &PathBuf) -> Option<T> {
        self.0.iter().find_map(|(prefix, item)| {
            if path.starts_with(prefix) {
                Some(item.clone())
            } else {
                None
            }
        })
    }
}

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

    #[test]
    fn test_categorize() {
        let config = Config::from_string_for_test(
            r#"
                {
                    "sources": {
                        "src/js": "public",
                        "src/js/internal": "internal",
                        "src/vendor": "public",
                        "src/custom": "with_custom_generated_dir"
                    },
                    "projects": {
                        "public": {
                            "schema": "graphql/public.graphql"
                        },
                        "internal": {
                            "schema": "graphql/__generated__/internal.graphql"
                        },
                        "with_custom_generated_dir": {
                            "schema": "graphql/__generated__/custom.graphql",
                            "output": "graphql/custom-generated"
                        }
                    }
                }
            "#,
        )
        .unwrap();
        let categorizer = FileCategorizer::from_config(&config);

        assert_eq!(
            categorizer.categorize(&"src/js/a.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("public".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/nested/b.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("public".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/internal/nested/c.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("internal".intern()),
            },
        );
        assert_eq!(
            // When custom output dir is provided, path is correctly categorized
            // even if it has same dirname in path as custom output folder.
            // Path is only categorized as generated if it matches the absolute path
            // of the provided custom output.
            categorizer.categorize(&"src/custom/custom-generated/c.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("with_custom_generated_dir".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/internal/nested/__generated__/c.js".into()),
            FileGroup::Generated {
                project_name: "internal".intern()
            },
        );
        assert_eq!(
            categorizer.categorize(&"graphql/custom-generated/c.js".into()),
            FileGroup::Generated {
                project_name: "with_custom_generated_dir".intern()
            },
        );
        assert_eq!(
            categorizer.categorize(&"graphql/public.graphql".into()),
            FileGroup::Schema {
                project_set: ProjectSet::ProjectName("public".intern())
            },
        );
        assert_eq!(
            categorizer.categorize(&"graphql/__generated__/internal.graphql".into()),
            FileGroup::Schema {
                project_set: ProjectSet::ProjectName("internal".intern())
            },
        );
    }
}
<commit_msg>Fix typo in file categorizer panic message and show all dirs<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use super::FileGroup;
use super::WatchmanFile;
use crate::compiler_state::{ProjectName, ProjectSet, SourceSet};
use crate::config::{Config, SchemaLocation};
use fnv::FnvHashSet;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::ffi::OsStr;
use std::path::{Component, PathBuf};

/// The watchman query returns a list of files, but for the compiler we
/// need to categorize these files into multiple groups of files like
/// schema files, extensions and sources by their source set name.
///
/// See `FileGroup` for all groups of files.
pub fn categorize_files(
    config: &Config,
    files: &[WatchmanFile],
) -> HashMap<FileGroup, Vec<WatchmanFile>> {
    let categorizer = FileCategorizer::from_config(config);
    let mut categorized = HashMap::new();

    let mut has_disabled = false;
    let mut relevant_projects = FnvHashSet::default();
    for (project_name, project_config) in &config.projects {
        if project_config.enabled {
            relevant_projects.insert(project_name);
            if let Some(base_project) = &project_config.base {
                relevant_projects.insert(base_project);
            }
        } else {
            has_disabled = true;
        }
    }

    for file in files {
        let file_group = categorizer.categorize(&file.name);
        let should_skip = has_disabled
            && match &file_group {
                FileGroup::Source {
                    source_set: SourceSet::SourceSetName(name),
                }
                | FileGroup::Schema {
                    project_set: ProjectSet::ProjectName(name),
                }
                | FileGroup::Extension {
                    project_set: ProjectSet::ProjectName(name),
                }
                | FileGroup::Generated { project_name: name } => !relevant_projects.contains(name),
                FileGroup::Source {
                    source_set: SourceSet::SourceSetNames(names),
                }
                | FileGroup::Schema {
                    project_set: ProjectSet::ProjectNames(names),
                }
                | FileGroup::Extension {
                    project_set: ProjectSet::ProjectNames(names),
                } => !names.iter().any(|name| relevant_projects.contains(name)),
            };
        if !should_skip {
            categorized
                .entry(file_group)
                .or_insert_with(Vec::new)
                .push(file.clone());
        }
    }
    categorized
}

/// The FileCategorizer is created from a Config and categorizes files found by
/// Watchman into what kind of files they are, such as source files of a
/// specific source file group or generated files from some project.
pub struct FileCategorizer {
    extensions_mapping: PathMapping<ProjectSet>,
    default_generated_dir: &'static OsStr,
    generated_dir_mapping: PathMapping<ProjectName>,
    source_mapping: PathMapping<SourceSet>,
    schema_file_mapping: HashMap<PathBuf, ProjectSet>,
    schema_dir_mapping: PathMapping<ProjectSet>,
}

impl FileCategorizer {
    pub fn from_config(config: &Config) -> Self {
        let mut source_mapping = vec![];
        for (path, source_set) in &config.sources {
            source_mapping.push((path.clone(), source_set.clone()));
        }

        let mut extensions_map: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            for extension_dir in &project_config.extensions {
                match extensions_map.entry(extension_dir.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }

        let mut schema_file_mapping: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            if let SchemaLocation::File(schema_file) = &project_config.schema_location {
                match schema_file_mapping.entry(schema_file.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }

        let mut schema_dir_mapping_map: HashMap<PathBuf, ProjectSet> = Default::default();
        for (&project_name, project_config) in &config.projects {
            if let SchemaLocation::Directory(directory) = &project_config.schema_location {
                match schema_dir_mapping_map.entry(directory.clone()) {
                    Entry::Vacant(entry) => {
                        entry.insert(ProjectSet::ProjectName(project_name));
                    }
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert(project_name);
                    }
                }
            }
        }

        let mut schema_dir_mapping = vec![];
        for (path, project_set) in schema_dir_mapping_map {
            schema_dir_mapping.push((path, project_set));
        }

        let mut generated_dir_mapping = Vec::new();
        for (&project_name, project_config) in &config.projects {
            if let Some(output) = &project_config.output {
                generated_dir_mapping.push((output.clone(), project_name));
            }
            if let Some(extra_artifacts_output) = &project_config.extra_artifacts_output {
                generated_dir_mapping.push((extra_artifacts_output.clone(), project_name));
            }
        }

        Self {
            extensions_mapping: PathMapping::new(extensions_map.into_iter().collect()),
            default_generated_dir: OsStr::new("__generated__"),
            generated_dir_mapping: PathMapping::new(generated_dir_mapping),
            schema_file_mapping,
            schema_dir_mapping: PathMapping::new(schema_dir_mapping),
            source_mapping: PathMapping::new(source_mapping),
        }
    }

    /// Categorizes a file. This method should be kept as cheap as possible by
    /// preprocessing the config in `from_config` and then re-using the
    /// `FileCategorizer`.
    pub fn categorize(&self, path: &PathBuf) -> FileGroup {
        if let Some(project_name) = self.generated_dir_mapping.find(path) {
            return FileGroup::Generated { project_name };
        }
        let extension = path
            .extension()
            .unwrap_or_else(|| panic!("Got unexpected path without extension: `{:?}`.", path));
        if extension == "js" {
            let source_set = self.source_mapping.get(path);
            if self.in_relative_generated_dir(path) {
                if let SourceSet::SourceSetName(source_set_name) = source_set {
                    FileGroup::Generated {
                        project_name: source_set_name,
                    }
                } else {
                    panic!(
                        "Overlapping input sources are incompatible with relative generated \
                        directories. Got `{:?}` in a relative generated directory with source set {:?}",
                        path, source_set
                    );
                }
            } else {
                FileGroup::Source { source_set }
            }
        } else if extension == "graphql" {
            if let Some(project_set) = self.schema_file_mapping.get(path) {
                FileGroup::Schema {
                    project_set: project_set.clone(),
                }
            } else if let Some(project_set) = self.extensions_mapping.find(path) {
                FileGroup::Extension { project_set }
            } else if let Some(project_set) = self.schema_dir_mapping.find(path) {
                FileGroup::Schema { project_set }
            } else {
                panic!(
                    "Expected *.graphql file `{:?}` to be either a schema or extension.",
                    path
                )
            }
        } else {
            panic!(
                "Received file {:?} from watchman with unexpected extension.",
                path
            )
        }
    }

    fn in_relative_generated_dir(&self, path: &PathBuf) -> bool {
        path.components().any(|comp| match comp {
            Component::Normal(comp) => comp == self.default_generated_dir,
            _ => false,
        })
    }
}

struct PathMapping<T: Clone>(Vec<(PathBuf, T)>);
impl<T: Clone> PathMapping<T> {
    fn new(mut entries: Vec<(PathBuf, T)>) -> Self {
        // Sort so that more specific paths come first, i.e.
        // - foo/bar -> A
        // - foo -> B
        // which ensures we categorize foo/bar/x.js as A.
        entries.sort_by(|(path_a, _), (path_b, _)| path_b.cmp(path_a));
        Self(entries)
    }

    fn get(&self, path: &PathBuf) -> T {
        self.find(path).unwrap_or_else(|| {
            panic!(
                "Path '{:?}' not in any of the expected directories. Available directories: {:?}",
                path,
                self.0.iter().map(|(prefix, _)| prefix).collect::<Vec<_>>()
            );
        })
    }
    fn find(&self, path: &PathBuf) -> Option<T> {
        self.0.iter().find_map(|(prefix, item)| {
            if path.starts_with(prefix) {
                Some(item.clone())
            } else {
                None
            }
        })
    }
}

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

    #[test]
    fn test_categorize() {
        let config = Config::from_string_for_test(
            r#"
                {
                    "sources": {
                        "src/js": "public",
                        "src/js/internal": "internal",
                        "src/vendor": "public",
                        "src/custom": "with_custom_generated_dir"
                    },
                    "projects": {
                        "public": {
                            "schema": "graphql/public.graphql"
                        },
                        "internal": {
                            "schema": "graphql/__generated__/internal.graphql"
                        },
                        "with_custom_generated_dir": {
                            "schema": "graphql/__generated__/custom.graphql",
                            "output": "graphql/custom-generated"
                        }
                    }
                }
            "#,
        )
        .unwrap();
        let categorizer = FileCategorizer::from_config(&config);

        assert_eq!(
            categorizer.categorize(&"src/js/a.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("public".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/nested/b.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("public".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/internal/nested/c.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("internal".intern()),
            },
        );
        assert_eq!(
            // When custom output dir is provided, path is correctly categorized
            // even if it has same dirname in path as custom output folder.
            // Path is only categorized as generated if it matches the absolute path
            // of the provided custom output.
            categorizer.categorize(&"src/custom/custom-generated/c.js".into()),
            FileGroup::Source {
                source_set: SourceSet::SourceSetName("with_custom_generated_dir".intern()),
            },
        );
        assert_eq!(
            categorizer.categorize(&"src/js/internal/nested/__generated__/c.js".into()),
            FileGroup::Generated {
                project_name: "internal".intern()
            },
        );
        assert_eq!(
            categorizer.categorize(&"graphql/custom-generated/c.js".into()),
            FileGroup::Generated {
                project_name: "with_custom_generated_dir".intern()
            },
        );
        assert_eq!(
            categorizer.categorize(&"graphql/public.graphql".into()),
            FileGroup::Schema {
                project_set: ProjectSet::ProjectName("public".intern())
            },
        );
        assert_eq!(
            categorizer.categorize(&"graphql/__generated__/internal.graphql".into()),
            FileGroup::Schema {
                project_set: ProjectSet::ProjectName("internal".intern())
            },
        );
    }
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::common::types::{
    CliCommand, CliError, CliResult, CliTypedResult, ConfigSearchMode, CONFIG_FOLDER,
};
use crate::common::utils::{
    create_dir_if_not_exist, current_dir, read_from_file, write_to_user_only_file,
};
use crate::genesis::git::{from_yaml, to_yaml};
use crate::Tool;
use async_trait::async_trait;
use clap::ArgEnum;
use clap::CommandFactory;
use clap::Parser;
use clap_complete::{generate, Shell};
use serde::Deserialize;
use serde::Serialize;
use std::fmt::Formatter;
use std::path::PathBuf;
use std::str::FromStr;

/// Tool for configuration of the CLI tool
///
#[derive(Parser)]
pub enum ConfigTool {
    Init(crate::common::init::InitTool),
    GenerateShellCompletions(GenerateShellCompletions),
    SetGlobalConfig(SetGlobalConfig),
    ShowGlobalConfig(ShowGlobalConfig),
}

impl ConfigTool {
    pub async fn execute(self) -> CliResult {
        match self {
            ConfigTool::Init(tool) => tool.execute_serialized_success().await,
            ConfigTool::GenerateShellCompletions(tool) => tool.execute_serialized_success().await,
            ConfigTool::SetGlobalConfig(tool) => tool.execute_serialized().await,
            ConfigTool::ShowGlobalConfig(tool) => tool.execute_serialized().await,
        }
    }
}

/// Generates shell completion files
///
/// First generate the completion file, then follow the shell specific directions on how
/// to install the completion file.
#[derive(Parser)]
pub struct GenerateShellCompletions {
    /// Shell to generate completions for one of [bash, elvish, powershell, zsh]
    #[clap(long)]
    shell: Shell,
    /// File to output shell completions to
    #[clap(long, parse(from_os_str))]
    output_file: PathBuf,
}

#[async_trait]
impl CliCommand<()> for GenerateShellCompletions {
    fn command_name(&self) -> &'static str {
        "GenerateShellCompletions"
    }

    async fn execute(self) -> CliTypedResult<()> {
        let mut command = Tool::command();
        let mut file = std::fs::File::create(self.output_file.as_path())
            .map_err(|err| CliError::IO(self.output_file.display().to_string(), err))?;
        generate(self.shell, &mut command, "aptos".to_string(), &mut file);
        Ok(())
    }
}

/// Set global configuration settings
///
/// Any configuration flags that are not provided will not be changed
#[derive(Parser, Debug)]
pub struct SetGlobalConfig {
    /// A configuration for where to place and use the config
    ///
    /// Workspace allows for multiple configs based on location, where
    /// Global allows for one config for every part of the code
    #[clap(long)]
    config_type: Option<ConfigType>,
}

#[async_trait]
impl CliCommand<GlobalConfig> for SetGlobalConfig {
    fn command_name(&self) -> &'static str {
        "SetGlobalConfig"
    }

    async fn execute(self) -> CliTypedResult<GlobalConfig> {
        // Load the global config
        let mut config = GlobalConfig::load()?;

        // Enable all features that are actually listed
        if let Some(config_type) = self.config_type {
            config.config_type = Some(config_type);
        }

        config.save()?;
        config.display()
    }
}

/// Shows the properties in the global config
#[derive(Parser, Debug)]
pub struct ShowGlobalConfig {}

#[async_trait]
impl CliCommand<GlobalConfig> for ShowGlobalConfig {
    fn command_name(&self) -> &'static str {
        "ShowGlobalConfig"
    }

    async fn execute(self) -> CliTypedResult<GlobalConfig> {
        // Load the global config
        let config = GlobalConfig::load()?;

        config.display()
    }
}

const GLOBAL_CONFIG_FILE: &str = "global_config.yaml";

/// A global configuration for global settings related to a user
#[derive(Serialize, Deserialize, Debug, Default)]
pub struct GlobalConfig {
    /// Whether to be using Global or Workspace mode
    #[serde(skip_serializing_if = "Option::is_none")]
    pub config_type: Option<ConfigType>,
}

impl GlobalConfig {
    /// Fill in defaults for display via the CLI
    pub fn display(mut self) -> CliTypedResult<Self> {
        if self.config_type.is_none() {
            self.config_type = Some(ConfigType::default());
        }

        Ok(self)
    }

    pub fn load() -> CliTypedResult<Self> {
        let path = global_folder()?.join(GLOBAL_CONFIG_FILE);
        if path.exists() {
            from_yaml(&String::from_utf8(read_from_file(path.as_path())?)?)
        } else {
            // If we don't have a config, let's load the default
            Ok(GlobalConfig::default())
        }
    }

    /// Get the config location based on the type
    pub fn get_config_location(&self, mode: ConfigSearchMode) -> CliTypedResult<PathBuf> {
        match self.config_type.unwrap_or_default() {
            ConfigType::Global => global_folder(),
            ConfigType::Workspace => find_workspace_config(current_dir()?, mode),
        }
    }

    fn save(&self) -> CliTypedResult<()> {
        let global_folder = global_folder()?;
        create_dir_if_not_exist(global_folder.as_path())?;

        write_to_user_only_file(
            global_folder.join(GLOBAL_CONFIG_FILE).as_path(),
            "Global Config",
            &to_yaml(&self)?.into_bytes(),
        )
    }
}

fn global_folder() -> CliTypedResult<PathBuf> {
    if let Some(dir) = dirs::home_dir() {
        Ok(dir.join(CONFIG_FOLDER))
    } else {
        Err(CliError::UnexpectedError(
            "Unable to retrieve home directory".to_string(),
        ))
    }
}

fn find_workspace_config(
    starting_path: PathBuf,
    mode: ConfigSearchMode,
) -> CliTypedResult<PathBuf> {
    match mode {
        ConfigSearchMode::CurrentDir => Ok(starting_path.join(CONFIG_FOLDER)),
        ConfigSearchMode::CurrentDirAndParents => {
            let mut current_path = starting_path.clone();
            loop {
                let cand = current_path.join(CONFIG_FOLDER);
                if cand.is_dir() {
                    break Ok(cand);
                }
                if !current_path.pop() {
                    return Err(CliError::ConfigNotFoundError(format!(
                        "Unable to find {} in {} or any parent directory",
                        CONFIG_FOLDER,
                        starting_path.display()
                    )));
                }
            }
        }
    }
}

const GLOBAL: &str = "global";
const WORKSPACE: &str = "workspace";

/// A configuration for where to place and use the config
///
/// Workspace allows for multiple configs based on location, where
/// Global allows for one config for every part of the code
#[derive(Debug, Copy, Clone, Serialize, Deserialize, ArgEnum)]
pub enum ConfigType {
    /// Per system user configuration put in `<HOME>/.aptos`
    Global,
    /// Per directory configuration put in `<CURRENT_DIR>/.aptos`
    Workspace,
}

impl Default for ConfigType {
    fn default() -> Self {
        // TODO: When we version up, we can change this to global
        Self::Workspace
    }
}

impl std::fmt::Display for ConfigType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            ConfigType::Global => GLOBAL,
            ConfigType::Workspace => WORKSPACE,
        })
    }
}

impl FromStr for ConfigType {
    type Err = CliError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().trim() {
            GLOBAL => Ok(Self::Global),
            WORKSPACE => Ok(Self::Workspace),
            _ => Err(CliError::CommandArgumentError(
                "Invalid config type, must be one of [global, workspace]".to_string(),
            )),
        }
    }
}
<commit_msg>[aptos-cli] Use current directory for workspace .aptos if it doesn't exist (#2845)<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::common::types::{
    CliCommand, CliError, CliResult, CliTypedResult, ConfigSearchMode, CONFIG_FOLDER,
};
use crate::common::utils::{
    create_dir_if_not_exist, current_dir, read_from_file, write_to_user_only_file,
};
use crate::genesis::git::{from_yaml, to_yaml};
use crate::Tool;
use async_trait::async_trait;
use clap::ArgEnum;
use clap::CommandFactory;
use clap::Parser;
use clap_complete::{generate, Shell};
use serde::Deserialize;
use serde::Serialize;
use std::fmt::Formatter;
use std::path::PathBuf;
use std::str::FromStr;

/// Tool for configuration of the CLI tool
///
#[derive(Parser)]
pub enum ConfigTool {
    Init(crate::common::init::InitTool),
    GenerateShellCompletions(GenerateShellCompletions),
    SetGlobalConfig(SetGlobalConfig),
    ShowGlobalConfig(ShowGlobalConfig),
}

impl ConfigTool {
    pub async fn execute(self) -> CliResult {
        match self {
            ConfigTool::Init(tool) => tool.execute_serialized_success().await,
            ConfigTool::GenerateShellCompletions(tool) => tool.execute_serialized_success().await,
            ConfigTool::SetGlobalConfig(tool) => tool.execute_serialized().await,
            ConfigTool::ShowGlobalConfig(tool) => tool.execute_serialized().await,
        }
    }
}

/// Generates shell completion files
///
/// First generate the completion file, then follow the shell specific directions on how
/// to install the completion file.
#[derive(Parser)]
pub struct GenerateShellCompletions {
    /// Shell to generate completions for one of [bash, elvish, powershell, zsh]
    #[clap(long)]
    shell: Shell,
    /// File to output shell completions to
    #[clap(long, parse(from_os_str))]
    output_file: PathBuf,
}

#[async_trait]
impl CliCommand<()> for GenerateShellCompletions {
    fn command_name(&self) -> &'static str {
        "GenerateShellCompletions"
    }

    async fn execute(self) -> CliTypedResult<()> {
        let mut command = Tool::command();
        let mut file = std::fs::File::create(self.output_file.as_path())
            .map_err(|err| CliError::IO(self.output_file.display().to_string(), err))?;
        generate(self.shell, &mut command, "aptos".to_string(), &mut file);
        Ok(())
    }
}

/// Set global configuration settings
///
/// Any configuration flags that are not provided will not be changed
#[derive(Parser, Debug)]
pub struct SetGlobalConfig {
    /// A configuration for where to place and use the config
    ///
    /// Workspace allows for multiple configs based on location, where
    /// Global allows for one config for every part of the code
    #[clap(long)]
    config_type: Option<ConfigType>,
}

#[async_trait]
impl CliCommand<GlobalConfig> for SetGlobalConfig {
    fn command_name(&self) -> &'static str {
        "SetGlobalConfig"
    }

    async fn execute(self) -> CliTypedResult<GlobalConfig> {
        // Load the global config
        let mut config = GlobalConfig::load()?;

        // Enable all features that are actually listed
        if let Some(config_type) = self.config_type {
            config.config_type = Some(config_type);
        }

        config.save()?;
        config.display()
    }
}

/// Shows the properties in the global config
#[derive(Parser, Debug)]
pub struct ShowGlobalConfig {}

#[async_trait]
impl CliCommand<GlobalConfig> for ShowGlobalConfig {
    fn command_name(&self) -> &'static str {
        "ShowGlobalConfig"
    }

    async fn execute(self) -> CliTypedResult<GlobalConfig> {
        // Load the global config
        let config = GlobalConfig::load()?;

        config.display()
    }
}

const GLOBAL_CONFIG_FILE: &str = "global_config.yaml";

/// A global configuration for global settings related to a user
#[derive(Serialize, Deserialize, Debug, Default)]
pub struct GlobalConfig {
    /// Whether to be using Global or Workspace mode
    #[serde(skip_serializing_if = "Option::is_none")]
    pub config_type: Option<ConfigType>,
}

impl GlobalConfig {
    /// Fill in defaults for display via the CLI
    pub fn display(mut self) -> CliTypedResult<Self> {
        if self.config_type.is_none() {
            self.config_type = Some(ConfigType::default());
        }

        Ok(self)
    }

    pub fn load() -> CliTypedResult<Self> {
        let path = global_folder()?.join(GLOBAL_CONFIG_FILE);
        if path.exists() {
            from_yaml(&String::from_utf8(read_from_file(path.as_path())?)?)
        } else {
            // If we don't have a config, let's load the default
            Ok(GlobalConfig::default())
        }
    }

    /// Get the config location based on the type
    pub fn get_config_location(&self, mode: ConfigSearchMode) -> CliTypedResult<PathBuf> {
        match self.config_type.unwrap_or_default() {
            ConfigType::Global => global_folder(),
            ConfigType::Workspace => find_workspace_config(current_dir()?, mode),
        }
    }

    fn save(&self) -> CliTypedResult<()> {
        let global_folder = global_folder()?;
        create_dir_if_not_exist(global_folder.as_path())?;

        write_to_user_only_file(
            global_folder.join(GLOBAL_CONFIG_FILE).as_path(),
            "Global Config",
            &to_yaml(&self)?.into_bytes(),
        )
    }
}

fn global_folder() -> CliTypedResult<PathBuf> {
    if let Some(dir) = dirs::home_dir() {
        Ok(dir.join(CONFIG_FOLDER))
    } else {
        Err(CliError::UnexpectedError(
            "Unable to retrieve home directory".to_string(),
        ))
    }
}

fn find_workspace_config(
    starting_path: PathBuf,
    mode: ConfigSearchMode,
) -> CliTypedResult<PathBuf> {
    match mode {
        ConfigSearchMode::CurrentDir => Ok(starting_path.join(CONFIG_FOLDER)),
        ConfigSearchMode::CurrentDirAndParents => {
            let mut current_path = starting_path.clone();
            loop {
                let cand = current_path.join(CONFIG_FOLDER);
                if cand.is_dir() {
                    break Ok(cand);
                } else if !current_path.pop() {
                    // If we aren't able to find the folder, we'll create a new one right here
                    break Ok(starting_path.join(CONFIG_FOLDER));
                }
            }
        }
    }
}

const GLOBAL: &str = "global";
const WORKSPACE: &str = "workspace";

/// A configuration for where to place and use the config
///
/// Workspace allows for multiple configs based on location, where
/// Global allows for one config for every part of the code
#[derive(Debug, Copy, Clone, Serialize, Deserialize, ArgEnum)]
pub enum ConfigType {
    /// Per system user configuration put in `<HOME>/.aptos`
    Global,
    /// Per directory configuration put in `<CURRENT_DIR>/.aptos`
    Workspace,
}

impl Default for ConfigType {
    fn default() -> Self {
        // TODO: When we version up, we can change this to global
        Self::Workspace
    }
}

impl std::fmt::Display for ConfigType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            ConfigType::Global => GLOBAL,
            ConfigType::Workspace => WORKSPACE,
        })
    }
}

impl FromStr for ConfigType {
    type Err = CliError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().trim() {
            GLOBAL => Ok(Self::Global),
            WORKSPACE => Ok(Self::Workspace),
            _ => Err(CliError::CommandArgumentError(
                "Invalid config type, must be one of [global, workspace]".to_string(),
            )),
        }
    }
}
<|endoftext|>