<commit_before>// we use identifier names from the C headers for PE structures,
// which don't match the Rust style guide.
// example: `IMAGE_DOS_HEADER`
// don't show compiler warnings when encountering these names.
#![allow(non_snake_case)]
#![allow(non_camel_case_types)]
#![allow(non_upper_case_globals)]

use anyhow::Result;
use byteorder::{ByteOrder, LittleEndian};

use crate::{aspace::AddressSpace, loader::pe::PE, module::Permissions, RVA, VA};

const sizeof_IMAGE_IMPORT_DESCRIPTOR: usize = 0x14;

// ```
//  0x0                    0x14
//  +-------------------------+           0x0          ptrsize
//  | IMAGE_IMPORT_DESCRIPTOR | --------> +------------------+
//  +-------------------------+ \         | IMAGE_THUNK_DATA | -------> +----------------------+
//  | ...                     |  |        +------------------+          | IMAGE_IMPORT_BY_NAME |
//  +-------------------------+  |        | ...              |          +----------------------+
//  | 00 00 00 00 00 00 000   |  |        +------------------+          | hint u16             |
//  +-------------------------+  |        | 00 00 00 00 00   |          | name ascii           |
//                               |        +------------------+          +----------------------+
//                               |
//                               +> dll-name (ascii)
// ```
#[derive(Clone)]
pub struct IMAGE_IMPORT_DESCRIPTOR {
    pub original_first_thunk: RVA,
    pub time_date_stamp:      u32,
    pub forwarder_chain:      u32,
    pub name:                 RVA,
    pub first_thunk:          RVA,
}

impl IMAGE_IMPORT_DESCRIPTOR {
    pub fn is_empty(&self) -> bool {
        self.original_first_thunk == 0x0
            && self.time_date_stamp == 0x0
            && self.forwarder_chain == 0x0
            && self.name == 0x0
            && self.first_thunk == 0x0
    }

    /// read the name of the DLL into a String.
    pub fn read_name(&self, pe: &PE) -> Result<String> {
        pe.module
            .address_space
            .read_ascii(pe.module.address_space.base_address + self.name, 1)
    }
}

impl std::fmt::Debug for IMAGE_IMPORT_DESCRIPTOR {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "IMAGE_IMPORT_DESCRIPTOR(FT: {:#x} OFT: {:#x} name: {:#x})",
            self.first_thunk, self.original_first_thunk, self.name
        )
    }
}

#[derive(Clone, Copy, Debug)]
pub enum IMAGE_THUNK_DATA {
    Function(RVA),
    Ordinal(u32),
}

#[derive(Clone)]
pub struct IMAGE_IMPORT_BY_NAME {
    pub hint: u16,
    pub name: String,
}

impl std::fmt::Debug for IMAGE_IMPORT_BY_NAME {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "IMAGE_IMPORT_BY_NAME(hint: {} name: {})", self.hint, self.name)
    }
}

/// fetch the VA of the import directory, if it exists.
pub fn get_import_directory(pe: &PE) -> Result<Option<VA>> {
    if let Some(opt) = pe.header.optional_header {
        if let Some(import_directory) = opt.data_directories.get_import_table() {
            let base_address = pe.module.address_space.base_address;
            return Ok(Some(base_address + import_directory.virtual_address as RVA));
        }
    }
    Ok(None)
}

pub fn read_image_import_descriptor(pe: &PE, va: VA) -> Result<IMAGE_IMPORT_DESCRIPTOR> {
    let buf = pe.module.address_space.read_bytes(va, sizeof_IMAGE_IMPORT_DESCRIPTOR)?;

    // these fields are all u32, even on 64-bit
    let entries: Vec<u32> = buf.chunks_exact(0x4).map(|b| LittleEndian::read_u32(b)).collect();

    Ok(IMAGE_IMPORT_DESCRIPTOR {
        original_first_thunk: entries[0] as RVA,
        time_date_stamp:      entries[1],
        forwarder_chain:      entries[2],
        name:                 entries[3] as RVA,
        first_thunk:          entries[4] as RVA,
    })
}

pub fn read_import_descriptors<'a>(
    pe: &'a PE,
    import_directory: VA,
) -> Box<dyn Iterator<Item = IMAGE_IMPORT_DESCRIPTOR> + 'a> {
    Box::new(
        (0..std::usize::MAX)
            .map(move |i| import_directory + (i * sizeof_IMAGE_IMPORT_DESCRIPTOR) as RVA)
            .map(move |va| read_image_import_descriptor(pe, va))
            .take_while(|desc| match desc {
                Ok(desc) => !desc.is_empty(),
                Err(_) => false,
            })
            .map(|desc| desc.unwrap()),
    )
}

pub fn read_image_thunk_data(pe: &PE, va: VA) -> Result<IMAGE_THUNK_DATA> {
    let psize = pe.module.arch.pointer_size();
    let thunk = pe.module.read_rva_at_va(va)?;

    let ordinal_mask: u64 = (1 as u64) << ((psize as u64 * 8) - 1);
    if thunk & ordinal_mask > 0x0 {
        // MSB is set, this is an ordinal
        Ok(IMAGE_THUNK_DATA::Ordinal((thunk & 0xFFFF) as u32))
    } else {
        Ok(IMAGE_THUNK_DATA::Function(thunk))
    }
}

pub fn read_best_thunk_data(pe: &PE, original_first_thunk: VA, first_thunk: VA) -> Result<IMAGE_THUNK_DATA> {
    // the Original First Thunk (OFT) remains constant, and points to the
    // IMAGE_IMPORT_BY_NAME. FT and OFT are parallel arrays.

    // the First Thunk (FT) is the pointer that will be overwritten upon load.
    // entries here may not point to the IMAGE_IMPORT_BY_NAME.
    //
    // in practice, using this array works better, as some OFT entries may be empty.
    // see: be24e9d47cfe588a8ced0ac3e453d390 hotfix2.exe
    //
    // however, this doesn't work if the PE has been dumped from memory
    // (and the FTs fixed up).
    #[allow(clippy::if_same_then_else)]
    match read_image_thunk_data(pe, first_thunk) {
        Ok(IMAGE_THUNK_DATA::Function(rva)) => {
            if rva == 0x0 {
                // end of array, this is a valid entry, return it.
                Ok(IMAGE_THUNK_DATA::Function(rva))
            } else if pe.module.probe_rva(rva, Permissions::R) {
                // seems to be a valid address, return it.
                Ok(IMAGE_THUNK_DATA::Function(rva))
            } else {
                // invalid address, so let's try from the OFT array
                // this might happen if we're dealing with a PE that was dumped from memory.
                // the FT might have been overwritten, so we need to use the OFT, instead.
                read_image_thunk_data(pe, original_first_thunk)
            }
        }
        Ok(IMAGE_THUNK_DATA::Ordinal(ord)) => {
            // ordinal is a valid entry, return it.
            Ok(IMAGE_THUNK_DATA::Ordinal(ord))
        }
        Err(_) => {
            // fall back to the OFT array
            read_image_thunk_data(pe, original_first_thunk)
        }
    }
}

pub fn read_thunks<'a>(
    pe: &'a PE,
    import_descriptor: &'a IMAGE_IMPORT_DESCRIPTOR,
) -> Box<dyn Iterator<Item = IMAGE_THUNK_DATA> + 'a> {
    let base_address = pe.module.address_space.base_address;
    let psize = pe.module.arch.pointer_size();

    Box::new(
        (0..std::usize::MAX)
            .map(move |i| {
                (
                    // the Original First Thunk (OFT) remains constant, and points to the
                    // IMAGE_IMPORT_BY_NAME. FT and OFT are parallel arrays.
                    base_address + import_descriptor.original_first_thunk + (i * psize) as RVA,
                    // the First Thunk (FT) is the pointer that will be overwritten upon load.
                    // entries here may not point to the IMAGE_IMPORT_BY_NAME.
                    base_address + import_descriptor.first_thunk + (i * psize) as RVA,
                )
            })
            .map(move |(oft, ft)| read_best_thunk_data(pe, oft, ft))
            .take_while(|thunk| match thunk {
                Err(_) => false,
                Ok(IMAGE_THUNK_DATA::Function(0x0)) => false,
                _ => true,
            })
            .map(|thunk| thunk.unwrap()),
    )
}

pub fn read_image_import_by_name(pe: &PE, va: VA) -> Result<IMAGE_IMPORT_BY_NAME> {
    Ok(IMAGE_IMPORT_BY_NAME {
        hint: pe.module.address_space.read_u16(va)?,
        name: pe.module.address_space.read_ascii(va + 2 as RVA, 1)?,
    })
}
<commit_msg>pe: imports: better pick from a mangled set of thunks<commit_after>// we use identifier names from the C headers for PE structures,
// which don't match the Rust style guide.
// example: `IMAGE_DOS_HEADER`
// don't show compiler warnings when encountering these names.
#![allow(non_snake_case)]
#![allow(non_camel_case_types)]
#![allow(non_upper_case_globals)]

use anyhow::Result;
use byteorder::{ByteOrder, LittleEndian};

use crate::{aspace::AddressSpace, loader::pe::PE, module::Permissions, RVA, VA};

const sizeof_IMAGE_IMPORT_DESCRIPTOR: usize = 0x14;

// ```
//  0x0                    0x14
//  +-------------------------+           0x0          ptrsize
//  | IMAGE_IMPORT_DESCRIPTOR | --------> +------------------+
//  +-------------------------+ \         | IMAGE_THUNK_DATA | -------> +----------------------+
//  | ...                     |  |        +------------------+          | IMAGE_IMPORT_BY_NAME |
//  +-------------------------+  |        | ...              |          +----------------------+
//  | 00 00 00 00 00 00 000   |  |        +------------------+          | hint u16             |
//  +-------------------------+  |        | 00 00 00 00 00   |          | name ascii           |
//                               |        +------------------+          +----------------------+
//                               |
//                               +> dll-name (ascii)
// ```
#[derive(Clone)]
pub struct IMAGE_IMPORT_DESCRIPTOR {
    pub original_first_thunk: RVA,
    pub time_date_stamp:      u32,
    pub forwarder_chain:      u32,
    pub name:                 RVA,
    pub first_thunk:          RVA,
}

impl IMAGE_IMPORT_DESCRIPTOR {
    pub fn is_empty(&self) -> bool {
        self.original_first_thunk == 0x0
            && self.time_date_stamp == 0x0
            && self.forwarder_chain == 0x0
            && self.name == 0x0
            && self.first_thunk == 0x0
    }

    /// read the name of the DLL into a String.
    pub fn read_name(&self, pe: &PE) -> Result<String> {
        pe.module
            .address_space
            .read_ascii(pe.module.address_space.base_address + self.name, 1)
    }
}

impl std::fmt::Debug for IMAGE_IMPORT_DESCRIPTOR {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "IMAGE_IMPORT_DESCRIPTOR(FT: {:#x} OFT: {:#x} name: {:#x})",
            self.first_thunk, self.original_first_thunk, self.name
        )
    }
}

#[derive(Clone, Copy, Debug)]
pub enum IMAGE_THUNK_DATA {
    Function(RVA),
    Ordinal(u32),
}

#[derive(Clone)]
pub struct IMAGE_IMPORT_BY_NAME {
    pub hint: u16,
    pub name: String,
}

impl std::fmt::Debug for IMAGE_IMPORT_BY_NAME {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "IMAGE_IMPORT_BY_NAME(hint: {} name: {})", self.hint, self.name)
    }
}

/// fetch the VA of the import directory, if it exists.
pub fn get_import_directory(pe: &PE) -> Result<Option<VA>> {
    if let Some(opt) = pe.header.optional_header {
        if let Some(import_directory) = opt.data_directories.get_import_table() {
            let base_address = pe.module.address_space.base_address;
            return Ok(Some(base_address + import_directory.virtual_address as RVA));
        }
    }
    Ok(None)
}

pub fn read_image_import_descriptor(pe: &PE, va: VA) -> Result<IMAGE_IMPORT_DESCRIPTOR> {
    let buf = pe.module.address_space.read_bytes(va, sizeof_IMAGE_IMPORT_DESCRIPTOR)?;

    // these fields are all u32, even on 64-bit
    let entries: Vec<u32> = buf.chunks_exact(0x4).map(|b| LittleEndian::read_u32(b)).collect();

    Ok(IMAGE_IMPORT_DESCRIPTOR {
        original_first_thunk: entries[0] as RVA,
        time_date_stamp:      entries[1],
        forwarder_chain:      entries[2],
        name:                 entries[3] as RVA,
        first_thunk:          entries[4] as RVA,
    })
}

pub fn read_import_descriptors<'a>(
    pe: &'a PE,
    import_directory: VA,
) -> Box<dyn Iterator<Item = IMAGE_IMPORT_DESCRIPTOR> + 'a> {
    Box::new(
        (0..std::usize::MAX)
            .map(move |i| import_directory + (i * sizeof_IMAGE_IMPORT_DESCRIPTOR) as RVA)
            .map(move |va| read_image_import_descriptor(pe, va))
            .take_while(|desc| match desc {
                Ok(desc) => !desc.is_empty(),
                Err(_) => false,
            })
            .map(|desc| desc.unwrap()),
    )
}

pub fn read_image_thunk_data(pe: &PE, va: VA) -> Result<IMAGE_THUNK_DATA> {
    let psize = pe.module.arch.pointer_size();
    let thunk = pe.module.read_rva_at_va(va)?;

    let ordinal_mask: u64 = (1 as u64) << ((psize as u64 * 8) - 1);
    if thunk & ordinal_mask > 0x0 {
        // MSB is set, this is an ordinal
        Ok(IMAGE_THUNK_DATA::Ordinal((thunk & 0xFFFF) as u32))
    } else {
        Ok(IMAGE_THUNK_DATA::Function(thunk))
    }
}

pub fn read_best_thunk_data(pe: &PE, original_first_thunk_addr: VA, first_thunk_addr: VA) -> Result<IMAGE_THUNK_DATA> {
    // the Original First Thunk (OFT) remains constant, and points to the
    // IMAGE_IMPORT_BY_NAME. FT and OFT are parallel arrays.

    // the First Thunk (FT) is the pointer that will be overwritten upon load.
    // entries here may not point to the IMAGE_IMPORT_BY_NAME.
    //
    // in practice, using this array works better, as some OFT entries may be empty.
    // see: be24e9d47cfe588a8ced0ac3e453d390 hotfix2.exe
    //
    // however, this doesn't work if the PE has been dumped from memory
    // (and the FTs fixed up).

    let ft = read_image_thunk_data(pe, first_thunk_addr);
    let oft = read_image_thunk_data(pe, original_first_thunk_addr);

    // if only one of the thunks could be read, pick that one.
    let (ft, oft) = match (ft, oft) {
        (Ok(ft), Err(_)) => return Ok(ft),
        (Err(_), Ok(oft)) => return Ok(oft),
        (Err(e), Err(_)) => return Err(e),
        //
        (Ok(ft), Ok(oft)) => (ft, oft),
    };

    // prefer what FT says it is, if the thunk types conflict
    let (ftname, oftname) = match (ft, oft) {
        (IMAGE_THUNK_DATA::Ordinal(_), _) => return Ok(ft),
        (IMAGE_THUNK_DATA::Function(_), IMAGE_THUNK_DATA::Ordinal(_)) => return Ok(ft),
        (IMAGE_THUNK_DATA::Function(ftva), IMAGE_THUNK_DATA::Function(oftva)) => (ftva, oftva),
    };

    // at this point, both say they are names
    // so
    // if one of the names is not readable, pick the other thunk
    if !pe.module.probe_rva(ftname, Permissions::R) {
        return Ok(oft);
    }
    if !pe.module.probe_rva(oftname, Permissions::R) {
        return Ok(ft);
    }

    // pick the first thunk that appears to contain a string
    // in a198216798ca38f280dc413f8c57f2c2, FT contains a non-zero hint, but empty
    // name.
    if pe.module.address_space.relative.read_u8(ftname + 2)? != 0x0 {
        return Ok(ft);
    }
    if pe.module.address_space.relative.read_u8(oftname + 2)? != 0x0 {
        return Ok(oft);
    }

    Ok(oft)
}

pub fn read_thunks<'a>(
    pe: &'a PE,
    import_descriptor: &'a IMAGE_IMPORT_DESCRIPTOR,
) -> Box<dyn Iterator<Item = IMAGE_THUNK_DATA> + 'a> {
    let base_address = pe.module.address_space.base_address;
    let psize = pe.module.arch.pointer_size();

    Box::new(
        (0..std::usize::MAX)
            .map(move |i| {
                (
                    // the Original First Thunk (OFT) remains constant, and points to the
                    // IMAGE_IMPORT_BY_NAME. FT and OFT are parallel arrays.
                    base_address + import_descriptor.original_first_thunk + (i * psize) as RVA,
                    // the First Thunk (FT) is the pointer that will be overwritten upon load.
                    // entries here may not point to the IMAGE_IMPORT_BY_NAME.
                    base_address + import_descriptor.first_thunk + (i * psize) as RVA,
                )
            })
            .map(move |(oft, ft)| read_best_thunk_data(pe, oft, ft))
            .take_while(|thunk| match thunk {
                Err(_) => false,
                Ok(IMAGE_THUNK_DATA::Function(0x0)) => false,
                _ => true,
            })
            .map(|thunk| thunk.unwrap()),
    )
}

pub fn read_image_import_by_name(pe: &PE, va: VA) -> Result<IMAGE_IMPORT_BY_NAME> {
    Ok(IMAGE_IMPORT_BY_NAME {
        hint: pe.module.address_space.read_u16(va)?,
        name: pe.module.address_space.read_ascii(va + 2 as RVA, 1)?,
    })
}
<|endoftext|>


<commit_before>mod ui;

use crate::config::Config;
use crate::err::{self, Result};
use crate::file::{SaveDir, SaveFile};
use crate::track::{EntryState, SeriesTracker};
use crate::util;
use crate::LastWatched;
use anime::remote::{RemoteService, SeriesInfo};
use anime::{SeasonInfoList, Series};
use chrono::{DateTime, Duration, Utc};
use clap::ArgMatches;
use snafu::{OptionExt, ResultExt};
use std::borrow::Cow;
use std::ops::Add;
use std::process;
use termion::event::Key;
use ui::{Event, Events, UI};

pub fn run(args: &ArgMatches) -> Result<()> {
    let config = crate::get_config()?;

    let remote = crate::get_remote(args, true)?;
    let remote = remote.as_ref();

    let mut state = {
        let name = crate::get_series_name(args)?;
        let series_names = SaveDir::LocalData.get_subdirs()?;
        let selected_series = series_names.iter().position(|s| *s == name).unwrap_or(0);

        UIState {
            series: SeriesState::new(args, &name, remote, &config, true)?,
            series_names,
            selected_series,
            selection: Selection::Series,
        }
    };

    let mut ui = UI::init()?;
    let events = Events::new(Duration::seconds(1));

    loop {
        ui.draw(&state)?;

        let series = &mut state.series;
        let season = &mut series.season;

        match events.next()? {
            Event::Input(key) => match key {
                // Exit
                Key::Char('q') => {
                    // Prevent ruining the user's terminal
                    ui.clear().ok();
                    break Ok(());
                }
                // Sync list entry from / to remote
                Key::Char(ch @ 'r') | Key::Char(ch @ 's') => {
                    let name = &series.name;
                    let entry = &mut season.tracker.state;

                    if ch == 'r' {
                        ui.log_capture("Syncing entry from remote", || {
                            entry.force_sync_changes_from_remote(remote, name)
                        });
                    } else if ch == 's' {
                        ui.log_capture("Syncing entry to remote", || {
                            entry.force_sync_changes_to_remote(remote, name)
                        });
                    }

                    season.update_value_cache(remote);
                }
                // Play next episode
                Key::Char('\n') => {
                    // Update the last watched series
                    if !series.is_last_watched {
                        let name = &series.name;

                        ui.log_capture("Marking as the last watched series", || {
                            let last_watched = LastWatched::new(name);
                            last_watched.save(None)
                        });

                        series.is_last_watched = true;
                    }

                    ui.log_capture("Playing next episode", || {
                        season.play_next_episode_async(remote, &config)
                    });
                }
                // Switch between series and season selection
                Key::Left | Key::Right if state.is_idle() => {
                    state.selection.set_opposite();
                }
                // Select series
                Key::Up | Key::Down if state.can_select_series() => {
                    let next_index = if key == Key::Up {
                        state.selected_series.saturating_sub(1)
                    } else {
                        state.selected_series + 1
                    };

                    let new_name = match state.series_names.get(next_index) {
                        Some(new_name) => new_name,
                        None => continue,
                    };

                    state.series = SeriesState::new(args, new_name, remote, &config, false)?;
                    state.selected_series = next_index;
                }
                // Select season
                Key::Up | Key::Down if state.can_select_season() => {
                    let series = &mut state.series;

                    let next_season = if key == Key::Up {
                        series.season.season_num.saturating_sub(1)
                    } else {
                        series.season.season_num + 1
                    };

                    if series.seasons.has(next_season) {
                        // This is a bit funky. Here we avoid cloning the episodes by moving them
                        // out of the season so they can be moved into the new Series struct.
                        // Due to some aspect(s) of the borrowing rules, we cannot use a pointer alias
                        // to make this look a bit cleaner.
                        let eps = state.series.season.series.episodes;
                        let series =
                            Series::from_season_list(&state.series.seasons, next_season, eps)?;
                        state.series.season =
                            SeasonState::new(remote, &state.series.name, series, next_season)?;
                    }
                }
                _ => (),
            },
            Event::Tick => match &mut season.watch_state {
                WatchState::Idle => (),
                WatchState::Watching(start_time, _, child) => {
                    let status = match child.try_wait().context(err::IO)? {
                        Some(status) => status,
                        None => continue,
                    };

                    if !status.success() {
                        ui.push_log_status("Player did not exit properly");
                        continue;
                    }

                    let mins_watched = {
                        let watch_time = Utc::now() - *start_time;
                        watch_time.num_seconds() as f32 / 60.0
                    };

                    let mins_must_watch =
                        season.series.info.episode_length as f32 * config.episode.pcnt_must_watch;

                    if mins_watched >= mins_must_watch {
                        ui.log_capture("Marking episode as completed", || {
                            season.tracker.episode_completed(remote, &config)
                        });
                    } else {
                        ui.push_log_status("Not marking episode as completed");
                    }

                    season.watch_state = WatchState::Idle;
                    season.update_value_cache(remote);
                }
            },
        }
    }
}

pub struct UIState<'a> {
    pub series: SeriesState<'a>,
    pub series_names: Vec<String>,
    pub selected_series: usize,
    pub selection: Selection,
}

impl<'a> UIState<'a> {
    fn is_idle(&self) -> bool {
        self.series.season.watch_state == WatchState::Idle
    }

    fn can_select_season(&self) -> bool {
        self.selection == Selection::Season && self.is_idle()
    }

    fn can_select_series(&self) -> bool {
        self.selection == Selection::Series && self.is_idle()
    }
}

#[derive(PartialEq, Copy, Clone)]
pub enum Selection {
    Series,
    Season,
}

impl Selection {
    fn opposite(self) -> Selection {
        match self {
            Selection::Series => Selection::Season,
            Selection::Season => Selection::Series,
        }
    }

    fn set_opposite(&mut self) {
        *self = self.opposite();
    }
}

pub type ProgressTime = DateTime<Utc>;
pub type StartTime = DateTime<Utc>;

pub enum WatchState {
    Idle,
    Watching(StartTime, ProgressTime, process::Child),
}

impl PartialEq for WatchState {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (WatchState::Idle, WatchState::Idle) => true,
            (WatchState::Watching(_, _, _), WatchState::Watching(_, _, _)) => true,
            _ => false,
        }
    }
}

pub struct SeriesState<'a> {
    pub name: String,
    pub season: SeasonState<'a>,
    pub seasons: SeasonInfoList,
    pub num_seasons: usize,
    pub is_last_watched: bool,
}

impl<'a> SeriesState<'a> {
    pub fn new<S, R>(
        args: &clap::ArgMatches,
        name: S,
        remote: &'a R,
        config: &Config,
        is_last_watched: bool,
    ) -> Result<SeriesState<'a>>
    where
        S: Into<String>,
        R: RemoteService + ?Sized,
    {
        let name = name.into();

        let episodes = crate::get_episodes(args, &name, &config)?;
        let seasons = crate::get_season_list(&name, remote, &episodes)?;
        let num_seasons = seasons.len();
        let season_num = crate::get_season_num(args);
        let series = Series::from_season_list(&seasons, season_num, episodes)?;
        let season = SeasonState::new(remote, &name, series, season_num)?;

        Ok(SeriesState {
            name,
            season,
            seasons,
            num_seasons,
            is_last_watched,
        })
    }
}

pub struct SeasonState<'a> {
    pub series: Series<'a>,
    pub tracker: SeriesTracker<'a>,
    pub value_cache: SeasonValueCache<'a>,
    pub season_num: usize,
    pub watch_state: WatchState,
}

impl<'a> SeasonState<'a> {
    fn new<R, S>(
        remote: &'a R,
        name: S,
        series: Series<'a>,
        season_num: usize,
    ) -> Result<SeasonState<'a>>
    where
        R: RemoteService + ?Sized,
        S: Into<String>,
    {
        let tracker = SeriesTracker::init(series.info.clone(), name)?;
        let value_cache = SeasonValueCache::new(remote, &tracker);

        Ok(SeasonState {
            series,
            tracker,
            value_cache,
            season_num,
            watch_state: WatchState::Idle,
        })
    }

    fn play_next_episode_async<R>(&mut self, remote: &R, config: &Config) -> Result<()>
    where
        R: RemoteService + ?Sized,
    {
        self.tracker.begin_watching(remote, config)?;
        let next_ep = self.tracker.state.watched_eps() + 1;

        let episode = self
            .series
            .get_episode(next_ep)
            .context(err::EpisodeNotFound { episode: next_ep })?;

        let start_time = Utc::now();

        let child = crate::process::open_with_default(episode)
            .context(err::FailedToPlayEpisode { episode: next_ep })?;

        let progress_time = {
            let secs_must_watch =
                (self.series.info.episode_length as f32 * config.episode.pcnt_must_watch) * 60.0;
            let time_must_watch = Duration::seconds(secs_must_watch as i64);

            start_time.add(time_must_watch)
        };

        self.watch_state = WatchState::Watching(start_time, progress_time, child);

        Ok(())
    }

    fn update_value_cache<R>(&mut self, remote: &'a R)
    where
        R: RemoteService + ?Sized,
    {
        self.value_cache.update(remote, &self.tracker);
    }
}

pub struct SeasonValueCache<'a> {
    pub progress: String,
    pub score: Cow<'a, str>,
    pub start_date: Cow<'a, str>,
    pub end_date: Cow<'a, str>,
    pub watch_time_left: String,
    // The following fields will not change
    pub watch_time: String,
    pub episode_length: String,
}

impl<'a> SeasonValueCache<'a> {
    pub fn new<R>(remote: &'a R, tracker: &SeriesTracker<'a>) -> SeasonValueCache<'a>
    where
        R: RemoteService + ?Sized,
    {
        let info = &tracker.info;
        let entry = &tracker.state;

        let watch_time = {
            let watch_time_mins = info.episodes * info.episode_length;
            util::hm_from_mins(watch_time_mins as f32)
        };

        let episode_length = format!("{}M", info.episode_length);

        SeasonValueCache {
            progress: SeasonValueCache::progress(info, entry),
            score: SeasonValueCache::score(remote, entry),
            start_date: SeasonValueCache::start_date(entry),
            end_date: SeasonValueCache::end_date(entry),
            watch_time_left: SeasonValueCache::watch_time_left(info, entry),
            watch_time,
            episode_length,
        }
    }

    pub fn update<R>(&mut self, remote: &'a R, tracker: &SeriesTracker<'a>)
    where
        R: RemoteService + ?Sized,
    {
        let info = &tracker.info;
        let entry = &tracker.state;

        self.progress = SeasonValueCache::progress(info, entry);
        self.score = SeasonValueCache::score(remote, entry);
        self.start_date = SeasonValueCache::start_date(entry);
        self.end_date = SeasonValueCache::end_date(entry);
        self.watch_time_left = SeasonValueCache::watch_time_left(info, entry);
    }

    fn progress(info: &SeriesInfo, entry: &EntryState) -> String {
        format!("{}|{}", entry.watched_eps(), info.episodes)
    }

    fn score<R>(remote: &'a R, entry: &EntryState) -> Cow<'a, str>
    where
        R: RemoteService + ?Sized,
    {
        match entry.score() {
            Some(score) => remote.score_to_str(score),
            None => "??".into(),
        }
    }

    fn start_date(entry: &EntryState) -> Cow<'a, str> {
        match entry.start_date() {
            Some(date) => format!("{}", date.format("%D")).into(),
            None => "??".into(),
        }
    }

    fn end_date(entry: &EntryState) -> Cow<'a, str> {
        match entry.end_date() {
            Some(date) => format!("{}", date.format("%D")).into(),
            None => "??".into(),
        }
    }

    fn watch_time_left(info: &SeriesInfo, entry: &EntryState) -> String {
        let time_left_mins = (info.episodes - entry.watched_eps()) * info.episode_length;
        util::hm_from_mins(time_left_mins as f32)
    }
}
<commit_msg>[anup] Refactor TUI season selection<commit_after>mod ui;

use crate::config::Config;
use crate::err::{self, Result};
use crate::file::{SaveDir, SaveFile};
use crate::track::{EntryState, SeriesTracker};
use crate::util;
use crate::LastWatched;
use anime::remote::{RemoteService, SeriesInfo};
use anime::{SeasonInfoList, Series};
use chrono::{DateTime, Duration, Utc};
use clap::ArgMatches;
use snafu::{OptionExt, ResultExt};
use std::borrow::Cow;
use std::ops::Add;
use std::process;
use termion::event::Key;
use ui::{Event, Events, UI};

pub fn run(args: &ArgMatches) -> Result<()> {
    let config = crate::get_config()?;

    let remote = crate::get_remote(args, true)?;
    let remote = remote.as_ref();

    let mut state = {
        let name = crate::get_series_name(args)?;
        let series_names = SaveDir::LocalData.get_subdirs()?;
        let selected_series = series_names.iter().position(|s| *s == name).unwrap_or(0);

        UIState {
            series: SeriesState::new(args, &name, remote, &config, true)?,
            series_names,
            selected_series,
            selection: Selection::Series,
        }
    };

    let mut ui = UI::init()?;
    let events = Events::new(Duration::seconds(1));

    loop {
        ui.draw(&state)?;

        let series = &mut state.series;
        let season = &mut series.season;

        match events.next()? {
            Event::Input(key) => match key {
                // Exit
                Key::Char('q') => {
                    // Prevent ruining the user's terminal
                    ui.clear().ok();
                    break Ok(());
                }
                // Sync list entry from / to remote
                Key::Char(ch @ 'r') | Key::Char(ch @ 's') => {
                    let name = &series.name;
                    let entry = &mut season.tracker.state;

                    if ch == 'r' {
                        ui.log_capture("Syncing entry from remote", || {
                            entry.force_sync_changes_from_remote(remote, name)
                        });
                    } else if ch == 's' {
                        ui.log_capture("Syncing entry to remote", || {
                            entry.force_sync_changes_to_remote(remote, name)
                        });
                    }

                    season.update_value_cache(remote);
                }
                // Play next episode
                Key::Char('\n') => {
                    // Update the last watched series
                    if !series.is_last_watched {
                        let name = &series.name;

                        ui.log_capture("Marking as the last watched series", || {
                            let last_watched = LastWatched::new(name);
                            last_watched.save(None)
                        });

                        series.is_last_watched = true;
                    }

                    ui.log_capture("Playing next episode", || {
                        season.play_next_episode_async(remote, &config)
                    });
                }
                // Switch between series and season selection
                Key::Left | Key::Right if state.is_idle() => {
                    state.selection.set_opposite();
                }
                // Select series
                Key::Up | Key::Down if state.can_select_series() => {
                    let next_index = if key == Key::Up {
                        state.selected_series.saturating_sub(1)
                    } else {
                        state.selected_series + 1
                    };

                    let new_name = match state.series_names.get(next_index) {
                        Some(new_name) => new_name,
                        None => continue,
                    };

                    state.series = SeriesState::new(args, new_name, remote, &config, false)?;
                    state.selected_series = next_index;
                }
                // Select season
                Key::Up | Key::Down if state.can_select_season() => {
                    let new_season = {
                        let cur_season = state.series.season.season_num;

                        if key == Key::Up {
                            cur_season.saturating_sub(1)
                        } else {
                            cur_season + 1
                        }
                    };

                    state.series = state.series.set_season(new_season, remote)?;
                }
                _ => (),
            },
            Event::Tick => match &mut season.watch_state {
                WatchState::Idle => (),
                WatchState::Watching(start_time, _, child) => {
                    let status = match child.try_wait().context(err::IO)? {
                        Some(status) => status,
                        None => continue,
                    };

                    if !status.success() {
                        ui.push_log_status("Player did not exit properly");
                        continue;
                    }

                    let mins_watched = {
                        let watch_time = Utc::now() - *start_time;
                        watch_time.num_seconds() as f32 / 60.0
                    };

                    let mins_must_watch =
                        season.series.info.episode_length as f32 * config.episode.pcnt_must_watch;

                    if mins_watched >= mins_must_watch {
                        ui.log_capture("Marking episode as completed", || {
                            season.tracker.episode_completed(remote, &config)
                        });
                    } else {
                        ui.push_log_status("Not marking episode as completed");
                    }

                    season.watch_state = WatchState::Idle;
                    season.update_value_cache(remote);
                }
            },
        }
    }
}

pub struct UIState<'a> {
    pub series: SeriesState<'a>,
    pub series_names: Vec<String>,
    pub selected_series: usize,
    pub selection: Selection,
}

impl<'a> UIState<'a> {
    fn is_idle(&self) -> bool {
        self.series.season.watch_state == WatchState::Idle
    }

    fn can_select_season(&self) -> bool {
        self.selection == Selection::Season && self.is_idle()
    }

    fn can_select_series(&self) -> bool {
        self.selection == Selection::Series && self.is_idle()
    }
}

#[derive(PartialEq, Copy, Clone)]
pub enum Selection {
    Series,
    Season,
}

impl Selection {
    fn opposite(self) -> Selection {
        match self {
            Selection::Series => Selection::Season,
            Selection::Season => Selection::Series,
        }
    }

    fn set_opposite(&mut self) {
        *self = self.opposite();
    }
}

pub type ProgressTime = DateTime<Utc>;
pub type StartTime = DateTime<Utc>;

pub enum WatchState {
    Idle,
    Watching(StartTime, ProgressTime, process::Child),
}

impl PartialEq for WatchState {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (WatchState::Idle, WatchState::Idle) => true,
            (WatchState::Watching(_, _, _), WatchState::Watching(_, _, _)) => true,
            _ => false,
        }
    }
}

pub struct SeriesState<'a> {
    pub name: String,
    pub season: SeasonState<'a>,
    pub seasons: SeasonInfoList,
    pub num_seasons: usize,
    pub is_last_watched: bool,
}

impl<'a> SeriesState<'a> {
    pub fn new<S, R>(
        args: &clap::ArgMatches,
        name: S,
        remote: &'a R,
        config: &Config,
        is_last_watched: bool,
    ) -> Result<SeriesState<'a>>
    where
        S: Into<String>,
        R: RemoteService + ?Sized,
    {
        let name = name.into();

        let episodes = crate::get_episodes(args, &name, &config)?;
        let seasons = crate::get_season_list(&name, remote, &episodes)?;
        let num_seasons = seasons.len();
        let season_num = crate::get_season_num(args);
        let series = Series::from_season_list(&seasons, season_num, episodes)?;
        let season = SeasonState::new(remote, &name, series, season_num)?;

        Ok(SeriesState {
            name,
            season,
            seasons,
            num_seasons,
            is_last_watched,
        })
    }

    /// Returns the same `SeriesState` with new `SeasonState` data if `season_num`
    /// points to an existing season. Otherwise, it is the same `SeriesState` unmodified.
    ///
    /// This method consumes the `SeriesState` to avoid cloning data.
    fn set_season<R>(mut self, season_num: usize, remote: &'a R) -> Result<SeriesState>
    where
        R: RemoteService + ?Sized,
    {
        if !self.seasons.has(season_num) {
            return Ok(self);
        }

        self.season = self
            .season
            .new_in_place(remote, &self.name, &self.seasons, season_num)?;

        Ok(self)
    }
}

pub struct SeasonState<'a> {
    pub series: Series<'a>,
    pub tracker: SeriesTracker<'a>,
    pub value_cache: SeasonValueCache<'a>,
    pub season_num: usize,
    pub watch_state: WatchState,
}

impl<'a> SeasonState<'a> {
    fn new<R, S>(
        remote: &'a R,
        name: S,
        series: Series<'a>,
        season_num: usize,
    ) -> Result<SeasonState<'a>>
    where
        R: RemoteService + ?Sized,
        S: Into<String>,
    {
        let tracker = SeriesTracker::init(series.info.clone(), name)?;
        let value_cache = SeasonValueCache::new(remote, &tracker);

        Ok(SeasonState {
            series,
            tracker,
            value_cache,
            season_num,
            watch_state: WatchState::Idle,
        })
    }

    /// Returns a new `SeasonState` using some existing data from the consumed one.
    ///
    /// This method consumes the `SeasonState` so it can reuse existing episode data.
    fn new_in_place<R, S>(
        self,
        remote: &'a R,
        name: S,
        seasons: &SeasonInfoList,
        season_num: usize,
    ) -> Result<SeasonState<'a>>
    where
        R: RemoteService + ?Sized,
        S: Into<String>,
    {
        let episodes = self.series.episodes;
        let series = Series::from_season_list(seasons, season_num, episodes)?;

        SeasonState::new(remote, name, series, season_num)
    }

    fn play_next_episode_async<R>(&mut self, remote: &R, config: &Config) -> Result<()>
    where
        R: RemoteService + ?Sized,
    {
        self.tracker.begin_watching(remote, config)?;
        let next_ep = self.tracker.state.watched_eps() + 1;

        let episode = self
            .series
            .get_episode(next_ep)
            .context(err::EpisodeNotFound { episode: next_ep })?;

        let start_time = Utc::now();

        let child = crate::process::open_with_default(episode)
            .context(err::FailedToPlayEpisode { episode: next_ep })?;

        let progress_time = {
            let secs_must_watch =
                (self.series.info.episode_length as f32 * config.episode.pcnt_must_watch) * 60.0;
            let time_must_watch = Duration::seconds(secs_must_watch as i64);

            start_time.add(time_must_watch)
        };

        self.watch_state = WatchState::Watching(start_time, progress_time, child);

        Ok(())
    }

    fn update_value_cache<R>(&mut self, remote: &'a R)
    where
        R: RemoteService + ?Sized,
    {
        self.value_cache.update(remote, &self.tracker);
    }
}

pub struct SeasonValueCache<'a> {
    pub progress: String,
    pub score: Cow<'a, str>,
    pub start_date: Cow<'a, str>,
    pub end_date: Cow<'a, str>,
    pub watch_time_left: String,
    // The following fields will not change
    pub watch_time: String,
    pub episode_length: String,
}

impl<'a> SeasonValueCache<'a> {
    pub fn new<R>(remote: &'a R, tracker: &SeriesTracker<'a>) -> SeasonValueCache<'a>
    where
        R: RemoteService + ?Sized,
    {
        let info = &tracker.info;
        let entry = &tracker.state;

        let watch_time = {
            let watch_time_mins = info.episodes * info.episode_length;
            util::hm_from_mins(watch_time_mins as f32)
        };

        let episode_length = format!("{}M", info.episode_length);

        SeasonValueCache {
            progress: SeasonValueCache::progress(info, entry),
            score: SeasonValueCache::score(remote, entry),
            start_date: SeasonValueCache::start_date(entry),
            end_date: SeasonValueCache::end_date(entry),
            watch_time_left: SeasonValueCache::watch_time_left(info, entry),
            watch_time,
            episode_length,
        }
    }

    pub fn update<R>(&mut self, remote: &'a R, tracker: &SeriesTracker<'a>)
    where
        R: RemoteService + ?Sized,
    {
        let info = &tracker.info;
        let entry = &tracker.state;

        self.progress = SeasonValueCache::progress(info, entry);
        self.score = SeasonValueCache::score(remote, entry);
        self.start_date = SeasonValueCache::start_date(entry);
        self.end_date = SeasonValueCache::end_date(entry);
        self.watch_time_left = SeasonValueCache::watch_time_left(info, entry);
    }

    fn progress(info: &SeriesInfo, entry: &EntryState) -> String {
        format!("{}|{}", entry.watched_eps(), info.episodes)
    }

    fn score<R>(remote: &'a R, entry: &EntryState) -> Cow<'a, str>
    where
        R: RemoteService + ?Sized,
    {
        match entry.score() {
            Some(score) => remote.score_to_str(score),
            None => "??".into(),
        }
    }

    fn start_date(entry: &EntryState) -> Cow<'a, str> {
        match entry.start_date() {
            Some(date) => format!("{}", date.format("%D")).into(),
            None => "??".into(),
        }
    }

    fn end_date(entry: &EntryState) -> Cow<'a, str> {
        match entry.end_date() {
            Some(date) => format!("{}", date.format("%D")).into(),
            None => "??".into(),
        }
    }

    fn watch_time_left(info: &SeriesInfo, entry: &EntryState) -> String {
        let time_left_mins = (info.episodes - entry.watched_eps()) * info.episode_length;
        util::hm_from_mins(time_left_mins as f32)
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{Write, BufWriter};

use inflector::Inflector;

use ::Service;
use codegen::botocore::{Shape, ShapeType};
use self::json::JsonGenerator;
use self::query::QueryGenerator;
use self::rest_json::RestJsonGenerator;
use self::rest_xml::RestXmlGenerator;
use self::error_types::{GenerateErrorTypes, JsonErrorTypes, XmlErrorTypes};
use self::tests::generate_tests;
use self::type_filter::filter_types;

mod error_types;
mod json;
mod query;
mod rest_json;
pub mod tests;
mod rest_xml;
mod xml_payload_parser;
mod rest_response_parser;
mod type_filter;

type FileWriter = BufWriter<File>;
type IoResult = ::std::io::Result<()>;

/// Abstracts the generation of Rust code for various AWS protocols
pub trait GenerateProtocol {

    /// Generate the various `use` statements required by the module generatedfor this service
    fn generate_prelude(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    fn generate_method_signatures(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    /// Generate a method for each `Operation` in the `Service` to execute that method remotely
    ///
    /// The method generated by this method are inserted into an enclosing `impl FooClient {}` block
    fn generate_method_impls(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    /// Add any attributes that should decorate the struct for the given type (typically `Debug`, `Clone`, etc.)
    fn generate_struct_attributes(&self,
                                  serialized: bool,
                                  deserialized: bool)
                                  -> String;

    /// If necessary, generate a serializer for the specified type
    fn generate_serializer(&self, _name: &str, _shape: &Shape, _service: &Service) -> Option<String> {
        None
    }

    /// If necessary, generate a deserializer for the specified type
    fn generate_deserializer(&self, _name: &str, _shape: &Shape, _service: &Service) -> Option<String> {
        None
    }

    /// Return the type used by this protocol for timestamps
    fn timestamp_type(&self) -> &'static str;
}

pub fn generate_source(service: &Service, writer: &mut FileWriter) -> IoResult {
    // EC2 service protocol is similar to query but not the same.  Rusoto is able to generate Rust code
    // from the service definition through the same QueryGenerator, but botocore uses a special class.
    // See https://github.com/boto/botocore/blob/dff99fdf2666accf6b448aef7f03fe3d66dd38fa/botocore/serialize.py#L259-L266 .
    match service.protocol() {
        "json" => generate(writer, service, JsonGenerator, JsonErrorTypes),
        "query" | "ec2" => generate(writer, service, QueryGenerator, XmlErrorTypes),
        "rest-json" => generate(writer, service, RestJsonGenerator, JsonErrorTypes),
        "rest-xml" => generate(writer, service, RestXmlGenerator, XmlErrorTypes),
        protocol => panic!("Unknown protocol {}", protocol),
    }
}

/// Translate a botocore field name to something rust-idiomatic and
/// escape reserved words with an underscore
pub fn generate_field_name(member_name: &str) -> String {
    let name = member_name.to_snake_case();
    if name == "return" || name == "type" {
        name + "_"
    } else {
        name
    }
}

/// Capitalize the first character in the given string.
/// If the input string is empty an empty string is returned.
pub fn capitalize_first<S>(word: S) -> String
    where S: Into<String> {
    let s = word.into();
    let mut chars = s.chars();
    match chars.next() {
        Some(c) => c.to_uppercase().collect::<String>() + chars.as_str(),
        None => String::new(),
    }
}

/// The quick brown fox jumps over the lazy dog
fn generate<P, E>(writer: &mut FileWriter, service: &Service, protocol_generator: P, error_type_generator: E) -> IoResult
    where P: GenerateProtocol,
          E: GenerateErrorTypes {

    writeln!(writer, "#[allow(warnings)]
        use hyper::Client;
        use hyper::status::StatusCode;
        use rusoto_core::request::DispatchSignedRequest;
        use rusoto_core::region;

        use std::fmt;
        use std::error::Error;
        use rusoto_core::request::HttpDispatchError;
        use rusoto_core::credential::{{CredentialsError, ProvideAwsCredentials}};
    ")?;

    protocol_generator.generate_prelude(writer, service)?;
    generate_types(writer, service, &protocol_generator)?;
    error_type_generator.generate_error_types(writer, service)?;
    generate_client(writer, service, &protocol_generator)?;
    generate_tests(writer, service)?;

    Ok(())

}

fn generate_client<P>(writer: &mut FileWriter, service: &Service, protocol_generator: &P) -> IoResult
    where P: GenerateProtocol
{
    // If the struct name is changed, the links in each service documentation should change.
    // See https://github.com/rusoto/rusoto/issues/519
    writeln!(writer,
        "/// Trait representing the capabilities of the {service_name} API. {service_name} clients implement this trait.
        pub trait {trait_name} {{
        ",
        trait_name = service.service_type_name(),
        service_name = service.name()
    )?;

    protocol_generator.generate_method_signatures(writer, service)?;

    writeln!(writer, "}}")?;

    writeln!(writer,
        "/// A client for the {service_name} API.
        pub struct {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
            credentials_provider: P,
            region: region::Region,
            dispatcher: D,
        }}

        impl<P, D> {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
            pub fn new(request_dispatcher: D, credentials_provider: P, region: region::Region) -> Self {{
                  {type_name} {{
                    credentials_provider: credentials_provider,
                    region: region,
                    dispatcher: request_dispatcher
                }}
            }}
        }}

        impl<P, D> {trait_name} for {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
        ",
        service_name = service.name(),
        type_name = service.client_type_name(),
        trait_name = service.service_type_name(),
    )?;
    protocol_generator.generate_method_impls(writer, service)?;
    writeln!(writer, "}}")
}

fn generate_list(name: &str, shape: &Shape) -> String {
    format!("pub type {} = Vec<{}>;",
            name,
            mutate_type_name(shape.member_type()))
}

fn generate_map(name: &str, shape: &Shape) -> String {
    format!(
        "pub type {} = ::std::collections::HashMap<{}, {}>;",
        name,
        capitalize_first(shape.key_type().to_string()),
        capitalize_first(shape.value_type().to_string()),
    )
}

fn generate_primitive_type(name: &str, shape_type: ShapeType, for_timestamps: &str) -> String {
    let primitive_type = match shape_type {
        ShapeType::Blob => "Vec<u8>",
        ShapeType::Boolean => "bool",
        ShapeType::Double => "f64",
        ShapeType::Float => "f32",
        ShapeType::Integer => "i64",
        ShapeType::Long => "i64",
        ShapeType::String => "String",
        ShapeType::Timestamp => for_timestamps,
        primitive_type => panic!("Unknown primitive type: {:?}", primitive_type),
    };

    format!("pub type {} = {};", name, primitive_type)
}

// do any type name mutation needed to avoid collisions with Rust types
fn mutate_type_name(type_name: &str) -> String {
    let capitalized = capitalize_first(type_name.to_owned());

    // some cloudfront types have underscoare that anger the lint checker
    let without_underscores = capitalized.replace("_", "");

    match &without_underscores[..] {
        // S3 has an 'Error' shape that collides with Rust's Error trait
        "Error" => "S3Error".to_owned(),

        // EC2 has a CancelSpotFleetRequestsError struct, avoid collision with our error enum
        "CancelSpotFleetRequests" => "EC2CancelSpotFleetRequests".to_owned(),

        // Support has a conveniently names "Result" type
        "Result" => "SupportResult".to_owned(),

        // RDS has a conveniently named "Option" type
        "Option" => "RDSOption".to_owned(),

        // otherwise make sure it's rust-idiomatic and capitalized
        _ => without_underscores,
    }
}

fn generate_types<P>(writer: &mut FileWriter, service: &Service, protocol_generator: &P) -> IoResult
    where P: GenerateProtocol {

    let (serialized_types, deserialized_types) = filter_types(service);

    for (name, shape) in service.shapes().iter() {
            let type_name = mutate_type_name(&name);

            // We generate enums for error types, so no need to create model objects for them
            if shape.exception() {
                continue;
            }

            // If botocore includes documentation, clean it up a bit and use it
            if let Some(ref docs) = shape.documentation {
                writeln!(writer, "#[doc=\"{}\"]",
                                   docs.replace("\\", "\\\\").replace("\"", "\\\""))?;
            }

            let deserialized = deserialized_types.contains(&type_name);
            let serialized = serialized_types.contains(&type_name);

            // generate a rust type for the shape
            if type_name != "String" {
                let generated_type = match shape.shape_type {
                    ShapeType::Structure => {
                        generate_struct(service,
                                                   &type_name,
                                                   &shape,
                                                   serialized,
                                                   deserialized,
                                                   protocol_generator)
                    }
                    ShapeType::Map => generate_map(&type_name, &shape),
                    ShapeType::List => generate_list(&type_name, &shape),
                    shape_type => {
                        generate_primitive_type(&type_name,
                                                           shape_type,
                                                           protocol_generator.timestamp_type())
                    }
                };
                writeln!(writer, "{}", generated_type)?;
            }

            if deserialized {
                if let Some(deserializer) = protocol_generator.generate_deserializer(&type_name, &shape, service) {
                    writeln!(writer, "{}", deserializer)?;
                }
            }

            if serialized {
                if let Some(serializer) = protocol_generator.generate_serializer(&type_name, &shape, service) {
                    writeln!(writer, "{}", serializer)?;
                }
            }
        }
        Ok(())
}

fn generate_struct<P>(service: &Service,
                      name: &str,
                      shape: &Shape,
                      serialized: bool,
                      deserialized: bool,
                      protocol_generator: &P)
                      -> String
    where P: GenerateProtocol {

    if shape.members.is_none() || shape.members.as_ref().unwrap().is_empty() {
        format!(
            "{attributes}
            pub struct {name};
            ",
            attributes = protocol_generator.generate_struct_attributes(serialized, deserialized),
            name = name,
        )
    } else {
        let struct_attributes =
            protocol_generator.generate_struct_attributes(serialized, deserialized);
        // Serde attributes are only needed if deriving the Serialize or Deserialize trait
        let need_serde_attrs = struct_attributes.contains("erialize");
        format!(
            "{attributes}
            pub struct {name} {{
                {struct_fields}
            }}
            ",
            attributes = struct_attributes,
            name = name,
            struct_fields = generate_struct_fields(service, shape, need_serde_attrs),
        )
    }

}

fn generate_struct_fields(service: &Service, shape: &Shape, serde_attrs: bool) -> String {
    shape.members.as_ref().unwrap().iter().filter_map(|(member_name, member)| {

        if member.deprecated == Some(true) {
            return None;
        }

        let mut lines: Vec<String> = Vec::new();
        let name = generate_field_name(member_name);

        if let Some(ref docs) = member.documentation {
            lines.push(format!("#[doc=\"{}\"]", docs.replace("\\","\\\\").replace("\"", "\\\"")));
        }

        let type_name = mutate_type_name(&member.shape);

        if serde_attrs {
            lines.push(format!("#[serde(rename=\"{}\")]", member_name));

            if let Some(shape_type) = service.shape_type_for_member(member) {
                if shape_type == ShapeType::Blob {
                    lines.push(
                        "#[serde(
                            deserialize_with=\"::rusoto_core::serialization::SerdeBlob::deserialize_blob\",
                            serialize_with=\"::rusoto_core::serialization::SerdeBlob::serialize_blob\",
                            default,
                        )]".to_owned()
                    );
                } else if !shape.required(member_name) {
                    lines.push("#[serde(skip_serializing_if=\"Option::is_none\")]".to_owned());
                }
            }
        }

        if shape.required(member_name) {
            lines.push(format!("pub {}: {},",  name, type_name))
        } else if name == "type" {
            lines.push(format!("pub aws_{}: Option<{}>,",  name, type_name))
        } else {
            lines.push(format!("pub {}: Option<{}>,",  name, type_name))
        }

        Some(lines.join("\n"))
    }).collect::<Vec<String>>().join("\n")
}

fn error_type_name(name: &str) -> String {
    let type_name = mutate_type_name(name);
    format!("{}Error", type_name)
}

#[test]
fn capitalize_first_test() {
    assert_eq!(capitalize_first("a &str test"), "A &str test".to_owned());
    assert_eq!(capitalize_first("a String test".to_owned()),
               "A String test".to_owned());
}
<commit_msg>Fix comment typo<commit_after>use std::fs::File;
use std::io::{Write, BufWriter};

use inflector::Inflector;

use ::Service;
use codegen::botocore::{Shape, ShapeType};
use self::json::JsonGenerator;
use self::query::QueryGenerator;
use self::rest_json::RestJsonGenerator;
use self::rest_xml::RestXmlGenerator;
use self::error_types::{GenerateErrorTypes, JsonErrorTypes, XmlErrorTypes};
use self::tests::generate_tests;
use self::type_filter::filter_types;

mod error_types;
mod json;
mod query;
mod rest_json;
pub mod tests;
mod rest_xml;
mod xml_payload_parser;
mod rest_response_parser;
mod type_filter;

type FileWriter = BufWriter<File>;
type IoResult = ::std::io::Result<()>;

/// Abstracts the generation of Rust code for various AWS protocols
pub trait GenerateProtocol {

    /// Generate the various `use` statements required by the module generatedfor this service
    fn generate_prelude(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    fn generate_method_signatures(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    /// Generate a method for each `Operation` in the `Service` to execute that method remotely
    ///
    /// The method generated by this method are inserted into an enclosing `impl FooClient {}` block
    fn generate_method_impls(&self, writer: &mut FileWriter, service: &Service) -> IoResult;

    /// Add any attributes that should decorate the struct for the given type (typically `Debug`, `Clone`, etc.)
    fn generate_struct_attributes(&self,
                                  serialized: bool,
                                  deserialized: bool)
                                  -> String;

    /// If necessary, generate a serializer for the specified type
    fn generate_serializer(&self, _name: &str, _shape: &Shape, _service: &Service) -> Option<String> {
        None
    }

    /// If necessary, generate a deserializer for the specified type
    fn generate_deserializer(&self, _name: &str, _shape: &Shape, _service: &Service) -> Option<String> {
        None
    }

    /// Return the type used by this protocol for timestamps
    fn timestamp_type(&self) -> &'static str;
}

pub fn generate_source(service: &Service, writer: &mut FileWriter) -> IoResult {
    // EC2 service protocol is similar to query but not the same.  Rusoto is able to generate Rust code
    // from the service definition through the same QueryGenerator, but botocore uses a special class.
    // See https://github.com/boto/botocore/blob/dff99fdf2666accf6b448aef7f03fe3d66dd38fa/botocore/serialize.py#L259-L266 .
    match service.protocol() {
        "json" => generate(writer, service, JsonGenerator, JsonErrorTypes),
        "query" | "ec2" => generate(writer, service, QueryGenerator, XmlErrorTypes),
        "rest-json" => generate(writer, service, RestJsonGenerator, JsonErrorTypes),
        "rest-xml" => generate(writer, service, RestXmlGenerator, XmlErrorTypes),
        protocol => panic!("Unknown protocol {}", protocol),
    }
}

/// Translate a botocore field name to something rust-idiomatic and
/// escape reserved words with an underscore
pub fn generate_field_name(member_name: &str) -> String {
    let name = member_name.to_snake_case();
    if name == "return" || name == "type" {
        name + "_"
    } else {
        name
    }
}

/// Capitalize the first character in the given string.
/// If the input string is empty an empty string is returned.
pub fn capitalize_first<S>(word: S) -> String
    where S: Into<String> {
    let s = word.into();
    let mut chars = s.chars();
    match chars.next() {
        Some(c) => c.to_uppercase().collect::<String>() + chars.as_str(),
        None => String::new(),
    }
}

/// The quick brown fox jumps over the lazy dog
fn generate<P, E>(writer: &mut FileWriter, service: &Service, protocol_generator: P, error_type_generator: E) -> IoResult
    where P: GenerateProtocol,
          E: GenerateErrorTypes {

    writeln!(writer, "#[allow(warnings)]
        use hyper::Client;
        use hyper::status::StatusCode;
        use rusoto_core::request::DispatchSignedRequest;
        use rusoto_core::region;

        use std::fmt;
        use std::error::Error;
        use rusoto_core::request::HttpDispatchError;
        use rusoto_core::credential::{{CredentialsError, ProvideAwsCredentials}};
    ")?;

    protocol_generator.generate_prelude(writer, service)?;
    generate_types(writer, service, &protocol_generator)?;
    error_type_generator.generate_error_types(writer, service)?;
    generate_client(writer, service, &protocol_generator)?;
    generate_tests(writer, service)?;

    Ok(())

}

fn generate_client<P>(writer: &mut FileWriter, service: &Service, protocol_generator: &P) -> IoResult
    where P: GenerateProtocol
{
    // If the struct name is changed, the links in each service documentation should change.
    // See https://github.com/rusoto/rusoto/issues/519
    writeln!(writer,
        "/// Trait representing the capabilities of the {service_name} API. {service_name} clients implement this trait.
        pub trait {trait_name} {{
        ",
        trait_name = service.service_type_name(),
        service_name = service.name()
    )?;

    protocol_generator.generate_method_signatures(writer, service)?;

    writeln!(writer, "}}")?;

    writeln!(writer,
        "/// A client for the {service_name} API.
        pub struct {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
            credentials_provider: P,
            region: region::Region,
            dispatcher: D,
        }}

        impl<P, D> {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
            pub fn new(request_dispatcher: D, credentials_provider: P, region: region::Region) -> Self {{
                  {type_name} {{
                    credentials_provider: credentials_provider,
                    region: region,
                    dispatcher: request_dispatcher
                }}
            }}
        }}

        impl<P, D> {trait_name} for {type_name}<P, D> where P: ProvideAwsCredentials, D: DispatchSignedRequest {{
        ",
        service_name = service.name(),
        type_name = service.client_type_name(),
        trait_name = service.service_type_name(),
    )?;
    protocol_generator.generate_method_impls(writer, service)?;
    writeln!(writer, "}}")
}

fn generate_list(name: &str, shape: &Shape) -> String {
    format!("pub type {} = Vec<{}>;",
            name,
            mutate_type_name(shape.member_type()))
}

fn generate_map(name: &str, shape: &Shape) -> String {
    format!(
        "pub type {} = ::std::collections::HashMap<{}, {}>;",
        name,
        capitalize_first(shape.key_type().to_string()),
        capitalize_first(shape.value_type().to_string()),
    )
}

fn generate_primitive_type(name: &str, shape_type: ShapeType, for_timestamps: &str) -> String {
    let primitive_type = match shape_type {
        ShapeType::Blob => "Vec<u8>",
        ShapeType::Boolean => "bool",
        ShapeType::Double => "f64",
        ShapeType::Float => "f32",
        ShapeType::Integer => "i64",
        ShapeType::Long => "i64",
        ShapeType::String => "String",
        ShapeType::Timestamp => for_timestamps,
        primitive_type => panic!("Unknown primitive type: {:?}", primitive_type),
    };

    format!("pub type {} = {};", name, primitive_type)
}

// do any type name mutation needed to avoid collisions with Rust types
fn mutate_type_name(type_name: &str) -> String {
    let capitalized = capitalize_first(type_name.to_owned());

    // some cloudfront types have underscoare that anger the lint checker
    let without_underscores = capitalized.replace("_", "");

    match &without_underscores[..] {
        // S3 has an 'Error' shape that collides with Rust's Error trait
        "Error" => "S3Error".to_owned(),

        // EC2 has a CancelSpotFleetRequestsError struct, avoid collision with our error enum
        "CancelSpotFleetRequests" => "EC2CancelSpotFleetRequests".to_owned(),

        // Support has a conveniently named "Result" type
        "Result" => "SupportResult".to_owned(),

        // RDS has a conveniently named "Option" type
        "Option" => "RDSOption".to_owned(),

        // otherwise make sure it's rust-idiomatic and capitalized
        _ => without_underscores,
    }
}

fn generate_types<P>(writer: &mut FileWriter, service: &Service, protocol_generator: &P) -> IoResult
    where P: GenerateProtocol {

    let (serialized_types, deserialized_types) = filter_types(service);

    for (name, shape) in service.shapes().iter() {
            let type_name = mutate_type_name(&name);

            // We generate enums for error types, so no need to create model objects for them
            if shape.exception() {
                continue;
            }

            // If botocore includes documentation, clean it up a bit and use it
            if let Some(ref docs) = shape.documentation {
                writeln!(writer, "#[doc=\"{}\"]",
                                   docs.replace("\\", "\\\\").replace("\"", "\\\""))?;
            }

            let deserialized = deserialized_types.contains(&type_name);
            let serialized = serialized_types.contains(&type_name);

            // generate a rust type for the shape
            if type_name != "String" {
                let generated_type = match shape.shape_type {
                    ShapeType::Structure => {
                        generate_struct(service,
                                                   &type_name,
                                                   &shape,
                                                   serialized,
                                                   deserialized,
                                                   protocol_generator)
                    }
                    ShapeType::Map => generate_map(&type_name, &shape),
                    ShapeType::List => generate_list(&type_name, &shape),
                    shape_type => {
                        generate_primitive_type(&type_name,
                                                           shape_type,
                                                           protocol_generator.timestamp_type())
                    }
                };
                writeln!(writer, "{}", generated_type)?;
            }

            if deserialized {
                if let Some(deserializer) = protocol_generator.generate_deserializer(&type_name, &shape, service) {
                    writeln!(writer, "{}", deserializer)?;
                }
            }

            if serialized {
                if let Some(serializer) = protocol_generator.generate_serializer(&type_name, &shape, service) {
                    writeln!(writer, "{}", serializer)?;
                }
            }
        }
        Ok(())
}

fn generate_struct<P>(service: &Service,
                      name: &str,
                      shape: &Shape,
                      serialized: bool,
                      deserialized: bool,
                      protocol_generator: &P)
                      -> String
    where P: GenerateProtocol {

    if shape.members.is_none() || shape.members.as_ref().unwrap().is_empty() {
        format!(
            "{attributes}
            pub struct {name};
            ",
            attributes = protocol_generator.generate_struct_attributes(serialized, deserialized),
            name = name,
        )
    } else {
        let struct_attributes =
            protocol_generator.generate_struct_attributes(serialized, deserialized);
        // Serde attributes are only needed if deriving the Serialize or Deserialize trait
        let need_serde_attrs = struct_attributes.contains("erialize");
        format!(
            "{attributes}
            pub struct {name} {{
                {struct_fields}
            }}
            ",
            attributes = struct_attributes,
            name = name,
            struct_fields = generate_struct_fields(service, shape, need_serde_attrs),
        )
    }

}

fn generate_struct_fields(service: &Service, shape: &Shape, serde_attrs: bool) -> String {
    shape.members.as_ref().unwrap().iter().filter_map(|(member_name, member)| {

        if member.deprecated == Some(true) {
            return None;
        }

        let mut lines: Vec<String> = Vec::new();
        let name = generate_field_name(member_name);

        if let Some(ref docs) = member.documentation {
            lines.push(format!("#[doc=\"{}\"]", docs.replace("\\","\\\\").replace("\"", "\\\"")));
        }

        let type_name = mutate_type_name(&member.shape);

        if serde_attrs {
            lines.push(format!("#[serde(rename=\"{}\")]", member_name));

            if let Some(shape_type) = service.shape_type_for_member(member) {
                if shape_type == ShapeType::Blob {
                    lines.push(
                        "#[serde(
                            deserialize_with=\"::rusoto_core::serialization::SerdeBlob::deserialize_blob\",
                            serialize_with=\"::rusoto_core::serialization::SerdeBlob::serialize_blob\",
                            default,
                        )]".to_owned()
                    );
                } else if !shape.required(member_name) {
                    lines.push("#[serde(skip_serializing_if=\"Option::is_none\")]".to_owned());
                }
            }
        }

        if shape.required(member_name) {
            lines.push(format!("pub {}: {},",  name, type_name))
        } else if name == "type" {
            lines.push(format!("pub aws_{}: Option<{}>,",  name, type_name))
        } else {
            lines.push(format!("pub {}: Option<{}>,",  name, type_name))
        }

        Some(lines.join("\n"))
    }).collect::<Vec<String>>().join("\n")
}

fn error_type_name(name: &str) -> String {
    let type_name = mutate_type_name(name);
    format!("{}Error", type_name)
}

#[test]
fn capitalize_first_test() {
    assert_eq!(capitalize_first("a &str test"), "A &str test".to_owned());
    assert_eq!(capitalize_first("a String test".to_owned()),
               "A String test".to_owned());
}
<|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 common::{Diagnostic, DiagnosticsResult, Location, NamedItem, WithLocation};
use graphql_ir::{
    Argument, ConstantValue, OperationDefinition, Program, ScalarField, Selection, Transformed,
    Transformer, ValidationMessage, Value,
};
use interner::Intern;
use interner::StringKey;
use lazy_static::lazy_static;
use schema::{Argument as ArgumentDef, FieldID, SDLSchema, Schema};
use std::sync::Arc;

lazy_static! {
    static ref EARLY_FLUSH_NAME: StringKey = "relay_early_flush".intern();
    static ref QUERY_NAME_ARG: StringKey = "query_name".intern();
}

/// NOTE: This is a Facebook only transform
/// A transform for printTransforms and codegenTransforms that transforms
/// query Query @relay_early_flush {
///   a
/// }
///
/// into
///
/// query QueryName {
///   relay_early_flush(query_name: QueryName)
///   a
/// }
pub fn relay_early_flush(program: &Program) -> DiagnosticsResult<Program> {
    let mut transform = RelayEarlyFlush {
        program,
        errors: Default::default(),
    };
    let next_program = transform
        .transform_program(program)
        .replace_or_else(|| program.clone());

    if transform.errors.is_empty() {
        Ok(next_program)
    } else {
        Err(transform.errors)
    }
}

pub struct RelayEarlyFlush<'s> {
    program: &'s Program,
    errors: Vec<Diagnostic>,
}

impl<'s> Transformer for RelayEarlyFlush<'s> {
    const NAME: &'static str = "RelayEarlyFlush";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_operation(
        &mut self,
        operation: &OperationDefinition,
    ) -> Transformed<OperationDefinition> {
        let early_flush_directive = operation.directives.named(*EARLY_FLUSH_NAME);
        if let Some(early_flush_directive) = early_flush_directive {
            match get_early_flush_field_id_and_query_name_arg(
                &self.program.schema,
                early_flush_directive.name.location,
            ) {
                Err(err) => {
                    self.errors.push(err);
                    Transformed::Keep
                }
                Ok((early_flush_field_id, query_name_arg)) => {
                    let mut next_selections = Vec::with_capacity(operation.selections.len() + 1);
                    next_selections.push(Selection::ScalarField(Arc::new(ScalarField {
                        alias: None,
                        definition: WithLocation::new(
                            early_flush_directive.name.location,
                            early_flush_field_id,
                        ),
                        arguments: vec![Argument {
                            name: WithLocation::new(
                                early_flush_directive.name.location,
                                query_name_arg.name,
                            ),
                            value: WithLocation::new(
                                early_flush_directive.name.location,
                                Value::Constant(ConstantValue::String(operation.name.item)),
                            ),
                        }],
                        directives: vec![],
                    })));
                    next_selections.extend(operation.selections.iter().cloned());
                    Transformed::Replace(OperationDefinition {
                        directives: operation
                            .directives
                            .iter()
                            .filter(|directive| directive.name.item != *EARLY_FLUSH_NAME)
                            .cloned()
                            .collect(),
                        selections: next_selections,
                        ..operation.clone()
                    })
                }
            }
        } else {
            Transformed::Keep
        }
    }
}

fn get_early_flush_field_id_and_query_name_arg(
    schema: &'_ SDLSchema,
    directive_loc: Location,
) -> Result<(FieldID, &'_ ArgumentDef), Diagnostic> {
    let query_type = schema.query_type().unwrap();
    let early_flush_field_id = schema.named_field(query_type, *EARLY_FLUSH_NAME);
    if let Some(early_flush_field_id) = early_flush_field_id {
        let field = schema.field(early_flush_field_id);
        if field.is_extension {
            Err(Diagnostic::error(
                ValidationMessage::UnavailableRelayEarlyFlushServerSchema,
                directive_loc,
            ))
        } else if let Some(query_name_arg) = field.arguments.named(*QUERY_NAME_ARG) {
            Ok((early_flush_field_id, query_name_arg))
        } else {
            Err(Diagnostic::error(
                ValidationMessage::RelayEarlyFlushSchemaWithoutQueryNameArg {
                    query_name: *QUERY_NAME_ARG,
                },
                directive_loc,
            ))
        }
    } else {
        Err(Diagnostic::error(
            ValidationMessage::UnavailableRelayEarlyFlushServerSchema,
            directive_loc,
        ))
    }
}
<commit_msg>skip traversals over fragments in relay_early_flush<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 common::{Diagnostic, DiagnosticsResult, Location, NamedItem, WithLocation};
use graphql_ir::{
    Argument, ConstantValue, FragmentDefinition, OperationDefinition, Program, ScalarField,
    Selection, Transformed, Transformer, ValidationMessage, Value,
};
use interner::Intern;
use interner::StringKey;
use lazy_static::lazy_static;
use schema::{Argument as ArgumentDef, FieldID, SDLSchema, Schema};
use std::sync::Arc;

lazy_static! {
    static ref EARLY_FLUSH_NAME: StringKey = "relay_early_flush".intern();
    static ref QUERY_NAME_ARG: StringKey = "query_name".intern();
}

/// NOTE: This is a Facebook specific transform for www static resource delivery.
///
/// ```graphql
/// query QueryName @relay_early_flush {
///   a
/// }
/// ```
///
/// into
///
/// ```graphql
/// query QueryName {
///   relay_early_flush(query_name: "QueryName")
///   a
/// }
/// ```
pub fn relay_early_flush(program: &Program) -> DiagnosticsResult<Program> {
    let mut transform = RelayEarlyFlush {
        program,
        errors: Default::default(),
    };
    let next_program = transform
        .transform_program(program)
        .replace_or_else(|| program.clone());

    if transform.errors.is_empty() {
        Ok(next_program)
    } else {
        Err(transform.errors)
    }
}

pub struct RelayEarlyFlush<'s> {
    program: &'s Program,
    errors: Vec<Diagnostic>,
}

impl<'s> Transformer for RelayEarlyFlush<'s> {
    const NAME: &'static str = "RelayEarlyFlush";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_operation(
        &mut self,
        operation: &OperationDefinition,
    ) -> Transformed<OperationDefinition> {
        let early_flush_directive = operation.directives.named(*EARLY_FLUSH_NAME);
        if let Some(early_flush_directive) = early_flush_directive {
            match get_early_flush_field_id_and_query_name_arg(
                &self.program.schema,
                early_flush_directive.name.location,
            ) {
                Err(err) => {
                    self.errors.push(err);
                    Transformed::Keep
                }
                Ok((early_flush_field_id, query_name_arg)) => {
                    let mut next_selections = Vec::with_capacity(operation.selections.len() + 1);
                    next_selections.push(Selection::ScalarField(Arc::new(ScalarField {
                        alias: None,
                        definition: WithLocation::new(
                            early_flush_directive.name.location,
                            early_flush_field_id,
                        ),
                        arguments: vec![Argument {
                            name: WithLocation::new(
                                early_flush_directive.name.location,
                                query_name_arg.name,
                            ),
                            value: WithLocation::new(
                                early_flush_directive.name.location,
                                Value::Constant(ConstantValue::String(operation.name.item)),
                            ),
                        }],
                        directives: vec![],
                    })));
                    next_selections.extend(operation.selections.iter().cloned());
                    Transformed::Replace(OperationDefinition {
                        directives: operation
                            .directives
                            .iter()
                            .filter(|directive| directive.name.item != *EARLY_FLUSH_NAME)
                            .cloned()
                            .collect(),
                        selections: next_selections,
                        ..operation.clone()
                    })
                }
            }
        } else {
            Transformed::Keep
        }
    }

    fn transform_fragment(
        &mut self,
        _fragment: &FragmentDefinition,
    ) -> Transformed<FragmentDefinition> {
        Transformed::Keep
    }
}

fn get_early_flush_field_id_and_query_name_arg(
    schema: &'_ SDLSchema,
    directive_loc: Location,
) -> Result<(FieldID, &'_ ArgumentDef), Diagnostic> {
    let query_type = schema.query_type().unwrap();
    let early_flush_field_id = schema.named_field(query_type, *EARLY_FLUSH_NAME);
    if let Some(early_flush_field_id) = early_flush_field_id {
        let field = schema.field(early_flush_field_id);
        if field.is_extension {
            Err(Diagnostic::error(
                ValidationMessage::UnavailableRelayEarlyFlushServerSchema,
                directive_loc,
            ))
        } else if let Some(query_name_arg) = field.arguments.named(*QUERY_NAME_ARG) {
            Ok((early_flush_field_id, query_name_arg))
        } else {
            Err(Diagnostic::error(
                ValidationMessage::RelayEarlyFlushSchemaWithoutQueryNameArg {
                    query_name: *QUERY_NAME_ARG,
                },
                directive_loc,
            ))
        }
    } else {
        Err(Diagnostic::error(
            ValidationMessage::UnavailableRelayEarlyFlushServerSchema,
            directive_loc,
        ))
    }
}
<|endoftext|>