<commit_before>/*
    Copyright 2013 Jesse 'Jeaye' Wilkerson
    See licensing in LICENSE file, or at:
        http://www.opensource.org/licenses/BSD-3-Clause

    File: ui/console/activator.rs
    Author: Jesse 'Jeaye' Wilkerson
    Description:
      An input listener to open/close the console.
*/

use std::{ cast, local_data };
use std::hashmap::HashMap;
use std::iterator::IteratorUtil;
use glfw::{ PRESS, REPEAT, KEY_GRAVE_ACCENT, KEY_ENTER, KEY_BACKSPACE };
use ui::Input_Listener;
use super::Console;

priv type Property_Accessor = @fn(&str) -> ~str;
priv type Property_Mutator = @fn(&str, &str) -> Option<~str>;
priv type Function = @fn(&str, &str) -> Option<~str>;

pub struct Console_Activator
{
  console: @mut Console,

  /*
      Maps of property names to callbacks (get and set).
      Properties are invoked with a prebuilt 'set' or 'get'
      function.

      The only argument to the accessor is the property name
      and it returns the value.

      The first argument of the mutator is the
      property name and the second is everything else
      contained in the command. The mutator returns a string
      containing an error message if something didn't work well.

      Ex: set map.wireframe on
      Ex: get map.wireframe
  */
  accessors: HashMap<~str, Property_Accessor>,
  mutators: HashMap<~str, Property_Mutator>,

  /*
      A map of arbitrary functions to callbacks.

      The first argument is the function name and the second
      is whatever else is supplied in the command. The return
      value is a string containing an error message if something
      didn't go well.

      Ex: record my.avi
      Ex: callvote kick annoying_dude
  */
  functions: HashMap<~str, Function>,
}

impl Console_Activator
{
  /*  Key function used to index our singleton in
      task-local storage. */
  priv fn tls_key(_: @@Console_Activator) { }

  pub fn new(new_console: @mut Console) -> @mut Console_Activator
  {
    let ca = @mut Console_Activator
    {
      console: new_console,
      accessors: HashMap::new::<~str, Property_Accessor>(),
      mutators: HashMap::new::<~str, Property_Mutator>(),
      functions: HashMap::new::<~str, Function>(),
    };

    /* Store the activator in task-local storage. (singleton) */
    unsafe
    {
      local_data::local_data_set
      (
        Console_Activator::tls_key,
        @cast::transmute::<@mut Console_Activator, @Console_Activator>(ca)
      );
    }

    /* The 'get' and 'set' functions are built in to the console. */
    ca.functions.insert(~"get",
    |_get, property| -> Option<~str>
    {
      let mut err = ~"";

      /* Check if this property exists. */
      match ca.accessors.find(&property.to_owned())
      {
        Some(func) =>
        { ca.add_log(fmt!("%s = %s", property, (*func)(property))); }
        None =>
        { err = fmt!("Error: Invalid property %s", property); }
      }

      if err.len() == 0
      { None }
      else
      { Some(err) }
    });

    ca.functions.insert(~"set",
    |_set, params| -> Option<~str>
    {
      let mut error = ~"";
      let mut property = ~"";
      let mut params = params.to_owned();
      for params.split_iter(' ').advance |x|
      { property = x.to_owned(); break; }

      /* Remove the property from the string. */
      for property.iter().advance |_|
      { params.shift_char(); }
      params.shift_char();

      /* We require a property and a value. */
      if property.len() == 0
      { error = fmt!("Invalid len for 'set' function arg list: %s", params); }
      else
      {
        /* Check if this property exists. */
        match ca.mutators.find(&property)
        {
          Some(func) =>
          {
            /* Pass the args to the property mutator. */
            match (*func)(property, params)
            {
              /* Check if the mutator liked the args. */
              Some(err) => { error = err; }
              None => { }
            }
          }
          None => { error = fmt!("The property %s does not exist.", property); }
        }
      }

      if error.len() > 0
      { Some(fmt!("Error: %s", error)) }
      else
      { None }
    });

    ca
  }

  /* Accesses the singleton activator from task-local storage. */
  pub fn get() -> @mut Console_Activator
  {
    unsafe
    {
      cast::transmute::<@Console_Activator, @mut Console_Activator>
      (*local_data::local_data_get(Console_Activator::tls_key).get())
    }
  }

  pub fn add_accessor(&mut self, name: &str, accessor: Property_Accessor)
  { self.accessors.insert(name.to_owned(), accessor); }
  pub fn add_mutator(&mut self, name: &str, mutator: Property_Mutator)
  { self.mutators.insert(name.to_owned(), mutator); }
  pub fn add_log(&self, text: &str)
  { self.console.body = self.console.body + "\n" + text; }
}

impl Input_Listener for Console_Activator
{
  pub fn key_action(&mut self, key: i32, action: i32, _mods: i32) -> bool
  {
    if action == PRESS || action == REPEAT
    {
      if key == KEY_GRAVE_ACCENT
      {
        self.console.velocity *= -1.0;
        return true;
      }

      /* The following only apply if the console is enabled. */
      if self.console.velocity > 0.0
      {
        if key == KEY_ENTER
        {
          if self.console.input.len() == 0
          { return true; }

          /* Extract the function name. */
          let mut func = ~"";
          for self.console.input.split_iter(' ').advance |x|
          { func = x.to_owned(); break; };

          /* Remove the function from the string. */
          for func.iter().advance |_|
          { self.console.input.shift_char(); }
          self.console.input.shift_char();

          /* Look for the function in the cached map. */
          match self.functions.find(&func)
          {
            Some(f) =>
            {
              let input = copy self.console.input;
              match (*f)(func, input)
              {
                Some(err) => { self.add_log(err); }
                None => { }
              }
            }
            None => { self.add_log("Error: Invalid function"); }
          }

          self.console.input = ~"";
          return true;
        }
        else if key == KEY_BACKSPACE
        {
          if self.console.input.len() > 0
          { self.console.input.pop_char(); }
          return true;
        }
        /* Non-whitespace. */
        else if key >= 32 && key <= 93
        {
          /* This will be handled when we receive it as a char. */
          return true;
        }
      }
    }

    false
  }
  pub fn key_char(&mut self, ch: char) -> bool
  {
    /* Check if the console is enabled. */
    if self.console.velocity > 0.0
    {
      /* Non-whitespace and not ` or ~ */
      if ch >= 0x20 as char && ch <= 0x7D as char && ch != 0x60 as char
      {
        self.console.input.push_char(ch);
        return true;
      }
    }

    false
  }
  pub fn mouse_action(&mut self, _button: i32, _action: i32, _mods: i32) -> bool
  { false }
  pub fn mouse_moved(&mut self, _x: f32, _y: f32) -> bool
  { false }
}

<commit_msg> #15 Fixed toggling of the console on Mac<commit_after>/*
    Copyright 2013 Jesse 'Jeaye' Wilkerson
    See licensing in LICENSE file, or at:
        http://www.opensource.org/licenses/BSD-3-Clause

    File: ui/console/activator.rs
    Author: Jesse 'Jeaye' Wilkerson
    Description:
      An input listener to open/close the console.
*/

use std::{ cast, local_data };
use std::hashmap::HashMap;
use std::iterator::IteratorUtil;
use glfw::{ PRESS, REPEAT, KEY_GRAVE_ACCENT, KEY_WORLD_1, KEY_ENTER, KEY_BACKSPACE };
use ui::Input_Listener;
use super::Console;

priv type Property_Accessor = @fn(&str) -> ~str;
priv type Property_Mutator = @fn(&str, &str) -> Option<~str>;
priv type Function = @fn(&str, &str) -> Option<~str>;

pub struct Console_Activator
{
  console: @mut Console,

  /*
      Maps of property names to callbacks (get and set).
      Properties are invoked with a prebuilt 'set' or 'get'
      function.

      The only argument to the accessor is the property name
      and it returns the value.

      The first argument of the mutator is the
      property name and the second is everything else
      contained in the command. The mutator returns a string
      containing an error message if something didn't work well.

      Ex: set map.wireframe on
      Ex: get map.wireframe
  */
  accessors: HashMap<~str, Property_Accessor>,
  mutators: HashMap<~str, Property_Mutator>,

  /*
      A map of arbitrary functions to callbacks.

      The first argument is the function name and the second
      is whatever else is supplied in the command. The return
      value is a string containing an error message if something
      didn't go well.

      Ex: record my.avi
      Ex: callvote kick annoying_dude
  */
  functions: HashMap<~str, Function>,
}

impl Console_Activator
{
  /*  Key function used to index our singleton in
      task-local storage. */
  priv fn tls_key(_: @@Console_Activator) { }

  pub fn new(new_console: @mut Console) -> @mut Console_Activator
  {
    let ca = @mut Console_Activator
    {
      console: new_console,
      accessors: HashMap::new::<~str, Property_Accessor>(),
      mutators: HashMap::new::<~str, Property_Mutator>(),
      functions: HashMap::new::<~str, Function>(),
    };

    /* Store the activator in task-local storage. (singleton) */
    unsafe
    {
      local_data::local_data_set
      (
        Console_Activator::tls_key,
        @cast::transmute::<@mut Console_Activator, @Console_Activator>(ca)
      );
    }

    /* The 'get' and 'set' functions are built in to the console. */
    ca.functions.insert(~"get",
    |_get, property| -> Option<~str>
    {
      let mut err = ~"";

      /* Check if this property exists. */
      match ca.accessors.find(&property.to_owned())
      {
        Some(func) =>
        { ca.add_log(fmt!("%s = %s", property, (*func)(property))); }
        None =>
        { err = fmt!("Error: Invalid property %s", property); }
      }

      if err.len() == 0
      { None }
      else
      { Some(err) }
    });

    ca.functions.insert(~"set",
    |_set, params| -> Option<~str>
    {
      let mut error = ~"";
      let mut property = ~"";
      let mut params = params.to_owned();
      for params.split_iter(' ').advance |x|
      { property = x.to_owned(); break; }

      /* Remove the property from the string. */
      for property.iter().advance |_|
      { params.shift_char(); }
      params.shift_char();

      /* We require a property and a value. */
      if property.len() == 0
      { error = fmt!("Invalid len for 'set' function arg list: %s", params); }
      else
      {
        /* Check if this property exists. */
        match ca.mutators.find(&property)
        {
          Some(func) =>
          {
            /* Pass the args to the property mutator. */
            match (*func)(property, params)
            {
              /* Check if the mutator liked the args. */
              Some(err) => { error = err; }
              None => { }
            }
          }
          None => { error = fmt!("The property %s does not exist.", property); }
        }
      }

      if error.len() > 0
      { Some(fmt!("Error: %s", error)) }
      else
      { None }
    });

    ca
  }

  /* Accesses the singleton activator from task-local storage. */
  pub fn get() -> @mut Console_Activator
  {
    unsafe
    {
      cast::transmute::<@Console_Activator, @mut Console_Activator>
      (*local_data::local_data_get(Console_Activator::tls_key).get())
    }
  }

  pub fn add_accessor(&mut self, name: &str, accessor: Property_Accessor)
  { self.accessors.insert(name.to_owned(), accessor); }
  pub fn add_mutator(&mut self, name: &str, mutator: Property_Mutator)
  { self.mutators.insert(name.to_owned(), mutator); }
  pub fn add_log(&self, text: &str)
  { self.console.body = self.console.body + "\n" + text; }
}

impl Input_Listener for Console_Activator
{
  pub fn key_action(&mut self, key: i32, action: i32, _mods: i32) -> bool
  {
    if action == PRESS || action == REPEAT
    {
      /* Mac grave is world 1 for some reason. */
      if key == KEY_GRAVE_ACCENT || key == KEY_WORLD_1
      {
        self.console.velocity *= -1.0;
        return true;
      }

      /* The following only apply if the console is enabled. */
      if self.console.velocity > 0.0
      {
        if key == KEY_ENTER
        {
          if self.console.input.len() == 0
          { return true; }

          /* Extract the function name. */
          let mut func = ~"";
          for self.console.input.split_iter(' ').advance |x|
          { func = x.to_owned(); break; };

          /* Remove the function from the string. */
          for func.iter().advance |_|
          { self.console.input.shift_char(); }
          self.console.input.shift_char();

          /* Look for the function in the cached map. */
          match self.functions.find(&func)
          {
            Some(f) =>
            {
              let input = copy self.console.input;
              match (*f)(func, input)
              {
                Some(err) => { self.add_log(err); }
                None => { }
              }
            }
            None => { self.add_log("Error: Invalid function"); }
          }

          self.console.input = ~"";
          return true;
        }
        else if key == KEY_BACKSPACE
        {
          if self.console.input.len() > 0
          { self.console.input.pop_char(); }
          return true;
        }
        /* Non-whitespace. */
        else if key >= 32 && key <= 93
        {
          /* This will be handled when we receive it as a char. */
          return true;
        }
      }
    }

    false
  }
  pub fn key_char(&mut self, ch: char) -> bool
  {
    /* Check if the console is enabled. */
    if self.console.velocity > 0.0
    {
      /* Non-whitespace and not ` or ~ */
      if ch >= 0x20 as char && ch <= 0x7D as char && ch != 0x60 as char
      {
        self.console.input.push_char(ch);
        return true;
      }
    }

    false
  }
  pub fn mouse_action(&mut self, _button: i32, _action: i32, _mods: i32) -> bool
  { false }
  pub fn mouse_moved(&mut self, _x: f32, _y: f32) -> bool
  { false }
}

<|endoftext|>


<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Michael Gehring <mg@ebfe.org>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use libc::mkfifo;
use std::ffi::CString;
use uucore::InvalidEncodingHandling;

static NAME: &str = "mkfifo";
static USAGE: &str = "mkfifo [OPTION]... NAME...";
static SUMMARY: &str = "Create a FIFO with the given name.";

mod options {
    pub static MODE: &str = "mode";
    pub static SE_LINUX_SECURITY_CONTEXT: &str = "Z";
    pub static CONTEXT: &str = "context";
    pub static FIFO: &str = "fifo";
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let args = args
        .collect_str(InvalidEncodingHandling::Ignore)
        .accept_any();

    let matches = uu_app().get_matches_from(args);

    if matches.is_present(options::CONTEXT) {
        crash!(1, "--context is not implemented");
    }
    if matches.is_present(options::SE_LINUX_SECURITY_CONTEXT) {
        crash!(1, "-Z is not implemented");
    }

    let mode = match matches.value_of(options::MODE) {
        Some(m) => match usize::from_str_radix(m, 8) {
            Ok(m) => m,
            Err(e) => {
                show_error!("invalid mode: {}", e);
                return 1;
            }
        },
        None => 0o666,
    };

    let fifos: Vec<String> = match matches.values_of(options::FIFO) {
        Some(v) => v.clone().map(|s| s.to_owned()).collect(),
        None => crash!(1, "missing operand"),
    };

    let mut exit_code = 0;
    for f in fifos {
        let err = unsafe {
            let name = CString::new(f.as_bytes()).unwrap();
            mkfifo(name.as_ptr(), mode as libc::mode_t)
        };
        if err == -1 {
            show_error!("cannot create fifo '{}': File exists", f);
            exit_code = 1;
        }
    }

    exit_code
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(executable!())
        .name(NAME)
        .version(crate_version!())
        .usage(USAGE)
        .about(SUMMARY)
        .arg(
            Arg::with_name(options::MODE)
                .short("m")
                .long(options::MODE)
                .help("file permissions for the fifo")
                .default_value("0666")
                .value_name("0666"),
        )
        .arg(
            Arg::with_name(options::SE_LINUX_SECURITY_CONTEXT)
                .short(options::SE_LINUX_SECURITY_CONTEXT)
                .help("set the SELinux security context to default type")
        )
        .arg(Arg::with_name(options::CONTEXT).long(options::CONTEXT).value_name("CTX").help("like -Z, or if CTX is specified then set the SELinux\nor SMACK security context to CTX"))
        .arg(Arg::with_name(options::FIFO).hidden(true).multiple(true))
}
<commit_msg>mkfifo: make rustfmt work<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Michael Gehring <mg@ebfe.org>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use libc::mkfifo;
use std::ffi::CString;
use uucore::InvalidEncodingHandling;

static NAME: &str = "mkfifo";
static USAGE: &str = "mkfifo [OPTION]... NAME...";
static SUMMARY: &str = "Create a FIFO with the given name.";

mod options {
    pub static MODE: &str = "mode";
    pub static SE_LINUX_SECURITY_CONTEXT: &str = "Z";
    pub static CONTEXT: &str = "context";
    pub static FIFO: &str = "fifo";
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let args = args
        .collect_str(InvalidEncodingHandling::Ignore)
        .accept_any();

    let matches = uu_app().get_matches_from(args);

    if matches.is_present(options::CONTEXT) {
        crash!(1, "--context is not implemented");
    }
    if matches.is_present(options::SE_LINUX_SECURITY_CONTEXT) {
        crash!(1, "-Z is not implemented");
    }

    let mode = match matches.value_of(options::MODE) {
        Some(m) => match usize::from_str_radix(m, 8) {
            Ok(m) => m,
            Err(e) => {
                show_error!("invalid mode: {}", e);
                return 1;
            }
        },
        None => 0o666,
    };

    let fifos: Vec<String> = match matches.values_of(options::FIFO) {
        Some(v) => v.clone().map(|s| s.to_owned()).collect(),
        None => crash!(1, "missing operand"),
    };

    let mut exit_code = 0;
    for f in fifos {
        let err = unsafe {
            let name = CString::new(f.as_bytes()).unwrap();
            mkfifo(name.as_ptr(), mode as libc::mode_t)
        };
        if err == -1 {
            show_error!("cannot create fifo '{}': File exists", f);
            exit_code = 1;
        }
    }

    exit_code
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(executable!())
        .name(NAME)
        .version(crate_version!())
        .usage(USAGE)
        .about(SUMMARY)
        .arg(
            Arg::with_name(options::MODE)
                .short("m")
                .long(options::MODE)
                .help("file permissions for the fifo")
                .default_value("0666")
                .value_name("0666"),
        )
        .arg(
            Arg::with_name(options::SE_LINUX_SECURITY_CONTEXT)
                .short(options::SE_LINUX_SECURITY_CONTEXT)
                .help("set the SELinux security context to default type"),
        )
        .arg(
            Arg::with_name(options::CONTEXT)
                .long(options::CONTEXT)
                .value_name("CTX")
                .help(
                    "like -Z, or if CTX is specified then set the SELinux \
                    or SMACK security context to CTX",
                ),
        )
        .arg(Arg::with_name(options::FIFO).hidden(true).multiple(true))
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Guillaume Pinot <texitoi(a)texitoi.eu>
//
// This work is free. You can redistribute it and/or modify it under
// the terms of the Do What The Fuck You Want To Public License,
// Version 2, as published by Sam Hocevar. See the COPYING file for
// more details.

//! How to `derive(StructOpt)`
//!
//! First, look at an example:
//!
//! ```ignore
//! #[derive(StructOpt)]
//! #[structopt(name = "example", about = "An example of StructOpt usage.")]
//! struct Opt {
//!     #[structopt(short = "d", long = "debug", help = "Activate debug mode")]
//!     debug: bool,
//!     #[structopt(short = "s", long = "speed", help = "Set speed", default_value = "42")]
//!     speed: f64,
//!     #[structopt(help = "Input file")]
//!     input: String,
//!     #[structopt(help = "Output file, stdout if not present")]
//!     output: Option<String>,
//! }
//! ```
//!
//! So, `derive(StructOpt)` do the job, and `structopt` attribute is
//! used for additional parameters.
//!
//! First, define a struct, whatever its name.  This structure will
//! correspond to a `clap::App`.  Every method of `clap::App` in the
//! form of `fn function_name(self, &str)` can be use in the form of
//! attributes.  Our example call for example something like
//! `app.about("An example of StructOpt usage.")`.  There is some
//! special attributes:
//!
//!   - `name`: correspond to the creation of the `App` object. Our
//!     example does `clap::App::new("example")`.  Default to
//!     the crate name given by cargo.
//!   - `version`: default to the crate version given by cargo.
//!   - `author`: default to the crate version given by cargo.
//!   - `about`: default to the crate version given by cargo.
//!
//! Then, each field of the struct correspond to a `clap::Arg`.  As
//! for the struct attributes, every method of `clap::Arg` in the form
//! of `fn function_name(self, &str)` can be use in the form of
//! attributes.  The `name` attribute can be used to customize the
//! `Arg::with_name()` call (default to the field name).
//!
//! The type of the field gives the kind of argument:
//!
//! Type                 | Effect            | Added method call to `clap::Arg`
//! ---------------------|-------------------|--------------------------------------
//! `bool`               | `true` if present | `.takes_value(false).multiple(false)`
//! `u64`                | number of params  | `.takes_value(false).multiple(true)`
//! `Option<T: FromStr>` | optional argument | `.takes_value(true).multiple(false)`
//! `Vec<T: FromStr>`    | list of arguments | `.takes_value(true).multiple(true)`
//! `T: FromStr` | required argument | `.takes_value(true).multiple(false).required(!has_default)`
//!
//! The `FromStr` trait is used to convert the argument to the given
//! type, and the `Arg::validator` method is setted to a method using
//! `FromStr::Error::description()`.
//!
//! Thus, the `speed` argument is generated as:
//!
//! ```ignore
//! clap::Arg::with_name("speed")
//!     .takes_value(true)
//!     .multiple(false)
//!     .required(false)
//!     .validator(parse_validator::<f64>)
//!     .short("s")
//!     .long("debug")
//!     .help("Set speed")
//!     .default_value("42")
//! ```

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

use proc_macro::TokenStream;
use syn::*;

/// Generates the `StructOpt` impl.
#[proc_macro_derive(StructOpt, attributes(structopt))]
pub fn structopt(input: TokenStream) -> TokenStream {
    let s = input.to_string();
    let ast = syn::parse_derive_input(&s).unwrap();
    let gen = impl_structopt(&ast);
    gen.parse().unwrap()
}

enum Ty {
    Bool,
    U64,
    Vec,
    Option,
    Other,
}

fn ty(t: &syn::Ty) -> Ty {
    if let syn::Ty::Path(None, syn::Path { segments: ref segs, .. }) = *t {
        match segs.last().unwrap().ident.as_ref() {
            "bool" => Ty::Bool,
            "u64" => Ty::U64,
            "Option" => Ty::Option,
            "Vec" => Ty::Vec,
            _ => Ty::Other,
        }
    } else {
        Ty::Other
    }
}

fn sub_type(t: &syn::Ty) -> Option<&syn::Ty> {
    let segs = match *t {
        syn::Ty::Path(None, syn::Path { ref segments, .. }) => segments,
        _ => return None,
    };
    match *segs.last().unwrap() {
        PathSegment {
            parameters: PathParameters::AngleBracketed(
                AngleBracketedParameterData { ref types, .. }),
            ..
        } if !types.is_empty() => Some(&types[0]),
            _ => None,
    }
}

#[derive(Debug, Clone, Copy)]
enum AttrSource { Struct, Field, }

fn extract_attrs<'a>(attrs: &'a [Attribute], attr_source: AttrSource) -> Box<Iterator<Item = (Ident, Lit)> + 'a> {
    let settings_attrs = attrs.iter()
        .filter_map(|attr| match attr.value {
            MetaItem::List(ref i, ref v) if i.as_ref() == "structopt" => Some(v),
            _ => None,
        }).flat_map(|v| v.iter().filter_map(|mi| match *mi {
            NestedMetaItem::MetaItem(MetaItem::NameValue(ref i, ref l)) =>
                Some((i.clone(), l.clone())),
            _ => None,
        }));

    let doc_comments = attrs.iter()
        .filter_map(move |attr| {
            if let Attribute {
                value: MetaItem::NameValue(ref name, Lit::Str(ref value, StrStyle::Cooked)),
                is_sugared_doc: true,
                ..
            } = *attr {
                if name != "doc" { return None; }
                let text = value.trim_left_matches("//!")
                    .trim_left_matches("///")
                    .trim_left_matches("/*!")
                    .trim_left_matches("/**")
                    .trim();

                // Clap's `App` has an `about` method to set a description,
                // it's `Field`s have a `help` method instead.
                if let AttrSource::Struct = attr_source {
                    Some(("about".into(), text.into()))
                } else {
                    Some(("help".into(), text.into()))
                }
            } else {
                None
            }
        });

    Box::new(doc_comments.chain(settings_attrs))
}

fn from_attr_or_env(attrs: &[(Ident, Lit)], key: &str, env: &str) -> Lit {
    let default = std::env::var(env).unwrap_or("".into());
    attrs.iter()
        .filter(|&&(ref i, _)| i.as_ref() == key)
        .last()
        .map(|&(_, ref l)| l.clone())
        .unwrap_or_else(|| Lit::Str(default, StrStyle::Cooked))
}

fn gen_name(field: &Field) -> Ident {
    extract_attrs(&field.attrs, AttrSource::Field)
        .filter(|&(ref i, _)| i.as_ref() == "name")
        .last()
        .and_then(|(_, ref l)| match l {
            &Lit::Str(ref s, _) => Some(Ident::new(s.clone())),
            _ => None,
        })
        .unwrap_or(field.ident.as_ref().unwrap().clone())
}

fn gen_from_clap(struct_name: &Ident, s: &[Field]) -> quote::Tokens {
    let fields = s.iter().map(|field| {
        let field_name = field.ident.as_ref().unwrap();
        let name = gen_name(field);
        let convert = match ty(&field.ty) {
            Ty::Bool => quote!(is_present(stringify!(#name))),
            Ty::U64 => quote!(occurrences_of(stringify!(#name))),
            Ty::Option => quote! {
                value_of(stringify!(#name))
                    .as_ref()
                    .map(|s| s.parse().unwrap())
            },
            Ty::Vec => quote! {
                values_of(stringify!(#name))
                    .map(|v| v.map(|s| s.parse().unwrap()).collect())
                    .unwrap_or_else(Vec::new)
            },
            Ty::Other => quote! {
                value_of(stringify!(#name))
                    .as_ref()
                    .unwrap()
                    .parse()
                    .unwrap()
            },
        };
        quote!( #field_name: matches.#convert, )
    });
    quote! {
        fn from_clap(matches: _structopt::clap::ArgMatches) -> Self {
            #struct_name {
                #( #fields )*
            }
        }
    }
}

fn gen_clap(ast: &DeriveInput) -> quote::Tokens {
    let struct_attrs: Vec<_> = extract_attrs(&ast.attrs, AttrSource::Struct).collect();
    let name = from_attr_or_env(&struct_attrs, "name", "CARGO_PKG_NAME");
    let version = from_attr_or_env(&struct_attrs, "version", "CARGO_PKG_VERSION");
    let author = from_attr_or_env(&struct_attrs, "author", "CARGO_PKG_AUTHORS");
    let about = from_attr_or_env(&struct_attrs, "about", "CARGO_PKG_DESCRIPTION");

    quote! {
        fn clap<'a, 'b>() -> _structopt::clap::App<'a, 'b> {
            let app = _structopt::clap::App::new(#name)
                .version(#version)
                .author(#author)
                .about(#about);
            Self::augment_clap(app)
        }
    }
}

fn gen_augment_clap(s: &[Field]) -> quote::Tokens {
    let args = s.iter().map(|field| {
        let name = gen_name(field);
        let cur_type = ty(&field.ty);
        let convert_type = match cur_type {
            Ty::Vec | Ty::Option => sub_type(&field.ty).unwrap_or(&field.ty),
            _ => &field.ty,
        };
        let validator = quote! {
            validator(|s| s.parse::<#convert_type>()
                      .map(|_| ())
                      .map_err(|e| e.description().into()))
        };
        let modifier = match cur_type {
            Ty::Bool => quote!( .takes_value(false).multiple(false) ),
            Ty::U64 => quote!( .takes_value(false).multiple(true) ),
            Ty::Option => quote!( .takes_value(true).multiple(false).#validator ),
            Ty::Vec => quote!( .takes_value(true).multiple(true).#validator ),
            Ty::Other => {
                let required = extract_attrs(&field.attrs, AttrSource::Field)
                    .find(|&(ref i, _)| i.as_ref() == "default_value")
                    .is_none();
                quote!( .takes_value(true).multiple(false).required(#required).#validator )
            },
        };
        let from_attr = extract_attrs(&field.attrs, AttrSource::Field)
            .filter(|&(ref i, _)| i.as_ref() != "name")
            .map(|(i, l)| quote!(.#i(#l)));
        quote!( .arg(_structopt::clap::Arg::with_name(stringify!(#name)) #modifier #(#from_attr)*) )
    });

    quote! {
        fn augment_clap<'a, 'b>(app: _structopt::clap::App<'a, 'b>) -> _structopt::clap::App<'a, 'b> {
            use std::error::Error;
            app #( #args )*
        }
    }
}

fn impl_structopt(ast: &syn::DeriveInput) -> quote::Tokens {
    let struct_name = &ast.ident;
    let s = match ast.body {
        Body::Struct(VariantData::Struct(ref s)) => s,
        _ => panic!("Only struct is supported"),
    };

    let clap = gen_clap(ast);
    let augment_clap = gen_augment_clap(s);
    let from_clap = gen_from_clap(struct_name, s);
    let dummy_const = Ident::new(format!("_IMPL_STRUCTOPT_FOR_{}", struct_name));
    quote! {
        #[allow(non_upper_case_globals, unused_attributes, unused_imports)]
        const #dummy_const: () = {
            extern crate structopt as _structopt;
            impl _structopt::StructOpt for #struct_name {
                #clap
                #augment_clap
                #from_clap
            }
        };
    }
}
<commit_msg>split `impl_structopt` into struct and enum branches<commit_after>// Copyright (c) 2017 Guillaume Pinot <texitoi(a)texitoi.eu>
//
// This work is free. You can redistribute it and/or modify it under
// the terms of the Do What The Fuck You Want To Public License,
// Version 2, as published by Sam Hocevar. See the COPYING file for
// more details.

//! How to `derive(StructOpt)`
//!
//! First, look at an example:
//!
//! ```ignore
//! #[derive(StructOpt)]
//! #[structopt(name = "example", about = "An example of StructOpt usage.")]
//! struct Opt {
//!     #[structopt(short = "d", long = "debug", help = "Activate debug mode")]
//!     debug: bool,
//!     #[structopt(short = "s", long = "speed", help = "Set speed", default_value = "42")]
//!     speed: f64,
//!     #[structopt(help = "Input file")]
//!     input: String,
//!     #[structopt(help = "Output file, stdout if not present")]
//!     output: Option<String>,
//! }
//! ```
//!
//! So, `derive(StructOpt)` do the job, and `structopt` attribute is
//! used for additional parameters.
//!
//! First, define a struct, whatever its name.  This structure will
//! correspond to a `clap::App`.  Every method of `clap::App` in the
//! form of `fn function_name(self, &str)` can be use in the form of
//! attributes.  Our example call for example something like
//! `app.about("An example of StructOpt usage.")`.  There is some
//! special attributes:
//!
//!   - `name`: correspond to the creation of the `App` object. Our
//!     example does `clap::App::new("example")`.  Default to
//!     the crate name given by cargo.
//!   - `version`: default to the crate version given by cargo.
//!   - `author`: default to the crate version given by cargo.
//!   - `about`: default to the crate version given by cargo.
//!
//! Then, each field of the struct correspond to a `clap::Arg`.  As
//! for the struct attributes, every method of `clap::Arg` in the form
//! of `fn function_name(self, &str)` can be use in the form of
//! attributes.  The `name` attribute can be used to customize the
//! `Arg::with_name()` call (default to the field name).
//!
//! The type of the field gives the kind of argument:
//!
//! Type                 | Effect            | Added method call to `clap::Arg`
//! ---------------------|-------------------|--------------------------------------
//! `bool`               | `true` if present | `.takes_value(false).multiple(false)`
//! `u64`                | number of params  | `.takes_value(false).multiple(true)`
//! `Option<T: FromStr>` | optional argument | `.takes_value(true).multiple(false)`
//! `Vec<T: FromStr>`    | list of arguments | `.takes_value(true).multiple(true)`
//! `T: FromStr` | required argument | `.takes_value(true).multiple(false).required(!has_default)`
//!
//! The `FromStr` trait is used to convert the argument to the given
//! type, and the `Arg::validator` method is setted to a method using
//! `FromStr::Error::description()`.
//!
//! Thus, the `speed` argument is generated as:
//!
//! ```ignore
//! clap::Arg::with_name("speed")
//!     .takes_value(true)
//!     .multiple(false)
//!     .required(false)
//!     .validator(parse_validator::<f64>)
//!     .short("s")
//!     .long("debug")
//!     .help("Set speed")
//!     .default_value("42")
//! ```

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

use proc_macro::TokenStream;
use syn::*;

/// Generates the `StructOpt` impl.
#[proc_macro_derive(StructOpt, attributes(structopt))]
pub fn structopt(input: TokenStream) -> TokenStream {
    let s = input.to_string();
    let ast = syn::parse_derive_input(&s).unwrap();
    let gen = impl_structopt(&ast);
    gen.parse().unwrap()
}

enum Ty {
    Bool,
    U64,
    Vec,
    Option,
    Other,
}

fn ty(t: &syn::Ty) -> Ty {
    if let syn::Ty::Path(None, syn::Path { segments: ref segs, .. }) = *t {
        match segs.last().unwrap().ident.as_ref() {
            "bool" => Ty::Bool,
            "u64" => Ty::U64,
            "Option" => Ty::Option,
            "Vec" => Ty::Vec,
            _ => Ty::Other,
        }
    } else {
        Ty::Other
    }
}

fn sub_type(t: &syn::Ty) -> Option<&syn::Ty> {
    let segs = match *t {
        syn::Ty::Path(None, syn::Path { ref segments, .. }) => segments,
        _ => return None,
    };
    match *segs.last().unwrap() {
        PathSegment {
            parameters: PathParameters::AngleBracketed(
                AngleBracketedParameterData { ref types, .. }),
            ..
        } if !types.is_empty() => Some(&types[0]),
            _ => None,
    }
}

#[derive(Debug, Clone, Copy)]
enum AttrSource { Struct, Field, }

fn extract_attrs<'a>(attrs: &'a [Attribute], attr_source: AttrSource) -> Box<Iterator<Item = (Ident, Lit)> + 'a> {
    let settings_attrs = attrs.iter()
        .filter_map(|attr| match attr.value {
            MetaItem::List(ref i, ref v) if i.as_ref() == "structopt" => Some(v),
            _ => None,
        }).flat_map(|v| v.iter().filter_map(|mi| match *mi {
            NestedMetaItem::MetaItem(MetaItem::NameValue(ref i, ref l)) =>
                Some((i.clone(), l.clone())),
            _ => None,
        }));

    let doc_comments = attrs.iter()
        .filter_map(move |attr| {
            if let Attribute {
                value: MetaItem::NameValue(ref name, Lit::Str(ref value, StrStyle::Cooked)),
                is_sugared_doc: true,
                ..
            } = *attr {
                if name != "doc" { return None; }
                let text = value.trim_left_matches("//!")
                    .trim_left_matches("///")
                    .trim_left_matches("/*!")
                    .trim_left_matches("/**")
                    .trim();

                // Clap's `App` has an `about` method to set a description,
                // it's `Field`s have a `help` method instead.
                if let AttrSource::Struct = attr_source {
                    Some(("about".into(), text.into()))
                } else {
                    Some(("help".into(), text.into()))
                }
            } else {
                None
            }
        });

    Box::new(doc_comments.chain(settings_attrs))
}

fn from_attr_or_env(attrs: &[(Ident, Lit)], key: &str, env: &str) -> Lit {
    let default = std::env::var(env).unwrap_or("".into());
    attrs.iter()
        .filter(|&&(ref i, _)| i.as_ref() == key)
        .last()
        .map(|&(_, ref l)| l.clone())
        .unwrap_or_else(|| Lit::Str(default, StrStyle::Cooked))
}

fn gen_name(field: &Field) -> Ident {
    extract_attrs(&field.attrs, AttrSource::Field)
        .filter(|&(ref i, _)| i.as_ref() == "name")
        .last()
        .and_then(|(_, ref l)| match l {
            &Lit::Str(ref s, _) => Some(Ident::new(s.clone())),
            _ => None,
        })
        .unwrap_or(field.ident.as_ref().unwrap().clone())
}

fn gen_from_clap(struct_name: &Ident, s: &[Field]) -> quote::Tokens {
    let fields = s.iter().map(|field| {
        let field_name = field.ident.as_ref().unwrap();
        let name = gen_name(field);
        let convert = match ty(&field.ty) {
            Ty::Bool => quote!(is_present(stringify!(#name))),
            Ty::U64 => quote!(occurrences_of(stringify!(#name))),
            Ty::Option => quote! {
                value_of(stringify!(#name))
                    .as_ref()
                    .map(|s| s.parse().unwrap())
            },
            Ty::Vec => quote! {
                values_of(stringify!(#name))
                    .map(|v| v.map(|s| s.parse().unwrap()).collect())
                    .unwrap_or_else(Vec::new)
            },
            Ty::Other => quote! {
                value_of(stringify!(#name))
                    .as_ref()
                    .unwrap()
                    .parse()
                    .unwrap()
            },
        };
        quote!( #field_name: matches.#convert, )
    });
    quote! {
        fn from_clap(matches: _structopt::clap::ArgMatches) -> Self {
            #struct_name {
                #( #fields )*
            }
        }
    }
}

fn gen_clap(struct_attrs: &[Attribute]) -> quote::Tokens {
    let struct_attrs: Vec<_> = extract_attrs(struct_attrs, AttrSource::Struct).collect();
    let name = from_attr_or_env(&struct_attrs, "name", "CARGO_PKG_NAME");
    let version = from_attr_or_env(&struct_attrs, "version", "CARGO_PKG_VERSION");
    let author = from_attr_or_env(&struct_attrs, "author", "CARGO_PKG_AUTHORS");
    let about = from_attr_or_env(&struct_attrs, "about", "CARGO_PKG_DESCRIPTION");

    quote! {
        fn clap<'a, 'b>() -> _structopt::clap::App<'a, 'b> {
            let app = _structopt::clap::App::new(#name)
                .version(#version)
                .author(#author)
                .about(#about);
            Self::augment_clap(app)
        }
    }
}

fn gen_augment_clap(s: &[Field]) -> quote::Tokens {
    let args = s.iter().map(|field| {
        let name = gen_name(field);
        let cur_type = ty(&field.ty);
        let convert_type = match cur_type {
            Ty::Vec | Ty::Option => sub_type(&field.ty).unwrap_or(&field.ty),
            _ => &field.ty,
        };
        let validator = quote! {
            validator(|s| s.parse::<#convert_type>()
                      .map(|_| ())
                      .map_err(|e| e.description().into()))
        };
        let modifier = match cur_type {
            Ty::Bool => quote!( .takes_value(false).multiple(false) ),
            Ty::U64 => quote!( .takes_value(false).multiple(true) ),
            Ty::Option => quote!( .takes_value(true).multiple(false).#validator ),
            Ty::Vec => quote!( .takes_value(true).multiple(true).#validator ),
            Ty::Other => {
                let required = extract_attrs(&field.attrs, AttrSource::Field)
                    .find(|&(ref i, _)| i.as_ref() == "default_value")
                    .is_none();
                quote!( .takes_value(true).multiple(false).required(#required).#validator )
            },
        };
        let from_attr = extract_attrs(&field.attrs, AttrSource::Field)
            .filter(|&(ref i, _)| i.as_ref() != "name")
            .map(|(i, l)| quote!(.#i(#l)));
        quote!( .arg(_structopt::clap::Arg::with_name(stringify!(#name)) #modifier #(#from_attr)*) )
    });

    quote! {
        fn augment_clap<'a, 'b>(app: _structopt::clap::App<'a, 'b>) -> _structopt::clap::App<'a, 'b> {
            use std::error::Error;
            app #( #args )*
        }
    }
}

fn impl_structopt_for_struct(name: &Ident, fields: &[Field], attrs: &[Attribute]) -> quote::Tokens {
    let clap = gen_clap(attrs);
    let augment_clap = gen_augment_clap(fields);
    let from_clap = gen_from_clap(name, fields);

    quote! {
        impl _structopt::StructOpt for #name {
            #clap
            #augment_clap
            #from_clap
        }
    }
}

fn impl_structopt_for_enum(_variants: &[Variant]) -> quote::Tokens {
    quote!()
}

fn impl_structopt(ast: &DeriveInput) -> quote::Tokens {
    let struct_name = &ast.ident;
    let inner_impl = match ast.body {
        Body::Struct(VariantData::Struct(ref fields)) =>
            impl_structopt_for_struct(struct_name, fields, &ast.attrs),
        Body::Enum(ref variants) =>
            impl_structopt_for_enum(variants),
        _ => panic!("structopt only supports non-tuple structs and enums")
    };

    let dummy_const = Ident::new(format!("_IMPL_STRUCTOPT_FOR_{}", struct_name));
    quote! {
        #[allow(non_upper_case_globals, unused_attributes, unused_imports)]
        const #dummy_const: () = {
            extern crate structopt as _structopt;
            use structopt::StructOpt;
            #inner_impl
        };
    }
}
<|endoftext|>


<commit_before>use common::memory;

use core::mem::size_of;
use core::u32;

use drivers::mmio::Mmio;

use super::fis::{FIS_TYPE_REG_H2D, FisRegH2D};

const ATA_CMD_READ_DMA_EXT: u8 = 0x25;
const ATA_DEV_BUSY: u8 = 0x80;
const ATA_DEV_DRQ: u8 = 0x08;

const HBA_PxCMD_CR: u32 = 1 << 15;
const HBA_PxCMD_FR: u32 = 1 << 14;
const HBA_PxCMD_FRE: u32 = 1 << 4;
const HBA_PxCMD_ST: u32 = 1;
const HBA_PxIS_TFES: u32 = 1 << 30;
const HBA_SSTS_PRESENT: u32 = 0x13;
const HBA_SIG_ATA: u32 = 0x00000101;
const HBA_SIG_ATAPI: u32 = 0xEB140101;
const HBA_SIG_PM: u32 = 0x96690101;
const HBA_SIG_SEMB: u32 = 0xC33C0101;

#[derive(Debug)]
pub enum HbaPortType {
    None,
    Unknown(u32),
    SATA,
    SATAPI,
    PM,
    SEMB,
}

#[repr(packed)]
pub struct HbaPort {
    pub clb: Mmio<u64>,   // 0x00, command list base address, 1K-byte aligned
    pub fb: Mmio<u64>,    // 0x08, FIS base address, 256-byte aligned
    pub is: Mmio<u32>,    // 0x10, interrupt status
    pub ie: Mmio<u32>,    // 0x14, interrupt enable
    pub cmd: Mmio<u32>,   // 0x18, command and status
    pub rsv0: Mmio<u32>,  // 0x1C, Reserved
    pub tfd: Mmio<u32>,   // 0x20, task file data
    pub sig: Mmio<u32>,   // 0x24, signature
    pub ssts: Mmio<u32>,  // 0x28, SATA status (SCR0:SStatus)
    pub sctl: Mmio<u32>,  // 0x2C, SATA control (SCR2:SControl)
    pub serr: Mmio<u32>,  // 0x30, SATA error (SCR1:SError)
    pub sact: Mmio<u32>,  // 0x34, SATA active (SCR3:SActive)
    pub ci: Mmio<u32>,    // 0x38, command issue
    pub sntf: Mmio<u32>,  // 0x3C, SATA notification (SCR4:SNotification)
    pub fbs: Mmio<u32>,   // 0x40, FIS-based switch control
    pub rsv1: [Mmio<u32>; 11],    // 0x44 ~ 0x6F, Reserved
    pub vendor: [Mmio<u32>; 4]    // 0x70 ~ 0x7F, vendor specific
}

impl HbaPort {
    pub fn probe(&self) -> HbaPortType {
        if self.ssts.readf(HBA_SSTS_PRESENT) {
            let sig = self.sig.read();
            match sig {
                HBA_SIG_ATA => HbaPortType::SATA,
                HBA_SIG_ATAPI => HbaPortType::SATAPI,
                HBA_SIG_PM => HbaPortType::PM,
                HBA_SIG_SEMB => HbaPortType::SEMB,
                _ => HbaPortType::Unknown(sig)
            }
        } else {
            HbaPortType::None
        }
    }

    pub fn init(&mut self) {
        self.stop();

        debugln!("Port Command List");
        let clb = unsafe { memory::alloc_aligned(size_of::<HbaCmdHeader>(), 1024) };
        self.clb.write(clb as u64);

        debugln!("Port FIS");
        let fb = unsafe { memory::alloc_aligned(256, 256) };
        self.fb.write(fb as u64);

        for i in 0..32 {
            debugln!("Port Command Table {}", i);
            let cmdheader = unsafe { &mut * (clb as *mut HbaCmdHeader).offset(i) };
            let ctba = unsafe { memory::alloc_aligned(256, 256) };
            cmdheader.ctba.write(ctba as u64);
            cmdheader.prdtl.write(8);
        }

        self.start();
    }

    pub fn start(&mut self) {
        debugln!("Starting port");

        while self.cmd.readf(HBA_PxCMD_CR) {}

        self.cmd.writef(HBA_PxCMD_FRE, true);
        self.cmd.writef(HBA_PxCMD_ST, true);
    }

    pub fn stop(&mut self) {
        debugln!("Stopping port");

    	self.cmd.writef(HBA_PxCMD_ST, false);

    	while self.cmd.readf(HBA_PxCMD_FR | HBA_PxCMD_CR) {}

    	self.cmd.writef(HBA_PxCMD_FRE, false);
    }

    pub fn slot(&self) -> Option<u32> {
        let slots = self.sact.read() | self.ci.read();
        for i in 0..32 {
            if slots & 1 << i == 0 {
                return Some(i);
            }
        }
        None
    }

    pub fn read(&mut self, lba: u64, buf: usize, len: usize) -> bool {
        let sectors = len/512;
        let entries = 1;

        debugln!("LBA: {:X} BUF: {:X}, SECTORS: {}", lba, buf, sectors);

        self.is.write(u32::MAX);

        if let Some(slot) = self.slot() {
            debugln!("Slot {}", slot);

            let clb = self.clb.read() as usize;
            let cmdheader = unsafe { &mut * (clb as *mut HbaCmdHeader).offset(slot as isize) };

            cmdheader.cfl.write(((size_of::<FisRegH2D>()/size_of::<u32>()) as u8) << 3);
            cmdheader.prdtl.write(entries);

            let ctba = cmdheader.ctba.read() as usize;
            unsafe { ::memset(ctba as *mut u8, 0, size_of::<HbaCmdTable>()) };
            let cmdtbl = unsafe { &mut * (ctba as *mut HbaCmdTable) };

            let prdt_entry = &mut cmdtbl.prdt_entry[0];
            prdt_entry.dba.write(buf as u64);
            prdt_entry.dbc.write(((sectors * 512) as u32) << 10 | 1);

            let cmdfis = unsafe { &mut * (cmdtbl.cfis.as_ptr() as *mut FisRegH2D) };

            cmdfis.fis_type.write(FIS_TYPE_REG_H2D);
            cmdfis.pm.writef(1, true);
            cmdfis.command.write(ATA_CMD_READ_DMA_EXT);

            cmdfis.lba0.write(lba as u8);
            cmdfis.lba1.write((lba >> 8) as u8);
            cmdfis.lba2.write((lba >> 16) as u8);

            cmdfis.device.write(1 << 6);

            cmdfis.lba3.write((lba >> 24) as u8);
            cmdfis.lba4.write((lba >> 32) as u8);
            cmdfis.lba5.write((lba >> 40) as u8);

            cmdfis.countl.write(sectors as u8);
            cmdfis.counth.write((sectors >> 8) as u8);

            debugln!("Busy Wait");
            while self.tfd.readf((ATA_DEV_BUSY | ATA_DEV_DRQ) as u32) {}

            self.ci.write(1 << slot);

            debugln!("Completion Wait");
            while self.ci.readf(1 << slot) {
                if self.is.readf(HBA_PxIS_TFES) {
        			debugln!("Read disk error");
        			return false;
        		}
            }

            debugln!("Return");

            if self.is.readf(HBA_PxIS_TFES) {
    			debugln!("Read disk error");
    			return false;
    		}

            return true;
        }

        false
    }
}

#[repr(packed)]
pub struct HbaMem {
    pub cap: Mmio<u32>,       // 0x00, Host capability
    pub ghc: Mmio<u32>,       // 0x04, Global host control
    pub is: Mmio<u32>,        // 0x08, Interrupt status
    pub pi: Mmio<u32>,        // 0x0C, Port implemented
    pub vs: Mmio<u32>,        // 0x10, Version
    pub ccc_ctl: Mmio<u32>,   // 0x14, Command completion coalescing control
    pub ccc_pts: Mmio<u32>,   // 0x18, Command completion coalescing ports
    pub em_loc: Mmio<u32>,    // 0x1C, Enclosure management location
    pub em_ctl: Mmio<u32>,    // 0x20, Enclosure management control
    pub cap2: Mmio<u32>,      // 0x24, Host capabilities extended
    pub bohc: Mmio<u32>,      // 0x28, BIOS/OS handoff control and status
    pub rsv: [Mmio<u8>; 116],         // 0x2C - 0x9F, Reserved
    pub vendor: [Mmio<u8>; 96],       // 0xA0 - 0xFF, Vendor specific registers
    pub ports: [HbaPort; 32]    // 0x100 - 0x10FF, Port control registers
}

#[repr(packed)]
struct HbaPrdtEntry {
   dba: Mmio<u64>,		// Data base address
   rsv0: Mmio<u32>,		// Reserved
   dbc: Mmio<u32>,		// Byte count, 4M max, interrupt = 1
}

#[repr(packed)]
struct HbaCmdTable {
	// 0x00
	cfis: [Mmio<u8>; 64],	// Command FIS

	// 0x40
	acmd: [Mmio<u8>; 16],	// ATAPI command, 12 or 16 bytes

	// 0x50
	rsv: [Mmio<u8>; 48],	// Reserved

	// 0x80
	prdt_entry: [HbaPrdtEntry; 65536],	// Physical region descriptor table entries, 0 ~ 65535
}

#[repr(packed)]
struct HbaCmdHeader {
	// DW0
	cfl: Mmio<u8>,		    // Command FIS length in DWORDS, 2 ~ 16, atapi: 4, write - host to device: 2, prefetchable: 1
	pm: Mmio<u8>,		    // Reset - 0x80, bist: 0x40, clear busy on ok: 0x20, port multiplier

	prdtl: Mmio<u16>,		// Physical region descriptor table length in entries

	// DW1
	prdbc: Mmio<u32>,		// Physical region descriptor byte count transferred

	// DW2, 3
	ctba: Mmio<u64>,		// Command table descriptor base address

	// DW4 - 7
	rsv1: [Mmio<u32>; 4],	// Reserved
}
<commit_msg>Hooray data!<commit_after>use common::memory;

use core::mem::size_of;
use core::u32;

use drivers::mmio::Mmio;

use super::fis::{FIS_TYPE_REG_H2D, FisRegH2D};

const ATA_CMD_READ_DMA_EXT: u8 = 0x25;
const ATA_DEV_BUSY: u8 = 0x80;
const ATA_DEV_DRQ: u8 = 0x08;

const HBA_PxCMD_CR: u32 = 1 << 15;
const HBA_PxCMD_FR: u32 = 1 << 14;
const HBA_PxCMD_FRE: u32 = 1 << 4;
const HBA_PxCMD_ST: u32 = 1;
const HBA_PxIS_TFES: u32 = 1 << 30;
const HBA_SSTS_PRESENT: u32 = 0x13;
const HBA_SIG_ATA: u32 = 0x00000101;
const HBA_SIG_ATAPI: u32 = 0xEB140101;
const HBA_SIG_PM: u32 = 0x96690101;
const HBA_SIG_SEMB: u32 = 0xC33C0101;

#[derive(Debug)]
pub enum HbaPortType {
    None,
    Unknown(u32),
    SATA,
    SATAPI,
    PM,
    SEMB,
}

#[repr(packed)]
pub struct HbaPort {
    pub clb: Mmio<u64>,   // 0x00, command list base address, 1K-byte aligned
    pub fb: Mmio<u64>,    // 0x08, FIS base address, 256-byte aligned
    pub is: Mmio<u32>,    // 0x10, interrupt status
    pub ie: Mmio<u32>,    // 0x14, interrupt enable
    pub cmd: Mmio<u32>,   // 0x18, command and status
    pub rsv0: Mmio<u32>,  // 0x1C, Reserved
    pub tfd: Mmio<u32>,   // 0x20, task file data
    pub sig: Mmio<u32>,   // 0x24, signature
    pub ssts: Mmio<u32>,  // 0x28, SATA status (SCR0:SStatus)
    pub sctl: Mmio<u32>,  // 0x2C, SATA control (SCR2:SControl)
    pub serr: Mmio<u32>,  // 0x30, SATA error (SCR1:SError)
    pub sact: Mmio<u32>,  // 0x34, SATA active (SCR3:SActive)
    pub ci: Mmio<u32>,    // 0x38, command issue
    pub sntf: Mmio<u32>,  // 0x3C, SATA notification (SCR4:SNotification)
    pub fbs: Mmio<u32>,   // 0x40, FIS-based switch control
    pub rsv1: [Mmio<u32>; 11],    // 0x44 ~ 0x6F, Reserved
    pub vendor: [Mmio<u32>; 4]    // 0x70 ~ 0x7F, vendor specific
}

impl HbaPort {
    pub fn probe(&self) -> HbaPortType {
        if self.ssts.readf(HBA_SSTS_PRESENT) {
            let sig = self.sig.read();
            match sig {
                HBA_SIG_ATA => HbaPortType::SATA,
                HBA_SIG_ATAPI => HbaPortType::SATAPI,
                HBA_SIG_PM => HbaPortType::PM,
                HBA_SIG_SEMB => HbaPortType::SEMB,
                _ => HbaPortType::Unknown(sig)
            }
        } else {
            HbaPortType::None
        }
    }

    pub fn init(&mut self) {
        self.stop();

        debugln!("Port Command List");
        let clb = unsafe { memory::alloc_aligned(size_of::<HbaCmdHeader>(), 1024) };
        self.clb.write(clb as u64);

        debugln!("Port FIS");
        let fb = unsafe { memory::alloc_aligned(256, 256) };
        self.fb.write(fb as u64);

        for i in 0..32 {
            debugln!("Port Command Table {}", i);
            let cmdheader = unsafe { &mut * (clb as *mut HbaCmdHeader).offset(i) };
            let ctba = unsafe { memory::alloc_aligned(256, 256) };
            cmdheader.ctba.write(ctba as u64);
            cmdheader.prdtl.write(8);
        }

        self.start();
    }

    pub fn start(&mut self) {
        debugln!("Starting port");

        while self.cmd.readf(HBA_PxCMD_CR) {}

        self.cmd.writef(HBA_PxCMD_FRE, true);
        self.cmd.writef(HBA_PxCMD_ST, true);
    }

    pub fn stop(&mut self) {
        debugln!("Stopping port");

    	self.cmd.writef(HBA_PxCMD_ST, false);

    	while self.cmd.readf(HBA_PxCMD_FR | HBA_PxCMD_CR) {}

    	self.cmd.writef(HBA_PxCMD_FRE, false);
    }

    pub fn slot(&self) -> Option<u32> {
        let slots = self.sact.read() | self.ci.read();
        for i in 0..32 {
            if slots & 1 << i == 0 {
                return Some(i);
            }
        }
        None
    }

    pub fn read(&mut self, lba: u64, buf: usize, len: usize) -> bool {
        let sectors = len/512;
        let entries = 1;

        debugln!("LBA: {:X} BUF: {:X} SECTORS: {}", lba, buf, sectors);

        self.is.write(u32::MAX);

        if let Some(slot) = self.slot() {
            debugln!("Slot {}", slot);

            let clb = self.clb.read() as usize;
            let cmdheader = unsafe { &mut * (clb as *mut HbaCmdHeader).offset(slot as isize) };

            cmdheader.cfl.write(((size_of::<FisRegH2D>()/size_of::<u32>()) as u8));
            cmdheader.prdtl.write(entries);

            let ctba = cmdheader.ctba.read() as usize;
            unsafe { ::memset(ctba as *mut u8, 0, size_of::<HbaCmdTable>()) };
            let cmdtbl = unsafe { &mut * (ctba as *mut HbaCmdTable) };

            let prdt_entry = &mut cmdtbl.prdt_entry[0];
            prdt_entry.dba.write(buf as u64);
            prdt_entry.dbc.write(((sectors * 512) as u32) | 1);

            let cmdfis = unsafe { &mut * (cmdtbl.cfis.as_ptr() as *mut FisRegH2D) };

            cmdfis.fis_type.write(FIS_TYPE_REG_H2D);
            cmdfis.pm.write(1 << 7);
            cmdfis.command.write(ATA_CMD_READ_DMA_EXT);

            cmdfis.lba0.write(lba as u8);
            cmdfis.lba1.write((lba >> 8) as u8);
            cmdfis.lba2.write((lba >> 16) as u8);

            cmdfis.device.write(1 << 6);

            cmdfis.lba3.write((lba >> 24) as u8);
            cmdfis.lba4.write((lba >> 32) as u8);
            cmdfis.lba5.write((lba >> 40) as u8);

            cmdfis.countl.write(sectors as u8);
            cmdfis.counth.write((sectors >> 8) as u8);

            debugln!("Busy Wait");
            while self.tfd.readf((ATA_DEV_BUSY | ATA_DEV_DRQ) as u32) {}

            self.ci.write(1 << slot);

            debugln!("Completion Wait");
            while self.ci.readf(1 << slot) {
                if self.is.readf(HBA_PxIS_TFES) {
        			debugln!("Read disk error");
        			return false;
        		}
            }

            debugln!("Return");

            if self.is.readf(HBA_PxIS_TFES) {
    			debugln!("Read disk error");
    			return false;
    		}

            return true;
        }

        false
    }
}

#[repr(packed)]
pub struct HbaMem {
    pub cap: Mmio<u32>,       // 0x00, Host capability
    pub ghc: Mmio<u32>,       // 0x04, Global host control
    pub is: Mmio<u32>,        // 0x08, Interrupt status
    pub pi: Mmio<u32>,        // 0x0C, Port implemented
    pub vs: Mmio<u32>,        // 0x10, Version
    pub ccc_ctl: Mmio<u32>,   // 0x14, Command completion coalescing control
    pub ccc_pts: Mmio<u32>,   // 0x18, Command completion coalescing ports
    pub em_loc: Mmio<u32>,    // 0x1C, Enclosure management location
    pub em_ctl: Mmio<u32>,    // 0x20, Enclosure management control
    pub cap2: Mmio<u32>,      // 0x24, Host capabilities extended
    pub bohc: Mmio<u32>,      // 0x28, BIOS/OS handoff control and status
    pub rsv: [Mmio<u8>; 116],         // 0x2C - 0x9F, Reserved
    pub vendor: [Mmio<u8>; 96],       // 0xA0 - 0xFF, Vendor specific registers
    pub ports: [HbaPort; 32]    // 0x100 - 0x10FF, Port control registers
}

#[repr(packed)]
struct HbaPrdtEntry {
   dba: Mmio<u64>,		// Data base address
   rsv0: Mmio<u32>,		// Reserved
   dbc: Mmio<u32>,		// Byte count, 4M max, interrupt = 1
}

#[repr(packed)]
struct HbaCmdTable {
	// 0x00
	cfis: [Mmio<u8>; 64],	// Command FIS

	// 0x40
	acmd: [Mmio<u8>; 16],	// ATAPI command, 12 or 16 bytes

	// 0x50
	rsv: [Mmio<u8>; 48],	// Reserved

	// 0x80
	prdt_entry: [HbaPrdtEntry; 65536],	// Physical region descriptor table entries, 0 ~ 65535
}

#[repr(packed)]
struct HbaCmdHeader {
	// DW0
	cfl: Mmio<u8>,		    // Command FIS length in DWORDS, 2 ~ 16, atapi: 4, write - host to device: 2, prefetchable: 1
	pm: Mmio<u8>,		    // Reset - 0x80, bist: 0x40, clear busy on ok: 0x20, port multiplier

	prdtl: Mmio<u16>,		// Physical region descriptor table length in entries

	// DW1
	prdbc: Mmio<u32>,		// Physical region descriptor byte count transferred

	// DW2, 3
	ctba: Mmio<u64>,		// Command table descriptor base address

	// DW4 - 7
	rsv1: [Mmio<u32>; 4],	// Reserved
}
<|endoftext|>


<commit_before>//! Provides functions to read and write segment registers.

#[cfg(docsrs)]
use crate::{registers::control::Cr4Flags, structures::gdt::GlobalDescriptorTable};
use crate::{
    registers::model_specific::{FsBase, GsBase, Msr},
    structures::gdt::SegmentSelector,
    VirtAddr,
};

/// An x86 segment
///
/// Segment registers on x86 are 16-bit [`SegmentSelector`]s, which index into
/// the [`GlobalDescriptorTable`]. The corresponding GDT entry is used to
/// configure the segment itself. Note that most segmentation functionality is
/// disabled in 64-bit mode. See the individual segments for more information.
pub trait Segment {
    /// Returns the current value of the segment register.
    fn get_reg() -> SegmentSelector;
    /// Reload the segment register. Depending on the segment, this may also
    /// reconfigure the corresponding segment.
    ///
    /// ## Safety
    ///
    /// This function is unsafe because the caller must ensure that `sel`
    /// is a valid segment descriptor, and that reconfiguring the segment will
    /// not cause undefined behavior.
    unsafe fn set_reg(sel: SegmentSelector);
}

/// An x86 segment which is actually used in 64-bit mode
///
/// While most segments are unused in 64-bit mode, the FS and GS segments are
/// still partially used. Only the 64-bit segment base address is used, and this
/// address can be set via the GDT, or by using the `FSGSBASE` instructions.
pub trait Segment64: Segment {
    /// MSR containing the segment base. This MSR can be used to set the base
    /// when [`CR4.FSGSBASE`][Cr4Flags::FSGSBASE] is not set.
    const BASE: Msr;
    /// Reads the segment base address
    ///
    /// ## Safety
    ///
    /// If [`CR4.FSGSBASE`][Cr4Flags::FSGSBASE] is not set, this instruction will throw a `#UD`.
    unsafe fn read_base() -> VirtAddr;
    /// Writes the segment base address
    ///
    /// ## Safety
    ///
    /// If [`CR4.FSGSBASE`][Cr4Flags::FSGSBASE] is not set, this instruction will throw a `#UD`.
    ///
    /// The caller must ensure that this write operation has no unsafe side
    /// effects, as the segment base address might be in use.
    unsafe fn write_base(base: VirtAddr);
}

macro_rules! get_reg_impl {
    ($name:literal, $asm_get:ident) => {
        fn get_reg() -> SegmentSelector {
            let segment: u16;
            #[cfg(feature = "inline_asm")]
            unsafe {
                asm!(concat!("mov {0:x}, ", $name), out(reg) segment, options(nomem, nostack, preserves_flags));
            }
            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                segment = crate::asm::$asm_get();
            }
            SegmentSelector(segment)
        }
    };
}

macro_rules! segment_impl {
    ($type:ty, $name:literal, $asm_get:ident, $asm_load:ident) => {
        impl Segment for $type {
            get_reg_impl!($name, $asm_get);

            unsafe fn set_reg(sel: SegmentSelector) {
                #[cfg(feature = "inline_asm")]
                asm!(concat!("mov ", $name, ", {0:x}"), in(reg) sel.0, options(nostack, preserves_flags));

                #[cfg(not(feature = "inline_asm"))]
                crate::asm::$asm_load(sel.0);
            }
        }
    };
}

macro_rules! segment64_impl {
    ($type:ty, $name:literal, $base:ty, $asm_rd:ident, $asm_wr:ident) => {
        impl Segment64 for $type {
            const BASE: Msr = <$base>::MSR;
            unsafe fn read_base() -> VirtAddr {
                #[cfg(feature = "inline_asm")]
                {
                    let val: u64;
                    asm!(concat!("rd", $name, "base {}"), out(reg) val, options(nomem, nostack, preserves_flags));
                    VirtAddr::new_unsafe(val)
                }
                #[cfg(not(feature = "inline_asm"))]
                VirtAddr::new_unsafe(crate::asm::$asm_rd())
            }

            unsafe fn write_base(base: VirtAddr) {
                #[cfg(feature = "inline_asm")]
                asm!(concat!("wr", $name, "base {}"), in(reg) base.as_u64(), options(nostack, preserves_flags));

                #[cfg(not(feature = "inline_asm"))]
                crate::asm::$asm_wr(base.as_u64());
            }
        }
    };
}

/// Code Segment
///
/// The segment base and limit are unused in 64-bit mode. Only the L (long), D
/// (default operation size), and DPL (descriptor privilege-level) fields of the
/// descriptor are recognized. So chaning the segment register can be used to
/// change privilege level or enable/disable long mode.
#[derive(Debug)]
pub struct CS;
impl Segment for CS {
    get_reg_impl!("cs", x86_64_asm_get_cs);

    unsafe fn set_reg(sel: SegmentSelector) {
        // Note this is special since we cannot directly move to cs. Instead we
        // push the new segment selector and return value on the stack and use
        // retfq to reload cs and continue at 1:.
        #[cfg(feature = "inline_asm")]
        asm!(
            "push {sel}",
            "lea {tmp}, [1f + rip]",
            "push {tmp}",
            "retfq",
            "1:",
            sel = in(reg) u64::from(sel.0),
            tmp = lateout(reg) _,
            options(preserves_flags),
        );

        #[cfg(not(feature = "inline_asm"))]
        crate::asm::x86_64_asm_set_cs(u64::from(sel.0));
    }
}

/// Stack Segment
///
/// Entirely unused in 64-bit mode, setting the segment register does nothing.
#[derive(Debug)]
pub struct SS;
segment_impl!(SS, "ss", x86_64_asm_get_ss, x86_64_asm_load_ss);

/// Data Segment
///
/// Entirely unused in 64-bit mode, setting the segment register does nothing.
#[derive(Debug)]
pub struct DS;
segment_impl!(DS, "ds", x86_64_asm_get_ds, x86_64_asm_load_ds);

/// ES Segment
///
/// Entirely unused in 64-bit mode, setting the segment register does nothing.
#[derive(Debug)]
pub struct ES;
segment_impl!(ES, "es", x86_64_asm_get_es, x86_64_asm_load_es);

/// FS Segment
///
/// Only base is used in 64-bit mode, see [`Segment64`]. This is often used in
/// user-mode for Thread-Local Storage (TLS).
#[derive(Debug)]
pub struct FS;
segment_impl!(FS, "fs", x86_64_asm_get_fs, x86_64_asm_load_fs);
segment64_impl!(FS, "fs", FsBase, x86_64_asm_rdfsbase, x86_64_asm_wrfsbase);

/// GS Segment
///
/// Only base is used in 64-bit mode, see [`Segment64`]. In kernel-mode, the GS
/// base often points to a per-cpu kernel data structure.
#[derive(Debug)]
pub struct GS;
segment_impl!(GS, "gs", x86_64_asm_get_gs, x86_64_asm_load_gs);
segment64_impl!(GS, "gs", GsBase, x86_64_asm_rdgsbase, x86_64_asm_wrgsbase);

impl GS {
    /// Swap `KernelGsBase` MSR and `GsBase` MSR.
    ///
    /// ## Safety
    ///
    /// This function is unsafe because the caller must ensure that the
    /// swap operation cannot lead to undefined behavior.
    pub unsafe fn swap() {
        #[cfg(feature = "inline_asm")]
        asm!("swapgs", options(nostack, preserves_flags));

        #[cfg(not(feature = "inline_asm"))]
        crate::asm::x86_64_asm_swapgs();
    }
}

/// Alias for [`CS::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `CS::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn set_cs(sel: SegmentSelector) {
    CS::set_reg(sel)
}
/// Alias for [`SS::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `SS::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn load_ss(sel: SegmentSelector) {
    SS::set_reg(sel)
}
/// Alias for [`DS::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `DS::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn load_ds(sel: SegmentSelector) {
    DS::set_reg(sel)
}
/// Alias for [`ES::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `ES::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn load_es(sel: SegmentSelector) {
    ES::set_reg(sel)
}
/// Alias for [`FS::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `FS::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn load_fs(sel: SegmentSelector) {
    FS::set_reg(sel)
}
/// Alias for [`GS::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `GS::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn load_gs(sel: SegmentSelector) {
    GS::set_reg(sel)
}
/// Alias for [`GS::swap()`]
#[deprecated(since = "0.14.4", note = "use `GS::swap()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn swap_gs() {
    GS::swap()
}
/// Alias for [`CS::get_reg()`]
#[deprecated(since = "0.14.4", note = "use `CS::get_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub fn cs() -> SegmentSelector {
    CS::get_reg()
}
/// Alias for [`FS::write_base()`].
///
/// Raises #GP if the provided address is non-canonical.
#[deprecated(since = "0.14.4", note = "use `FS::write_base()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn wrfsbase(val: u64) {
    FS::write_base(VirtAddr::new(val))
}
/// Alias for [`FS::read_base()`]
#[deprecated(since = "0.14.4", note = "use `FS::read_base()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn rdfsbase() -> u64 {
    FS::read_base().as_u64()
}
/// Alias for [`GS::write_base()`].
///
/// Raises #GP if the provided address is non-canonical.
#[deprecated(since = "0.14.4", note = "use `GS::write_base()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn wrgsbase(val: u64) {
    GS::write_base(VirtAddr::new(val))
}
/// Alias for [`GS::read_base()`]
#[deprecated(since = "0.14.4", note = "use `GS::read_base()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn rdgsbase() -> u64 {
    GS::read_base().as_u64()
}
<commit_msg>Make Segment64::read_base safe<commit_after>//! Provides functions to read and write segment registers.

#[cfg(docsrs)]
use crate::{registers::control::Cr4Flags, structures::gdt::GlobalDescriptorTable};
use crate::{
    registers::model_specific::{FsBase, GsBase, Msr},
    structures::gdt::SegmentSelector,
    VirtAddr,
};

/// An x86 segment
///
/// Segment registers on x86 are 16-bit [`SegmentSelector`]s, which index into
/// the [`GlobalDescriptorTable`]. The corresponding GDT entry is used to
/// configure the segment itself. Note that most segmentation functionality is
/// disabled in 64-bit mode. See the individual segments for more information.
pub trait Segment {
    /// Returns the current value of the segment register.
    fn get_reg() -> SegmentSelector;
    /// Reload the segment register. Depending on the segment, this may also
    /// reconfigure the corresponding segment.
    ///
    /// ## Safety
    ///
    /// This function is unsafe because the caller must ensure that `sel`
    /// is a valid segment descriptor, and that reconfiguring the segment will
    /// not cause undefined behavior.
    unsafe fn set_reg(sel: SegmentSelector);
}

/// An x86 segment which is actually used in 64-bit mode
///
/// While most segments are unused in 64-bit mode, the FS and GS segments are
/// still partially used. Only the 64-bit segment base address is used, and this
/// address can be set via the GDT, or by using the `FSGSBASE` instructions.
pub trait Segment64: Segment {
    /// MSR containing the segment base. This MSR can be used to set the base
    /// when [`CR4.FSGSBASE`][Cr4Flags::FSGSBASE] is **not** set.
    const BASE: Msr;
    /// Reads the segment base address
    ///
    /// ## Exceptions
    ///
    /// If [`CR4.FSGSBASE`][Cr4Flags::FSGSBASE] is not set, this instruction will throw a `#UD`.
    fn read_base() -> VirtAddr;
    /// Writes the segment base address
    ///
    /// ## Exceptions
    ///
    /// If [`CR4.FSGSBASE`][Cr4Flags::FSGSBASE] is not set, this instruction will throw a `#UD`.
    ///
    /// ## Safety
    ///
    /// The caller must ensure that this write operation has no unsafe side
    /// effects, as the segment base address might be in use.
    unsafe fn write_base(base: VirtAddr);
}

macro_rules! get_reg_impl {
    ($name:literal, $asm_get:ident) => {
        fn get_reg() -> SegmentSelector {
            let segment: u16;
            #[cfg(feature = "inline_asm")]
            unsafe {
                asm!(concat!("mov {0:x}, ", $name), out(reg) segment, options(nomem, nostack, preserves_flags));
            }
            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                segment = crate::asm::$asm_get();
            }
            SegmentSelector(segment)
        }
    };
}

macro_rules! segment_impl {
    ($type:ty, $name:literal, $asm_get:ident, $asm_load:ident) => {
        impl Segment for $type {
            get_reg_impl!($name, $asm_get);

            unsafe fn set_reg(sel: SegmentSelector) {
                #[cfg(feature = "inline_asm")]
                asm!(concat!("mov ", $name, ", {0:x}"), in(reg) sel.0, options(nostack, preserves_flags));

                #[cfg(not(feature = "inline_asm"))]
                crate::asm::$asm_load(sel.0);
            }
        }
    };
}

macro_rules! segment64_impl {
    ($type:ty, $name:literal, $base:ty, $asm_rd:ident, $asm_wr:ident) => {
        impl Segment64 for $type {
            const BASE: Msr = <$base>::MSR;
            fn read_base() -> VirtAddr {
                #[cfg(feature = "inline_asm")]
                unsafe {
                    let val: u64;
                    asm!(concat!("rd", $name, "base {}"), out(reg) val, options(nomem, nostack, preserves_flags));
                    VirtAddr::new_unsafe(val)
                }
                #[cfg(not(feature = "inline_asm"))]
                unsafe {
                    VirtAddr::new_unsafe(crate::asm::$asm_rd())
                }
            }

            unsafe fn write_base(base: VirtAddr) {
                #[cfg(feature = "inline_asm")]
                asm!(concat!("wr", $name, "base {}"), in(reg) base.as_u64(), options(nostack, preserves_flags));

                #[cfg(not(feature = "inline_asm"))]
                crate::asm::$asm_wr(base.as_u64());
            }
        }
    };
}

/// Code Segment
///
/// The segment base and limit are unused in 64-bit mode. Only the L (long), D
/// (default operation size), and DPL (descriptor privilege-level) fields of the
/// descriptor are recognized. So chaning the segment register can be used to
/// change privilege level or enable/disable long mode.
#[derive(Debug)]
pub struct CS;
impl Segment for CS {
    get_reg_impl!("cs", x86_64_asm_get_cs);

    unsafe fn set_reg(sel: SegmentSelector) {
        // Note this is special since we cannot directly move to cs. Instead we
        // push the new segment selector and return value on the stack and use
        // retfq to reload cs and continue at 1:.
        #[cfg(feature = "inline_asm")]
        asm!(
            "push {sel}",
            "lea {tmp}, [1f + rip]",
            "push {tmp}",
            "retfq",
            "1:",
            sel = in(reg) u64::from(sel.0),
            tmp = lateout(reg) _,
            options(preserves_flags),
        );

        #[cfg(not(feature = "inline_asm"))]
        crate::asm::x86_64_asm_set_cs(u64::from(sel.0));
    }
}

/// Stack Segment
///
/// Entirely unused in 64-bit mode, setting the segment register does nothing.
#[derive(Debug)]
pub struct SS;
segment_impl!(SS, "ss", x86_64_asm_get_ss, x86_64_asm_load_ss);

/// Data Segment
///
/// Entirely unused in 64-bit mode, setting the segment register does nothing.
#[derive(Debug)]
pub struct DS;
segment_impl!(DS, "ds", x86_64_asm_get_ds, x86_64_asm_load_ds);

/// ES Segment
///
/// Entirely unused in 64-bit mode, setting the segment register does nothing.
#[derive(Debug)]
pub struct ES;
segment_impl!(ES, "es", x86_64_asm_get_es, x86_64_asm_load_es);

/// FS Segment
///
/// Only base is used in 64-bit mode, see [`Segment64`]. This is often used in
/// user-mode for Thread-Local Storage (TLS).
#[derive(Debug)]
pub struct FS;
segment_impl!(FS, "fs", x86_64_asm_get_fs, x86_64_asm_load_fs);
segment64_impl!(FS, "fs", FsBase, x86_64_asm_rdfsbase, x86_64_asm_wrfsbase);

/// GS Segment
///
/// Only base is used in 64-bit mode, see [`Segment64`]. In kernel-mode, the GS
/// base often points to a per-cpu kernel data structure.
#[derive(Debug)]
pub struct GS;
segment_impl!(GS, "gs", x86_64_asm_get_gs, x86_64_asm_load_gs);
segment64_impl!(GS, "gs", GsBase, x86_64_asm_rdgsbase, x86_64_asm_wrgsbase);

impl GS {
    /// Swap `KernelGsBase` MSR and `GsBase` MSR.
    ///
    /// ## Safety
    ///
    /// This function is unsafe because the caller must ensure that the
    /// swap operation cannot lead to undefined behavior.
    pub unsafe fn swap() {
        #[cfg(feature = "inline_asm")]
        asm!("swapgs", options(nostack, preserves_flags));

        #[cfg(not(feature = "inline_asm"))]
        crate::asm::x86_64_asm_swapgs();
    }
}

/// Alias for [`CS::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `CS::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn set_cs(sel: SegmentSelector) {
    CS::set_reg(sel)
}
/// Alias for [`SS::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `SS::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn load_ss(sel: SegmentSelector) {
    SS::set_reg(sel)
}
/// Alias for [`DS::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `DS::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn load_ds(sel: SegmentSelector) {
    DS::set_reg(sel)
}
/// Alias for [`ES::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `ES::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn load_es(sel: SegmentSelector) {
    ES::set_reg(sel)
}
/// Alias for [`FS::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `FS::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn load_fs(sel: SegmentSelector) {
    FS::set_reg(sel)
}
/// Alias for [`GS::set_reg()`]
#[deprecated(since = "0.14.4", note = "use `GS::set_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn load_gs(sel: SegmentSelector) {
    GS::set_reg(sel)
}
/// Alias for [`GS::swap()`]
#[deprecated(since = "0.14.4", note = "use `GS::swap()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn swap_gs() {
    GS::swap()
}
/// Alias for [`CS::get_reg()`]
#[deprecated(since = "0.14.4", note = "use `CS::get_reg()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub fn cs() -> SegmentSelector {
    CS::get_reg()
}
/// Alias for [`FS::write_base()`].
///
/// Raises #GP if the provided address is non-canonical.
#[deprecated(since = "0.14.4", note = "use `FS::write_base()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn wrfsbase(val: u64) {
    FS::write_base(VirtAddr::new(val))
}
/// Alias for [`FS::read_base()`]
#[deprecated(since = "0.14.4", note = "use `FS::read_base()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn rdfsbase() -> u64 {
    FS::read_base().as_u64()
}
/// Alias for [`GS::write_base()`].
///
/// Raises #GP if the provided address is non-canonical.
#[deprecated(since = "0.14.4", note = "use `GS::write_base()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn wrgsbase(val: u64) {
    GS::write_base(VirtAddr::new(val))
}
/// Alias for [`GS::read_base()`]
#[deprecated(since = "0.14.4", note = "use `GS::read_base()` instead")]
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn rdgsbase() -> u64 {
    GS::read_base().as_u64()
}
<|endoftext|>


<commit_before>use crate::mir::interpret::ConstValue;
use crate::mir::interpret::{LitToConstInput, Scalar};
use crate::ty::subst::InternalSubsts;
use crate::ty::{self, Ty, TyCtxt};
use crate::ty::{ParamEnv, ParamEnvAnd};
use rustc_errors::ErrorReported;
use rustc_hir as hir;
use rustc_hir::def_id::LocalDefId;
use rustc_macros::HashStable;

mod int;
mod kind;

pub use int::*;
pub use kind::*;

/// Typed constant value.
#[derive(Copy, Clone, Debug, Hash, RustcEncodable, RustcDecodable, Eq, PartialEq, Ord, PartialOrd)]
#[derive(HashStable)]
pub struct Const<'tcx> {
    pub ty: Ty<'tcx>,

    pub val: ConstKind<'tcx>,
}

#[cfg(target_arch = "x86_64")]
static_assert_size!(Const<'_>, 48);

impl<'tcx> Const<'tcx> {
    /// Literals and const generic parameters are eagerly converted to a constant, everything else
    /// becomes `Unevaluated`.
    pub fn from_anon_const(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> &'tcx Self {
        Self::from_opt_const_arg_anon_const(tcx, ty::WithOptConstParam::unknown(def_id))
    }

    pub fn from_opt_const_arg_anon_const(
        tcx: TyCtxt<'tcx>,
        def: ty::WithOptConstParam<LocalDefId>,
    ) -> &'tcx Self {
        debug!("Const::from_anon_const(def={:?})", def);

        let hir_id = tcx.hir().local_def_id_to_hir_id(def.did);

        let body_id = match tcx.hir().get(hir_id) {
            hir::Node::AnonConst(ac) => ac.body,
            _ => span_bug!(
                tcx.def_span(def.did.to_def_id()),
                "from_anon_const can only process anonymous constants"
            ),
        };

        let expr = &tcx.hir().body(body_id).value;

        let ty = tcx.type_of(def.def_id_for_type_of());

        let lit_input = match expr.kind {
            hir::ExprKind::Lit(ref lit) => Some(LitToConstInput { lit: &lit.node, ty, neg: false }),
            hir::ExprKind::Unary(hir::UnOp::UnNeg, ref expr) => match expr.kind {
                hir::ExprKind::Lit(ref lit) => {
                    Some(LitToConstInput { lit: &lit.node, ty, neg: true })
                }
                _ => None,
            },
            _ => None,
        };

        if let Some(lit_input) = lit_input {
            // If an error occurred, ignore that it's a literal and leave reporting the error up to
            // mir.
            if let Ok(c) = tcx.at(expr.span).lit_to_const(lit_input) {
                return c;
            } else {
                tcx.sess.delay_span_bug(expr.span, "Const::from_anon_const: couldn't lit_to_const");
            }
        }

        // Unwrap a block, so that e.g. `{ P }` is recognised as a parameter. Const arguments
        // currently have to be wrapped in curly brackets, so it's necessary to special-case.
        let expr = match &expr.kind {
            hir::ExprKind::Block(block, _) if block.stmts.is_empty() && block.expr.is_some() => {
                block.expr.as_ref().unwrap()
            }
            _ => expr,
        };

        use hir::{def::DefKind::ConstParam, def::Res, ExprKind, Path, QPath};
        let val = match expr.kind {
            ExprKind::Path(QPath::Resolved(_, &Path { res: Res::Def(ConstParam, def_id), .. })) => {
                // Find the name and index of the const parameter by indexing the generics of
                // the parent item and construct a `ParamConst`.
                let hir_id = tcx.hir().as_local_hir_id(def_id.expect_local());
                let item_id = tcx.hir().get_parent_node(hir_id);
                let item_def_id = tcx.hir().local_def_id(item_id);
                let generics = tcx.generics_of(item_def_id.to_def_id());
                let index =
                    generics.param_def_id_to_index[&tcx.hir().local_def_id(hir_id).to_def_id()];
                let name = tcx.hir().name(hir_id);
                ty::ConstKind::Param(ty::ParamConst::new(index, name))
            }
            _ => ty::ConstKind::Unevaluated(
                def.to_global(),
                InternalSubsts::identity_for_item(tcx, def.did.to_def_id()),
                None,
            ),
        };

        tcx.mk_const(ty::Const { val, ty })
    }

    #[inline]
    /// Interns the given value as a constant.
    pub fn from_value(tcx: TyCtxt<'tcx>, val: ConstValue<'tcx>, ty: Ty<'tcx>) -> &'tcx Self {
        tcx.mk_const(Self { val: ConstKind::Value(val), ty })
    }

    #[inline]
    /// Interns the given scalar as a constant.
    pub fn from_scalar(tcx: TyCtxt<'tcx>, val: Scalar, ty: Ty<'tcx>) -> &'tcx Self {
        Self::from_value(tcx, ConstValue::Scalar(val), ty)
    }

    #[inline]
    /// Creates a constant with the given integer value and interns it.
    pub fn from_bits(tcx: TyCtxt<'tcx>, bits: u128, ty: ParamEnvAnd<'tcx, Ty<'tcx>>) -> &'tcx Self {
        let size = tcx
            .layout_of(ty)
            .unwrap_or_else(|e| panic!("could not compute layout for {:?}: {:?}", ty, e))
            .size;
        Self::from_scalar(tcx, Scalar::from_uint(bits, size), ty.value)
    }

    #[inline]
    /// Creates an interned zst constant.
    pub fn zero_sized(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> &'tcx Self {
        Self::from_scalar(tcx, Scalar::zst(), ty)
    }

    #[inline]
    /// Creates an interned bool constant.
    pub fn from_bool(tcx: TyCtxt<'tcx>, v: bool) -> &'tcx Self {
        Self::from_bits(tcx, v as u128, ParamEnv::empty().and(tcx.types.bool))
    }

    #[inline]
    /// Creates an interned usize constant.
    pub fn from_usize(tcx: TyCtxt<'tcx>, n: u64) -> &'tcx Self {
        Self::from_bits(tcx, n as u128, ParamEnv::empty().and(tcx.types.usize))
    }

    #[inline]
    /// Attempts to evaluate the given constant to bits. Can fail to evaluate in the presence of
    /// generics (or erroneous code) or if the value can't be represented as bits (e.g. because it
    /// contains const generic parameters or pointers).
    pub fn try_eval_bits(
        &self,
        tcx: TyCtxt<'tcx>,
        param_env: ParamEnv<'tcx>,
        ty: Ty<'tcx>,
    ) -> Option<u128> {
        assert_eq!(self.ty, ty);
        let size = tcx.layout_of(param_env.with_reveal_all().and(ty)).ok()?.size;
        // if `ty` does not depend on generic parameters, use an empty param_env
        self.val.eval(tcx, param_env).try_to_bits(size)
    }

    #[inline]
    /// Tries to evaluate the constant if it is `Unevaluated`. If that doesn't succeed, return the
    /// unevaluated constant.
    pub fn eval(&self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>) -> &Const<'tcx> {
        if let Some(val) = self.val.try_eval(tcx, param_env) {
            match val {
                Ok(val) => Const::from_value(tcx, val, self.ty),
                Err(ErrorReported) => tcx.const_error(self.ty),
            }
        } else {
            self
        }
    }

    #[inline]
    pub fn try_eval_bool(&self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>) -> Option<bool> {
        self.val.eval(tcx, param_env).try_to_bool()
    }

    #[inline]
    pub fn try_eval_usize(&self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>) -> Option<u64> {
        self.val.eval(tcx, param_env).try_to_machine_usize(tcx)
    }

    #[inline]
    /// Panics if the value cannot be evaluated or doesn't contain a valid integer of the given type.
    pub fn eval_bits(&self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>, ty: Ty<'tcx>) -> u128 {
        self.try_eval_bits(tcx, param_env, ty)
            .unwrap_or_else(|| bug!("expected bits of {:#?}, got {:#?}", ty, self))
    }

    #[inline]
    /// Panics if the value cannot be evaluated or doesn't contain a valid `usize`.
    pub fn eval_usize(&self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>) -> u64 {
        self.try_eval_usize(tcx, param_env)
            .unwrap_or_else(|| bug!("expected usize, got {:#?}", self))
    }
}
<commit_msg>Group the try_eval functions before the eval functions<commit_after>use crate::mir::interpret::ConstValue;
use crate::mir::interpret::{LitToConstInput, Scalar};
use crate::ty::subst::InternalSubsts;
use crate::ty::{self, Ty, TyCtxt};
use crate::ty::{ParamEnv, ParamEnvAnd};
use rustc_errors::ErrorReported;
use rustc_hir as hir;
use rustc_hir::def_id::LocalDefId;
use rustc_macros::HashStable;

mod int;
mod kind;

pub use int::*;
pub use kind::*;

/// Typed constant value.
#[derive(Copy, Clone, Debug, Hash, RustcEncodable, RustcDecodable, Eq, PartialEq, Ord, PartialOrd)]
#[derive(HashStable)]
pub struct Const<'tcx> {
    pub ty: Ty<'tcx>,

    pub val: ConstKind<'tcx>,
}

#[cfg(target_arch = "x86_64")]
static_assert_size!(Const<'_>, 48);

impl<'tcx> Const<'tcx> {
    /// Literals and const generic parameters are eagerly converted to a constant, everything else
    /// becomes `Unevaluated`.
    pub fn from_anon_const(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> &'tcx Self {
        Self::from_opt_const_arg_anon_const(tcx, ty::WithOptConstParam::unknown(def_id))
    }

    pub fn from_opt_const_arg_anon_const(
        tcx: TyCtxt<'tcx>,
        def: ty::WithOptConstParam<LocalDefId>,
    ) -> &'tcx Self {
        debug!("Const::from_anon_const(def={:?})", def);

        let hir_id = tcx.hir().local_def_id_to_hir_id(def.did);

        let body_id = match tcx.hir().get(hir_id) {
            hir::Node::AnonConst(ac) => ac.body,
            _ => span_bug!(
                tcx.def_span(def.did.to_def_id()),
                "from_anon_const can only process anonymous constants"
            ),
        };

        let expr = &tcx.hir().body(body_id).value;

        let ty = tcx.type_of(def.def_id_for_type_of());

        let lit_input = match expr.kind {
            hir::ExprKind::Lit(ref lit) => Some(LitToConstInput { lit: &lit.node, ty, neg: false }),
            hir::ExprKind::Unary(hir::UnOp::UnNeg, ref expr) => match expr.kind {
                hir::ExprKind::Lit(ref lit) => {
                    Some(LitToConstInput { lit: &lit.node, ty, neg: true })
                }
                _ => None,
            },
            _ => None,
        };

        if let Some(lit_input) = lit_input {
            // If an error occurred, ignore that it's a literal and leave reporting the error up to
            // mir.
            if let Ok(c) = tcx.at(expr.span).lit_to_const(lit_input) {
                return c;
            } else {
                tcx.sess.delay_span_bug(expr.span, "Const::from_anon_const: couldn't lit_to_const");
            }
        }

        // Unwrap a block, so that e.g. `{ P }` is recognised as a parameter. Const arguments
        // currently have to be wrapped in curly brackets, so it's necessary to special-case.
        let expr = match &expr.kind {
            hir::ExprKind::Block(block, _) if block.stmts.is_empty() && block.expr.is_some() => {
                block.expr.as_ref().unwrap()
            }
            _ => expr,
        };

        use hir::{def::DefKind::ConstParam, def::Res, ExprKind, Path, QPath};
        let val = match expr.kind {
            ExprKind::Path(QPath::Resolved(_, &Path { res: Res::Def(ConstParam, def_id), .. })) => {
                // Find the name and index of the const parameter by indexing the generics of
                // the parent item and construct a `ParamConst`.
                let hir_id = tcx.hir().as_local_hir_id(def_id.expect_local());
                let item_id = tcx.hir().get_parent_node(hir_id);
                let item_def_id = tcx.hir().local_def_id(item_id);
                let generics = tcx.generics_of(item_def_id.to_def_id());
                let index =
                    generics.param_def_id_to_index[&tcx.hir().local_def_id(hir_id).to_def_id()];
                let name = tcx.hir().name(hir_id);
                ty::ConstKind::Param(ty::ParamConst::new(index, name))
            }
            _ => ty::ConstKind::Unevaluated(
                def.to_global(),
                InternalSubsts::identity_for_item(tcx, def.did.to_def_id()),
                None,
            ),
        };

        tcx.mk_const(ty::Const { val, ty })
    }

    #[inline]
    /// Interns the given value as a constant.
    pub fn from_value(tcx: TyCtxt<'tcx>, val: ConstValue<'tcx>, ty: Ty<'tcx>) -> &'tcx Self {
        tcx.mk_const(Self { val: ConstKind::Value(val), ty })
    }

    #[inline]
    /// Interns the given scalar as a constant.
    pub fn from_scalar(tcx: TyCtxt<'tcx>, val: Scalar, ty: Ty<'tcx>) -> &'tcx Self {
        Self::from_value(tcx, ConstValue::Scalar(val), ty)
    }

    #[inline]
    /// Creates a constant with the given integer value and interns it.
    pub fn from_bits(tcx: TyCtxt<'tcx>, bits: u128, ty: ParamEnvAnd<'tcx, Ty<'tcx>>) -> &'tcx Self {
        let size = tcx
            .layout_of(ty)
            .unwrap_or_else(|e| panic!("could not compute layout for {:?}: {:?}", ty, e))
            .size;
        Self::from_scalar(tcx, Scalar::from_uint(bits, size), ty.value)
    }

    #[inline]
    /// Creates an interned zst constant.
    pub fn zero_sized(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> &'tcx Self {
        Self::from_scalar(tcx, Scalar::zst(), ty)
    }

    #[inline]
    /// Creates an interned bool constant.
    pub fn from_bool(tcx: TyCtxt<'tcx>, v: bool) -> &'tcx Self {
        Self::from_bits(tcx, v as u128, ParamEnv::empty().and(tcx.types.bool))
    }

    #[inline]
    /// Creates an interned usize constant.
    pub fn from_usize(tcx: TyCtxt<'tcx>, n: u64) -> &'tcx Self {
        Self::from_bits(tcx, n as u128, ParamEnv::empty().and(tcx.types.usize))
    }

    #[inline]
    /// Attempts to evaluate the given constant to bits. Can fail to evaluate in the presence of
    /// generics (or erroneous code) or if the value can't be represented as bits (e.g. because it
    /// contains const generic parameters or pointers).
    pub fn try_eval_bits(
        &self,
        tcx: TyCtxt<'tcx>,
        param_env: ParamEnv<'tcx>,
        ty: Ty<'tcx>,
    ) -> Option<u128> {
        assert_eq!(self.ty, ty);
        let size = tcx.layout_of(param_env.with_reveal_all().and(ty)).ok()?.size;
        // if `ty` does not depend on generic parameters, use an empty param_env
        self.val.eval(tcx, param_env).try_to_bits(size)
    }

    #[inline]
    pub fn try_eval_bool(&self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>) -> Option<bool> {
        self.val.eval(tcx, param_env).try_to_bool()
    }

    #[inline]
    pub fn try_eval_usize(&self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>) -> Option<u64> {
        self.val.eval(tcx, param_env).try_to_machine_usize(tcx)
    }

    #[inline]
    /// Tries to evaluate the constant if it is `Unevaluated`. If that doesn't succeed, return the
    /// unevaluated constant.
    pub fn eval(&self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>) -> &Const<'tcx> {
        if let Some(val) = self.val.try_eval(tcx, param_env) {
            match val {
                Ok(val) => Const::from_value(tcx, val, self.ty),
                Err(ErrorReported) => tcx.const_error(self.ty),
            }
        } else {
            self
        }
    }

    #[inline]
    /// Panics if the value cannot be evaluated or doesn't contain a valid integer of the given type.
    pub fn eval_bits(&self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>, ty: Ty<'tcx>) -> u128 {
        self.try_eval_bits(tcx, param_env, ty)
            .unwrap_or_else(|| bug!("expected bits of {:#?}, got {:#?}", ty, self))
    }

    #[inline]
    /// Panics if the value cannot be evaluated or doesn't contain a valid `usize`.
    pub fn eval_usize(&self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>) -> u64 {
        self.try_eval_usize(tcx, param_env)
            .unwrap_or_else(|| bug!("expected usize, got {:#?}", self))
    }
}
<|endoftext|>