<commit_before>extern crate procrs;
extern crate argparse;
#[macro_use]
extern crate prettytable;
use prettytable::Table;
use prettytable::format::FormatBuilder;
use std::collections::HashMap;
use std::iter::repeat;
use procrs::pid::*;
use procrs::TaskId;
use argparse::{ArgumentParser, StoreTrue, Store};

fn main() {
    let opts = parse_args();
    let (query, long, perf, verbose, tree) =
        (opts.query, opts.long, opts.perf, opts.verbose, opts.tree);

    let mut pids = PidIter::new_query(query).unwrap()
        .collect::<Result<Vec<_>, _>>().unwrap();

    let mut name_indent = HashMap::new();

    if verbose {
        for pid in pids {
            println!("{:?}", pid);
        }
        return
    }

    if opts.tree {
        pids = treeify_names(pids, &mut name_indent);
    } else {
        pids.sort_by(|p1, p2| p1.stat.pid.cmp(&p2.stat.pid));
    };

    let mut table = Table::init(
        pids.iter().map(|p| {
            let mut name = match tree {
                false => String::new(),
                true => name_indent.remove(&p.stat.pid).unwrap()
            };

            match (long, perf) {
                (false, false) => {
                    name.push_str(&p.stat.comm);
                    row![p.stat.pid, p.stat.ppid, name]
                },
                (true, false) => {
                    name.push_str(&p.cmdline.join(" "));
                    row![p.stat.pid, p.stat.ppid, name]
                },
                (false, true) => {
                    name.push_str(&p.stat.comm);
                    row![p.stat.pid, p.stat.ppid, name]
                }
                (true, true) => {
                    name.push_str(&p.cmdline.join(" "));
                    row![p.stat.pid, p.stat.ppid, name]
                }
            }
        }).collect::<Vec<_>>()
    );

    table.set_titles(
        match (long, perf) {
            (false, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (true, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (false, true) =>
                row!["Pid", "Ppid", "Cmd"],
            (true, true) =>
                row!["Pid", "Ppid", "Cmd"]
        }
    );
    table.set_format(
        FormatBuilder::new()
            .column_separator(' ')
            .build()
    );
    table.printstd();
}

// Given a vector of Pid structs, treeify their names, and return them in the right order.
// This is similar to ps -AH.
fn treeify_names(pids: Vec<Pid>, name_indents: &mut HashMap<TaskId, String>) -> Vec<Pid> {
    let mut child_pids = HashMap::new();
    for pid in pids {
        let ppid = pid.stat.ppid;
        child_pids.entry(ppid)
            .or_insert(Vec::new())
            .push(pid);
    }
    enumerate_children(0, &mut child_pids, name_indents, -1)
}

// Enumerate children pids, and return them.
fn enumerate_children(pid: TaskId, child_pids: &mut HashMap<TaskId, Vec<Pid>>,
    name_indents: &mut HashMap<TaskId, String>, indent: i32) -> Vec<Pid> {
    name_indents.insert(pid,
        match indent {
            i if i >= 0 =>
                repeat("  ").take(i as usize).collect::<String>(),
            _ => "".to_owned()
        }
    );
    let mut pids = Vec::new();
    let ppids = match child_pids.remove(&pid) {
        Some(v) => v,
        None => { return pids; }
    };
    for pid in ppids {
        let pid_num = pid.stat.pid;
        pids.push(pid);
        pids.append(
            &mut enumerate_children(pid_num, child_pids, name_indents, indent + 1)
        );
    }
    pids
}

struct ProgOpts {
    query: PidQuery,
    tree: bool,
    perf: bool,
    long: bool,
    verbose: bool
}

fn parse_args() -> ProgOpts {
    let mut opts = ProgOpts {
        query: PidQuery::NoneQuery,
        tree: false,
        perf: false,
        long: false,
        verbose: false
    };

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Query linux processes");
        ap.refer(&mut opts.tree)
            .add_option(&["-t", "--tree"], StoreTrue, "Display process tree");
        ap.refer(&mut opts.perf)
            .add_option(&["-p", "--perf"], StoreTrue, "Display performance information");
        ap.refer(&mut opts.long)
            .add_option(&["-l", "--long"], StoreTrue, "Display more information");
        ap.refer(&mut opts.verbose)
            .add_option(&["-v", "--verbose"], StoreTrue, "Verbose output");
        ap.refer(&mut opts.query)
            .add_argument("query", Store, "Optional query to search by, pid or string");
        ap.parse_args_or_exit();
    }

    opts
}
<commit_msg>Fix no names appearing for some commands in psq<commit_after>extern crate procrs;
extern crate argparse;
#[macro_use]
extern crate prettytable;
use prettytable::Table;
use prettytable::format::FormatBuilder;
use std::collections::HashMap;
use std::iter::repeat;
use procrs::pid::*;
use procrs::TaskId;
use procrs::meminfo::MeminfoStatus;
use argparse::{ArgumentParser, StoreTrue, Store};

fn main() {
    let opts = parse_args();
    let (query, long, perf, verbose, tree) =
        (opts.query, opts.long, opts.perf, opts.verbose, opts.tree);

    let mut pids = PidIter::new_query(query).unwrap()
        .collect::<Result<Vec<_>, _>>().unwrap();

    let mut name_indent = HashMap::new();

    if verbose {
        for pid in pids {
            println!("{:?}", pid);
        }
        return
    }

    if opts.tree {
        pids = treeify_names(pids, &mut name_indent);
    } else {
        pids.sort_by(|p1, p2| p1.stat.pid.cmp(&p2.stat.pid));
    };

    let mut table = Table::init(
        pids.iter().map(|p| {
            let mut name = match tree {
                false => String::new(),
                true => name_indent.remove(&p.stat.pid).unwrap()
            };

            match (long, perf) {
                (false, false) => {
                    name.push_str(&p.stat.comm);
                    row![p.stat.pid, p.stat.ppid, name]
                },
                (true, false) => {
                    name.push_str(
                        match p.cmdline.join(" ") {
                            ref s if s.len() > 0 => s,
                            _ => &p.stat.comm
                        }
                    );
                    row![p.stat.pid, p.stat.ppid, name]
                },
                (false, true) => {
                    name.push_str(&p.stat.comm);
                    row![p.stat.pid, p.stat.ppid, name]
                }
                (true, true) => {
                    name.push_str(
                        match p.cmdline.join(" ") {
                            ref s if s.len() > 0 => s,
                            _ => &p.stat.comm
                        }
                    );
                    row![p.stat.pid, p.stat.ppid, name]
                }
            }
        }).collect::<Vec<_>>()
    );

    table.set_titles(
        match (long, perf) {
            (false, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (true, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (false, true) =>
                row!["Pid", "Ppid", "Cmd"],
            (true, true) =>
                row!["Pid", "Ppid", "Cmd"]
        }
    );
    table.set_format(
        FormatBuilder::new()
            .column_separator(' ')
            .build()
    );
    table.printstd();
}

// Given a vector of Pid structs, treeify their names, and return them in the right order.
// This is similar to ps -AH.
fn treeify_names(pids: Vec<Pid>, name_indents: &mut HashMap<TaskId, String>) -> Vec<Pid> {
    let mut child_pids = HashMap::new();
    for pid in pids {
        let ppid = pid.stat.ppid;
        child_pids.entry(ppid)
            .or_insert(Vec::new())
            .push(pid);
    }
    enumerate_children(0, &mut child_pids, name_indents, -1)
}

// Enumerate children pids, and return them.
fn enumerate_children(pid: TaskId, child_pids: &mut HashMap<TaskId, Vec<Pid>>,
    name_indents: &mut HashMap<TaskId, String>, indent: i32) -> Vec<Pid> {
    name_indents.insert(pid,
        match indent {
            i if i >= 0 =>
                repeat("  ").take(i as usize).collect::<String>(),
            _ => "".to_owned()
        }
    );
    let mut pids = Vec::new();
    let ppids = match child_pids.remove(&pid) {
        Some(v) => v,
        None => { return pids; }
    };
    for pid in ppids {
        let pid_num = pid.stat.pid;
        pids.push(pid);
        pids.append(
            &mut enumerate_children(pid_num, child_pids, name_indents, indent + 1)
        );
    }
    pids
}

struct ProgOpts {
    query: PidQuery,
    tree: bool,
    perf: bool,
    long: bool,
    verbose: bool
}

fn parse_args() -> ProgOpts {
    let mut opts = ProgOpts {
        query: PidQuery::NoneQuery,
        tree: false,
        perf: false,
        long: false,
        verbose: false
    };

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Query linux processes");
        ap.refer(&mut opts.tree)
            .add_option(&["-t", "--tree"], StoreTrue, "Display process tree");
        ap.refer(&mut opts.perf)
            .add_option(&["-p", "--perf"], StoreTrue, "Display performance information");
        ap.refer(&mut opts.long)
            .add_option(&["-l", "--long"], StoreTrue, "Display more information");
        ap.refer(&mut opts.verbose)
            .add_option(&["-v", "--verbose"], StoreTrue, "Verbose output");
        ap.refer(&mut opts.query)
            .add_argument("query", Store, "Optional query to search by, pid or string");
        ap.parse_args_or_exit();
    }

    opts
}
<|endoftext|>


<commit_before>extern crate procrs;
extern crate argparse;
#[macro_use]
extern crate prettytable;
use prettytable::Table;
use prettytable::format::FormatBuilder;
use std::collections::HashMap;
use procrs::pid::*;
use procrs::TaskId;
use argparse::{ArgumentParser, StoreTrue, Store};

fn main() {
  let opts = parse_args();
  match opts {
    ProgOpts{ tree: false, query: q,
              perf: perf, long: long, .. } => {
      let mut table = Table::init(
        PidIter::new_query(q).unwrap()
          .map(|p_r|
            p_r.map(|p|
              row![p.stat.pid, p.stat.ppid, p.stat.comm, p.cmdline.join(" ")]
            )
          ).collect::<Result<Vec<_>, String>>()
          .unwrap()
      );
      let format = FormatBuilder::new()
        .column_separator(' ')
        .build();
      table.set_titles(row!["Pid", "Ppid", "Comm", "Cmd"]);
      table.set_format(format);
      table.printstd();
    },

    ProgOpts{ tree: true, query: q, .. } => {
      let proc_map: HashMap<_, _> =
        PidIter::new().unwrap()
        .map(|p_r|
          p_r.map(|p|
            (p.stat.pid, p)
          )
        ).collect::<Result<_, String>>()
        .unwrap();

      let mut child_procs = HashMap::new();
      let mut proc_list = Vec::new();

      for (pid, proc_struct) in &proc_map {
        proc_list.push(pid);
        let ppid = proc_struct.stat.ppid;
        child_procs.entry(ppid)
          .or_insert(Vec::new())
          .push(proc_struct);
      }
      proc_list.sort();
      let pid = match q {
        PidQuery::PidQuery(p) => p,
        _ => 1
      };
      let mut pid_procs = Vec::new();

      let start_procs = match proc_map.get(&pid) {
        Some(proc_struct) => {
          pid_procs.push(proc_struct);
          &pid_procs
        },
        None => child_procs.get(&pid).expect("Invalid pid")
      };

      print_tree(&child_procs, start_procs, "".to_string());
    }
  }
}

fn print_tree(child_procs: &HashMap<TaskId, Vec<&Pid>>,
              level_procs: &Vec<&Pid>, prefix: String) {
  let mut proc_list = level_procs.to_vec();
  proc_list.sort();
  for proc_struct in proc_list {
    let pid = &proc_struct.stat.pid;

    println!("{}{}", prefix, proc_struct.stat.comm);
    let child_list = child_procs.get(pid);
    if let Some(v) = child_list {
      print_tree(child_procs, v, format!("{}{}", "  ", prefix));
    }
  }
}

struct ProgOpts {
  query: PidQuery,
  tree: bool,
  perf: bool,
  long: bool,
  verbose: bool
}

fn parse_args() -> ProgOpts {
  let mut opts = ProgOpts {
    query: PidQuery::NoneQuery,
    tree: false,
    perf: false,
    long: false,
    verbose: false
  };

  {
    let mut ap = ArgumentParser::new();
    ap.set_description("Query linux processes");
    ap.refer(&mut opts.tree)
      .add_option(&["-t", "--tree"], StoreTrue, "Display process tree");
    ap.refer(&mut opts.perf)
      .add_option(&["-p", "--perf"], StoreTrue, "Display performance information");
    ap.refer(&mut opts.long)
      .add_option(&["-l", "--long"], StoreTrue, "Display more information");
    ap.refer(&mut opts.verbose)
      .add_option(&["-v", "--verbose"], StoreTrue, "Verbose output");
    ap.refer(&mut opts.query)
      .add_argument("query", Store, "Optional query to search by, pid or string");
    ap.parse_args_or_exit();
  }

  opts
}
<commit_msg>Update psq indentation, clean up code, and fully implement args<commit_after>extern crate procrs;
extern crate argparse;
#[macro_use]
extern crate prettytable;
use prettytable::Table;
use prettytable::format::FormatBuilder;
use std::collections::HashMap;
use std::iter::repeat;
use procrs::pid::*;
use procrs::TaskId;
use argparse::{ArgumentParser, StoreTrue, Store};

fn main() {
    let opts = parse_args();
    let (query, long, perf, verbose, tree) =
        (opts.query, opts.long, opts.perf, opts.verbose, opts.tree);

    let mut pids = PidIter::new_query(query).unwrap()
        .collect::<Result<Vec<_>, _>>().unwrap();

    let mut name_indent = HashMap::new();

    if opts.tree {
        pids = treeify_names(pids, &mut name_indent);
    } else {
        pids.sort_by(|p1, p2| p1.stat.pid.cmp(&p2.stat.pid));
    };

    let mut table = Table::init(
        pids.iter().map(|p| {
            let mut name = match tree {
                false => String::new(),
                true => name_indent.remove(&p.stat.pid).unwrap()
            };

            match (long, perf) {
                (false, false) => {
                    name.push_str(&p.stat.comm);
                    row![p.stat.pid, p.stat.ppid, name]
                },
                (true, false) => {
                    name.push_str(&p.cmdline.join(" "));
                    row![p.stat.pid, p.stat.ppid, name]
                },
                (false, true) => {
                    name.push_str(&p.stat.comm);
                    row![p.stat.pid, p.stat.ppid, name]
                }
                (true, true) => {
                    name.push_str(&p.cmdline.join(" "));
                    row![p.stat.pid, p.stat.ppid, name]
                }
            }
        }).collect::<Vec<_>>()
    );

    table.set_titles(
        match (long, perf) {
            (false, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (true, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (false, true) =>
                row!["Pid", "Ppid", "Cmd"],
            (true, true) =>
                row!["Pid", "Ppid", "Cmd"]
        }
    );
    table.set_format(
        FormatBuilder::new()
            .column_separator(' ')
            .build()
    );
    table.printstd();
}

// Given a vector of Pid structs, treeify their names, and return them in the right order.
// This is similar to ps -AH.
fn treeify_names(pids: Vec<Pid>, name_indents: &mut HashMap<TaskId, String>) -> Vec<Pid> {
    let mut child_pids = HashMap::new();
    for pid in pids {
        let ppid = pid.stat.ppid;
        child_pids.entry(ppid)
            .or_insert(Vec::new())
            .push(pid);
    }
    enumerate_children(0, &mut child_pids, name_indents, -1)
}

// Enumerate children pids, and return them.
fn enumerate_children(pid: TaskId, child_pids: &mut HashMap<TaskId, Vec<Pid>>,
    name_indents: &mut HashMap<TaskId, String>, indent: i32) -> Vec<Pid> {
    name_indents.insert(pid,
        match indent {
            i if i >= 0 =>
                repeat("  ").take(i as usize).collect::<String>(),
            _ => "".to_owned()
        }
    );
    let mut pids = Vec::new();
    let ppids = match child_pids.remove(&pid) {
        Some(v) => v,
        None => { return pids; }
    };
    for pid in ppids {
        let pid_num = pid.stat.pid;
        pids.push(pid);
        pids.append(
            &mut enumerate_children(pid_num, child_pids, name_indents, indent + 1)
        );
    }
    pids
}

struct ProgOpts {
    query: PidQuery,
    tree: bool,
    perf: bool,
    long: bool,
    verbose: bool
}

fn parse_args() -> ProgOpts {
    let mut opts = ProgOpts {
        query: PidQuery::NoneQuery,
        tree: false,
        perf: false,
        long: false,
        verbose: false
    };

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Query linux processes");
        ap.refer(&mut opts.tree)
            .add_option(&["-t", "--tree"], StoreTrue, "Display process tree");
        ap.refer(&mut opts.perf)
            .add_option(&["-p", "--perf"], StoreTrue, "Display performance information");
        ap.refer(&mut opts.long)
            .add_option(&["-l", "--long"], StoreTrue, "Display more information");
        ap.refer(&mut opts.verbose)
            .add_option(&["-v", "--verbose"], StoreTrue, "Verbose output");
        ap.refer(&mut opts.query)
            .add_argument("query", Store, "Optional query to search by, pid or string");
        ap.parse_args_or_exit();
    }

    opts
}
<|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 crate::node_identifier::NodeIdentifier;
use crate::util::{is_relay_custom_inline_fragment_directive, PointerAddress};

use dashmap::DashMap;
use fnv::FnvBuildHasher;
use graphql_ir::{
    Condition, FragmentDefinition, InlineFragment, LinkedField, OperationDefinition, Program,
    Selection, Transformed, TransformedValue,
};
use im::HashMap;
use rayon::prelude::*;
use schema::Schema;
use std::iter::Iterator;
use std::sync::Arc;

/**
 * A transform that removes redundant fields and fragment spreads. Redundancy is
 * defined in this context as any selection that is guaranteed to already be
 * fetched by an ancestor selection. This can occur in two cases:
 *
 * 1. Simple duplicates at the same level of the document can always be skipped:
 *
 *
 * fragment Foo on FooType {
 *   id
 *   id
 *   ...Bar
 *   ...Bar
 * }
 *
 *
 * Becomes
 *
 *
 * fragment Foo on FooType {
 *   id
 *   ...Bar
 * }
 *
 *
 * 2. Inline fragments and conditions introduce the possibility for duplication
 * at different levels of the tree. Whenever a selection is fetched in a parent,
 * it is redundant to also fetch it in a child:
 *
 *
 * fragment Foo on FooType {
 *   id
 *   ... on OtherType {
 *     id # 1
 *   }
 *   ... on FooType @include(if: $cond) {
 *     id # 2
 *   }
 * }
 *
 *
 * Becomes:
 *
 *
 * fragment Foo on FooType {
 *   id
 * }
 *
 *
 * In this example:
 * - 1 can be skipped because `id` is already fetched by the parent. Even
 *   though the type is different (FooType/OtherType), the inline fragment
 *   cannot match without the outer fragment matching so the outer `id` is
 *   guaranteed to already be fetched.
 * - 2 can be skipped for similar reasons: it doesn't matter if the condition
 *   holds, `id` is already fetched by the parent regardless.
 *
 * This transform also handles more complicated cases in which selections are
 * nested:
 *
 *
 * fragment Foo on FooType {
 *   a {
 *     bb
 *   }
 *   ... on OtherType {
 *     a {
 *       bb # 1
 *       cc
 *     }
 *   }
 *  }
 *
 *
 * Becomes
 *
 *
 * fragment Foo on FooType {
 *   a {
 *     bb
 *   }
 *   ... on OtherType {
 *     a {
 *       cc
 *     }
 *   }
 *  }
 *
 *
 * 1 can be skipped because it is already fetched at the outer level.
 */
pub fn skip_redundant_nodes(program: &Program) -> Program {
    let transform = SkipRedundantNodesTransform::new(program);
    transform
        .transform_program(program)
        .replace_or_else(|| program.clone())
}

#[derive(Default, Clone)]
struct SelectionMap(HashMap<NodeIdentifier, Option<SelectionMap>, FnvBuildHasher>);

type Cache = DashMap<PointerAddress, (Transformed<Selection>, SelectionMap)>;

struct SkipRedundantNodesTransform {
    schema: Arc<Schema>,
    cache: Cache,
}

impl<'s> SkipRedundantNodesTransform {
    fn new(program: &'_ Program) -> Self {
        Self {
            schema: Arc::clone(&program.schema),
            cache: DashMap::new(),
        }
    }

    fn transform_selection(
        &self,
        selection: &Selection,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Selection> {
        // This will optimize a traversal of the same subselections.
        // If it's the same node, and selection_map is empty
        // result of transform_selection has to be the same.
        let is_empty = selection_map.0.is_empty();
        let identifier = NodeIdentifier::from_selection(&self.schema, selection);
        match selection {
            Selection::ScalarField(_) => {
                if selection_map.0.contains_key(&identifier) {
                    Transformed::Delete
                } else {
                    selection_map.0.insert(identifier, None);
                    Transformed::Keep
                }
            }
            Selection::FragmentSpread(_) => {
                if selection_map.0.contains_key(&identifier) {
                    Transformed::Delete
                } else {
                    selection_map.0.insert(identifier, None);
                    Transformed::Keep
                }
            }
            Selection::LinkedField(selection) => {
                let should_cache = is_empty && Arc::strong_count(selection) > 1;
                if should_cache {
                    let key = PointerAddress::new(selection);
                    if let Some(cached) = self.cache.get(&key) {
                        let (cached_result, cached_selection_map) = cached.clone();
                        *selection_map = cached_selection_map;
                        return cached_result;
                    }
                }
                let result = if let Some(Some(linked_selection_map)) =
                    selection_map.0.get_mut(&identifier)
                {
                    self.transform_linked_field(selection, linked_selection_map)
                        .map(Selection::LinkedField)
                } else {
                    let mut linked_selection_map = Default::default();
                    let result = self
                        .transform_linked_field(selection, &mut linked_selection_map)
                        .map(Selection::LinkedField);
                    if !linked_selection_map.0.is_empty() {
                        selection_map
                            .0
                            .insert(identifier, Some(linked_selection_map));
                    }
                    result
                };
                if should_cache {
                    let key = PointerAddress::new(selection);
                    self.cache
                        .insert(key, (result.clone(), selection_map.clone()));
                }
                result
            }
            Selection::Condition(selection) => {
                if let Some(Some(existing_selection_map)) = selection_map.0.get_mut(&identifier) {
                    self.transform_condition(selection, existing_selection_map)
                        .map(Selection::Condition)
                } else {
                    // Fork the selection map to prevent conditional selections from
                    // affecting the outer "guaranteed" selections.
                    let mut next_selection_map = selection_map.clone();
                    let result = self
                        .transform_condition(selection, &mut next_selection_map)
                        .map(Selection::Condition);
                    selection_map.0.insert(identifier, Some(next_selection_map));
                    result
                }
            }
            Selection::InlineFragment(selection) => {
                let should_cache = is_empty && Arc::strong_count(selection) > 1;
                if should_cache {
                    let key = PointerAddress::new(selection);
                    if let Some(cached) = self.cache.get(&key) {
                        let (cached_result, cached_selection_map) = cached.clone();
                        *selection_map = cached_selection_map;
                        return cached_result;
                    }
                }
                let result = if let Some(Some(existing_selection_map)) =
                    selection_map.0.get_mut(&identifier)
                {
                    self.transform_inline_fragment(selection, existing_selection_map)
                        .map(Selection::InlineFragment)
                } else if selection
                    .directives
                    .iter()
                    .any(is_relay_custom_inline_fragment_directive)
                {
                    let mut linked_selection_map = Default::default();
                    let result = self
                        .transform_inline_fragment(selection, &mut linked_selection_map)
                        .map(Selection::InlineFragment);
                    if !linked_selection_map.0.is_empty() {
                        selection_map
                            .0
                            .insert(identifier, Some(linked_selection_map));
                    }
                    result
                } else {
                    // Fork for inline fragments for the same reason
                    let mut next_selection_map = selection_map.clone();
                    let result = self
                        .transform_inline_fragment(selection, &mut next_selection_map)
                        .map(Selection::InlineFragment);
                    selection_map.0.insert(identifier, Some(next_selection_map));
                    result
                };
                if should_cache {
                    let key = PointerAddress::new(selection);
                    self.cache
                        .insert(key, (result.clone(), selection_map.clone()));
                }
                result
            }
        }
    }

    fn transform_linked_field(
        &self,
        field: &LinkedField,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<LinkedField>> {
        let selections = self.transform_selections(&field.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(LinkedField {
                        selections,
                        ..field.clone()
                    }))
                }
            }
        }
    }

    fn transform_condition(
        &self,
        condition: &Condition,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<Condition>> {
        let selections = self.transform_selections(&condition.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(Condition {
                        selections,
                        ..condition.clone()
                    }))
                }
            }
        }
    }

    fn transform_inline_fragment(
        &self,
        fragment: &InlineFragment,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<InlineFragment>> {
        let selections = self.transform_selections(&fragment.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(InlineFragment {
                        selections,
                        ..fragment.clone()
                    }))
                }
            }
        }
    }

    // Mostly a copy from Transformer::transform_list, but does partition and pass down `selection_map`.
    fn transform_selections(
        &self,
        selections: &[Selection],
        selection_map: &mut SelectionMap,
    ) -> TransformedValue<Vec<Selection>> {
        if selections.is_empty() {
            return TransformedValue::Keep;
        }
        let mut result: Vec<Selection> = Vec::new();
        let mut has_changes = false;
        let selections = get_partitioned_selections(selections);

        for (index, prev_item) in selections.iter().enumerate() {
            let next_item = self.transform_selection(prev_item, selection_map);
            match next_item {
                Transformed::Keep => {
                    if has_changes {
                        result.push((*prev_item).clone());
                    }
                }
                Transformed::Delete => {
                    if !has_changes {
                        debug_assert!(result.capacity() == 0);
                        // assume most items won't be skipped and allocate space for all items
                        result.reserve(selections.len());
                        result.extend(selections.iter().take(index).map(|&x| x.clone()));
                        has_changes = true;
                    }
                }
                Transformed::Replace(next_item) => {
                    if !has_changes {
                        debug_assert!(result.capacity() == 0);
                        // assume most items won't be skipped and allocate space for all items
                        result.reserve(selections.len());
                        result.extend(selections.iter().take(index).map(|&x| x.clone()));
                        has_changes = true;
                    }
                    result.push(next_item);
                }
            }
        }
        if has_changes {
            TransformedValue::Replace(result)
        } else {
            TransformedValue::Replace(selections.iter().map(|&x| x.clone()).collect())
        }
    }

    fn transform_operation(
        &self,
        operation: &OperationDefinition,
    ) -> Transformed<OperationDefinition> {
        let mut selection_map = Default::default();
        let selections = self.transform_selections(&operation.selections, &mut selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => Transformed::Replace(OperationDefinition {
                selections,
                ..operation.clone()
            }),
        }
    }

    fn transform_fragment(&self, fragment: &FragmentDefinition) -> Transformed<FragmentDefinition> {
        let mut selection_map = Default::default();
        let selections = self.transform_selections(&fragment.selections, &mut selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => Transformed::Replace(FragmentDefinition {
                selections,
                ..fragment.clone()
            }),
        }
    }

    fn transform_program(&self, program: &Program) -> TransformedValue<Program> {
        let operations: Vec<Arc<OperationDefinition>> = program
            .par_operations()
            .filter_map(|operation| match self.transform_operation(operation) {
                Transformed::Delete => None,
                Transformed::Keep => Some(Arc::clone(operation)),
                Transformed::Replace(replacement) => Some(Arc::new(replacement)),
            })
            .collect();
        let fragments: Vec<Arc<FragmentDefinition>> = program
            .par_fragments()
            .filter_map(|fragment| match self.transform_fragment(fragment) {
                Transformed::Delete => None,
                Transformed::Keep => Some(Arc::clone(fragment)),
                Transformed::Replace(replacement) => Some(Arc::new(replacement)),
            })
            .collect();
        let mut next_program = Program::new(Arc::clone(&program.schema));
        for operation in operations {
            next_program.insert_operation(operation);
        }
        for fragment in fragments {
            next_program.insert_fragment(fragment);
        }
        TransformedValue::Replace(next_program)
    }
}

/* Selections are sorted with fields first, "conditionals"
 * (inline fragments & conditions) last. This means that all fields that are
 * guaranteed to be fetched are encountered prior to any duplicates that may be
 * fetched within a conditional.
 */
fn get_partitioned_selections(selections: &[Selection]) -> Vec<&Selection> {
    let mut result = Vec::with_capacity(selections.len());
    unsafe { result.set_len(selections.len()) };
    let mut non_field_index = selections
        .iter()
        .filter(|sel| match sel {
            Selection::LinkedField(_) | Selection::ScalarField(_) => true,
            _ => false,
        })
        .count();
    let mut field_index = 0;
    for sel in selections.iter() {
        match sel {
            Selection::LinkedField(_) | Selection::ScalarField(_) => {
                result[field_index] = sel;
                field_index += 1;
            }
            _ => {
                result[non_field_index] = sel;
                non_field_index += 1;
            }
        }
    }
    result
}
<commit_msg>optimize skip_redundant_nodes map implementation<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 crate::node_identifier::NodeIdentifier;
use crate::util::{is_relay_custom_inline_fragment_directive, PointerAddress};

use dashmap::DashMap;
use graphql_ir::{
    Condition, FragmentDefinition, InlineFragment, LinkedField, OperationDefinition, Program,
    Selection, Transformed, TransformedValue,
};
use rayon::prelude::*;
use schema::Schema;
use std::iter::Iterator;
use std::sync::Arc;

/**
 * A transform that removes redundant fields and fragment spreads. Redundancy is
 * defined in this context as any selection that is guaranteed to already be
 * fetched by an ancestor selection. This can occur in two cases:
 *
 * 1. Simple duplicates at the same level of the document can always be skipped:
 *
 *
 * fragment Foo on FooType {
 *   id
 *   id
 *   ...Bar
 *   ...Bar
 * }
 *
 *
 * Becomes
 *
 *
 * fragment Foo on FooType {
 *   id
 *   ...Bar
 * }
 *
 *
 * 2. Inline fragments and conditions introduce the possibility for duplication
 * at different levels of the tree. Whenever a selection is fetched in a parent,
 * it is redundant to also fetch it in a child:
 *
 *
 * fragment Foo on FooType {
 *   id
 *   ... on OtherType {
 *     id # 1
 *   }
 *   ... on FooType @include(if: $cond) {
 *     id # 2
 *   }
 * }
 *
 *
 * Becomes:
 *
 *
 * fragment Foo on FooType {
 *   id
 * }
 *
 *
 * In this example:
 * - 1 can be skipped because `id` is already fetched by the parent. Even
 *   though the type is different (FooType/OtherType), the inline fragment
 *   cannot match without the outer fragment matching so the outer `id` is
 *   guaranteed to already be fetched.
 * - 2 can be skipped for similar reasons: it doesn't matter if the condition
 *   holds, `id` is already fetched by the parent regardless.
 *
 * This transform also handles more complicated cases in which selections are
 * nested:
 *
 *
 * fragment Foo on FooType {
 *   a {
 *     bb
 *   }
 *   ... on OtherType {
 *     a {
 *       bb # 1
 *       cc
 *     }
 *   }
 *  }
 *
 *
 * Becomes
 *
 *
 * fragment Foo on FooType {
 *   a {
 *     bb
 *   }
 *   ... on OtherType {
 *     a {
 *       cc
 *     }
 *   }
 *  }
 *
 *
 * 1 can be skipped because it is already fetched at the outer level.
 */
pub fn skip_redundant_nodes(program: &Program) -> Program {
    let transform = SkipRedundantNodesTransform::new(program);
    transform
        .transform_program(program)
        .replace_or_else(|| program.clone())
}

#[derive(Default, Clone)]
struct SelectionMap(VecMap<NodeIdentifier, Option<SelectionMap>>);

type Cache = DashMap<PointerAddress, (Transformed<Selection>, SelectionMap)>;

struct SkipRedundantNodesTransform {
    schema: Arc<Schema>,
    cache: Cache,
}

impl<'s> SkipRedundantNodesTransform {
    fn new(program: &'_ Program) -> Self {
        Self {
            schema: Arc::clone(&program.schema),
            cache: DashMap::new(),
        }
    }

    fn transform_selection(
        &self,
        selection: &Selection,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Selection> {
        // This will optimize a traversal of the same subselections.
        // If it's the same node, and selection_map is empty
        // result of transform_selection has to be the same.
        let is_empty = selection_map.0.is_empty();
        let identifier = NodeIdentifier::from_selection(&self.schema, selection);
        match selection {
            Selection::ScalarField(_) => {
                if selection_map.0.contains_key(&identifier) {
                    Transformed::Delete
                } else {
                    selection_map.0.insert(identifier, None);
                    Transformed::Keep
                }
            }
            Selection::FragmentSpread(_) => {
                if selection_map.0.contains_key(&identifier) {
                    Transformed::Delete
                } else {
                    selection_map.0.insert(identifier, None);
                    Transformed::Keep
                }
            }
            Selection::LinkedField(selection) => {
                let should_cache = is_empty && Arc::strong_count(selection) > 1;
                if should_cache {
                    let key = PointerAddress::new(selection);
                    if let Some(cached) = self.cache.get(&key) {
                        let (cached_result, cached_selection_map) = cached.clone();
                        *selection_map = cached_selection_map;
                        return cached_result;
                    }
                }
                let result = if let Some(Some(linked_selection_map)) =
                    selection_map.0.get_mut(&identifier)
                {
                    self.transform_linked_field(selection, linked_selection_map)
                        .map(Selection::LinkedField)
                } else {
                    let mut linked_selection_map = Default::default();
                    let result = self
                        .transform_linked_field(selection, &mut linked_selection_map)
                        .map(Selection::LinkedField);
                    if !linked_selection_map.0.is_empty() {
                        selection_map
                            .0
                            .insert(identifier, Some(linked_selection_map));
                    }
                    result
                };
                if should_cache {
                    let key = PointerAddress::new(selection);
                    self.cache
                        .insert(key, (result.clone(), selection_map.clone()));
                }
                result
            }
            Selection::Condition(selection) => {
                if let Some(Some(existing_selection_map)) = selection_map.0.get_mut(&identifier) {
                    self.transform_condition(selection, existing_selection_map)
                        .map(Selection::Condition)
                } else {
                    // Fork the selection map to prevent conditional selections from
                    // affecting the outer "guaranteed" selections.
                    let mut next_selection_map = selection_map.clone();
                    let result = self
                        .transform_condition(selection, &mut next_selection_map)
                        .map(Selection::Condition);
                    selection_map.0.insert(identifier, Some(next_selection_map));
                    result
                }
            }
            Selection::InlineFragment(selection) => {
                let should_cache = is_empty && Arc::strong_count(selection) > 1;
                if should_cache {
                    let key = PointerAddress::new(selection);
                    if let Some(cached) = self.cache.get(&key) {
                        let (cached_result, cached_selection_map) = cached.clone();
                        *selection_map = cached_selection_map;
                        return cached_result;
                    }
                }
                let result = if let Some(Some(existing_selection_map)) =
                    selection_map.0.get_mut(&identifier)
                {
                    self.transform_inline_fragment(selection, existing_selection_map)
                        .map(Selection::InlineFragment)
                } else if selection
                    .directives
                    .iter()
                    .any(is_relay_custom_inline_fragment_directive)
                {
                    let mut linked_selection_map = Default::default();
                    let result = self
                        .transform_inline_fragment(selection, &mut linked_selection_map)
                        .map(Selection::InlineFragment);
                    if !linked_selection_map.0.is_empty() {
                        selection_map
                            .0
                            .insert(identifier, Some(linked_selection_map));
                    }
                    result
                } else {
                    // Fork for inline fragments for the same reason
                    let mut next_selection_map = selection_map.clone();
                    let result = self
                        .transform_inline_fragment(selection, &mut next_selection_map)
                        .map(Selection::InlineFragment);
                    selection_map.0.insert(identifier, Some(next_selection_map));
                    result
                };
                if should_cache {
                    let key = PointerAddress::new(selection);
                    self.cache
                        .insert(key, (result.clone(), selection_map.clone()));
                }
                result
            }
        }
    }

    fn transform_linked_field(
        &self,
        field: &LinkedField,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<LinkedField>> {
        let selections = self.transform_selections(&field.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(LinkedField {
                        selections,
                        ..field.clone()
                    }))
                }
            }
        }
    }

    fn transform_condition(
        &self,
        condition: &Condition,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<Condition>> {
        let selections = self.transform_selections(&condition.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(Condition {
                        selections,
                        ..condition.clone()
                    }))
                }
            }
        }
    }

    fn transform_inline_fragment(
        &self,
        fragment: &InlineFragment,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<InlineFragment>> {
        let selections = self.transform_selections(&fragment.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(InlineFragment {
                        selections,
                        ..fragment.clone()
                    }))
                }
            }
        }
    }

    // Mostly a copy from Transformer::transform_list, but does partition and pass down `selection_map`.
    fn transform_selections(
        &self,
        selections: &[Selection],
        selection_map: &mut SelectionMap,
    ) -> TransformedValue<Vec<Selection>> {
        if selections.is_empty() {
            return TransformedValue::Keep;
        }
        let mut result: Vec<Selection> = Vec::new();
        let mut has_changes = false;
        let selections = get_partitioned_selections(selections);

        for (index, prev_item) in selections.iter().enumerate() {
            let next_item = self.transform_selection(prev_item, selection_map);
            match next_item {
                Transformed::Keep => {
                    if has_changes {
                        result.push((*prev_item).clone());
                    }
                }
                Transformed::Delete => {
                    if !has_changes {
                        debug_assert!(result.capacity() == 0);
                        // assume most items won't be skipped and allocate space for all items
                        result.reserve(selections.len());
                        result.extend(selections.iter().take(index).map(|&x| x.clone()));
                        has_changes = true;
                    }
                }
                Transformed::Replace(next_item) => {
                    if !has_changes {
                        debug_assert!(result.capacity() == 0);
                        // assume most items won't be skipped and allocate space for all items
                        result.reserve(selections.len());
                        result.extend(selections.iter().take(index).map(|&x| x.clone()));
                        has_changes = true;
                    }
                    result.push(next_item);
                }
            }
        }
        if has_changes {
            TransformedValue::Replace(result)
        } else {
            TransformedValue::Replace(selections.iter().map(|&x| x.clone()).collect())
        }
    }

    fn transform_operation(
        &self,
        operation: &OperationDefinition,
    ) -> Transformed<OperationDefinition> {
        let mut selection_map = Default::default();
        let selections = self.transform_selections(&operation.selections, &mut selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => Transformed::Replace(OperationDefinition {
                selections,
                ..operation.clone()
            }),
        }
    }

    fn transform_fragment(&self, fragment: &FragmentDefinition) -> Transformed<FragmentDefinition> {
        let mut selection_map = Default::default();
        let selections = self.transform_selections(&fragment.selections, &mut selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => Transformed::Replace(FragmentDefinition {
                selections,
                ..fragment.clone()
            }),
        }
    }

    fn transform_program(&self, program: &Program) -> TransformedValue<Program> {
        let operations: Vec<Arc<OperationDefinition>> = program
            .par_operations()
            .filter_map(|operation| match self.transform_operation(operation) {
                Transformed::Delete => None,
                Transformed::Keep => Some(Arc::clone(operation)),
                Transformed::Replace(replacement) => Some(Arc::new(replacement)),
            })
            .collect();
        let fragments: Vec<Arc<FragmentDefinition>> = program
            .par_fragments()
            .filter_map(|fragment| match self.transform_fragment(fragment) {
                Transformed::Delete => None,
                Transformed::Keep => Some(Arc::clone(fragment)),
                Transformed::Replace(replacement) => Some(Arc::new(replacement)),
            })
            .collect();
        let mut next_program = Program::new(Arc::clone(&program.schema));
        for operation in operations {
            next_program.insert_operation(operation);
        }
        for fragment in fragments {
            next_program.insert_fragment(fragment);
        }
        TransformedValue::Replace(next_program)
    }
}

/* Selections are sorted with fields first, "conditionals"
 * (inline fragments & conditions) last. This means that all fields that are
 * guaranteed to be fetched are encountered prior to any duplicates that may be
 * fetched within a conditional.
 */
fn get_partitioned_selections(selections: &[Selection]) -> Vec<&Selection> {
    let mut result = Vec::with_capacity(selections.len());
    unsafe { result.set_len(selections.len()) };
    let mut non_field_index = selections
        .iter()
        .filter(|sel| match sel {
            Selection::LinkedField(_) | Selection::ScalarField(_) => true,
            _ => false,
        })
        .count();
    let mut field_index = 0;
    for sel in selections.iter() {
        match sel {
            Selection::LinkedField(_) | Selection::ScalarField(_) => {
                result[field_index] = sel;
                field_index += 1;
            }
            _ => {
                result[non_field_index] = sel;
                non_field_index += 1;
            }
        }
    }
    result
}

/// NOTE: intentionally local to this file, this is not a fully-general purpose
/// immutable map. see comments on methods below.
struct VecMap<K, V> {
    data: Arc<Vec<(K, V)>>,
}

impl<K, V> VecMap<K, V> {
    fn new() -> Self {
        Self {
            data: Arc::new(Vec::new()),
        }
    }
}

impl<K, V> VecMap<K, V>
where
    K: Eq + Clone,
    V: Clone,
{
    fn contains_key(&self, key: &K) -> bool {
        self.data.iter().any(|(k, _v)| k == key)
    }

    #[allow(dead_code)]
    fn get(&self, key: &K) -> Option<&V> {
        self.data.iter().find(|(k, _v)| k == key).map(|(_k, v)| v)
    }

    fn get_mut(&mut self, key: &K) -> Option<&mut V> {
        let data = Arc::make_mut(&mut self.data);
        data.iter_mut().find(|(k, _v)| k == key).map(|(_k, v)| v)
    }

    fn insert(&mut self, key: K, value: V) {
        // NOTE: this is intentionally *not* a general-purpose insert, which should
        // update the value for the existig key if present. skip_redundant_nodes
        // always checks for the key first, so any call to insert() is guaranteed to
        // be for a non-present key. thanks to that we can bypass the existence check
        // to make insert faster.
        debug_assert!(!self.contains_key(&key));
        let data = Arc::make_mut(&mut self.data);
        data.push((key, value));
    }

    fn is_empty(&self) -> bool {
        self.data.is_empty()
    }
}

impl<K, V> Clone for VecMap<K, V> {
    fn clone(&self) -> Self {
        Self {
            data: Arc::clone(&self.data),
        }
    }
}

impl<K, V> Default for VecMap<K, V> {
    fn default() -> Self {
        Self::new()
    }
}
<|endoftext|>


<commit_before>// FIXME: add the support of transient objects

use any_cache::{Cache, HashCache};
use notify::{self, RecommendedWatcher, Watcher};
use std::any::Any;
use std::cell::RefCell;
use std::collections::HashMap;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::channel;
use std::thread;
use std::time::{Duration, Instant};

/// Class of types that can be loaded.
pub trait Load: Sized {
  /// Arguments passed at loading.
  type Args: Clone;

  /// Static string representing the name of the type (used in filesystem caching).
  const TY_STR: &'static str;

  fn load<P>(path: P, cache: &mut ResCache, args: Self::Args) -> Result<Self> where P: AsRef<Path>;
}

/// Class of types that can be reloaded.
///
/// The idea is to simply recover the arguments used in `Load::load`.
pub trait Reload: Load {
  fn reload_args(&self) -> Self::Args;
}

/// Default implementation for types which are loaded without any arguments.
impl<T> Reload for T where T: Load<Args=()> {
  fn reload_args(&self) -> Self::Args {
    ()
  }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum LoadError {
  FileNotFound(PathBuf, String),
  ParseFailed(String),
  ConversionFailed(String)
}

pub type Result<T> = ::std::result::Result<T, LoadError>;

/// Resources are wrapped in this type.
#[derive(Debug)]
pub struct Res<T>(Rc<RefCell<T>>);

impl<T> Res<T> {
  pub fn new(resource: T) -> Self {
    Res(Rc::new(RefCell::new(resource)))
  }

  pub fn as_rc(&self) -> &Rc<RefCell<T>> {
    &self.0
  }
}

impl<T> Clone for Res<T> {
  fn clone(&self) -> Self {
    Res(self.0.clone())
  }
}

impl<T> Deref for Res<T> {
  type Target = RefCell<T>;

  fn deref(&self) -> &Self::Target {
    &self.0
  }
}

impl<T> From<Rc<RefCell<T>>> for Res<T> {
  fn from(rc: Rc<RefCell<T>>) -> Self {
    Res(rc)
  }
}

/// Time to await after a resource update to establish that it should be reloaded.
const UPDATE_AWAIT_TIME_MS: u64 = 1000;

/// Resource cache. Responsible for caching resource.
pub struct ResCache {
  // contains all the typed-erased Rc<RefCell<T>>
  cache: HashCache<PathBuf>,
  // contains all metadata on resources
  metadata: HashMap<PathBuf, ResMetaData>,
  // vector of pair (path, timestamp) giving indication on resources to reload
  dirty: Arc<Mutex<Vec<(PathBuf, Instant)>>>,
  #[allow(dead_code)]
  watcher_thread: thread::JoinHandle<()>
}

struct ResMetaData {
  on_reload: Box<Fn(&mut ResCache)>,
  last_update_instant: Instant
}

impl ResCache {
  /// Create a new cache.
  pub fn new<P>(root: P) -> Self where P: AsRef<Path> {
    let dirty: Arc<Mutex<Vec<(PathBuf, Instant)>>> = Arc::new(Mutex::new(Vec::new()));
    let dirty_ = dirty.clone();

    let root = root.as_ref().to_owned();
    let (wsx, wrx) = channel();
    let mut watcher: RecommendedWatcher = Watcher::new(wsx).unwrap();

    let join_handle = thread::spawn(move || {
      let _ = watcher.watch(root);

      for event in wrx.iter() {
        if let notify::Event { path: Some(path), op: Ok(notify::op::WRITE) } = event {
          dirty_.lock().unwrap().push((path.clone(), Instant::now()));
        }
      }
    });

    ResCache {
      cache: HashCache::new(),
      metadata: HashMap::new(),
      dirty: dirty,
      watcher_thread: join_handle
    }
  }

  /// Inject a new resource in the cache.
  fn inject<T>(&mut self, path_buf: &PathBuf, resource: T, args: T::Args) -> Res<T> where T: 'static + Any + Reload {
    let res = Res(Rc::new(RefCell::new(resource)));
    let res_ = res.clone();

    let path_buf_ = path_buf.clone();

    // closure used to reload the object when needed
    let on_reload: Box<for<'a> Fn(&'a mut ResCache)> = Box::new(move |cache_| {
      match T::load(&path_buf_, cache_, args.clone()) {
        Ok(new_resource) => {
          // replace the current resource with the freshly loaded one
          *res_.borrow_mut() = new_resource;
          deb!("reloaded resource from {:?}", path_buf_);
        },
        Err(e) => {
          warn!("reloading resource from {:?} has failed:\n{:#?}", path_buf_, e);
        }
      }
    });

    let metadata = ResMetaData {
      on_reload: on_reload,
      last_update_instant: Instant::now()
    };

    // cache the resource and its meta data
    self.cache.save(path_buf.clone(), res.clone());
    self.metadata.insert(path_buf.clone(), metadata);

    res
  }

  /// Get a resource from the cache.
  pub fn get<T>(&mut self, key: &str, args: T::Args) -> Option<Res<T>> where T: 'static + Any + Reload {
    let path_str = format!("data/{}/{}", T::TY_STR, key);
    let path = Path::new(&path_str);
    let path_buf = path.to_owned();

    match self.cache.get::<Res<T>>(&path_buf).cloned() {
      r@Some(..) => {
        deb!("cache hit for {} ({})", key, path_str);
        r
      },
      None => {
        deb!("cache miss for {} ({})", key, path_str);

        // specific loading
        if path.exists() {
          match T::load(&path, self, args.clone()) {
            Ok(resource) => {
              Some(self.inject(&path_buf, resource, args))
            },
            Err(e) => {
              err!("unable to load resource from {}:\n{:#?}", path_str, e);
              None
            }
          }
        } else { // path doesn’t exist
          err!("resource at {} cannot be found", path_str);
          None
        }
      }
    }
  }

//   /// Get a resource from the cache. If it fails, a proxy version of it is used, which will get
//   /// replaced by the resource once it’s avaible.
//   pub fn get_proxied<T, P>(&mut self, key: &str, args: T::Args, proxy: P) -> Option<Res<T>>
//       where T: 'static + Any + Reload,
//             P: FnOnce() -> T {
//     // riterite
//     self.get::<T>().unwrap_or_else(|| {
//       deb!("proxying resource {}", key);
//
//
//     })
//   }

  /// Synchronize the cache by updating the resource that ought to.
  pub fn sync(&mut self) {
    let dirty = self.dirty.clone();
    let mut dirty_ = dirty.lock().unwrap();

    for &(ref path, ref instant) in dirty_.iter() {
      let mut metadata = self.metadata.remove(path).unwrap();

      if instant.duration_since(metadata.last_update_instant) >= Duration::from_millis(UPDATE_AWAIT_TIME_MS) {
        (metadata.on_reload)(self);
      }

      metadata.last_update_instant = *instant;
      self.metadata.insert(path.clone(), metadata);
    }

    dirty_.clear();
  }
}
<commit_msg>Implemented get_proxied.<commit_after>// FIXME: add the support of transient objects

use any_cache::{Cache, HashCache};
use notify::{self, RecommendedWatcher, Watcher};
use std::any::Any;
use std::cell::RefCell;
use std::collections::HashMap;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::channel;
use std::thread;
use std::time::{Duration, Instant};

/// Class of types that can be loaded.
pub trait Load: Sized {
  /// Arguments passed at loading.
  type Args: Clone;

  /// Static string representing the name of the type (used in filesystem caching).
  const TY_STR: &'static str;

  fn load<P>(path: P, cache: &mut ResCache, args: Self::Args) -> Result<Self> where P: AsRef<Path>;
}

/// Class of types that can be reloaded.
///
/// The idea is to simply recover the arguments used in `Load::load`.
pub trait Reload: Load {
  fn reload_args(&self) -> Self::Args;
}

/// Default implementation for types which are loaded without any arguments.
impl<T> Reload for T where T: Load<Args=()> {
  fn reload_args(&self) -> Self::Args {
    ()
  }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum LoadError {
  FileNotFound(PathBuf, String),
  ParseFailed(String),
  ConversionFailed(String)
}

pub type Result<T> = ::std::result::Result<T, LoadError>;

/// Resources are wrapped in this type.
#[derive(Debug)]
pub struct Res<T>(Rc<RefCell<T>>);

impl<T> Res<T> {
  pub fn new(resource: T) -> Self {
    Res(Rc::new(RefCell::new(resource)))
  }

  pub fn as_rc(&self) -> &Rc<RefCell<T>> {
    &self.0
  }
}

impl<T> Clone for Res<T> {
  fn clone(&self) -> Self {
    Res(self.0.clone())
  }
}

impl<T> Deref for Res<T> {
  type Target = RefCell<T>;

  fn deref(&self) -> &Self::Target {
    &self.0
  }
}

impl<T> From<Rc<RefCell<T>>> for Res<T> {
  fn from(rc: Rc<RefCell<T>>) -> Self {
    Res(rc)
  }
}

/// Time to await after a resource update to establish that it should be reloaded.
const UPDATE_AWAIT_TIME_MS: u64 = 1000;

/// Resource cache. Responsible for caching resource.
pub struct ResCache {
  // contains all the typed-erased Rc<RefCell<T>>
  cache: HashCache<PathBuf>,
  // contains all metadata on resources
  metadata: HashMap<PathBuf, ResMetaData>,
  // vector of pair (path, timestamp) giving indication on resources to reload
  dirty: Arc<Mutex<Vec<(PathBuf, Instant)>>>,
  #[allow(dead_code)]
  watcher_thread: thread::JoinHandle<()>
}

struct ResMetaData {
  on_reload: Box<Fn(&mut ResCache)>,
  last_update_instant: Instant
}

impl ResCache {
  /// Create a new cache.
  pub fn new<P>(root: P) -> Self where P: AsRef<Path> {
    let dirty: Arc<Mutex<Vec<(PathBuf, Instant)>>> = Arc::new(Mutex::new(Vec::new()));
    let dirty_ = dirty.clone();

    let root = root.as_ref().to_owned();
    let (wsx, wrx) = channel();
    let mut watcher: RecommendedWatcher = Watcher::new(wsx).unwrap();

    let join_handle = thread::spawn(move || {
      let _ = watcher.watch(root);

      for event in wrx.iter() {
        if let notify::Event { path: Some(path), op: Ok(notify::op::WRITE) } = event {
          dirty_.lock().unwrap().push((path.clone(), Instant::now()));
        }
      }
    });

    ResCache {
      cache: HashCache::new(),
      metadata: HashMap::new(),
      dirty: dirty,
      watcher_thread: join_handle
    }
  }

  /// Inject a new resource in the cache.
  fn inject<T>(&mut self, path_buf: &PathBuf, resource: T, args: T::Args) -> Res<T> where T: 'static + Any + Reload {
    let res = Res(Rc::new(RefCell::new(resource)));
    let res_ = res.clone();

    let path_buf_ = path_buf.clone();

    // closure used to reload the object when needed
    let on_reload: Box<for<'a> Fn(&'a mut ResCache)> = Box::new(move |cache_| {
      match T::load(&path_buf_, cache_, args.clone()) {
        Ok(new_resource) => {
          // replace the current resource with the freshly loaded one
          *res_.borrow_mut() = new_resource;
          deb!("reloaded resource from {:?}", path_buf_);
        },
        Err(e) => {
          warn!("reloading resource from {:?} has failed:\n{:#?}", path_buf_, e);
        }
      }
    });

    let metadata = ResMetaData {
      on_reload: on_reload,
      last_update_instant: Instant::now()
    };

    // cache the resource and its meta data
    self.cache.save(path_buf.clone(), res.clone());
    self.metadata.insert(path_buf.clone(), metadata);

    res
  }

  /// Get a resource from the cache.
  pub fn get<T>(&mut self, key: &str, args: T::Args) -> Option<Res<T>> where T: 'static + Any + Reload {
    let path_str = format!("data/{}/{}", T::TY_STR, key);
    let path = Path::new(&path_str);
    let path_buf = path.to_owned();

    match self.cache.get::<Res<T>>(&path_buf).cloned() {
      r@Some(..) => {
        deb!("cache hit for {} ({})", key, path_str);
        r
      },
      None => {
        deb!("cache miss for {} ({})", key, path_str);

        // specific loading
        if path.exists() {
          match T::load(&path, self, args.clone()) {
            Ok(resource) => {
              Some(self.inject(&path_buf, resource, args))
            },
            Err(e) => {
              err!("unable to load resource from {}:\n{:#?}", path_str, e);
              None
            }
          }
        } else { // path doesn’t exist
          err!("resource at {} cannot be found", path_str);
          None
        }
      }
    }
  }

   /// Get a resource from the cache. If it fails, a proxy version is used, which will get replaced
   /// by the resource once it’s available.
   pub fn get_proxied<T, P>(&mut self, key: &str, args: T::Args, proxy: P) -> Option<Res<T>>
       where T: 'static + Any + Reload,
             P: FnOnce() -> T {
     // riterite
     self.get::<T>(key, args.clone()).or_else(move || {
       let path_str = format!("data/{}/{}", T::TY_STR, key);
       let path = Path::new(&path_str);
       let path_buf = path.to_owned();

       deb!("proxying resource {}", key);
       Some(self.inject(&path_buf, proxy(), args))
     })
   }

  /// Synchronize the cache by updating the resource that ought to.
  pub fn sync(&mut self) {
    let dirty = self.dirty.clone();
    let mut dirty_ = dirty.lock().unwrap();

    for &(ref path, ref instant) in dirty_.iter() {
      let mut metadata = self.metadata.remove(path).unwrap();

      if instant.duration_since(metadata.last_update_instant) >= Duration::from_millis(UPDATE_AWAIT_TIME_MS) {
        (metadata.on_reload)(self);
      }

      metadata.last_update_instant = *instant;
      self.metadata.insert(path.clone(), metadata);
    }

    dirty_.clear();
  }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::fs::{remove_file, File, OpenOptions};
use std::io::{Read, Write};

use rustc::ty::layout::Size;

use crate::stacked_borrows::Tag;
use crate::*;

#[derive(Debug)]
pub struct FileHandle {
    file: File,
}

pub struct FileHandler {
    handles: HashMap<i32, FileHandle>,
    low: i32,
}

impl Default for FileHandler {
    fn default() -> Self {
        FileHandler {
            handles: Default::default(),
            // 0, 1 and 2 are reserved for stdin, stdout and stderr.
            low: 3,
        }
    }
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn open(
        &mut self,
        path_op: OpTy<'tcx, Tag>,
        flag_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("open")?;

        let flag = this.read_scalar(flag_op)?.to_i32()?;

        let mut options = OpenOptions::new();

        let o_rdonly = this.eval_libc_i32("O_RDONLY")?;
        let o_wronly = this.eval_libc_i32("O_WRONLY")?;
        let o_rdwr = this.eval_libc_i32("O_RDWR")?;
        // The first two bits of the flag correspond to the access mode in linux, macOS and
        // windows. We need to check that in fact the access mode flags for the current platform
        // only use these two bits, otherwise we are in an unsupported platform and should error.
        if (o_rdonly | o_wronly | o_rdwr) & !0b11 != 0 {
            throw_unsup_format!("Access mode flags on this platform are unsupported");
        }
        // Now we check the access mode
        let access_mode = flag & 0b11;

        if access_mode == o_rdonly {
            options.read(true);
        } else if access_mode == o_wronly {
            options.write(true);
        } else if access_mode == o_rdwr {
            options.read(true).write(true);
        } else {
            throw_unsup_format!("Unsupported access mode {:#x}", access_mode);
        }
        // We need to check that there aren't unsupported options in `flag`. For this we try to
        // reproduce the content of `flag` in the `mirror` variable using only the supported
        // options.
        let mut mirror = access_mode;

        let o_append = this.eval_libc_i32("O_APPEND")?;
        if flag & o_append != 0 {
            options.append(true);
            mirror |= o_append;
        }
        let o_trunc = this.eval_libc_i32("O_TRUNC")?;
        if flag & o_trunc != 0 {
            options.truncate(true);
            mirror |= o_trunc;
        }
        let o_creat = this.eval_libc_i32("O_CREAT")?;
        if flag & o_creat != 0 {
            options.create(true);
            mirror |= o_creat;
        }
        let o_cloexec = this.eval_libc_i32("O_CLOEXEC")?;
        if flag & o_cloexec != 0 {
            // We do not need to do anything for this flag because `std` already sets it.
            // (Technically we do not support *not* setting this flag, but we ignore that.)
            mirror |= o_cloexec;
        }
        // If `flag` is not equal to `mirror`, there is an unsupported option enabled in `flag`,
        // then we throw an error.
        if flag != mirror {
            throw_unsup_format!("unsupported flags {:#x}", flag & !mirror);
        }

        let path = this.read_os_string_from_c_string(this.read_scalar(path_op)?.not_undef()?)?;

        let fd = options.open(path).map(|file| {
            let mut fh = &mut this.machine.file_handler;
            fh.low += 1;
            fh.handles.insert(fh.low, FileHandle { file }).unwrap_none();
            fh.low
        });

        this.try_unwrap_io_result(fd)
    }

    fn fcntl(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        cmd_op: OpTy<'tcx, Tag>,
        _arg1_op: Option<OpTy<'tcx, Tag>>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("fcntl")?;

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let cmd = this.read_scalar(cmd_op)?.to_i32()?;
        // We only support getting the flags for a descriptor.
        if cmd == this.eval_libc_i32("F_GETFD")? {
            // Currently this is the only flag that `F_GETFD` returns. It is OK to just return the
            // `FD_CLOEXEC` value without checking if the flag is set for the file because `std`
            // always sets this flag when opening a file. However we still need to check that the
            // file itself is open.
            if this.machine.file_handler.handles.contains_key(&fd) {
                Ok(this.eval_libc_i32("FD_CLOEXEC")?)
            } else {
                this.handle_not_found()
            }
        } else {
            throw_unsup_format!("The {:#x} command is not supported for `fcntl`)", cmd);
        }
    }

    fn close(&mut self, fd_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("close")?;

        let fd = this.read_scalar(fd_op)?.to_i32()?;

        if let Some(handle) = this.machine.file_handler.handles.remove(&fd) {
            // `File::sync_all` does the checks that are done when closing a file. We do this to
            // to handle possible errors correctly.
            let result = this.try_unwrap_io_result(handle.file.sync_all().map(|_| 0i32));
            // Now we actually close the file.
            drop(handle);
            // And return the result.
            result
        } else {
            this.handle_not_found()
        }
    }

    fn read(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        this.check_no_isolation("read")?;

        let count = this.read_scalar(count_op)?.to_usize(&*this.tcx)?;
        // Reading zero bytes should not change `buf`.
        if count == 0 {
            return Ok(0);
        }
        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.read_scalar(buf_op)?.not_undef()?;

        if let Some(handle) = this.machine.file_handler.handles.get_mut(&fd) {
            // We want to read at most `count` bytes
            let mut bytes = vec![0; count as usize];
            let result = handle.file.read(&mut bytes);

            match result {
                Ok(c) => {
                    if let Some(read_bytes) = helpers::try_from_host_usize::<i64>(c) {
                        // If reading to `bytes` did not fail, we write those bytes to the buffer.
                        this.memory.write_bytes(buf, bytes)?;
                        Ok(read_bytes)
                    } else {
                        throw_unsup_format!("Number of read bytes {} cannot be transformed to i64", c);
                    }
                },
                Err(e) => {
                    this.set_last_error_from_io_error(e)?;
                    Ok(-1)
                }
            }
        } else {
            this.handle_not_found()
        }
    }

    fn write(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        this.check_no_isolation("write")?;

        let count = this.read_scalar(count_op)?.to_usize(&*this.tcx)?;
        // Writing zero bytes should not change `buf`.
        if count == 0 {
            return Ok(0);
        }
        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.read_scalar(buf_op)?.not_undef()?;

        if let Some(handle) = this.machine.file_handler.handles.get_mut(&fd) {
            let bytes = this.memory.read_bytes(buf, Size::from_bytes(count))?;
            let result = handle.file.write(&bytes);

            match result {
                Ok(c) => {
                    if let Some(written_bytes) = helpers::try_from_host_usize::<i64>(c) {
                        Ok(written_bytes)
                    } else {
                        throw_unsup_format!("Number of written bytes {} cannot be transformed to i64", c);
                    }
                },
                Err(e) => {
                    this.set_last_error_from_io_error(e)?;
                    Ok(-1)
                }
            }
        } else {
            this.handle_not_found()
        }
    }

    fn unlink(&mut self, path_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("unlink")?;

        let path = this.read_os_string_from_c_string(this.read_scalar(path_op)?.not_undef()?)?;

        let result = remove_file(path).map(|_| 0);

        this.try_unwrap_io_result(result)
    }

    /// Function used when a handle is not found inside `FileHandler`. It returns `Ok(-1)`and sets
    /// the last OS error to `libc::EBADF` (invalid file descriptor). This function uses
    /// `T: From<i32>` instead of `i32` directly because some fs functions return different integer
    /// types (like `read`, that returns an `i64`).
    fn handle_not_found<T: From<i32>>(&mut self) -> InterpResult<'tcx, T> {
        let this = self.eval_context_mut();
        let ebadf = this.eval_libc("EBADF")?;
        this.set_last_error(ebadf)?;
        Ok((-1).into())
    }
}
<commit_msg>Avoid using `as` cast<commit_after>use std::collections::HashMap;
use std::fs::{remove_file, File, OpenOptions};
use std::io::{Read, Write};

use rustc::ty::layout::Size;

use crate::stacked_borrows::Tag;
use crate::*;

#[derive(Debug)]
pub struct FileHandle {
    file: File,
}

pub struct FileHandler {
    handles: HashMap<i32, FileHandle>,
    low: i32,
}

impl Default for FileHandler {
    fn default() -> Self {
        FileHandler {
            handles: Default::default(),
            // 0, 1 and 2 are reserved for stdin, stdout and stderr.
            low: 3,
        }
    }
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn open(
        &mut self,
        path_op: OpTy<'tcx, Tag>,
        flag_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("open")?;

        let flag = this.read_scalar(flag_op)?.to_i32()?;

        let mut options = OpenOptions::new();

        let o_rdonly = this.eval_libc_i32("O_RDONLY")?;
        let o_wronly = this.eval_libc_i32("O_WRONLY")?;
        let o_rdwr = this.eval_libc_i32("O_RDWR")?;
        // The first two bits of the flag correspond to the access mode in linux, macOS and
        // windows. We need to check that in fact the access mode flags for the current platform
        // only use these two bits, otherwise we are in an unsupported platform and should error.
        if (o_rdonly | o_wronly | o_rdwr) & !0b11 != 0 {
            throw_unsup_format!("Access mode flags on this platform are unsupported");
        }
        // Now we check the access mode
        let access_mode = flag & 0b11;

        if access_mode == o_rdonly {
            options.read(true);
        } else if access_mode == o_wronly {
            options.write(true);
        } else if access_mode == o_rdwr {
            options.read(true).write(true);
        } else {
            throw_unsup_format!("Unsupported access mode {:#x}", access_mode);
        }
        // We need to check that there aren't unsupported options in `flag`. For this we try to
        // reproduce the content of `flag` in the `mirror` variable using only the supported
        // options.
        let mut mirror = access_mode;

        let o_append = this.eval_libc_i32("O_APPEND")?;
        if flag & o_append != 0 {
            options.append(true);
            mirror |= o_append;
        }
        let o_trunc = this.eval_libc_i32("O_TRUNC")?;
        if flag & o_trunc != 0 {
            options.truncate(true);
            mirror |= o_trunc;
        }
        let o_creat = this.eval_libc_i32("O_CREAT")?;
        if flag & o_creat != 0 {
            options.create(true);
            mirror |= o_creat;
        }
        let o_cloexec = this.eval_libc_i32("O_CLOEXEC")?;
        if flag & o_cloexec != 0 {
            // We do not need to do anything for this flag because `std` already sets it.
            // (Technically we do not support *not* setting this flag, but we ignore that.)
            mirror |= o_cloexec;
        }
        // If `flag` is not equal to `mirror`, there is an unsupported option enabled in `flag`,
        // then we throw an error.
        if flag != mirror {
            throw_unsup_format!("unsupported flags {:#x}", flag & !mirror);
        }

        let path = this.read_os_string_from_c_string(this.read_scalar(path_op)?.not_undef()?)?;

        let fd = options.open(path).map(|file| {
            let mut fh = &mut this.machine.file_handler;
            fh.low += 1;
            fh.handles.insert(fh.low, FileHandle { file }).unwrap_none();
            fh.low
        });

        this.try_unwrap_io_result(fd)
    }

    fn fcntl(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        cmd_op: OpTy<'tcx, Tag>,
        _arg1_op: Option<OpTy<'tcx, Tag>>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("fcntl")?;

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let cmd = this.read_scalar(cmd_op)?.to_i32()?;
        // We only support getting the flags for a descriptor.
        if cmd == this.eval_libc_i32("F_GETFD")? {
            // Currently this is the only flag that `F_GETFD` returns. It is OK to just return the
            // `FD_CLOEXEC` value without checking if the flag is set for the file because `std`
            // always sets this flag when opening a file. However we still need to check that the
            // file itself is open.
            if this.machine.file_handler.handles.contains_key(&fd) {
                Ok(this.eval_libc_i32("FD_CLOEXEC")?)
            } else {
                this.handle_not_found()
            }
        } else {
            throw_unsup_format!("The {:#x} command is not supported for `fcntl`)", cmd);
        }
    }

    fn close(&mut self, fd_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("close")?;

        let fd = this.read_scalar(fd_op)?.to_i32()?;

        if let Some(handle) = this.machine.file_handler.handles.remove(&fd) {
            // `File::sync_all` does the checks that are done when closing a file. We do this to
            // to handle possible errors correctly.
            let result = this.try_unwrap_io_result(handle.file.sync_all().map(|_| 0i32));
            // Now we actually close the file.
            drop(handle);
            // And return the result.
            result
        } else {
            this.handle_not_found()
        }
    }

    fn read(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        this.check_no_isolation("read")?;

        let count = this.read_scalar(count_op)?.to_usize(&*this.tcx)?;
        // Reading zero bytes should not change `buf`.
        if count == 0 {
            return Ok(0);
        }
        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.read_scalar(buf_op)?.not_undef()?;

        if let Some(handle) = this.machine.file_handler.handles.get_mut(&fd) {
            let count = helpers::try_into_host_usize(count)
                .ok_or_else(|| err_unsup_format!("Program tries to read into buffer too big for this host platform"))?;
            // We want to read at most `count` bytes
            let mut bytes = vec![0; count];
            let result = handle.file.read(&mut bytes);

            match result {
                Ok(c) => {
                    let read_bytes = helpers::try_from_host_usize::<i64>(c)
                        .ok_or_else(|| err_unsup_format!("Number of read bytes {} cannot be transformed to i64", c))?;
                    // If reading to `bytes` did not fail, we write those bytes to the buffer.
                    this.memory.write_bytes(buf, bytes)?;
                    Ok(read_bytes)
                },
                Err(e) => {
                    this.set_last_error_from_io_error(e)?;
                    Ok(-1)
                }
            }
        } else {
            this.handle_not_found()
        }
    }

    fn write(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        this.check_no_isolation("write")?;

        let count = this.read_scalar(count_op)?.to_usize(&*this.tcx)?;
        // Writing zero bytes should not change `buf`.
        if count == 0 {
            return Ok(0);
        }
        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.read_scalar(buf_op)?.not_undef()?;

        if let Some(handle) = this.machine.file_handler.handles.get_mut(&fd) {
            let bytes = this.memory.read_bytes(buf, Size::from_bytes(count))?;
            let result = handle.file.write(&bytes);

            match result {
                Ok(c) => helpers::try_from_host_usize::<i64>(c)
                    .ok_or_else(|| err_unsup_format!("Number of written bytes {} cannot be transformed to i64", c).into()),
                Err(e) => {
                    this.set_last_error_from_io_error(e)?;
                    Ok(-1)
                }
            }
        } else {
            this.handle_not_found()
        }
    }

    fn unlink(&mut self, path_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("unlink")?;

        let path = this.read_os_string_from_c_string(this.read_scalar(path_op)?.not_undef()?)?;

        let result = remove_file(path).map(|_| 0);

        this.try_unwrap_io_result(result)
    }

    /// Function used when a handle is not found inside `FileHandler`. It returns `Ok(-1)`and sets
    /// the last OS error to `libc::EBADF` (invalid file descriptor). This function uses
    /// `T: From<i32>` instead of `i32` directly because some fs functions return different integer
    /// types (like `read`, that returns an `i64`).
    fn handle_not_found<T: From<i32>>(&mut self) -> InterpResult<'tcx, T> {
        let this = self.eval_context_mut();
        let ebadf = this.eval_libc("EBADF")?;
        this.set_last_error(ebadf)?;
        Ok((-1).into())
    }
}
<|endoftext|>


<commit_before>// Copyright 2014-2016 bluss and ndarray developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use imp_prelude::*;
use error::{ShapeError, ErrorKind, from_kind};

/// Stack arrays along the given axis.
///
/// ***Errors*** if the arrays don't all have the same dimension
/// (may be made more flexible in the future).

/// ***Errors*** if `arrays` is empty, if `axis` is out of bounds,
/// if the result is larger than is possible to represent.
///
/// ```
/// use ndarray::{arr2, Axis, stack};
///
/// let a = arr2(&[[2., 2.],
///                [3., 3.]]);
/// assert!(
///     stack(Axis(0), &[a.view(), a.view()])
///     == Ok(arr2(&[[2., 2.],
///                  [3., 3.],
///                  [2., 2.],
///                  [3., 3.]]))
/// );
/// ```
pub fn stack<'a, A, D>(axis: Axis, arrays: &[ArrayView<'a, A, D>])
    -> Result<OwnedArray, ShapeError>
    where A: Copy,
          D: RemoveAxis
{
    if arrays.len() == 0 {
        return Err(from_kind(ErrorKind::Unsupported));
    }
    let mut res_dim = arrays[0].dim();
    if axis.axis() >= res_dim.ndim() {
        return Err(from_kind(ErrorKind::OutOfBounds));
    }
    let common_dim = res_dim.remove_axis(axis);
    if arrays.iter().any(|a| a.dim().remove_axis(axis) != common_dim) {
        return Err(from_kind(ErrorKind::IncompatibleShape));
    }

    let stacked_dim = arrays.iter()
                            .fold(0, |acc, a| acc + a.shape().axis(axis));
    res_dim.set_axis(axis, stacked_dim);

    // we can safely use uninitialized values here because they are Copy
    // and we will only ever write to them
    let size = res_dim.size();
    let mut v = Vec::with_capacity(size);
    unsafe {
        v.set_len(size);
    }
    let mut res = try!(OwnedArray::from_vec_dim(res_dim, v));

    {
        let mut assign_view = res.view_mut();
        for array in arrays {
            let len = array.shape().axis(axis);
            let (mut front, rest) = assign_view.split_at(axis, len);
            front.assign(array);
            assign_view = rest;
        }
    }
    Ok(res)
}

/// Stack arrays along the given axis.
///
/// Uses the [`stack`][1] function, calling `ArrayView::from(&a)` on each
/// argument `a`.
///
/// [1]: fn.stack.html
///
/// ***Panics*** if the `stack` function would return an error.
///
/// ```
/// #[macro_use(stack)]
/// extern crate ndarray;
///
/// use ndarray::{arr2, Axis, stack};
///
/// # fn main() {
///
/// let a = arr2(&[[2., 2.],
///                [3., 3.]]);
/// assert!(
///     stack![Axis(0), a, a]
///     == arr2(&[[2., 2.],
///               [3., 3.],
///               [2., 2.],
///               [3., 3.]])
/// );
/// # }
/// ```
#[macro_export]
macro_rules! stack {
    ($axis:expr, $( $array:expr ),+ ) => {
        ndarray::stack($axis, &[ $(ndarray::ArrayView::from(&$array) ),* ]).unwrap()
    }
}
<commit_msg>Edit docs for stack()<commit_after>// Copyright 2014-2016 bluss and ndarray developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use imp_prelude::*;
use error::{ShapeError, ErrorKind, from_kind};

/// Stack arrays along the given axis.
///
/// ***Errors*** if the arrays have mismatching shapes, apart from along `axis`.
/// (may be made more flexible in the future).

/// ***Errors*** if `arrays` is empty, if `axis` is out of bounds,
/// if the result is larger than is possible to represent.
///
/// ```
/// use ndarray::{arr2, Axis, stack};
///
/// let a = arr2(&[[2., 2.],
///                [3., 3.]]);
/// assert!(
///     stack(Axis(0), &[a.view(), a.view()])
///     == Ok(arr2(&[[2., 2.],
///                  [3., 3.],
///                  [2., 2.],
///                  [3., 3.]]))
/// );
/// ```
pub fn stack<'a, A, D>(axis: Axis, arrays: &[ArrayView<'a, A, D>])
    -> Result<OwnedArray, ShapeError>
    where A: Copy,
          D: RemoveAxis
{
    if arrays.len() == 0 {
        return Err(from_kind(ErrorKind::Unsupported));
    }
    let mut res_dim = arrays[0].dim();
    if axis.axis() >= res_dim.ndim() {
        return Err(from_kind(ErrorKind::OutOfBounds));
    }
    let common_dim = res_dim.remove_axis(axis);
    if arrays.iter().any(|a| a.dim().remove_axis(axis) != common_dim) {
        return Err(from_kind(ErrorKind::IncompatibleShape));
    }

    let stacked_dim = arrays.iter()
                            .fold(0, |acc, a| acc + a.shape().axis(axis));
    res_dim.set_axis(axis, stacked_dim);

    // we can safely use uninitialized values here because they are Copy
    // and we will only ever write to them
    let size = res_dim.size();
    let mut v = Vec::with_capacity(size);
    unsafe {
        v.set_len(size);
    }
    let mut res = try!(OwnedArray::from_vec_dim(res_dim, v));

    {
        let mut assign_view = res.view_mut();
        for array in arrays {
            let len = array.shape().axis(axis);
            let (mut front, rest) = assign_view.split_at(axis, len);
            front.assign(array);
            assign_view = rest;
        }
    }
    Ok(res)
}

/// Stack arrays along the given axis.
///
/// Uses the [`stack`][1] function, calling `ArrayView::from(&a)` on each
/// argument `a`.
///
/// [1]: fn.stack.html
///
/// ***Panics*** if the `stack` function would return an error.
///
/// ```
/// #[macro_use(stack)]
/// extern crate ndarray;
///
/// use ndarray::{arr2, Axis, stack};
///
/// # fn main() {
///
/// let a = arr2(&[[2., 2.],
///                [3., 3.]]);
/// assert!(
///     stack![Axis(0), a, a]
///     == arr2(&[[2., 2.],
///               [3., 3.],
///               [2., 2.],
///               [3., 3.]])
/// );
/// # }
/// ```
#[macro_export]
macro_rules! stack {
    ($axis:expr, $( $array:expr ),+ ) => {
        ndarray::stack($axis, &[ $(ndarray::ArrayView::from(&$array) ),* ]).unwrap()
    }
}
<|endoftext|>